diff --git a/.buildkite/scripts/hardware_ci/run-xpu-test.sh b/.buildkite/scripts/hardware_ci/run-xpu-test.sh
index ab8c2fd5d..fa4f00e29 100755
--- a/.buildkite/scripts/hardware_ci/run-xpu-test.sh
+++ b/.buildkite/scripts/hardware_ci/run-xpu-test.sh
@@ -52,6 +52,7 @@ time timeout -k 30 30m docker run \
     pip install tblib==3.1.0
     cd /workspace/vllm-omni
     pytest -v -s -m "core_model and xpu and B60"
+    pytest -v -s tests/diffusion/quantization/test_mxfp8_config.py
     export VLLM_XPU_USE_SAMPLER_KERNEL=0    # NOTE: Remove this after vLLM v0.21.1 is merged. Fixes Qwen2-5 omni-expansion tests.
     pytest -v -s -m "advanced_model and xpu and B60"
 '
diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml
index 9e5027272..318b872b5 100644
--- a/.buildkite/test-merge.yml
+++ b/.buildkite/test-merge.yml
@@ -248,6 +248,47 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: "TTS · GLM-TTS Test"
+        timeout_in_minutes: 20
+        commands:
+          - |
+            timeout 20m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              pytest -s -v tests/e2e/online_serving/test_glm_tts.py tests/e2e/offline_inference/test_glm_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+            "
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 1
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
 
       - label: "Omni · Qwen3-Omni Test"
         timeout_in_minutes: 30
@@ -348,6 +389,31 @@ steps:
               volumes:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
+      - label: "Omni · Higgs-Audio V2 Test"
+        timeout_in_minutes: 30
+        commands:
+          - |
+            timeout 30m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              export VLLM_USE_DEEP_GEMM=0
+              export VLLM_MOE_USE_DEEP_GEMM=0
+              pytest -s -v tests/e2e/online_serving/test_higgs_audio_v2.py tests/e2e/offline_inference/test_higgs_audio_v2.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+            "
+        agents:
+          queue: "gpu_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
       - label: "Diffusion · Qwen Image Edit Test"
         timeout_in_minutes: 20
         commands:
diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml
index ae93d2353..d23080cbe 100644
--- a/.buildkite/test-nightly.yml
+++ b/.buildkite/test-nightly.yml
@@ -305,6 +305,51 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: ":full_moon: Omni · Perf Test · Multi-Replica"
+        key: nightly-omni-performance-multi-replicas
+        timeout_in_minutes: 180
+        commands:
+          - export BENCHMARK_DIR=tests/dfx/perf/results
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
+            exit $$EXIT
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 3
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
 
       - label: ":full_moon: Omni · Multi-Replica Startup Test with 4x H100"
         timeout_in_minutes: 45
@@ -414,6 +459,52 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: ":full_moon: VoxCPM2 · Perf Test"
+        key: nightly-voxcpm2-performance
+        timeout_in_minutes: 180
+        commands:
+          - export BENCHMARK_DIR=tests/dfx/perf/results
+          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_voxcpm2.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
+            exit $$EXIT
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 1
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
   # Diffusion X2I suite: x2i / x2a / x2t and related non-video paths; x2v is only in "Diffusion X2V Model Test" below.
   - group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
     key: nightly-diffusion-x2iat-group
@@ -525,6 +616,7 @@ steps:
             tests/e2e/online_serving/test_longcat_image_expansion.py
             tests/e2e/online_serving/test_longcat_image_edit_expansion.py
             tests/e2e/online_serving/test_flux_2_dev_expansion.py
+            tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py
             -m "full_model and diffusion and H100 and distributed_cuda"
             --run-level "full_model"
         agents:
@@ -578,7 +670,7 @@ steps:
                   - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
                     resources:
                       limits:
-                        nvidia.com/gpu: 3
+                        nvidia.com/gpu: 2
                     volumeMounts:
                       - name: devshm
                         mountPath: /dev/shm
@@ -797,8 +889,8 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
-      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2509"
-        key: nightly-diffusion-x2iat-performance-qwen-image-edit-2509
+      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
+        key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
         timeout_in_minutes: 180
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
@@ -806,7 +898,7 @@ steps:
           - 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_2509_vllm_omni.json
+            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
             EXIT=$$?
             buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
             buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
@@ -893,26 +985,20 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
-      - label: ":full_moon: Diffusion X2I(&A&T) · HunyuanImage3 · DiT Perf Test"
-        key: nightly-hunyuan-image3-performance
-        soft_fail: true
-        timeout_in_minutes: 120
-        if: build.env("RUN_HUNYUAN_IMAGE3_PERF") == "1"
+      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · BAGEL"
+        key: nightly-diffusion-x2iat-performance-bagel
+        timeout_in_minutes: 180
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
           - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
           - 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_hunyuan_image_tp4_fp8.json
-            EXIT1=$$?
-            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_sp2.json
-            EXIT2=$$?
-            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_cfgp2.json
-            EXIT3=$$?
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
+            EXIT=$$?
             buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
             buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
-            exit $$((EXIT1 | EXIT2 | EXIT3))
+            exit $$EXIT
         agents:
           queue: "mithril-h100-pool"
         plugins:
@@ -922,7 +1008,7 @@ steps:
                   - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
                     resources:
                       limits:
-                        nvidia.com/gpu: 4
+                        nvidia.com/gpu: 1
                     volumeMounts:
                       - name: devshm
                         mountPath: /dev/shm
@@ -998,6 +1084,8 @@ steps:
         timeout_in_minutes: 180
         commands:
           - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m full_model --run-level full_model
+          - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model
+          - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
           - pytest -s -v tests/e2e/accuracy/test_ltx2_3_video_similarity.py -m full_model --run-level full_model
           - pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -m full_model --run-level full_model
         agents:
@@ -1086,9 +1174,8 @@ steps:
     if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "diffusion-x2iat-test"
     steps:
       - label: ":full_moon: Quantization Test with H100"
-        timeout_in_minutes: 90
         commands:
-          - timeout 30m pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model"
+          - timeout 60m pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model"
         agents:
           queue: "mithril-h100-pool"
         plugins:
@@ -1189,14 +1276,16 @@ steps:
       - nightly-omni-performance-no-async-chunk
       - nightly-omni-performance-async-chunk
       - nightly-omni-performance-vllm-text
+      - nightly-omni-performance-multi-replicas
       - nightly-tts-performance
       - nightly-diffusion-x2iat-performance-qwen-image
       - nightly-diffusion-x2iat-performance-qwen-image-edit
-      - nightly-diffusion-x2iat-performance-qwen-image-edit-2509
+      - nightly-diffusion-x2iat-performance-qwen-image-edit-2511
       - nightly-diffusion-x2iat-performance-qwen-image-layered
+      - nightly-diffusion-x2iat-performance-bagel
       - nightly-diffusion-x2v-performance
       - nightly-testcase-statistics
-    if: build.env("NIGHTLY") == "1"
+    if: build.env("NIGHTLY") == "1" && build.env("EMAIL_DISTRIBUTION") == "1"
     commands:
       - pip install openpyxl
       - export DEFAULT_INPUT_DIR=tests/dfx/perf/results
@@ -1205,10 +1294,12 @@ steps:
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-vllm-text
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-multi-replicas
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit
-      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2509
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2511
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-layered
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-bagel
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance
       - buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-testcase-statistics
       - python tools/nightly/generate_nightly_perf_excel.py
diff --git a/.claude/skills/diffusion-perf-opt/SKILL.md b/.claude/skills/diffusion-perf-opt/SKILL.md
index 559134935..0d1b714fd 100644
--- a/.claude/skills/diffusion-perf-opt/SKILL.md
+++ b/.claude/skills/diffusion-perf-opt/SKILL.md
@@ -21,6 +21,9 @@ Before proposing changes, ask for the optimization scene if it is not already kn
 - Current enabled strategies: USP/SP, CFG parallel, HSDP/FSDP, VAE patch parallel, torch.compile, profiler options.
 - Optimization target: latency, throughput, memory, cost, or quality-preserving speed.
 - Precision/quality tolerance: bf16/fp8/quantization/sparsity/approximate attention allowed or not.
+- Quality validation method: exact-output comparison when feasible, or metrics
+  such as SSIM, PSNR, LPIPS, cosine similarity, MAE, MSE, temporal flicker
+  checks, baseline self-run variance, and reviewable output artifacts.
 
 ## Workflow
 
@@ -79,7 +82,19 @@ Before proposing changes, ask for the optimization scene if it is not already kn
    - Record latency, stage timings, memory, output quality, and logs.
    - Report comparison tables in ms. Include at least end-to-end client time, server `inference_time`, server stage generation time if available, `vae.encode`, `diffuse`, `vae.decode`, and peak memory.
 
-6. **Collect diagnostic trace only after narrowing hypotheses.**
+6. **Enforce the quality and precision gate for every optimization.**
+   - Do not mark an optimization ready until quality is checked on the same
+     workload: model, prompt/media, seed, shape, steps, scheduler, dtype,
+     backend, parallelism, and output encoding.
+   - For math-preserving changes, require exact/near-exact agreement or show
+     that differences are within baseline self-run variance.
+   - For precision, quantization, approximate attention, sparse/custom kernels,
+     or backend changes, include reviewable artifacts and metrics such as SSIM,
+     PSNR, LPIPS, cosine similarity, MAE, MSE, and temporal flicker checks.
+   - If quality tolerance is not stated, default to quality-preserving behavior.
+     Failed or inconclusive quality validation blocks a ready-to-merge claim.
+
+7. **Collect diagnostic trace only after narrowing hypotheses.**
    - Use torch profiler for a small number of requests.
    - Run two separate diagnostic traces instead of mixing concerns:
      - **Operator/shape trace**: enable `torch_profiler_record_shapes=True` and keep stack collection disabled. Use this to rank CUDA kernels, NCCL collectives, attention/MLP/norm/RoPE work, and shape-specific hot operators.
@@ -91,15 +106,16 @@ Before proposing changes, ask for the optimization scene if it is not already kn
    - Diagnostic reports must be written to disk and preserve: server command, profiler config, warmup/request/polling commands, trace artifact paths, rank analyzed, analyzer output or summary, and the decision about whether additional ranks are necessary.
    - Analyze both rank-level balance and device-level free bubbles when additional ranks are opened.
 
-7. **Analyze host, communication, and operators.**
+8. **Analyze host, communication, and operators.**
    - Find GPU idle/free intervals and map each large gap to the enclosing CPU/Python code.
    - Separate real GPU idle from profiler overhead such as CUPTI `Command Buffer Full`.
    - Compare NCCL kernel time to user annotations; annotations can overcount nested intervals.
    - Rank operator work by total CUDA time and by repeated small-kernel launch count.
 
-8. **Produce an optimization plan.**
+9. **Produce an optimization plan.**
    - Classify candidates as P0/P1/P2.
    - For each candidate, state necessity, expected benefit, implementation path, validation plan, and quality risk.
+   - Include a concrete quality/precision validation gate for each candidate.
    - Do not implement high-risk operator rewrites before proving the operator is a bottleneck for the target shapes.
    - End the plan with a user-facing candidate selection table. The assistant
      should not automatically choose a risky optimization just because it is
@@ -124,7 +140,13 @@ Before proposing changes, ask for the optimization scene if it is not already kn
 - **P1:** meaningful code changes with contained risk. Examples: cross-attention KV caching, VAE gather/broadcast reduction, AdaLayerNorm/RMSNorm/RoPE fusion after trace evidence.
 - **P2:** high implementation or quality risk. Examples: FA to LA replacement, custom Triton/CUDA fused kernels, FP8/quantization, sparsity/Rainfusion-style acceleration.
 
-Every implemented optimization needs A/B validation and quality regression. A/B means same workload and hardware before/after. Quality regression means checking generated image/video stability, artifacts, temporal flicker, and seed behavior when precision or approximate kernels change.
+Every implemented optimization needs A/B validation and a passing
+quality/precision gate, including math-preserving changes such as padding trim,
+layout cleanup, cache reuse, or host/runtime cleanup. Use objective metrics
+when available, keep reviewable artifacts, and compare against baseline
+self-run variance for precision, quantization, approximate kernels, or backend
+changes. If quality validation fails or is inconclusive, do not present the
+optimization as ready to merge.
 
 ## Optimization Layers
 
diff --git a/.readthedocs.yml b/.readthedocs.yml
index 5a06c6634..2b5b81bbc 100644
--- a/.readthedocs.yml
+++ b/.readthedocs.yml
@@ -10,6 +10,11 @@ build:
   jobs:
     post_checkout:
       - git fetch --unshallow || true
+    pre_install:
+      # Pre-install CPU torch so the subsequent `pip install .[docs]` does not
+      # pull ~2 GB of CUDA 13 libs (cuBLAS/cuDNN/NCCL/...) that exceed RTD's
+      # 1 GB memory cap and time out the mkdocs build.
+      - python -m pip install --index-url https://download.pytorch.org/whl/cpu torch
 
 mkdocs:
   configuration: mkdocs.yml
diff --git a/benchmarks/README.md b/benchmarks/README.md
index 68ffd40ef..37a6264a9 100644
--- a/benchmarks/README.md
+++ b/benchmarks/README.md
@@ -1,44 +1,60 @@
-# Benchmarks Overview and Architecture
+# Benchmarks
 
-This document explains the benchmark architecture across all benchmark assets in this repo. It describes what we measure, and where to find or plug in new scenarios. Per-task details remain in subfolder READMEs (e.g., `benchmarks/<model>/README.md`).
+This directory contains benchmark suites for evaluating different model families and infrastructure components in vLLM-Omni. Each subfolder targets a different benchmark family with its own scripts, configs, and metrics. See the per-subfolder READMEs for detailed usage.
 
-## Scope and goals
-- Establish repeatable latency/throughput measurements for multimodal LLM pipelines.
-- Provide both HF Transformers (offline) and vLLM-Omni (multi-stage/pipeline) baselines.
-- Make it easy to plug in new datasets and models with minimal changes to the runner scripts.
+## Benchmark families
 
-## Dataset and inputs
-- Default example: SeedTTS top-100 prompts (`benchmarks/build_dataset/top100.txt`) via `benchmarks/build_dataset/`.
-- Extensible: drop in new prompt files or modality-aligned payloads; keep the expected format for the consuming scripts (e.g., one prompt per line).
-- If you add a new dataset, document it under `benchmarks/<model>/README.md` and point scripts to your data path.
+### [TTS](tts/README.md) — Text-to-Speech
 
-## Directory layout
-- `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
-- `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
-- `benchmarks/accuracy/` — accuracy benchmark integrations that adapt external
-  benchmark suites to vLLM-Omni serving and evaluation flows.
-- Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
+Model-agnostic serving benchmarks for TTS models, including Qwen3-TTS and VoxCPM2.
 
-## Reference workflows
-- **HF Transformers (offline, single process)**  
-  Script (example): `benchmarks/<model>/transformers/eval_qwen3_moe_omni_transformers.sh`  
-  Outputs: `benchmark_results/perf_stats.json`, `benchmark_results/results.json`, `benchmark_results/audio/` (if audio is produced).
+- **Layout**: `tts/bench_tts.py` (serving benchmark driver), `tts/model_configs.yaml` (model registry), `tts/plot_results.py` (result plotting)
+- **Dataset**: Seed-TTS full or text-only datasets, plus bundled smoke/design prompts under `build_dataset/`
+- **Key metrics**: TTFP (time to first audio packet), E2E latency, RTF (real-time factor), throughput (audio seconds / wall-clock second)
 
-- **vLLM-Omni end-to-end pipeline**  
-  Script (example): `benchmarks/<model>/vllm_omni/eval_qwen3_moe_omni.sh`  
-  Outputs: `vllm_omni/logs/*.stats.jsonl` (per-stage/overall latency & TPS), `vllm_omni/logs/stage*.log`, `vllm_omni/outputs/` (text/audio artifacts).
+### [Diffusion](diffusion/README.md) — Image and Video Generation
 
-- **Adding a new task/model**  
-  1) Create `benchmarks/<model>/transformers/` and/or `benchmarks/<model>/vllm_omni/` with scripts referencing your model and dataset.  
-  2) Add a task README describing dataset, configs, and expected outputs.  
-  3) Keep the output/log structure similar for easy comparison (perf_stats/results/audio or text outputs; stats.jsonl/logs for pipeline).
+Online-serving benchmark for diffusion image/video models, sending requests to the configured vLLM serving endpoint (`/v1/chat/completions`, `/v1/images/generations`, or `/v1/videos`, depending on backend/task).
 
-## Metrics to watch
-- **Throughput**: `overall_tps`, `*_tps_avg` per stage.
-- **Latency distribution**: look for long tails in `*.stats.jsonl`.
-- **Quality/completeness**: missing outputs or errors in stage logs indicate pipeline failures or misconfigurations.
+- **Tasks**: text-to-image, text-to-video, image-to-image, image-to-video, text+image-to-image, text+image-to-video
+- **Datasets**: `vbench`, `trace`, `random`
+- **Key metrics**: request throughput, latency percentiles, optional SLO attainment
 
-## Troubleshooting
-- Verify GPU/driver/FlashAttention2 requirements for your chosen model/config.
-- Ensure network access for dataset/model downloads (Google Drive, Hugging Face, etc.).
-- If outputs are missing or slow, inspect per-stage logs and `*.stats.jsonl` for errors, stragglers, or contention.
+### [GLM-Image](glm_image/README.md) — Text-to-Image and Image-to-Image
+
+Benchmarks for GLM-Image performance across HuggingFace baseline, vLLM-Omni offline inference, and vLLM-Omni online serving.
+
+- **Layout**: `glm_image/huggingface/` (HF baseline), `glm_image/vllm-omni/` (offline inference), `glm_image/benchmark_glm_image.py` (online serving)
+- **Tasks**: text-to-image and image-to-image
+- **Key metrics**: request/image throughput, latency percentiles, optional per-stage pipeline timings
+
+### [Distributed](distributed/omni_connectors/README.md) — RDMA Connector Testing
+
+RDMA environment setup and transfer tests for `MooncakeTransferEngineConnector`, including pytest-based single-node checks and manual cross-node benchmarks.
+
+- **Transfer modes**: `copy`, `zerocopy`, `gpu` (GPUDirect)
+- **Supports**: single-node pytest suites and manual multi-node/cross-node transfer testing
+
+### [Accuracy](accuracy/README.md) — Image Generation and Editing Quality
+
+Accuracy benchmarks for image generation/editing models, adapting external suites to vLLM-Omni serving and local judge-evaluation flows.
+
+- **Layout**: `accuracy/text_to_image/` (GEBench), `accuracy/image_to_image/` (GEdit-Bench)
+- **Method**: generation and judge scoring both run through local `vllm-omni serve` endpoints
+
+### Common serving metrics framework
+
+`vllm_omni/benchmarks/` extends `vllm bench serve --omni` with Omni-specific datasets, backends, and multimodal metrics. Key metrics include:
+
+- **Text output**: TTFT (time to first token), TPOT (time per output token), ITL (inter-token latency)
+- **Audio output**: TTFP (time to first audio packet), E2E latency, RTF (real-time factor)
+- **Throughput**: request throughput, output token throughput, total token throughput, audio throughput
+
+See `vllm_omni/benchmarks/serve.py` for the `vllm bench serve --omni` runner wrapper and `vllm_omni/benchmarks/metrics/` for Omni metric definitions.
+
+## Adding a new benchmark
+
+1. Create a subfolder under `benchmarks/<name>/` with scripts, configs if needed, and a `README.md`.
+2. If comparing against another runtime, use clear backend subfolders where applicable, such as `huggingface/` and `vllm-omni/`, or follow the shared TTS serving benchmark pattern in `tts/`.
+3. Add reusable dataset or prompt-building utilities to `build_dataset/` if applicable.
+4. Update this README with a link to the new benchmark family.
diff --git a/benchmarks/diffusion/bench_attention_backends.py b/benchmarks/diffusion/bench_attention_backends.py
new file mode 100644
index 000000000..0d34b2192
--- /dev/null
+++ b/benchmarks/diffusion/bench_attention_backends.py
@@ -0,0 +1,328 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Standalone diagnostic for the diffusion attention backends added in #3079.
+
+Exercises the same synthetic attention shape through each backend and each
+torch SDPA sub-kernel, with and without an attn_mask, so we can pinpoint why
+CUDNN_ATTN underperforms on some SKUs. Two typical causes:
+
+  (a) cuDNN lacks a tuned kernel for this (SM, head_dim, seq) combination
+      and silently falls back to MATH.
+  (b) attn_mask is non-None and the cuDNN SDPA dispatch rejects it,
+      walking CUDNN -> FLASH -> MATH.
+
+Run:
+    python benchmarks/diffusion/bench_attention_backends.py --preset hv15
+    python benchmarks/diffusion/bench_attention_backends.py --preset wan22
+    python benchmarks/diffusion/bench_attention_backends.py \
+        --batch 1 --heads 24 --seq 14336 --head-dim 128
+
+Optional extras:
+    pip install --pre flash-attn-4        # FA4 is currently pre-release only
+    pip install -U flashinfer              # latest FlashInfer (0.6.9)
+
+The table at the end is the data we want on the PR — surface the row where a
+backend is >1.5x the SDPA baseline, that's the one to gate off the auto-route.
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+
+import torch
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+# Presets mirror the DiT hot-path attention call (single-stream block) at the
+# shapes we run in the PR's validation configs.
+_PRESETS = {
+    # HunyuanVideo-1.5 480p/33f: (30 * 52 * 9) latent + ~256 text tokens.
+    "hv15": {"batch": 1, "heads": 24, "seq": 14336, "head_dim": 128},
+    # Wan 2.2 480p/33f rough estimate; adjust when we measure real shapes.
+    "wan22": {"batch": 1, "heads": 40, "seq": 16384, "head_dim": 128},
+    # Smaller image-gen shape for quick smoke tests.
+    "flux": {"batch": 1, "heads": 24, "seq": 4096, "head_dim": 128},
+}
+
+_SDPA_BACKENDS = [
+    ("CUDNN_ATTENTION", [SDPBackend.CUDNN_ATTENTION]),
+    ("FLASH_ATTENTION", [SDPBackend.FLASH_ATTENTION]),
+    ("EFFICIENT_ATTENTION", [SDPBackend.EFFICIENT_ATTENTION]),
+    ("MATH", [SDPBackend.MATH]),
+    # The PR's CUDNN_ATTN impl uses this priority chain; keep it here to show
+    # what torch actually picks when multiple are allowed.
+    ("CUDNN_ATTN_CHAIN", [SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]),
+]
+
+
+def _env_report() -> None:
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA not available; this diagnostic needs a GPU.")
+    cc = torch.cuda.get_device_capability()
+    print("=" * 72)
+    print(f"GPU           : {torch.cuda.get_device_name()}")
+    print(f"SM capability : sm_{cc[0]}{cc[1]}")
+    print(f"torch         : {torch.__version__}")
+    print(f"cuDNN version : {torch.backends.cudnn.version()}")
+    print("flashinfer    : ", end="")
+    try:
+        import flashinfer
+
+        print(getattr(flashinfer, "__version__", "present"))
+    except Exception as e:
+        print(f"not installed ({type(e).__name__})")
+    print("=" * 72)
+
+
+def _make_qkv(batch: int, heads: int, seq: int, head_dim: int, device: str, dtype: torch.dtype):
+    # Layout: (B, S, H, D) — what the PR's backends accept as input.
+    q = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
+    k = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
+    v = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
+    return q, k, v
+
+
+def _make_mask(batch: int, seq: int, device: str, dtype: torch.dtype, pad_tokens: int = 128) -> torch.Tensor:
+    # Mirrors a text-encoder padding mask: last `pad_tokens` positions masked.
+    mask = torch.zeros(batch, 1, seq, seq, device=device, dtype=dtype)
+    if pad_tokens > 0:
+        mask[..., -pad_tokens:] = float("-inf")
+    return mask
+
+
+def _time_call(fn, *args, warmup: int = 3, iters: int = 10) -> tuple[float, str]:
+    """Return (median ms, error string). Skips on any exception so that a
+    backend that rejects our args (wrong dtype, missing JIT module, unsupported
+    kwarg value) doesn't abort the whole sweep."""
+    try:
+        for _ in range(warmup):
+            fn(*args)
+        torch.accelerator.synchronize()
+        times = []
+        for _ in range(iters):
+            torch.accelerator.synchronize()
+            t0 = time.perf_counter()
+            fn(*args)
+            torch.accelerator.synchronize()
+            times.append((time.perf_counter() - t0) * 1000.0)
+        times.sort()
+        return times[len(times) // 2], ""
+    except Exception as e:  # noqa: BLE001 — probe script, keep the whole table going
+        msg = str(e).split("\n", 1)[0][:60]
+        return float("nan"), f"{type(e).__name__}: {msg}" if msg else type(e).__name__
+
+
+def _run_sdpa_variants(q, k, v, attn_mask, scale: float) -> list[tuple[str, float, str]]:
+    # q/k/v come in (B, S, H, D); F.sdpa wants (B, H, S, D).
+    q_t = q.permute(0, 2, 1, 3)
+    k_t = k.permute(0, 2, 1, 3)
+    v_t = v.permute(0, 2, 1, 3)
+    rows: list[tuple[str, float, str]] = []
+    for name, backends in _SDPA_BACKENDS:
+
+        def _call(mask=attn_mask):
+            with sdpa_kernel(backends):
+                return torch.nn.functional.scaled_dot_product_attention(
+                    q_t, k_t, v_t, attn_mask=mask, dropout_p=0.0, is_causal=False, scale=scale
+                )
+
+        ms, err = _time_call(_call)
+        rows.append((name, ms, err))
+    return rows
+
+
+def _run_flashinfer(q, k, v, scale: float, backend: str | None = None, attn_mask=None) -> tuple[float, str]:
+    """Call FlashInfer's dense single-prefill.
+
+    ``backend`` hints at cutlass/fa3/trtllm if the installed version exposes
+    that kwarg. ``attn_mask`` is the same additive 4D mask we pass to SDPA —
+    we convert it to the 2D boolean ``custom_mask`` form FlashInfer accepts.
+    Per FlashInfer docs, ``custom_mask`` only applies when ``causal=False``.
+    """
+    try:
+        import inspect
+
+        from flashinfer.prefill import single_prefill_with_kv_cache
+    except Exception as e:
+        return float("nan"), f"import-{type(e).__name__}"
+
+    kwargs: dict = {"sm_scale": scale, "causal": False, "return_lse": False}
+    if backend is not None:
+        sig = inspect.signature(single_prefill_with_kv_cache)
+        if "backend" not in sig.parameters:
+            return float("nan"), "no-backend-kwarg"
+        kwargs["backend"] = backend
+
+    if attn_mask is not None:
+        # Collapse (B, 1, S, S) additive float mask to (S, S) boolean.
+        # FlashInfer expects True = keep, False = masked out.
+        mask_2d = attn_mask[0, 0]
+        kwargs["custom_mask"] = mask_2d != float("-inf")
+
+    def _call():
+        out = single_prefill_with_kv_cache(q[0], k[0], v[0], **kwargs)
+        return out.unsqueeze(0)
+
+    return _time_call(_call)
+
+
+def _run_fa4(q, k, v, scale: float) -> tuple[float, str]:
+    """Call FlashAttention-4 directly (``pip install flash-attn-4``). FA4
+    ships a Blackwell-native kernel via CuTe-DSL; on sm_120 it should beat
+    cuDNN by ~20%. API lives under ``flash_attn.cute``; the older
+    ``flash_attn.flash_attn_func`` path is FA2/FA3 only."""
+    try:
+        from flash_attn.cute import flash_attn_func
+    except Exception as e:
+        return float("nan"), f"import-{type(e).__name__}"
+
+    def _call():
+        # FA4 accepts (B, S, H, D) directly — same layout the PR's backends use.
+        return flash_attn_func(q, k, v, softmax_scale=scale, causal=False)
+
+    return _time_call(_call)
+
+
+def _run_flashinfer_cudnn_batch(q, k, v, scale: float) -> tuple[float, str]:
+    """Call FlashInfer's direct cuDNN wrapper. Bypasses PyTorch SDPA dispatch
+    (which has a few hundred ns of overhead per call) and talks to cuDNN FMHA
+    straight. Useful as a ceiling for 'pure cuDNN, no SDPA' on Blackwell."""
+    try:
+        from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
+    except Exception as e:
+        return float("nan"), f"import-{type(e).__name__}"
+
+    b, s, h, d = q.shape
+    qo_indptr = torch.tensor([0, s], dtype=torch.int32, device=q.device)
+    kv_indptr = torch.tensor([0, s], dtype=torch.int32, device=q.device)
+
+    def _call():
+        return cudnn_batch_prefill_with_kv_cache(
+            q.reshape(b * s, h, d),
+            k.reshape(b * s, h, d),
+            v.reshape(b * s, h, d),
+            qo_indptr=qo_indptr,
+            kv_indptr=kv_indptr,
+            max_qo_len=s,
+            max_kv_len=s,
+            sm_scale=scale,
+            causal=False,
+        )
+
+    return _time_call(_call)
+
+
+def _print_table(title: str, rows: list[tuple[str, float, str]], baseline_name: str | None = None) -> None:
+    print(f"\n{title}")
+    print("-" * 88)
+    header = f"{'backend':<24} {'median (ms)':>14}    {'vs baseline':>12}    status"
+    print(header)
+    print("-" * 88)
+    baseline_ms = None
+    if baseline_name is not None:
+        for name, ms, _ in rows:
+            if name == baseline_name and ms == ms:  # not NaN
+                baseline_ms = ms
+                break
+    for name, ms, err in rows:
+        ms_str = f"{ms:>14.3f}" if ms == ms else f"{'n/a':>14}"
+        if baseline_ms is not None and ms == ms:
+            ratio = baseline_ms / ms
+            ratio_str = f"{ratio:>11.2f}x"
+        else:
+            ratio_str = f"{'—':>12}"
+        status = "ok" if not err else f"FAILED ({err})"
+        print(f"{name:<24} {ms_str}    {ratio_str}    {status}")
+
+
+def _pick_winner(rows: list[tuple[str, float, str]]) -> tuple[str, float] | None:
+    """Return (backend, ms) of the fastest non-failing row, or None."""
+    ok_rows = [(n, m) for n, m, err in rows if not err and m == m]
+    if not ok_rows:
+        return None
+    return min(ok_rows, key=lambda x: x[1])
+
+
+def _bench_one_shape(shape: dict, dtype: torch.dtype, device: str) -> tuple[list, list]:
+    torch.manual_seed(0)
+    q, k, v = _make_qkv(**shape, device=device, dtype=dtype)
+    scale = 1.0 / (shape["head_dim"] ** 0.5)
+
+    rows_nomask = _run_sdpa_variants(q, k, v, attn_mask=None, scale=scale)
+    rows_nomask.append(("FLASHINFER (default)", *_run_flashinfer(q, k, v, scale)))
+    # `trtllm-gen` has no sm_120 cubins (NVIDIA/TensorRT-LLM#11799) — skipped.
+    for fi_backend in ("fa2", "fa3", "cutlass", "auto"):
+        rows_nomask.append((f"FLASHINFER ({fi_backend})", *_run_flashinfer(q, k, v, scale, backend=fi_backend)))
+    rows_nomask.append(("FLASHINFER (cudnn-batch)", *_run_flashinfer_cudnn_batch(q, k, v, scale)))
+    rows_nomask.append(("FA4 (direct)", *_run_fa4(q, k, v, scale)))
+
+    mask = _make_mask(shape["batch"], shape["seq"], device=device, dtype=dtype, pad_tokens=256)
+    rows_mask = _run_sdpa_variants(q, k, v, attn_mask=mask, scale=scale)
+    rows_mask.append(("FLASHINFER (dense)", *_run_flashinfer(q, k, v, scale, attn_mask=mask)))
+    rows_mask.append(("FA4 (direct)", float("nan"), "mask-not-supported"))
+
+    return rows_nomask, rows_mask
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--preset", choices=list(_PRESETS.keys()), default="hv15")
+    parser.add_argument("--sweep", action="store_true", help="Run all presets and print a ranking")
+    parser.add_argument("--batch", type=int, default=None)
+    parser.add_argument("--heads", type=int, default=None)
+    parser.add_argument("--seq", type=int, default=None)
+    parser.add_argument("--head-dim", type=int, default=None)
+    parser.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16")
+    parser.add_argument("--device", default="cuda:0")
+    args = parser.parse_args()
+
+    dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
+    _env_report()
+    print(f"dtype         : {args.dtype}")
+
+    presets = list(_PRESETS.keys()) if args.sweep else [args.preset]
+    summary: dict[str, tuple[str, float] | None] = {}
+
+    for preset_name in presets:
+        shape = dict(_PRESETS[preset_name])
+        if not args.sweep:
+            for k in ("batch", "heads", "seq", "head_dim"):
+                v = getattr(args, k if k != "head_dim" else "head_dim")
+                if v is not None:
+                    shape[k] = v
+
+        print("\n" + "=" * 88)
+        print(
+            f"Preset: {preset_name}  |  batch={shape['batch']} heads={shape['heads']} "
+            f"seq={shape['seq']} head_dim={shape['head_dim']}"
+        )
+        print("=" * 88)
+
+        rows_nomask, rows_mask = _bench_one_shape(shape, dtype, args.device)
+        _print_table(f"[{preset_name}] No attention mask", rows_nomask, baseline_name="CUDNN_ATTENTION")
+        _print_table(
+            f"[{preset_name}] With attention mask (pad 256 tokens)", rows_mask, baseline_name="CUDNN_ATTENTION"
+        )
+
+        summary[preset_name] = _pick_winner(rows_nomask)
+
+    if args.sweep or len(presets) > 1:
+        print("\n" + "=" * 88)
+        print("Winners per preset (no-mask path)")
+        print("=" * 88)
+        for preset_name, winner in summary.items():
+            if winner is None:
+                print(f"  {preset_name:<10} — no successful backend")
+            else:
+                name, ms = winner
+                print(f"  {preset_name:<10} {name:<24} {ms:>8.3f} ms")
+
+    print("\nNotes:")
+    print("  * Ratios are relative to CUDNN_ATTENTION. >1.0x means faster than cuDNN.")
+    print("  * Mask-path winner inherits CUDNN_ATTN's fallback in the PR's backends.")
+    print("  * `trtllm-gen` and FA4 4.0.0b10 are known-broken on sm_120 as of Apr 2026.")
+    print("  * For e2e timings, run text_to_video.py / text_to_image.py while varying DIFFUSION_ATTENTION_BACKEND.")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/diffusion/bench_sage_attn3_matrix.sh b/benchmarks/diffusion/bench_sage_attn3_matrix.sh
new file mode 100644
index 000000000..f7a3f4816
--- /dev/null
+++ b/benchmarks/diffusion/bench_sage_attn3_matrix.sh
@@ -0,0 +1,108 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# 2x2 timing matrix (backend x compile mode) for SAGE_ATTN_3 vs TORCH_SDPA,
+# on HunyuanVideo-1.5 480p T2V and Wan 2.2 TI2V-5B. Run from the repo root:
+#
+#     bash benchmarks/diffusion/bench_sage_attn3_matrix.sh
+#
+# Requires: sageattn3 installed (for SAGE_ATTN_3 rows) and a Blackwell GPU.
+
+set -u
+
+STEPS=${STEPS:-50}
+SEED=${SEED:-42}
+PROMPT=${PROMPT:-"A dog running across a field of golden wheat."}
+OUT_DIR=${OUT_DIR:-/tmp/bench_sage3}
+mkdir -p "$OUT_DIR"
+
+declare -A TOTAL PER_STEP TEXT_ENC TRANSFORMER VAE
+
+run() {
+    local name="$1"; shift
+    local model="$1"; shift
+    local h="$1"; shift
+    local w="$1"; shift
+    local f="$1"; shift
+    local g="$1"; shift
+    local log="$OUT_DIR/${name}.log"
+    echo "[$(date +%T)] Running $name ..."
+    env "$@" python examples/offline_inference/text_to_video/text_to_video.py \
+        --model "$model" \
+        --prompt "$PROMPT" \
+        --height "$h" --width "$w" --num-frames "$f" \
+        --num-inference-steps "$STEPS" --seed "$SEED" --guidance-scale "$g" \
+        --enable-diffusion-pipeline-profiler \
+        --output "$OUT_DIR/${name}.mp4" > "$log" 2>&1 || {
+            echo "  FAILED — see $log"
+            TOTAL[$name]="FAIL"; PER_STEP[$name]="FAIL"
+            TEXT_ENC[$name]="-"; TRANSFORMER[$name]="-"; VAE[$name]="-"
+            return
+        }
+
+    local wait_ms
+    wait_ms=$(grep "add_req_and_wait=" "$log" | tail -1 | sed -nE 's/.*add_req_and_wait=([0-9.]+) ms.*/\1/p')
+    local total_s
+    total_s=$(grep -oE "Total generation time: [0-9.]+ seconds" "$log" | tail -1 | sed -nE 's/.*: ([0-9.]+) seconds.*/\1/p')
+    local text_enc_s
+    text_enc_s=$(grep -oE "text_encoder.forward took [0-9.]+s" "$log" | awk '{sub(/s$/,"",$NF); sum+=$NF} END {printf "%.3f", sum+0}')
+    local vae_s
+    vae_s=$(grep -oE "vae.decode took [0-9.]+s" "$log" | awk '{sub(/s$/,"",$NF); sum+=$NF} END {printf "%.3f", sum+0}')
+    local pipeline_s
+    pipeline_s=$(grep -oE "Pipeline\.forward took [0-9.]+s" "$log" | tail -1 | sed -nE 's/.*took ([0-9.]+)s.*/\1/p')
+    local transformer_s
+    transformer_s=$(awk "BEGIN {printf \"%.3f\", ${pipeline_s:-0} - ${vae_s:-0} - ${text_enc_s:-0}}")
+
+    TOTAL[$name]="${total_s:-0}"
+    PER_STEP[$name]=$(awk "BEGIN {printf \"%.3f\", ${wait_ms:-0} / 1000 / $STEPS}")
+    TEXT_ENC[$name]="$text_enc_s"
+    TRANSFORMER[$name]="$transformer_s"
+    VAE[$name]="$vae_s"
+
+    echo "  total=${TOTAL[$name]}s  per-step=${PER_STEP[$name]}s/it  transformer=${TRANSFORMER[$name]}s  vae=${VAE[$name]}s"
+}
+
+# ========== HunyuanVideo-1.5 480p T2V (480x832x33, CFG=6.0) ==========
+HV_MODEL=hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v
+run hv15_sdpa_eager     "$HV_MODEL" 480 832 33 6.0 TORCH_COMPILE_DISABLE=1
+run hv15_sdpa_compiled  "$HV_MODEL" 480 832 33 6.0
+run hv15_sage3_eager    "$HV_MODEL" 480 832 33 6.0 TORCH_COMPILE_DISABLE=1 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
+run hv15_sage3_compiled "$HV_MODEL" 480 832 33 6.0                         DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
+
+# ========== Wan 2.2 TI2V-5B (704x1280x33, CFG=5.0) ==========
+WAN_MODEL=Wan-AI/Wan2.2-TI2V-5B-Diffusers
+run wan22_sdpa_eager     "$WAN_MODEL" 704 1280 33 5.0 TORCH_COMPILE_DISABLE=1
+run wan22_sdpa_compiled  "$WAN_MODEL" 704 1280 33 5.0
+run wan22_sage3_eager    "$WAN_MODEL" 704 1280 33 5.0 TORCH_COMPILE_DISABLE=1 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
+run wan22_sage3_compiled "$WAN_MODEL" 704 1280 33 5.0                         DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
+
+print_table() {
+    local title="$1"; shift
+    echo
+    echo "$title (${STEPS} steps, seed=${SEED})"
+    echo
+    printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
+        "Config" "Total (s)" "Per-step (s/it)" "Text Enc (s)" "Transformer (s)" "VAE Decode (s)"
+    printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
+        "-----------------------" "---------" "---------------" "------------" "---------------" "--------------"
+    for row in "$@"; do
+        case "$row" in
+            *_sdpa_eager)     label="SDPA + Eager" ;;
+            *_sdpa_compiled)  label="SDPA + Compiled" ;;
+            *_sage3_eager)    label="SAGE_ATTN_3 + Eager" ;;
+            *_sage3_compiled) label="SAGE_ATTN_3 + Compiled" ;;
+        esac
+        printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
+            "$label" "${TOTAL[$row]}" "${PER_STEP[$row]}" "${TEXT_ENC[$row]}" "${TRANSFORMER[$row]}" "${VAE[$row]}"
+    done
+}
+
+print_table "HunyuanVideo-1.5 480p T2V, 832x480x33" \
+    hv15_sdpa_eager hv15_sdpa_compiled hv15_sage3_eager hv15_sage3_compiled
+
+print_table "Wan 2.2 TI2V-5B, 1280x704x33" \
+    wan22_sdpa_eager wan22_sdpa_compiled wan22_sage3_eager wan22_sage3_compiled
+
+echo
+echo "Logs: $OUT_DIR/*.log"
diff --git a/benchmarks/diffusion/diffusion_benchmark_serving.py b/benchmarks/diffusion/diffusion_benchmark_serving.py
index 7e8ea97e1..2f385b538 100644
--- a/benchmarks/diffusion/diffusion_benchmark_serving.py
+++ b/benchmarks/diffusion/diffusion_benchmark_serving.py
@@ -47,6 +47,13 @@ Usage:
         --endpoint /v1/images/edits --dataset random --task ti2i --num-prompts 10 \
         --bot-task think
 
+    python benchmarks/diffusion/diffusion_benchmark_serving.py \
+      --endpoint /v1/images/edits \
+      --dataset custom \
+      --dataset-path custom_requests.jsonl \
+      --task ti2i \
+      --bot-task think
+
     i2i (chat-based models such as Qwen-Image-Edit):
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
         --endpoint /v1/chat/completions --dataset vbench --task i2i --num-prompts 10
@@ -566,6 +573,135 @@ class TraceDataset(BaseDataset):
         return [self[i] for i in range(len(self))]
 
 
+class CustomDataset(BaseDataset):
+    """
+    Custom dataset that loads requests from a JSONL file.
+
+    Each line in the JSONL file should be a JSON object with the following fields:
+    - prompt (required): The text prompt for the request
+    - width (optional): Image/video width
+    - height (optional): Image height
+    - num_inference_steps (optional): Number of diffusion steps
+    - seed (optional): Random seed
+    - image_paths (optional): List of input image paths for i2i/i2v/ti2i tasks
+    - image_urls (optional): List of input image URLs (alternative to image_paths)
+
+    Example JSONL for ti2i:
+    {"prompt": "Add sunset lighting", "width": 1024, "height": 1024, "image_urls": ["https://example.com/image.jpg"]}
+    """
+
+    def __init__(self, args, api_url: str, model: str):
+        super().__init__(args, api_url, model)
+        self.dataset_path = args.dataset_path
+        if not self.dataset_path:
+            raise ValueError("--dataset-path must be provided when using 'custom' dataset")
+        self.load_data()
+
+    def load_data(self) -> None:
+        """Load data from JSONL file."""
+        import pandas as pd
+
+        if not self.dataset_path.endswith(".jsonl"):
+            raise ValueError("Custom dataset must be a JSONL file")
+
+        try:
+            df = pd.read_json(path_or_buf=self.dataset_path, lines=True)
+        except Exception as e:
+            raise ValueError(f"Failed to load JSONL file {self.dataset_path}: {e}")
+
+        if "prompt" not in df.columns:
+            raise ValueError("JSONL file must contain a 'prompt' column")
+
+        self.data = df.to_dict("records")
+        print(f"Loaded {len(self.data)} requests from {self.dataset_path}")
+
+    def _resolve_image_paths(self, item: dict) -> list[str] | None:
+        # Handle local file paths
+        if "image_paths" in item and item["image_paths"]:
+            paths = item["image_paths"]
+            if isinstance(paths, str):
+                paths = [paths]
+
+            # Verify all paths exist
+            valid_paths = []
+            for path in paths:
+                if os.path.exists(path):
+                    valid_paths.append(path)
+                else:
+                    raise ValueError(f"Image file not found: {path}")
+
+            return valid_paths if valid_paths else None
+
+        # Handle URLs - download to temp files
+        if "image_urls" in item and item["image_urls"]:
+            urls = item["image_urls"]
+            if isinstance(urls, str):
+                urls = [urls]
+
+            downloaded_paths = []
+            for url in urls:
+                try:
+                    response = requests.get(url, timeout=30)
+                    response.raise_for_status()
+
+                    # Create temp file with appropriate extension
+                    suffix = os.path.splitext(url)[1] or ".png"
+                    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
+                        tmp.write(response.content)
+                        downloaded_paths.append(tmp.name)
+                except Exception as e:
+                    raise ValueError(f"Failed to download image from {url}: {e}")
+
+            return downloaded_paths if downloaded_paths else None
+
+        return None
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def __getitem__(self, idx: int) -> RequestFuncInput:
+        item = self.data[idx]
+
+        # Extract fields with fallback to args defaults
+        prompt = item["prompt"]
+        width = item.get("width", self.args.width)
+        height = item.get("height", self.args.height)
+        num_inference_steps = item.get("num_inference_steps", self.args.num_inference_steps)
+        seed = item.get("seed", self.args.seed)
+        reserved_keys = {"prompt", "width", "height", "num_inference_steps", "seed", "image_paths", "image_urls"}
+        extra_body = {k: v for k, v in item.items() if k not in reserved_keys}
+
+        # Handle image paths/URLs
+        image_paths = self._resolve_image_paths(item)
+
+        return RequestFuncInput(
+            prompt=prompt,
+            api_url=self.api_url,
+            model=self.model,
+            seed=seed,
+            image_paths=image_paths,
+            extra_body=extra_body,
+            width=width,
+            height=height,
+            num_inference_steps=num_inference_steps,
+        )
+
+    def get_requests(self) -> list[RequestFuncInput]:
+        """Get all requests, cycling through data if num_prompts > dataset size."""
+        num_requests = getattr(self.args, "num_prompts", len(self.data))
+
+        if num_requests <= len(self.data):
+            return [self[i] for i in range(num_requests)]
+
+        # Cycle through the dataset to reach num_prompts
+        requests = []
+        for i in range(num_requests):
+            idx = i % len(self.data)
+            requests.append(self[idx])
+
+        return requests
+
+
 class RandomDataset(BaseDataset):
     def __init__(self, args, api_url: str, model: str, enable_negative_prompt: bool = False):
         super().__init__(args, api_url, model)
@@ -950,6 +1086,8 @@ async def benchmark(args):
         dataset = TraceDataset(args, api_url, args.model)
     elif args.dataset == "random":
         dataset = RandomDataset(args, api_url, args.model, args.enable_negative_prompt)
+    elif args.dataset == "custom":
+        dataset = CustomDataset(args, api_url, args.model)
     else:
         raise ValueError(f"Unknown dataset: {args.dataset}")
 
@@ -1097,7 +1235,7 @@ if __name__ == "__main__":
         "--dataset",
         type=str,
         default="vbench",
-        choices=["vbench", "trace", "random"],
+        choices=["vbench", "trace", "random", "custom"],
         help="Dataset to use.",
     )
     parser.add_argument(
diff --git a/benchmarks/tts/bench_tts.py b/benchmarks/tts/bench_tts.py
index 8d2eb99c8..78608e844 100644
--- a/benchmarks/tts/bench_tts.py
+++ b/benchmarks/tts/bench_tts.py
@@ -130,6 +130,9 @@ def build_bench_args(
     if task_extra_body:
         cmd += ["--extra-body", json.dumps(task_extra_body, separators=(",", ":"))]
 
+    if model_cfg.get("trust_remote_code"):
+        cmd.append("--trust-remote-code")
+
     if concurrency is not None:
         cmd += ["--max-concurrency", str(concurrency), "--request-rate", "inf"]
 
diff --git a/benchmarks/tts/model_configs.yaml b/benchmarks/tts/model_configs.yaml
index 83b253705..17254e37c 100644
--- a/benchmarks/tts/model_configs.yaml
+++ b/benchmarks/tts/model_configs.yaml
@@ -32,8 +32,11 @@ models:
         task_type: Base
 
   openbmb/VoxCPM2:
-    supported_tasks: [voice_clone]
+    supported_tasks: [voice_clone, default_voice]
+    trust_remote_code: true
     backend: openai-audio-speech
     endpoint: /v1/audio/speech
     task_extra_body:
       voice_clone: {}
+      default_voice:
+        voice: default
diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda
index 78f64f6a5..6c94689e6 100644
--- a/docker/Dockerfile.cuda
+++ b/docker/Dockerfile.cuda
@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai:v0.20.0
+ARG BASE_IMAGE=vllm/vllm-openai:v0.21.0
 FROM ${BASE_IMAGE}
 
 ARG COMMON_WORKDIR=/app
diff --git a/docs/assets/WeChat.jpg b/docs/assets/WeChat.jpg
index 18a0c53ce..c4e836f68 100644
Binary files a/docs/assets/WeChat.jpg and b/docs/assets/WeChat.jpg differ
diff --git a/docs/design/module/dit_module.md b/docs/design/module/dit_module.md
index 4390793e9..009df77ff 100644
--- a/docs/design/module/dit_module.md
+++ b/docs/design/module/dit_module.md
@@ -169,15 +169,13 @@ class DiffusionRequestStatus(enum.IntEnum):
 
 @dataclass
 class DiffusionRequestState:
-    sched_req_id: str
+    request_id: str
     req: OmniDiffusionRequest
     status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
 ```
 
 **Design Features**:
 
-- **Scheduler-owned ID**: Each `OmniDiffusionRequest` is tracked by an internal `sched_req_id`, separated from public `request_id` values.
-
 - **Explicit lifecycle**: Requests move through waiting, running, optional preemption, and terminal states.
 
 - **Centralized error handling**: Completion, abort, and error states are all normalized in the scheduler layer.
@@ -188,7 +186,6 @@ class DiffusionRequestState:
 class _BaseScheduler(SchedulerInterface):
     def __init__(self) -> None:
         self._request_states = {}
-        self._request_id_to_sched_req_id = {}
         self._waiting = deque()
         self._running = []
         self._finished_req_ids = set()
@@ -199,7 +196,8 @@ class _BaseScheduler(SchedulerInterface):
 
 - **Common state storage**: Shared request maps and waiting/running sets live in the base class.
 
-- **Shared cleanup logic**: Request-id registration, finish handling, and state removal are centralized instead of duplicated in each policy.
+- **Shared cleanup logic**: Duplicate-request checks, finish handling, and state removal are centralized instead of
+  duplicated in each policy.
 
 - **Current constraint boundary**: `_BaseScheduler` derives `max_num_running_reqs` from `max_num_seqs`, but request-mode diffusion is still clamped back to `1` by the engine. The step-wise path can keep this above `1` for compatible-request batching.
 
@@ -219,12 +217,12 @@ class RequestScheduler(_BaseScheduler):
 
 - **Single-request admission**: `RequestScheduler` still admits one active request at a time because request-mode execution completes a whole request per dispatch.
 
-- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished scheduler ids.
+- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished request ids.
 
 #### 2.5 Engine-Driven Execution Loop
 
 ```python
-sched_req_id = scheduler.add_request(request)
+request_id = scheduler.add_request(request)
 while True:
     sched_output = scheduler.schedule()
     output = executor.add_req(req)
diff --git a/docs/getting_started/installation/gpu/cuda.inc.md b/docs/getting_started/installation/gpu/cuda.inc.md
index 69fc05351..9f3c86843 100644
--- a/docs/getting_started/installation/gpu/cuda.inc.md
+++ b/docs/getting_started/installation/gpu/cuda.inc.md
@@ -103,10 +103,10 @@ docker run --runtime nvidia --gpus 2 \
     -p 8091:8091 \
     --ipc=host \
     vllm/vllm-omni:v0.21.0 \
-    --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
+    vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
 !!! tip
-    You can use this docker image to serve models the same way you would with in vLLM! To do so, make sure you overwrite the default entrypoint (`vllm serve --omni`) which works only for models supported in the vLLM-Omni project.
+    The CUDA image does not define a default entrypoint, so include `vllm serve ... --omni` after the image name.
 
 # --8<-- [end:pre-built-images]
diff --git a/docs/mkdocs/hooks/generate_examples.py b/docs/mkdocs/hooks/generate_examples.py
index b4c928936..81f2bac24 100644
--- a/docs/mkdocs/hooks/generate_examples.py
+++ b/docs/mkdocs/hooks/generate_examples.py
@@ -16,6 +16,7 @@ ROOT_DIR_RELATIVE = "../../../../.."
 EXAMPLE_DIR = ROOT_DIR / "examples"
 EXAMPLE_DOC_DIR = ROOT_DIR / "docs/user_guide/examples"
 NAV_FILE = ROOT_DIR / "docs/.nav.yml"
+MAX_INLINE_MATERIAL_SIZE = 8 * 1024
 
 
 def fix_case(text: str) -> str:
@@ -191,18 +192,21 @@ class Example:
 
         return re.sub(link_pattern, replace_link, content)
 
+    def github_url(self, path: Path) -> str:
+        url = "https://github.com/vllm-project/vllm-omni/"
+        url += "tree/main" if path.is_dir() else "blob/main"
+        return f"{url}/{path.relative_to(ROOT_DIR).as_posix()}"
+
     def generate(self) -> str:
         content = f"# {self.title}\n\n"
-        url = "https://github.com/vllm-project/vllm-omni/"
-        url += "tree/main" if self.path.is_dir() else "blob/main"
-        content += f"Source <{url}/{self.path.relative_to(ROOT_DIR)}>.\n\n"
+        content += f"Source <{self.github_url(self.path)}>.\n\n"
 
         # Use long code fence to avoid issues with
         # included files containing code fences too
         code_fence = "``````"
 
         if self.is_code:
-            main_file_rel = self.main_file.relative_to(ROOT_DIR)
+            main_file_rel = self.main_file.relative_to(ROOT_DIR).as_posix()
             content += f'{code_fence}{self.main_file.suffix[1:]}\n--8<-- "{main_file_rel}"\n{code_fence}\n'
         else:
             with open(self.main_file, encoding="utf-8") as f:
@@ -216,10 +220,15 @@ class Example:
 
         content += "## Example materials\n\n"
         for file in sorted(self.other_files):
-            content += f'??? abstract "{file.relative_to(self.path)}"\n'
+            content += f'??? abstract "{file.relative_to(self.path).as_posix()}"\n'
+            if file.stat().st_size > MAX_INLINE_MATERIAL_SIZE:
+                content += (
+                    f"    Large file omitted from the rendered docs. View it on GitHub: <{self.github_url(file)}>.\n\n"
+                )
+                continue
             if file.suffix != ".md":
                 content += f"    {code_fence}{file.suffix[1:]}\n"
-            content += f'    --8<-- "{file.relative_to(ROOT_DIR)}"\n'
+            content += f'    --8<-- "{file.relative_to(ROOT_DIR).as_posix()}"\n'
             if file.suffix != ".md":
                 content += f"    {code_fence}\n"
 
@@ -291,7 +300,7 @@ def update_nav_file(examples: list[Example]):
         for example in category_examples:
             doc_path = EXAMPLE_DOC_DIR / example.category / f"{example.path.stem}.md"
             rel_path = doc_path.relative_to(ROOT_DIR / "docs")
-            category_items.append({example.title: str(rel_path)})
+            category_items.append({example.title: rel_path.as_posix()})
 
         if category_items:
             # Format category name (e.g., "offline_inference" -> "Offline Inference")
diff --git a/docs/mkdocs/hooks/search_index.py b/docs/mkdocs/hooks/search_index.py
new file mode 100644
index 000000000..1b9e268d9
--- /dev/null
+++ b/docs/mkdocs/hooks/search_index.py
@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Reduce browser search-index memory for generated API reference pages."""
+
+from material.plugins.search.plugin import SearchIndex
+
+_add_entry_from_context = SearchIndex.add_entry_from_context
+
+
+def add_entry_from_context(self: SearchIndex, page) -> None:
+    search = page.meta.get("search") or {}
+    if search.get("exclude"):
+        return
+
+    if str(page.url).startswith("api/"):
+        entry = {
+            "location": page.url,
+            "title": str(page.meta.get("title", page.title)),
+            "text": "",
+        }
+        if "boost" in search:
+            entry["boost"] = search["boost"]
+        self.entries.append(entry)
+        return
+
+    _add_entry_from_context(self, page)
+
+
+SearchIndex.add_entry_from_context = add_entry_from_context
diff --git a/docs/mkdocs/javascript/mathjax.js b/docs/mkdocs/javascript/mathjax.js
index eb89ace06..2238e32e4 100644
--- a/docs/mkdocs/javascript/mathjax.js
+++ b/docs/mkdocs/javascript/mathjax.js
@@ -1,20 +1,44 @@
-// Enables MathJax rendering
+// Enables MathJax rendering only on pages that contain math.
 window.MathJax = {
-    tex: {
-      inlineMath: [["\\(", "\\)"]],
-      displayMath: [["\\[", "\\]"]],
-      processEscapes: true,
-      processEnvironments: true
-    },
-    options: {
-      ignoreHtmlClass: ".*|",
-      processHtmlClass: "arithmatex"
-    }
-  };
+  tex: {
+    inlineMath: [["\\(", "\\)"]],
+    displayMath: [["\\[", "\\]"]],
+    processEscapes: true,
+    processEnvironments: true,
+  },
+  options: {
+    ignoreHtmlClass: ".*|",
+    processHtmlClass: "arithmatex",
+  },
+};
 
-  document$.subscribe(() => {
-    MathJax.startup.output.clearCache()
-    MathJax.typesetClear()
-    MathJax.texReset()
-    MathJax.typesetPromise()
-  })
+let mathJaxLoading;
+
+function loadMathJax() {
+  if (window.MathJax.typesetPromise) {
+    return Promise.resolve();
+  }
+  if (!mathJaxLoading) {
+    mathJaxLoading = new Promise((resolve, reject) => {
+      const script = document.createElement("script");
+      script.src = "https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js";
+      script.async = true;
+      script.onload = resolve;
+      script.onerror = reject;
+      document.head.appendChild(script);
+    });
+  }
+  return mathJaxLoading;
+}
+
+document$.subscribe(() => {
+  if (!document.querySelector(".arithmatex")) {
+    return;
+  }
+  loadMathJax().then(() => {
+    MathJax.startup.output.clearCache();
+    MathJax.typesetClear();
+    MathJax.texReset();
+    MathJax.typesetPromise();
+  });
+});
diff --git a/docs/mkdocs/javascript/memory.js b/docs/mkdocs/javascript/memory.js
new file mode 100644
index 000000000..1cc22a8ff
--- /dev/null
+++ b/docs/mkdocs/javascript/memory.js
@@ -0,0 +1,2 @@
+// Opt out of bfcache so large generated API pages are released after navigation.
+window.addEventListener("unload", () => {});
diff --git a/docs/mkdocs/javascript/mermaid.js b/docs/mkdocs/javascript/mermaid.js
index 63676de9d..f8fae2d47 100644
--- a/docs/mkdocs/javascript/mermaid.js
+++ b/docs/mkdocs/javascript/mermaid.js
@@ -1,21 +1,42 @@
-// Initialize Mermaid for diagram rendering
-mermaid.initialize({
-  startOnLoad: false,
-  theme: 'default',
-  securityLevel: 'loose',
-  flowchart: {
-    useMaxWidth: true,
-    htmlLabels: true
-  }
-});
+// Load Mermaid only on pages that contain diagrams.
+let mermaidLoading;
 
-// Render Mermaid diagrams when page content is ready
-document$.subscribe(() => {
-  const mermaidElements = document.querySelectorAll('.mermaid');
-  if (mermaidElements.length > 0) {
-    mermaid.run({
-      querySelector: '.mermaid',
-      nodes: mermaidElements
+function loadMermaid() {
+  if (window.mermaid) {
+    return Promise.resolve();
+  }
+  if (!mermaidLoading) {
+    mermaidLoading = new Promise((resolve, reject) => {
+      const script = document.createElement("script");
+      script.src = "https://unpkg.com/mermaid@10/dist/mermaid.min.js";
+      script.async = true;
+      script.onload = () => {
+        mermaid.initialize({
+          startOnLoad: false,
+          theme: "default",
+          securityLevel: "loose",
+          flowchart: {
+            useMaxWidth: true,
+            htmlLabels: true,
+          },
+        });
+        resolve();
+      };
+      script.onerror = reject;
+      document.head.appendChild(script);
     });
   }
+  return mermaidLoading;
+}
+
+document$.subscribe(() => {
+  if (!document.querySelector(".mermaid")) {
+    return;
+  }
+  loadMermaid().then(() => {
+    mermaid.run({
+      querySelector: ".mermaid",
+      nodes: document.querySelectorAll(".mermaid"),
+    });
+  });
 });
diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index f30f34758..f58123317 100644
--- a/docs/models/supported_models.md
+++ b/docs/models/supported_models.md
@@ -19,6 +19,7 @@ th {
 |--------------|--------|-------------------|------------|---------|-----|-----------|
 | `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ |   |   |   |
 | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ |
 | `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | |
 | `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
@@ -46,7 +47,7 @@ th {
 | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
-| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | |
+| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ |
 | `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
 | `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |
@@ -58,6 +59,7 @@ th {
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
 | `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
 | `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
 | `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
@@ -71,4 +73,6 @@ th {
 | `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
 |`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
 | `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
+
 ✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
diff --git a/docs/serving/image_edit_api.md b/docs/serving/image_edit_api.md
index 79303e1a6..f64d37f44 100644
--- a/docs/serving/image_edit_api.md
+++ b/docs/serving/image_edit_api.md
@@ -13,7 +13,7 @@ For example...
 ```bash
 # Qwen-Image
 vllm serve Qwen/Qwen-Image-Edit-2511 --omni --port 8000
-
+```
 
 ### Generate Images
 
@@ -53,7 +53,7 @@ result = client.images.edit(
     output_format='jpeg',
     # url格式
     extra_body={
-        "url": [input_image_url1,input_image_url],
+        "url": [input_image_url,input_image_url],
         "num_inference_steps": 50,
         "guidance_scale": 1,
         "seed": 777,
@@ -92,7 +92,8 @@ Content-Type: multipart/form-data
 | `user` | string | null | User identifier for tracking |
 | `output_format` | string | "png" | The format in which the generated images are returned. Must be one of "png", "jpg", "jpeg", "webp". |
 | `output_compression` | integer | 100 | The compression level (0-100%) for the generated images. |
-| `background` | string or null | "auto" | Allows to set transparency for the background of the generated image(s).
+| `background` | string or null | "auto" | Allows to set transparency for the background of the generated image(s). |
+| `stream` | boolean | false | Return Server-Sent Events instead of a single JSON response. Currently supported only by multi-stage image edit pipelines that produce AR text before the final image, such as HunyuanImage3 IT2I AR+DiT. |
 
 #### vllm-omni Extension Parameters
 
@@ -109,6 +110,8 @@ Content-Type: multipart/form-data
 
 ### Response Format
 
+When `stream=false` or omitted, the endpoint returns the standard image edit response:
+
 ```json
 {
   "created": 1701234567,
@@ -120,10 +123,56 @@ Content-Type: multipart/form-data
     }
   ],
   "output_format": null,
-  "size": null,
+  "size": null
 }
 ```
 
+### Streaming Response Format
+
+Set `stream=true` when you want to receive the HunyuanImage3 IT2I AR
+recaption text before the final edited image is ready. Streaming is only
+available for multi-stage image edit pipelines where stage 0 produces AR text
+and a later diffusion stage produces the image. Single-stage diffusion image
+edit pipelines reject `stream=true` with HTTP 400.
+
+The response uses Server-Sent Events with `Content-Type: text/event-stream`.
+Each event is sent as one `data:` line. The stream order is:
+
+1. One or more AR text delta chunks.
+2. One final image chunk.
+3. `data: [DONE]`.
+
+AR text chunks expose only the generated text delta and the AR completion
+`index`. They do not expose token ids, ratio tokens, KV cache data, or internal
+prompt token ids.
+
+Example AR delta event:
+
+```text
+data: {"object":"image.edit.chunk","type":"ar_delta","delta":"A close-up product photo...","index":0,"created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct"}
+```
+
+Example final image event:
+
+```text
+data: {"object":"image.edit.chunk","type":"image","data":[{"b64_json":"<base64-encoded PNG>","url":null,"revised_prompt":null}],"output_format":"png","size":"1024x1024","created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct"}
+```
+
+Terminal event:
+
+```text
+data: [DONE]
+```
+
+If the engine fails after the SSE response has started, the stream emits an
+error chunk followed by `[DONE]`:
+
+```text
+data: {"object":"error","created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct","error":{"message":"<error message>","type":"server_error","code":500}}
+
+data: [DONE]
+```
+
 ## Examples
 
 ### Multiple Images input
@@ -134,12 +183,28 @@ curl -s -D >(grep -i x-request-id >&2) \
   -X POST "http://localhost:8000/v1/images/edits" \
   -F "model=xxx" \
   -F "image=@xx.png" \
-  -F "image=@xx.png"
+  -F "image=@xx.png" \
   -F "prompt='this bear is wearing sportwear. holding a basketball, and bending one leg.'" \
   -F "size=1024x1024" \
   -F "output_format=png"
 ```
 
+### Streaming HunyuanImage3 IT2I AR Text
+
+```bash
+curl -N -X POST "http://localhost:8000/v1/images/edits" \
+  -F "model=tencent/HunyuanImage-3.0-Instruct" \
+  -F "image=@./input.png" \
+  -F "prompt=Turn the product into a clean studio advertisement." \
+  -F "size=1024x1024" \
+  -F "output_format=png" \
+  -F "stream=true"
+```
+
+Use the `ar_delta` events for progressive display of the AR-generated
+recaption. Decode the `b64_json` field from the final `image` event to get the
+edited image.
+
 
 ## Parameter Handling
 
diff --git a/docs/user_guide/diffusion/attention_backends.md b/docs/user_guide/diffusion/attention_backends.md
index 4a19ad8c4..abc63ebed 100644
--- a/docs/user_guide/diffusion/attention_backends.md
+++ b/docs/user_guide/diffusion/attention_backends.md
@@ -1,28 +1,26 @@
 # Diffusion Attention Backends
 
-This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, and how to use SageAttention.
+This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, the per-platform defaults, and how to use SageAttention.
 
 ## Overview
 
 Diffusion attention backend selection is resolved in `vllm_omni.diffusion.attention.selector`. It looks up the backend from a structured `AttentionConfig` carried on `OmniDiffusionConfig` and falls back to the platform default when nothing is configured.
 
-This backend is used by diffusion attention layers such as the DiT attention in video and image generation models.
+This backend is used by diffusion attention layers such as the DiT attention in video and image generation models. It does **not** affect autoregressive (LLM) attention paths — those go through vLLM's own attention backend selector.
 
-On CUDA, the practical choices today are:
-
-- `FLASH_ATTN`: FlashAttention backend. This is the default on supported CUDA systems when FlashAttention is installed.
-- `TORCH_SDPA`: PyTorch `scaled_dot_product_attention`.
-- `SAGE_ATTN`: SageAttention backend, if `sageattention` is installed.
-
-If no attention backend is configured, vLLM-Omni asks the current platform to choose the default. On CUDA, that normally means `FLASH_ATTN` when available, otherwise `TORCH_SDPA`.
+The full set of backends and their platform defaults is in the **Backend Options** and **Platform Defaults** sections below. If no attention backend is configured, vLLM-Omni asks the current platform to choose the default.
 
 ## Backend Options
 
 | Value | Notes |
 |---|---|
-| `FLASH_ATTN` | Default on CUDA when FlashAttention is available. Good default for most diffusion workloads. |
-| `TORCH_SDPA` | Most conservative fallback. Useful for debugging or compatibility. |
-| `SAGE_ATTN` | Requires `sageattention`. Can improve performance on some workloads, but output quality must be validated model-by-model. |
+| `FLASH_ATTN` | Wraps FlashAttention 2. Default on Hopper / Ada / Ampere when `flash-attn` is installed. |
+| `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
+| `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
+| `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
+| `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
+| `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. |
+
 
 ## Configuration
 
@@ -95,6 +93,80 @@ config = OmniDiffusionConfig(
 
 A plain dict is also accepted and normalized to `AttentionConfig`.
 
+## Platform Defaults
+
+### Blackwell (sm_100 / sm_103 / sm_120 / sm_121)
+
+Auto-route preference, in order:
+
+1. `CUDNN_ATTN` — when cuDNN ≥ 9.5 is available (ships in PyTorch 2.5+ wheels)
+2. `FLASHINFER_ATTN` — when `flashinfer` is installed but cuDNN < 9.5
+3. `FLASH_ATTN` — when `flash-attn` is installed with the Blackwell CUTE kernel
+4. `TORCH_SDPA` — last resort
+
+The startup log line `Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell sm_120, cuDNN 91002)` confirms the route.
+
+**Why CUDNN_ATTN by default**: on mask-heavy diffusion models (HunyuanVideo-1.5, Qwen-Image), cuDNN's pinned FMHA kernel sidesteps a PyTorch SDPA dispatch quirk where the unpinned dispatcher picks `EFFICIENT_ATTENTION` (~25 ms) for masked calls instead of cuDNN (~11 ms). The pin gives 2× e2e on HV-1.5 with no regression on lighter models.
+
+### Hopper (sm_90) / Ada (sm_89) / Ampere (sm_80–sm_86)
+
+Auto-route preference:
+
+1. `FLASH_ATTN` — when `flash-attn` is installed
+2. `TORCH_SDPA` — fallback
+
+`CUDNN_ATTN` and `FLASHINFER_ATTN` are still selectable via env var on these GPUs but are not in the auto-route — FlashAttention 2 is the well-tuned path on pre-Blackwell hardware.
+
+## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
+
+Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`.
+
+| Model | Shape | TORCH_SDPA | CUDNN_ATTN | FLASHINFER_ATTN |
+|---|---|---|---|---|
+| HunyuanVideo-1.5 (T2V) | 480p / 33f / 50 steps | 147.05 s | **73.02 s** | 127.84 s |
+| Wan 2.2 14B (T2V) | 480p / 33f / 40 steps | 117.75 s | 117.17 s | **115.07 s** |
+| Qwen-Image (T2I) | 1024² / 50 steps | 17.41 s | **15.14 s** | 16.02 s |
+| FLUX.2-dev (T2I) | 1024² / 50 steps, TP=2 | 53.62 s | **53.30 s** | 54.94 s |
+
+Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise.
+
+## Known Limitations
+
+### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile
+
+LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121).
+
+**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0:
+
+```bash
+DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
+    --model Lightricks/LTX-2 ...
+```
+
+### FA4 not yet integrated
+
+FlashAttention-4 (released March 2026) targets Blackwell natively and reportedly beats cuDNN by ~20% on B200. As of this writing the `flash-attn-4 4.0.0b10` wheel crashes with `AttributeError: 'NoneType' object has no attribute '_trait'` during JIT on sm_120. Not yet wired into vLLM-Omni; revisit when stable lands.
+
+## Choosing a Backend Manually
+
+### When to override the default
+
+- **Quality validation**: compare a new backend against `TORCH_SDPA` as the reference, since SDPA's default dispatcher is the most extensively tested.
+- **Lossy speedup hunting**: try `SAGE_ATTN` (INT8 quantized) on diffusion outputs — typically indistinguishable visually but always validate.
+- **Workaround for known issues**: see Known Limitations above.
+
+### Verifying which backend is in use
+
+The startup log prints one of:
+
+```
+Using diffusion attention backend 'CUDNN_ATTN'           # explicit override
+Defaulting to diffusion attention backend CUDNN_ATTN ... # auto-route
+Defaulting to diffusion attention backend SDPA           # nothing else available
+```
+
+If you don't see one of these, the model didn't reach diffusion stage init — check earlier logs for failures.
+
 ## SageAttention Installation
 
 vLLM-Omni expects SageAttention to be installed into the same Python environment as vLLM-Omni.
@@ -115,11 +187,56 @@ Quick check:
 python -c "import sageattention; print(sageattention.__file__)"
 ```
 
-## Usage
+## SageAttention3 Installation
 
-### Enable SageAttention
+vLLM-Omni expects SageAttention3 to be installed into the same Python environment as vLLM-Omni.
 
-Example: HunyuanVideo-1.5 text-to-video
+Build from source:
+
+```bash
+git clone https://github.com/thu-ml/SageAttention.git
+cd SageAttention/sageattention3_blackwell
+python setup.py install
+```
+
+Quick check:
+
+```bash
+python -c "import sageattn3; print(sageattn3.__file__)"
+```
+
+Notes:
+
+- `SAGE_ATTN_3` is only selected on CUDA when `sageattn3` is importable and the GPU is Blackwell-class.
+- SageAttention3's Blackwell kernel assumes `Hq == Hkv`. In vLLM-Omni, GQA/MQA diffusion requests fall back to PyTorch SDPA for correctness.
+
+## Usage Examples
+
+### Default (auto-route)
+
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+    --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
+    --prompt "A dog running across a field of golden wheat." \
+    --height 480 --width 832 --num-frames 33 \
+    --num-inference-steps 50 --seed 42 --guidance-scale 6.0 \
+    --output hv15.mp4
+```
+
+On Blackwell this picks `CUDNN_ATTN` automatically. Check the log for the `Defaulting to ...` line.
+
+### Explicit backend selection
+
+```bash
+DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
+    --model Lightricks/LTX-2 \
+    --prompt "A dog running across a field of golden wheat." \
+    --height 480 --width 832 --num-frames 33 \
+    --num-inference-steps 40 --seed 42 --guidance-scale 4.0 \
+    --output ltx2.mp4
+```
+
+### SageAttention (lossy)
 
 ```bash
 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
@@ -128,7 +245,7 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_
     --height 480 --width 832 --num-frames 33 \
     --num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
     --tensor-parallel-size 2 \
-    --output ../tmp/hv15_modelopt_sage.mp4
+    --output hv15_sage.mp4
 ```
 
 Example: Wan2.2 TI2V 5B
@@ -143,6 +260,20 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_
     --output outputs/wan22_sage.mp4
 ```
 
+### Enable SageAttention3
+
+Example:
+
+```bash
+DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3 python examples/offline_inference/text_to_video/text_to_video.py \
+    --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
+    --prompt "A dog running across a field of golden wheat." \
+    --height 480 --width 832 --num-frames 33 \
+    --num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
+    --tensor-parallel-size 2 \
+    --output outputs/hv15_sage3.mp4
+```
+
 ### Mixed backends across roles
 
 Use `FLASH_ATTN` for self-attention and `TORCH_SDPA` for cross-attention:
@@ -173,13 +304,13 @@ python examples/offline_inference/text_to_video/text_to_video.py \
 
 ## Validation Guidance
 
-Do not assume that a faster attention backend is numerically interchangeable with `FLASH_ATTN`.
+Don't assume a faster attention backend is numerically interchangeable with `TORCH_SDPA`.
 
 Always compare:
 
 - End-to-end runtime
-- DiT / diffusion stage runtime
-- Output quality against a known-good baseline
+- Diffusion-stage runtime (`add_req_and_wait` line in DiffusionEngine.step breakdown)
+- Output quality against a known-good baseline (CLIP similarity, frame-level diff, or visual review)
 
 At minimum, keep the same:
 
@@ -187,6 +318,19 @@ At minimum, keep the same:
 - prompt
 - seed
 - resolution
-- frame count
-- inference steps
-- parallel config
+- frame count / step count
+- parallel config (TP / CFG-parallel / Ulysses degrees)
+
+## Reproducing the Benchmark Table
+
+The end-to-end numbers above were collected by running `text_to_video.py` /
+`text_to_image.py` with the same prompt and seed while varying
+`DIFFUSION_ATTENTION_BACKEND`. For a quick kernel-level comparison of the
+backends without loading a model:
+
+```bash
+python benchmarks/diffusion/bench_attention_backends.py --preset hv15
+```
+
+It runs all three BF16 backends on representative DiT attention shapes and
+prints a ranking table at the end.
diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md
index 61694ac9d..4e62cd665 100644
--- a/docs/user_guide/diffusion_features.md
+++ b/docs/user_guide/diffusion_features.md
@@ -113,8 +113,9 @@ The following tables show which models support each feature:
 | **FLUX.1-schnell**       |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
 | **FLUX.2-klein**         |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
 | **FLUX.1-Kontext-dev**   |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
-| **FLUX.2-dev**           |     ✅     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **FLUX.2-dev**           |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **GLM-Image**            |     ❌     |     ❌      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **Hidream-I1-Full**        |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **HunyuanImage3**        |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ✅        |        ❌         |
 | **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LongCat-Image-Edit**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
@@ -128,7 +129,7 @@ The following tables show which models support each feature:
 | **Qwen-Image-Edit**      |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Qwen-Image-Edit-2509** |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |        ✅ (decode)         |          ✅           |       ❌        |        ❌         |
 | **Qwen-Image-Layered**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
-| **SenseNova-U1**         |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **SenseNova-U1**         |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Stable-Diffusion3.5**  |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Z-Image**              |     ✅     |     ✅      |           ✅           |       ❓        |   ✅ (TP=2 only)   |          ❌          |   ✅    |             ❌             |      ✅ (decode)      |       ✅        |        ❌         |
 | **ERNIE-Image**          |     ❌     |     ✅      |           ✅           |       ❓        |         ✅         |          ❌          |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
@@ -144,9 +145,9 @@ The following tables show which models support each feature:
 | **Wan2.2**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ✅         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ❌        |        ❌         |
 | **Wan2.1-VACE**              |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2**                    |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ❌        |         ✅         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ❌         |
+| **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 
 **Frame Interpolation Support**
diff --git a/docs/user_guide/examples/offline_inference/text_to_speech.md b/docs/user_guide/examples/offline_inference/text_to_speech.md
index 7ad3d2d40..ad7777014 100644
--- a/docs/user_guide/examples/offline_inference/text_to_speech.md
+++ b/docs/user_guide/examples/offline_inference/text_to_speech.md
@@ -19,9 +19,11 @@ list of supported architectures across all modalities, see
 | VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | — | continuation (`--ref-audio` + `--ref-text`) | 48 kHz |
 | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ (`async_chunk: true` default) | — | 22.05 kHz |
 | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ (`--streaming`) | — | 44.1 kHz |
+| GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
 | OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design (`--instruct`), language (`--lang`) | 24 kHz |
 | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants (`--query-type`) | 24 kHz |
 | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets (`--voice`) | 24 kHz |
+| higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` (+ `k2-fsa/OmniVoice/audio_tokenizer/` codec) | 2 (talker + code2wav, DualFFN) | ✓ (`--ref-audio` + `--ref-text`) | — | — | 24 kHz |
 
 ## Common Quick Start
 
@@ -158,6 +160,35 @@ Streaming requires `async_chunk: true` in the stage config.
 
 ---
 
+## GLM-TTS
+
+2-stage TTS pipeline (AR + DiT flow-matching) at 24 kHz. Every request requires reference audio and its transcript for zero-shot voice cloning.
+
+### Quick start
+```bash
+python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
+    --model zai-org/GLM-TTS \
+    --text "你好，这是语音合成测试。" \
+    --ref-audio /path/to/reference.wav \
+    --ref-text "这是参考音频的文本内容。" \
+    --output-dir ./output
+```
+
+### Architecture
+```
+Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 kHz)
+        (Llama-based)    (32k vocab)      (Flow Matching)
+```
+
+### Notes
+- `--ref-audio` and `--ref-text` are **required** together; GLM-TTS does not support text-only synthesis.
+- Reference audio should be 3-10 seconds.
+- First run may be slow due to lazy loading of WhisperVQ tokenizer and CampPlus ONNX speaker embedder.
+- Default sampling: temperature=1.0, top_k=25, top_p=0.8 (RAS method).
+- The `--model` path should point to the repository root (not `llm/` subdirectory).
+
+---
+
 ## OmniVoice
 
 Zero-shot multilingual TTS supporting 600+ languages, with three modes (auto / clone / design).
@@ -303,6 +334,46 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
 - `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`.
 - Run `--help` for the full argument surface.
 
+---
+
+## higgs-audio v2
+
+2-stage TTS at 24 kHz: a vLLM-native Llama-3.2-3B talker with a DualFFN audio expert (Stage 0) feeding a HiggsAudio codec decoder (Stage 1) over the shared-memory connector. Stage 1 builds on the HiggsAudio decoder kernel at `vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py`.
+
+### Prerequisites
+
+Voice clone uses HF's `HiggsAudioV2TokenizerModel`, instantiated from `k2-fsa/OmniVoice/audio_tokenizer/` — the boson-ai standalone tokenizer Hub repo's `model.safetensors` is actually the 3B talker LM, so we point HF at k2's repackaged codec weights instead. Only the `audio_tokenizer/` subdir (~806 MB) is downloaded.
+
+```bash
+pip install -U "transformers>=5.3.0"
+```
+
+### Quick start (plain TTS)
+
+```bash
+python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
+    --texts "Hello world." "The quick brown fox jumps over the lazy dog." \
+    --output-dir results/wavs
+```
+
+### Voice cloning
+
+Pass both `--ref-audio` and `--ref-text` together:
+
+```bash
+python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
+    --texts "Hello, this is a cloned voice." \
+    --ref-audio /path/to/reference.wav \
+    --ref-text  "Exact transcript spoken in reference.wav." \
+    --output-dir results/wavs
+```
+
+### Notes
+
+- Deploy config auto-loads from `vllm_omni/deploy/higgs_audio_v2.yaml`.
+- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
+- Out of scope (rejected with 4xx by the request validator): multi-speaker `[SPEAKERn]` dialogue, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.
+
 ## Example materials
 
 ??? abstract "cosyvoice3/end2end.py"
@@ -313,6 +384,13 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
     ``````py
     --8<-- "examples/offline_inference/text_to_speech/fish_speech/end2end.py"
     ``````
+??? abstract "higgs_audio_v2/end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py"
+??? abstract "glm_tts/end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/text_to_speech/glm_tts/end2end.py"
+    ``````
 ??? abstract "omnivoice/end2end.py"
     ``````py
     --8<-- "examples/offline_inference/text_to_speech/omnivoice/end2end.py"
diff --git a/docs/user_guide/examples/online_serving/text_to_speech.md b/docs/user_guide/examples/online_serving/text_to_speech.md
index ff22dd0fd..f971cced1 100644
--- a/docs/user_guide/examples/online_serving/text_to_speech.md
+++ b/docs/user_guide/examples/online_serving/text_to_speech.md
@@ -19,6 +19,8 @@ For the full list of supported architectures across all modalities, see
 |---|---|---|---|---|---|
 | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | — |
 | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
+| higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✓ (`ref_audio`+`ref_text`) | ✓ (codec_streaming) | — | — |
+| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
 | OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — |
 | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
 | VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ |
@@ -192,6 +194,84 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091  # if server
 
 ---
 
+## higgs-audio v2
+
+2-stage TTS at 24 kHz: a vLLM-native Llama-3.2-3B talker with a DualFFN audio expert (Stage 0) feeding a HiggsAudio codec decoder (Stage 1) that streams chunks back to the client.
+
+### Prerequisites
+
+Voice clone uses HF's `HiggsAudioV2TokenizerModel` loaded from `k2-fsa/OmniVoice/audio_tokenizer/` (~806 MB subdir; the boson-ai standalone tokenizer Hub repo's `model.safetensors` is the 3B talker LM, not the codec):
+
+```bash
+pip install -U "transformers>=5.3.0"
+```
+
+### Launch
+```bash
+GPUS=6,7 PORT=8094 bash examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh
+```
+Deploy config auto-loads from `vllm_omni/deploy/higgs_audio_v2.yaml`.
+
+### Sending requests
+```bash
+# Plain TTS
+python higgs_audio_v2/batch_speech_client.py \
+    --base-url http://localhost:8094 \
+    --output-dir /tmp/higgs_out \
+    --prompts "Hello world." "The quick brown fox jumps over the lazy dog."
+
+# Voice cloning — pass a reference clip and its transcript together
+python higgs_audio_v2/batch_speech_client.py \
+    --base-url http://localhost:8094 \
+    --output-dir /tmp/higgs_clone \
+    --ref-audio /path/to/reference.wav \
+    --ref-text  "Exact transcript spoken in reference.wav." \
+    --prompts "Hello, this is a cloned voice."
+```
+
+### Notes
+- Output: 24 kHz mono.
+- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
+- Out of scope (rejected with explicit 4xx): multi-speaker `[SPEAKERn]` tags inside `input`, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.
+## GLM-TTS
+
+2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`.
+
+### Launch
+```bash
+vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
+# or:
+bash examples/online_serving/text_to_speech/glm_tts/run_server.sh /path/to/GLM-TTS
+```
+
+### Sending requests
+```bash
+# Voice cloning (required)
+python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
+    --text "你好，这是语音克隆测试。" \
+    --ref-audio file:///path/to/ref.wav \
+    --ref-text "这是参考音频的文本内容。"
+
+# Custom format
+python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
+    --text "Hello, this is a voice cloning test." \
+    --ref-audio file:///path/to/ref.wav \
+    --ref-text "Transcript of the reference audio." \
+    --response-format mp3 -o output.mp3
+```
+
+### Gradio demo
+```bash
+bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
+```
+
+### Notes
+- Output: 24 kHz mono WAV via HiFT vocoder.
+- `ref_audio` + `ref_text` are **required** together on every request. Reference audio should be 3-10 seconds.
+- Voice cloning feature extraction (WhisperVQ, CampPlus, mel) runs on the model side — no external dependency on the serving layer.
+
+---
+
 ## OmniVoice
 
 Zero-shot multilingual TTS (600+ languages). Online serving currently exposes **auto voice** only; voice cloning and voice design are available offline.
@@ -249,6 +329,16 @@ Single-GPU serves now default to the uniproc executor (lower IPC overhead, the B
 
 To opt out of chunked streaming, pass `--no-async-chunk` — the pipeline auto-dispatches to the end-to-end codec processor.
 
+### Tuning stage 1 `max_num_seqs` per task type
+The bundled `qwen3_tts.yaml` ships stage 1 (Code2Wav) at `max_num_seqs: 10`, tuned for Base voice cloning: stage-1 lifetimes are long (~3 s/req), so admitting up to 10 concurrent codec sequences lets requests progress in parallel in the scheduler — ~2× TTFA p95 at c=4 / c=8 (1× H100, 1.7B-Base, seed-tts) at an 8–12 % audio-throughput cost.
+
+CustomVoice / VoiceDesign have much shorter stage-1 lifetimes (~50–200 ms) and are TTFA-optimal at `max_num_seqs: 1`. Override the default when serving those task types:
+
+```bash
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni \
+    --stage-overrides '{"1": {"max_num_seqs": 1}}'
+```
+
 ### Sending requests
 ```bash
 # CustomVoice with a predefined speaker
@@ -419,6 +509,29 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
     ``````py
     --8<-- "examples/online_serving/text_to_speech/fish_speech/speech_client.py"
     ``````
+??? abstract "higgs_audio_v2/batch_speech_client.py"
+    ``````py
+    --8<-- "examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py"
+    ``````
+??? abstract "higgs_audio_v2/run_server.sh"
+    ``````sh
+    --8<-- "examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh"
+??? abstract "glm_tts/gradio_demo.py"
+    ``````py
+    --8<-- "examples/online_serving/text_to_speech/glm_tts/gradio_demo.py"
+    ``````
+??? abstract "glm_tts/openai_speech_client.py"
+    ``````py
+    --8<-- "examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py"
+    ``````
+??? abstract "glm_tts/run_gradio_demo.sh"
+    ``````sh
+    --8<-- "examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh"
+    ``````
+??? abstract "glm_tts/run_server.sh"
+    ``````sh
+    --8<-- "examples/online_serving/text_to_speech/glm_tts/run_server.sh"
+    ``````
 ??? abstract "omnivoice/run_server.sh"
     ``````sh
     --8<-- "examples/online_serving/text_to_speech/omnivoice/run_server.sh"
diff --git a/docs/user_guide/quantization/autoround.md b/docs/user_guide/quantization/autoround.md
index 88fed3b62..3afc8b062 100644
--- a/docs/user_guide/quantization/autoround.md
+++ b/docs/user_guide/quantization/autoround.md
@@ -52,8 +52,8 @@ vLLM-Omni's runtime module names.
 
 | Model | Scope | Status | Notes |
 |-------|-------|--------|-------|
+| GLM-Image | Diffusion transformer  | ✅ | `Intel/GLM-Image-int4-AutoRound` |
 | BAGEL | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
-| GLM-Image | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
 
 ## Configuration
 
diff --git a/docs/user_guide/quantization/modelopt.md b/docs/user_guide/quantization/modelopt.md
index 8a4389015..66e8427bb 100644
--- a/docs/user_guide/quantization/modelopt.md
+++ b/docs/user_guide/quantization/modelopt.md
@@ -6,34 +6,46 @@ ModelOpt quantization loads checkpoints produced by NVIDIA ModelOpt. The
 quantized weights and scale tensors are generated before serving, so inference
 does not run online calibration or convert a BF16 checkpoint at startup.
 
-vLLM-Omni currently validates the ModelOpt FP8 checkpoint path for diffusion
-transformers. The loader auto-detects supported ModelOpt FP8 checkpoint configs
-and keeps non-transformer components, such as the tokenizer, scheduler, text
+vLLM-Omni validates ModelOpt FP8, ModelOpt NVFP4, and ModelOpt mixed
+FP8/NVFP4 checkpoint loading for diffusion transformer stages. The loader
+auto-detects supported ModelOpt checkpoint configs and keeps non-transformer
+components, such as the tokenizer, scheduler, text encoder, vision/audio
 encoder, and VAE, on the base checkpoint unless a model-specific guide says
 otherwise.
 
 !!! note
-    `--force-cutlass-fp8` is an explicit runtime override for diffusion
-    checkpoints that already carry a supported ModelOpt FP8 config. It does not
-    quantize BF16 checkpoints and it does not apply to online `--quantization
-    fp8`. The flag only takes effect for ModelOpt FP8 diffusion stages on CUDA
-    SM89+ devices; other platforms and non-ModelOpt FP8 paths fall back to the
-    normal vLLM kernel selection.
+    ModelOpt checkpoints are pre-quantized checkpoints. Do not pass
+    `--quantization fp8` for these checkpoints. The checkpoint
+    `quantization_config` selects the ModelOpt path.
+
+!!! note
+    `--force-cutlass-fp8`, `--linear-backend cutlass`, and
+    `--moe-backend cutlass` are runtime backend selections for checkpoints that
+    already carry supported ModelOpt quantized weights and scales. They do not
+    quantize BF16 checkpoints at startup.
 
 ## Supported ModelOpt Checkpoint Formats
 
 vLLM-Omni treats ModelOpt checkpoints as pre-quantized checkpoints. The
 checkpoint config must identify ModelOpt as the quantization method or producer,
-and the quantization algorithm must be one of the validated FP8 algorithms.
+and the quantization algorithm must be one of the validated algorithms below.
 
 | Checkpoint field | Supported value |
 |------------------|-----------------|
-| `method` / `quant_method` | `modelopt` |
+| `method` / `quant_method` | `modelopt`, `modelopt_fp4`, `modelopt_mixed` |
 | `producer.name` | `modelopt` |
-| `quant_algo` | `FP8`, `FP8_PER_CHANNEL_PER_TOKEN` |
+| `quant_algo` | `FP8`, `FP8_PER_CHANNEL_PER_TOKEN`, `NVFP4`, `MIXED_PRECISION` |
 
-Other ModelOpt algorithms, such as NVFP4, are not enabled by this diffusion
-FP8 path until they have separate model and quality validation.
+| `quant_algo` | Runtime method | Typical use |
+|--------------|----------------|-------------|
+| `FP8`, `FP8_PER_CHANNEL_PER_TOKEN` | `modelopt` | FP8 diffusion transformer checkpoints |
+| `NVFP4` | `modelopt_fp4` | NVFP4 diffusion transformer checkpoints |
+| `MIXED_PRECISION` | `modelopt_mixed` | Mixed FP8/NVFP4 checkpoints with a ModelOpt per-layer policy |
+
+For multi-component diffusion or omni models, only the checkpoint component
+that contains ModelOpt quantized weights should use the ModelOpt quantization
+method. Encoders, decoders, tokenizers, schedulers, and other BF16 components
+stay unquantized unless the model-specific recipe validates otherwise.
 
 ## Hardware Support
 
@@ -47,7 +59,10 @@ FP8 path until they have separate model and quality validation.
 | Ascend NPU | ❌ |
 
 Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
-The optional CUTLASS FP8 runtime override requires CUDA SM89+.
+The optional CUTLASS FP8 runtime override requires CUDA SM89+. ModelOpt NVFP4
+and mixed FP8/NVFP4 diffusion checkpoints are currently validated on Blackwell
+CUDA systems in the recipes below; other CUDA generations require separate
+backend and quality validation.
 
 ## Model Type Support
 
@@ -56,12 +71,19 @@ The optional CUTLASS FP8 runtime override requires CUDA SM89+.
 | Model | HF checkpoint | Scope | Status |
 |-------|---------------|-------|--------|
 | Qwen-Image 2512 | `feizhai123/qwen-image-2512-modelopt-fp8-dynamic-all` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
+| Qwen-Image 2512 | `feizhai123/qwen-image-2512-modelopt-mixed-fp8-sensitive-nvfp4-heavy` | Diffusion transformer | Validated for ModelOpt mixed FP8/NVFP4 checkpoints |
 | Z-Image | `feizhai123/z-image-modelopt-fp8-conservative` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
 | FLUX.2-dev | `feizhai123/flux2-dev-modelopt-fp8` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
 | FLUX.2-klein 4B | `feizhai123/flux2-klein-4b-modelopt-fp8` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
 | HunyuanImage-3.0 | `feizhai123/hunyuan-image3-modelopt-fp8` | MoE diffusion transformer | Validated for ModelOpt FP8 checkpoints |
+| HunyuanImage-3.0 | `feizhai123/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8` | MoE diffusion transformer | Validated for ModelOpt mixed FP8/NVFP4 checkpoints |
 | Wan2.2 | Not available | Diffusion transformer | Not validated |
 
+For full serving commands and benchmark context, see
+[`recipes/Qwen/Qwen-Image.md`](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image.md)
+and
+[`recipes/Tencent/HunyuanImage-3.0-Instruct.md`](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md).
+
 ### Multi-Stage Omni/TTS Model
 
 | Model | Scope | Status |
@@ -80,15 +102,28 @@ ModelOpt targets yet.
 
 ## Configuration
 
-For pre-quantized ModelOpt FP8 checkpoints, no `--quantization fp8` flag is
-needed. The checkpoint config selects the ModelOpt path.
+For pre-quantized ModelOpt checkpoints, no `--quantization fp8` flag is needed.
+The checkpoint config selects the ModelOpt path.
 
 Online serving:
 
 ```bash
-vllm serve <modelopt-fp8-checkpoint> \
+vllm serve <modelopt-checkpoint> \
   --omni \
   --tensor-parallel-size <N> \
+  --linear-backend cutlass \
+  --force-cutlass-fp8
+```
+
+For mixed FP8/NVFP4 MoE checkpoints, also select the validated MoE backend:
+
+```bash
+vllm serve <modelopt-mixed-moe-checkpoint> \
+  --omni \
+  --tensor-parallel-size <N> \
+  --enable-expert-parallel \
+  --linear-backend cutlass \
+  --moe-backend cutlass \
   --force-cutlass-fp8
 ```
 
@@ -96,14 +131,14 @@ Offline inference:
 
 ```bash
 python examples/offline_inference/text_to_image/text_to_image.py \
-  --model <modelopt-fp8-checkpoint> \
+  --model <modelopt-checkpoint> \
   --tensor-parallel-size <N> \
   --prompt "a red ceramic teapot on a wooden table" \
   --height 1024 \
   --width 1024 \
   --num-inference-steps 20 \
   --seed 42 \
-  --output outputs/modelopt_fp8.png
+  --output outputs/modelopt.png
 ```
 
 Python API:
@@ -112,7 +147,7 @@ Python API:
 from vllm_omni import Omni
 
 omni = Omni(
-    model="<modelopt-fp8-checkpoint>",
+    model="<modelopt-checkpoint>",
     tensor_parallel_size=2,
     force_cutlass_fp8=True,
 )
@@ -123,16 +158,34 @@ omni = Omni(
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
 | `force_cutlass_fp8` / `--force-cutlass-fp8` | bool | `False` | Force CUTLASS FP8 linear kernels for supported ModelOpt FP8 diffusion stages on CUDA SM89+ |
+| `--linear-backend cutlass` | str | auto | Select the validated CUTLASS linear backend for supported ModelOpt NVFP4 or mixed FP8/NVFP4 diffusion stages |
+| `--moe-backend cutlass` | str | auto | Select the validated CUTLASS MoE backend for supported ModelOpt mixed MoE checkpoints |
 
 ## Validation and Notes
 
-1. Compare the ModelOpt FP8 checkpoint against the BF16 baseline with the same
+1. Compare the ModelOpt checkpoint against the BF16 baseline with the same
    prompt, resolution, seed, and inference steps.
 2. Use `tests/diffusion/quantization/test_quantization_quality.py` with
    `VLLM_OMNI_QUALITY_CONFIGS` to validate local baseline and quantized model
    paths.
-3. Report LPIPS, PSNR, MAE, throughput, latency, and peak memory when adding a
-   new validated ModelOpt diffusion checkpoint.
-4. Keep `--quantization fp8` for online FP8 from BF16 checkpoints; use this
-   ModelOpt path only when the checkpoint already contains ModelOpt FP8 weights
-   and scales.
+3. For HunyuanImage-3.0 quantized DiT checkpoints, the opt-in accuracy check is:
+
+   ```bash
+   CUDA_VISIBLE_DEVICES=2,3 \
+   HUNYUAN_IMAGE3_RUN_QUANT_ACCURACY=1 \
+   HUNYUAN_IMAGE3_QUANT_DEVICES=0,1 \
+   HUNYUAN_IMAGE3_QUANT_TP=2 \
+   HUNYUAN_IMAGE3_BF16_MODEL=/path/to/hunyuan-image3-bf16 \
+   HUNYUAN_IMAGE3_FP8_MODEL=/path/to/hunyuan-image3-modelopt-fp8 \
+   HUNYUAN_IMAGE3_NVFP4_MODEL=/path/to/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8 \
+   PYTHONPATH=/path/to/vllm-omni:${PYTHONPATH:-} \
+   python -m pytest -s -v \
+     tests/e2e/accuracy/test_hunyuan_image3.py \
+     -k quantized_dit_matches_bf16_accuracy
+   ```
+
+4. Report CLIP score deltas, SSIM, PSNR, throughput, latency, and peak memory
+   when adding a new validated ModelOpt diffusion checkpoint.
+5. Keep `--quantization fp8` for online FP8 from BF16 checkpoints; use this
+   ModelOpt path only when the checkpoint already contains ModelOpt quantized
+   weights and scales.
diff --git a/docs/user_guide/quantization/mxfp8.md b/docs/user_guide/quantization/mxfp8.md
index 418aad369..823e898fd 100644
--- a/docs/user_guide/quantization/mxfp8.md
+++ b/docs/user_guide/quantization/mxfp8.md
@@ -7,26 +7,28 @@ using the OCP MX format: groups of 32 K-dimension elements share a single
 `float8_e8m0fnu` exponent scale. This gives better accuracy than channel-wise
 FP8 while keeping the same 8-bit weight footprint.
 
-This method supports two modes:
+This method supports three modes:
 
 | Mode | Description |
 |------|-------------|
 | **Online** | BF16 weights are quantized to MXFP8 at load time — no pre-processing needed |
-| **Offline** | msModelSlim-exported MXFP8 weights converted to diffusers format via `merge_mxfp8_checkpoint.py` — weights and scales are loaded directly from the preprocessed checkpoint |
+| **Offline (Native)** | msModelSlim-exported MXFP8 weights converted to diffusers format via `merge_mxfp8_checkpoint.py` — weights and scales are loaded directly from the preprocessed checkpoint |
+| **Offline (AutoRound)** | AutoRound MXFP8 checkpoints with `data_type="mx_fp"` — auto-detected from `config.json` |
 
 ## Hardware Support
 
-| Device | Support |
-|--------|---------|
-| NVIDIA Blackwell GPU (SM 100+) | ⭕ |
-| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ |
-| NVIDIA Ampere GPU (SM 80+) | ⭕ |
-| AMD ROCm | ⭕ |
-| Intel XPU | ⭕ |
-| Ascend NPU (Atlas 950 A5) | ✅ |
+| Device | Online | Offline (Native) | Offline (AutoRound) |
+|--------|--------|------------------|---------------------|
+| NVIDIA Blackwell GPU (SM 100+) | ⭕ | ⭕ | ⭕ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ | ⭕ | ⭕ |
+| NVIDIA Ampere GPU (SM 80+) | ⭕ | ⭕ | ⭕ |
+| AMD ROCm | ⭕ | ⭕ | ⭕ |
+| Intel XPU | ✅ | ❌ | ✅ |
+| Ascend NPU (Atlas 950 A5) | ✅ | ✅ | ⭕ |
 
-Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
-guide.
+Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
+
+**Note**: Intel XPU only supports AutoRound MXFP8 for offline mode. Use AutoRound quantized checkpoints or online mode for XPU.
 
 ## Model Type Support
 
@@ -82,9 +84,9 @@ python text_to_video.py --model <your-model> --quantization mxfp8
 vllm serve <your-model> --omni --quantization mxfp8
 ```
 
-### Offline Mode
+### Offline Mode (Native)
 
-Offline mode loads a pre-quantized checkpoint from msModelSlim. A preprocessing
+Native offline mode loads a pre-quantized checkpoint from msModelSlim. A preprocessing
 step converts the raw quantized output to the diffusers format expected by
 vLLM-Omni and injects the quantization config into `transformer/config.json` so
 that vLLM-Omni auto-detects the offline path without a `--quantization` flag.
@@ -159,10 +161,41 @@ omni = Omni(model="/path/to/Wan2.2-TI2V-5B-MXFP8")
 ```
 
 !!! note
-    No `--quantization` flag is needed for offline mode. The preprocessing
+    No `--quantization` flag is needed for native offline mode. The preprocessing
     script injects `quantization_config` into each `transformer/config.json`,
     which vLLM-Omni reads automatically to activate the offline MXFP8 method.
 
+### Offline Mode (AutoRound)
+
+AutoRound MXFP8 checkpoints declare `quant_method="auto-round"` with `data_type="mx_fp"` in their `config.json`. These are automatically detected and use the `IncMxfp8OfflineLinearMethod` backend.
+
+To use an AutoRound MXFP8 checkpoint:
+
+```bash
+python text_to_video.py --model <autoround-mxfp8-model>
+
+# Online serving
+vllm serve <autoround-mxfp8-model> --omni
+```
+
+Python API:
+
+```python
+omni = Omni(model="<autoround-mxfp8-model>")
+```
+
+!!! note
+    AutoRound MXFP8 checkpoints are auto-detected from `config.json` and do not require a `--quantization` flag. The config must include:
+    ```json
+    {
+      "quantization_config": {
+        "quant_method": "auto-round",
+        "data_type": "mx_fp",
+        ...
+      }
+    }
+    ```
+
 ## Parameters
 
 | Parameter | Type | Default | Description |
diff --git a/docs/user_guide/quantization/online.md b/docs/user_guide/quantization/online.md
index b9a6861d1..6103be9a0 100644
--- a/docs/user_guide/quantization/online.md
+++ b/docs/user_guide/quantization/online.md
@@ -21,7 +21,7 @@ by msModelSlim and the merge tools.
 | NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ⭕ | ⭕ |
 | NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ |
 | AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ |
-| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ |
+| Intel XPU | ⭕ | ⭕ | ✅ | ⭕ |
 | Ascend NPU | ❌ | ✅ | ✅ | ✅ |
 
 Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
@@ -36,7 +36,7 @@ MXFP4 are documented for the Ascend NPU path.
 |--------|-------|----------------|--------|
 | FP8 W8A8 | [FP8](fp8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
 | Int8 W8A8 | [Int8](int8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
-| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B, Wan2.2-TI2V-5B | Ascend NPU only |
+| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B, Wan2.2-TI2V-5B | Validated on Ascend NPU and Intel XPU |
 | MXFP4 W4A4 | [MXFP4](mxfp4.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B | Ascend NPU only; TI2V-5B is not supported |
 
 ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
diff --git a/examples/offline_inference/hunyuan_image3/end2end.py b/examples/offline_inference/hunyuan_image3/end2end.py
index aed06bcb9..e8cd8a36f 100644
--- a/examples/offline_inference/hunyuan_image3/end2end.py
+++ b/examples/offline_inference/hunyuan_image3/end2end.py
@@ -17,7 +17,7 @@ from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniPromptType
 
 _REPO_ROOT = Path(__file__).resolve().parents[3]
-_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3.yaml")
+_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image_3_moe.yaml")
 _DEFAULT_AR_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3_ar.yaml")
 
 _MODALITY_TASK_MAP: dict[str, tuple[str, str | None]] = {
@@ -61,8 +61,8 @@ def parse_args():
     parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.")
     parser.add_argument("--guidance-scale", type=float, default=5.0, help="Classifier-free guidance scale.")
     parser.add_argument("--seed", type=int, default=42, help="Random seed.")
-    parser.add_argument("--height", type=int, default=1024, help="Output image height.")
-    parser.add_argument("--width", type=int, default=1024, help="Output image width.")
+    parser.add_argument("--height", type=int, default=None, help="Output image height.")
+    parser.add_argument("--width", type=int, default=None, help="Output image width.")
     parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.")
     parser.add_argument(
         "--bot-task",
@@ -229,7 +229,21 @@ def main():
 
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
-    ar_stop_token_ids = resolve_stop_token_ids(task=task, bot_task=bot_task, tokenizer=tokenizer)
+    if (args.height is None) != (args.width is None):
+        raise ValueError("--height and --width must both be specified or both omitted.")
+    user_specified_size = args.height is not None and args.width is not None
+    if args.modality in ("img2text", "text2text"):
+        ar_image_size = "auto"
+    elif user_specified_size:
+        ar_image_size = f"{args.width}x{args.height}"
+    else:
+        ar_image_size = None
+    ar_stop_token_ids = resolve_stop_token_ids(
+        task=task, bot_task=bot_task, tokenizer=tokenizer, image_size=ar_image_size
+    )
+    print(
+        f"[AR Config] task={task}, bot_task={bot_task}, image_size={ar_image_size}, stop_token_ids={ar_stop_token_ids}"
+    )
     for sp in params_list:
         if isinstance(sp, OmniDiffusionSamplingParams):
             sp.num_inference_steps = args.steps
diff --git a/examples/offline_inference/internvla_a1/end2end.py b/examples/offline_inference/internvla_a1/end2end.py
index 6bc69f4e0..72cdbff0c 100755
--- a/examples/offline_inference/internvla_a1/end2end.py
+++ b/examples/offline_inference/internvla_a1/end2end.py
@@ -127,7 +127,12 @@ def run_one_path(
             ),
             args.device,
         )
-        pred = run_pipeline_forward(pipeline, batch_inputs, noise)
+        pred = run_pipeline_forward(
+            pipeline,
+            batch_inputs,
+            noise,
+            request_id=f"internvla-a1-sample-{index}",
+        )
         pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
         results.append(
             {
@@ -150,6 +155,7 @@ def run_pipeline_forward(
     pipeline,
     batch_inputs: dict[str, torch.Tensor],
     noise: torch.Tensor,
+    request_id: str,
 ) -> torch.Tensor:
     output = pipeline.forward(
         OmniDiffusionRequest(
@@ -161,6 +167,7 @@ def run_pipeline_forward(
                     "decode_image": False,
                 }
             ),
+            request_id=request_id,
         )
     )
     if output.error:
@@ -212,23 +219,38 @@ def benchmark_forward(
 
     _synchronize(args.device)
     cold_start_begin = time.perf_counter()
-    pred = run_pipeline_forward(pipeline, batch_inputs, noise)
+    pred = run_pipeline_forward(
+        pipeline,
+        batch_inputs,
+        noise,
+        request_id=f"internvla-a1-benchmark-{index}-cold",
+    )
     _synchronize(args.device)
     cold_start_ms = (time.perf_counter() - cold_start_begin) * 1000.0
 
     warmup_ms: list[float] = []
-    for _ in range(args.warmup_iters):
+    for iter_idx in range(args.warmup_iters):
         _synchronize(args.device)
         begin = time.perf_counter()
-        _ = run_pipeline_forward(pipeline, batch_inputs, noise)
+        _ = run_pipeline_forward(
+            pipeline,
+            batch_inputs,
+            noise,
+            request_id=f"internvla-a1-benchmark-{index}-warmup-{iter_idx}",
+        )
         _synchronize(args.device)
         warmup_ms.append((time.perf_counter() - begin) * 1000.0)
 
     benchmark_ms: list[float] = []
-    for _ in range(args.benchmark_iters):
+    for iter_idx in range(args.benchmark_iters):
         _synchronize(args.device)
         begin = time.perf_counter()
-        _ = run_pipeline_forward(pipeline, batch_inputs, noise)
+        _ = run_pipeline_forward(
+            pipeline,
+            batch_inputs,
+            noise,
+            request_id=f"internvla-a1-benchmark-{index}-iter-{iter_idx}",
+        )
         _synchronize(args.device)
         benchmark_ms.append((time.perf_counter() - begin) * 1000.0)
 
@@ -295,7 +317,12 @@ def main() -> None:
             dataset=dataset,
             train_meta=train_meta,
             collate_samples=collate_open_loop_samples,
-            run_sample_actions=lambda policy, batch_inputs, noise: run_pipeline_forward(policy, batch_inputs, noise),
+            run_sample_actions=lambda policy, batch_inputs, noise: run_pipeline_forward(
+                policy,
+                batch_inputs,
+                noise,
+                request_id="internvla-a1-open-loop",
+            ),
             num_episodes=args.num_episodes,
             seed=args.seed,
             device=args.device,
diff --git a/examples/offline_inference/lance/end2end.py b/examples/offline_inference/lance/end2end.py
index ec7985304..7d2e5a6d9 100644
--- a/examples/offline_inference/lance/end2end.py
+++ b/examples/offline_inference/lance/end2end.py
@@ -77,12 +77,10 @@ def parse_args():
             "caption-style default and the model describes instead of answering)."
         ),
     )
-    p.add_argument(
-        "--deploy-config",
-        type=str,
-        default="vllm_omni/deploy/lance.yaml",
-        help="Deploy YAML (single-stage Lance).",
-    )
+    # Lance is single-stage diffusion — no deploy YAML needed.  Required
+    # engine knobs (``pipeline``, ``enforce_eager``, ``trust_remote_code``,
+    # ``max_num_seqs=1`` …) are passed as flat kwargs to ``Omni`` below
+    # and ``create_default_diffusion`` materializes the stage config.
 
     from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 
@@ -235,8 +233,15 @@ def main():
                 resolved_model = candidate
 
     omni_kwargs = vars(args).copy()
-    omni_kwargs["deploy_config"] = args.deploy_config
     omni_kwargs["model"] = resolved_model
+    # Lance single-stage defaults (formerly in vllm_omni/deploy/lance.yaml):
+    omni_kwargs.setdefault("pipeline", "lance")
+    omni_kwargs.setdefault("max_num_batched_tokens", 32768)
+    omni_kwargs.setdefault("max_num_seqs", 1)
+    omni_kwargs.setdefault("enforce_eager", True)
+    omni_kwargs.setdefault("trust_remote_code", True)
+    omni_kwargs.setdefault("enable_prefix_caching", False)
+    omni_kwargs.setdefault("async_chunk", False)
     omni = Omni(**omni_kwargs)
 
     if args.modality == "img2text":
diff --git a/examples/offline_inference/lance/gradio_demo.py b/examples/offline_inference/lance/gradio_demo.py
index 3d6b240d3..a8297288a 100644
--- a/examples/offline_inference/lance/gradio_demo.py
+++ b/examples/offline_inference/lance/gradio_demo.py
@@ -141,7 +141,6 @@ _PARAMS_VIDEO = None
 def parse_args():
     p = argparse.ArgumentParser()
     p.add_argument("--model", required=True, help="Parent snapshot dir containing Lance_3B and Lance_3B_Video subdirs.")
-    p.add_argument("--deploy-config", default="vllm_omni/deploy/lance.yaml", help="Lance deploy YAML.")
     p.add_argument("--host", default="0.0.0.0")
     p.add_argument("--port", type=int, default=7860)
     p.add_argument("--share", action="store_true")
@@ -198,7 +197,6 @@ def _resolve_ckpts(snapshot: str):
 
 def _init_models(
     snapshot: str,
-    deploy_config: str,
     ulysses_degree: int = 1,
     replicas_per_omni: int = 1,
 ):
@@ -290,7 +288,20 @@ def _init_models(
         f"  replicas={replicas_per_omni}...",
         flush=True,
     )
-    img_kwargs: dict = {"model": img_ckpt, "deploy_config": deploy_config}
+    # Lance is single-stage diffusion — no deploy YAML.  The defaults
+    # below match the (now-removed) ``vllm_omni/deploy/lance.yaml`` so
+    # ``create_default_diffusion`` produces the same stage config.
+    lance_engine_defaults = dict(
+        pipeline="lance",
+        max_num_batched_tokens=32768,
+        max_num_seqs=1,
+        enforce_eager=True,
+        trust_remote_code=True,
+        enable_prefix_caching=False,
+        async_chunk=False,
+    )
+
+    img_kwargs: dict = {"model": img_ckpt, **lance_engine_defaults}
     if img_devices != "0":
         img_kwargs["stage_0_devices"] = img_devices
     if use_sp:
@@ -309,7 +320,7 @@ def _init_models(
         f"  replicas={replicas_per_omni}...",
         flush=True,
     )
-    vid_kwargs: dict = {"model": vid_ckpt, "deploy_config": deploy_config}
+    vid_kwargs: dict = {"model": vid_ckpt, **lance_engine_defaults}
     vid_kwargs["stage_0_devices"] = vid_devices
     if use_sp:
         vid_kwargs["ulysses_degree"] = ulysses_degree
@@ -1240,7 +1251,6 @@ def main():
     print("=" * 60, flush=True)
     _init_models(
         args.model,
-        args.deploy_config,
         ulysses_degree=args.ulysses_degree,
         replicas_per_omni=args.replicas_per_omni,
     )
diff --git a/examples/offline_inference/ming_flash_omni/README.md b/examples/offline_inference/ming_flash_omni/README.md
index 4b798619e..54e993cab 100644
--- a/examples/offline_inference/ming_flash_omni/README.md
+++ b/examples/offline_inference/ming_flash_omni/README.md
@@ -2,12 +2,13 @@
 
 [Ming-flash-omni-2.0](https://github.com/inclusionAI/Ming) is an omni-modal model supporting text, image, video, and audio understanding, with text and speech outputs.
 
-vLLM-Omni supports two deployment modes:
+vLLM-Omni supports three deployment modes:
 
 | Mode | Deploy config | Output |
 |------|--------------|--------|
 | Thinker + Talker (omni-speech, default) | `vllm_omni/deploy/ming_flash_omni.yaml` | Text + Audio |
 | Thinker only (multimodal understanding) | `vllm_omni/deploy/ming_flash_omni_thinker_only.yaml` | Text |
+| Thinker + Imagegen (text-to-image / img2img) | `vllm_omni/deploy/ming_flash_omni_image.yaml` | Image (online-serving only at the moment) |
 
 For standalone TTS (talker only), see the [Ming-flash-omni-TTS section in the Text-To-Speech hub](../text_to_speech/README.md#ming-flash-omni-tts).
 
@@ -87,6 +88,17 @@ The default deploy YAML allocates thinker on GPUs 0–3 and talker on GPU 3 for
 
 Pass `--deploy-config /path/to/your_deploy.yaml` to any of the commands above to override the bundled deploy config.
 
+### Image generation (text-to-image / img2img)
+
+The diffusion-stage image-generation path is currently only wired through
+the online OpenAI-compatible chat endpoint (`/v1/chat/completions` with
+`"modalities": ["image"]`). For payloads, optional knobs
+(`steps`/`cfg`/`seed`/`byte5_text`/`negative_prompt`), and the img2img
+reference-image flow, see the
+[image-gen section in the online-serving README](../../online_serving/ming_flash_omni/README.md#image-generation-text-to-image).
+`end2end.py` does not yet exercise the imagegen stage; that is tracked as
+follow-up work.
+
 ## Online serving
 
 For online serving via the OpenAI-compatible API, see [examples/online_serving/ming_flash_omni/README.md](../../online_serving/ming_flash_omni/README.md).
diff --git a/examples/offline_inference/sensenova_u1/README.md b/examples/offline_inference/sensenova_u1/README.md
index 8052c90d3..3943b8d4e 100644
--- a/examples/offline_inference/sensenova_u1/README.md
+++ b/examples/offline_inference/sensenova_u1/README.md
@@ -37,6 +37,10 @@ python end2end.py --modality text2text \
 # Custom resolution (image generation)
 python end2end.py --prompt "A futuristic cityscape at sunset" \
                   --width 2048 --height 1024 --think
+
+# Text-to-image with Cache-DiT acceleration
+python end2end.py --prompt "A cute cat" \
+                  --cache-backend cache_dit \
 ```
 
 > **Note**: Default configuration works on a single **NVIDIA A100 (80GB)** or **H100** GPU.
@@ -135,6 +139,8 @@ python end2end.py \
 | `--tensor-parallel-size` | 1 | Number of GPUs for tensor parallelism |
 | `--enforce-eager` | False | Disable torch.compile |
 | `--enable-cpu-offload` | False | Enable module-wise (sequential) CPU offload to reduce peak VRAM |
+| `--cache-backend` | None | Set to `cache_dit` for Cache-DiT acceleration |
+| `--enable-cache-dit-summary` | False | Print Cache-DiT cache statistics after generation |
 
 ## Reducing GPU Memory Usage
 
diff --git a/examples/offline_inference/sensenova_u1/end2end.py b/examples/offline_inference/sensenova_u1/end2end.py
index 12f28da6f..ebb0d7f24 100644
--- a/examples/offline_inference/sensenova_u1/end2end.py
+++ b/examples/offline_inference/sensenova_u1/end2end.py
@@ -155,6 +155,18 @@ def parse_args():
             "of extra CPU<->GPU transfers."
         ),
     )
+    parser.add_argument(
+        "--cache-backend",
+        type=str,
+        default=None,
+        choices=["cache_dit"],
+        help="Cache backend to use for image generation acceleration.",
+    )
+    parser.add_argument(
+        "--enable-cache-dit-summary",
+        action="store_true",
+        help="Enable cache-dit summary logging after diffusion forward passes.",
+    )
 
     nullify_stage_engine_defaults(parser)
     return parser.parse_args()
@@ -182,6 +194,8 @@ def main():
         tensor_parallel_size=args.tensor_parallel_size,
         enforce_eager=args.enforce_eager,
         enable_cpu_offload=args.enable_cpu_offload,
+        cache_backend=args.cache_backend,
+        enable_cache_dit_summary=args.enable_cache_dit_summary,
     )
 
     extra_args = {
@@ -225,6 +239,7 @@ def main():
     print(f"  Seed           : {args.seed}")
     print(f"  Think mode     : {args.think}")
     print(f"  TP size        : {args.tensor_parallel_size}")
+    print(f"  Cache backend  : {args.cache_backend or 'none'}")
     print(f"{'=' * 60}\n")
 
     # Build prompt dict
diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md
index c71773972..02538b40f 100644
--- a/examples/offline_inference/text_to_image/README.md
+++ b/examples/offline_inference/text_to_image/README.md
@@ -35,6 +35,7 @@ This folder provides several entrypoints for experimenting with text-to-image di
 | `black-forest-labs/FLUX.2-klein-9B` | 1024 x 1024 | 37.1 | 32.3 |
 | `black-forest-labs/FLUX.2-dev` | 1024 x 1024 | 65.7 | >80 (CPU offload required) |
 | `HunyuanImage-3.0` | 1024 x 1024 | 80.0 (TP≥3)  | 160 |
+| `HiDream-I1-Full` | 1024 x 1024 | 63.7  | 57.7 |
 
 !!! info
 *Peak VRAM:  based on basic single-card usage, batch size =1, without any acceleration/optimization features. FLUX.2-dev requires `--enable-cpu-offload` on a single 80 GiB GPU.
@@ -94,6 +95,7 @@ python text_to_image.py \
 | `--lora-scale` | float | `1.0` | Scale factor for LoRA weights |
 | `--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). |
 
 **NextStep-1.1 specific arguments:**
 
@@ -177,6 +179,24 @@ python examples/offline_inference/text_to_image/text_to_image.py \
   --output flux2-dev.png
 ```
 
+
+### HiDream-I1-Full Models
+
+The `--auxiliary-text-encoder` parameter is required when running HiDream‑I1‑Full:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model HiDream-ai/HiDream-I1-Full \
+  --prompt "The setting sun of late autumn dyes the riverside with a warm orange hue" \
+  --seed 42 \
+  --guidance-scale 5.0 \
+  --tensor-parallel-size 1 \
+  --num-images-per-prompt 1 \
+  --num-inference-steps 50 \
+  --auxiliary-text-encoder meta-llama/Meta-Llama-3.1-8B-Instruct \
+  --output /output.png
+```
+
 ### Batch Requests (Multiple Prompts)
 
 You can pass multiple prompts in a single `generate` call.
diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py
index c0fd337bd..443b21acc 100644
--- a/examples/offline_inference/text_to_image/text_to_image.py
+++ b/examples/offline_inference/text_to_image/text_to_image.py
@@ -300,6 +300,12 @@ def parse_args() -> argparse.Namespace:
         default=None,
         help=("Custom system prompt. Used when --use-system-prompt is custom. "),
     )
+    parser.add_argument(
+        "--auxiliary-text-encoder",
+        type=str,
+        default=None,
+        help="Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).",
+    )
     current_omni_platform.pre_register_and_update(parser)
     from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 
@@ -392,6 +398,7 @@ def main():
         "profiler_config": args.profiler_config,
         "init_timeout": args.init_timeout,
         "stage_init_timeout": args.stage_init_timeout,
+        "auxiliary_text_encoder": args.auxiliary_text_encoder,
         **lora_args,
         **quant_kwargs,
     }
diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md
index 722b9eeb3..a0e9b529b 100644
--- a/examples/offline_inference/text_to_speech/README.md
+++ b/examples/offline_inference/text_to_speech/README.md
@@ -16,6 +16,7 @@ list of supported architectures across all modalities, see
 |---|---|---|---|---|---|---|
 | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz |
 | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz |
+| GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
 | Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz |
 | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | single (AR + codec) | ✓ (required) | ✓ | voice_clone, continuation | 48 kHz |
 | OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
@@ -92,6 +93,35 @@ Streaming is enabled by default via `async_chunk: true` in `vllm_omni/deploy/cos
 
 ---
 
+## GLM-TTS
+
+2-stage TTS pipeline (AR + DiT flow-matching) at 24 kHz. Every request requires reference audio and its transcript for zero-shot voice cloning.
+
+### Quick start
+```bash
+python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
+    --model zai-org/GLM-TTS \
+    --text "你好，这是语音合成测试。" \
+    --ref-audio /path/to/reference.wav \
+    --ref-text "这是参考音频的文本内容。" \
+    --output-dir ./output
+```
+
+### Architecture
+```
+Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 kHz)
+        (Llama-based)    (32k vocab)      (Flow Matching)
+```
+
+### Notes
+- `--ref-audio` and `--ref-text` are **required** together; GLM-TTS does not support text-only synthesis.
+- Reference audio should be 3-10 seconds.
+- First run may be slow due to lazy loading of WhisperVQ tokenizer and CampPlus ONNX speaker embedder.
+- Default sampling: temperature=1.0, top_k=25, top_p=0.8 (RAS method).
+- The `--model` path should point to the repository root (not `llm/` subdirectory).
+
+---
+
 ## Fish Speech S2 Pro
 
 4B dual-AR text-to-speech model from FishAudio with the DAC codec at 44.1 kHz.
@@ -241,6 +271,14 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
     --lang zh
 ```
 
+### Seed for Reproducibility
+```bash
+python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
+    --model k2-fsa/OmniVoice \
+    --text "Hello, this is a test." \
+    --seed 42
+```
+
 ### Notes
 - Stage 0 (Generator): Qwen3-0.6B with 32-step iterative unmasking.
 - Stage 1 (Decoder): HiggsAudioV2 RVQ + DAC at 24 kHz.
diff --git a/examples/offline_inference/text_to_speech/glm_tts/end2end.py b/examples/offline_inference/text_to_speech/glm_tts/end2end.py
new file mode 100644
index 000000000..d1f19f8a5
--- /dev/null
+++ b/examples/offline_inference/text_to_speech/glm_tts/end2end.py
@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GLM-TTS End-to-End Offline Inference Example.
+
+GLM-TTS is a two-stage TTS system:
+  - Stage 0 (AR): Llama-based model generates speech tokens from text
+  - Stage 1 (DiT): Flow matching model converts speech tokens to audio
+
+Usage:
+    # Sync two-stage (default)
+    python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
+        --model /path/to/GLM-TTS \
+        --text "你好，这是一个语音合成测试。" \
+        --ref-audio /path/to/reference.wav \
+        --ref-text "参考音频的转录文本。" \
+        --output-dir ./output
+
+    # Async chunk mode (streaming DiT)
+    python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
+        --model /path/to/GLM-TTS --async-chunk \
+        --text "你好，这是一个语音合成测试。" \
+        --ref-audio /path/to/reference.wav \
+        --ref-text "参考音频的转录文本。" \
+        --output-dir ./output
+"""
+
+import base64
+import io
+import logging
+import os
+import tempfile
+import time
+from typing import Any
+from urllib.request import urlopen
+
+import numpy as np
+import soundfile as sf
+import torch
+import yaml
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni import Omni
+from vllm_omni.model_executor.models.glm_tts.glm_tts import build_glm_tts_prefill_metadata
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_DEPLOY_CONFIG = os.path.join(
+    os.path.dirname(__file__),
+    "..",
+    "..",
+    "..",
+    "..",
+    "vllm_omni",
+    "deploy",
+    "glm_tts.yaml",
+)
+SAMPLE_RATE = 24000
+
+
+def _load_ref_audio(ref_audio: str) -> tuple[torch.Tensor, int]:
+    """Load reference audio from file path, URL, or data URI."""
+    if ref_audio.startswith(("http://", "https://")):
+        with urlopen(ref_audio, timeout=60) as response:
+            audio_obj: Any = io.BytesIO(response.read())
+    elif ref_audio.startswith("data:"):
+        _, _, encoded = ref_audio.partition(",")
+        audio_obj = io.BytesIO(base64.b64decode(encoded))
+    else:
+        audio_obj = ref_audio
+    wav_np, sr = sf.read(audio_obj, dtype="float32")
+    if wav_np.ndim > 1:
+        wav_np = wav_np.mean(axis=1)
+    return torch.from_numpy(wav_np), int(sr)
+
+
+def _concat_audio(audio_val: Any) -> np.ndarray:
+    """Concatenate audio tensors from multimodal output."""
+    if isinstance(audio_val, list):
+        tensors = [torch.as_tensor(t).float().reshape(-1) for t in audio_val if t is not None]
+        if not tensors:
+            return np.zeros((0,), dtype=np.float32)
+        return torch.cat(tensors, dim=-1).cpu().numpy().astype(np.float32, copy=False)
+    if isinstance(audio_val, torch.Tensor):
+        return audio_val.float().cpu().numpy().reshape(-1)
+    return np.asarray(audio_val, dtype=np.float32).reshape(-1)
+
+
+def _extract_sample_rate(audio_mm: dict) -> int:
+    """Extract sample rate from multimodal output dict."""
+    sr_raw = audio_mm.get("sr", SAMPLE_RATE)
+    if isinstance(sr_raw, list):
+        sr_raw = sr_raw[-1] if sr_raw else SAMPLE_RATE
+    if hasattr(sr_raw, "item"):
+        return int(sr_raw.item())
+    return int(sr_raw)
+
+
+def _modify_deploy_config(base_path: str, async_chunk: bool) -> str:
+    """Build deploy config with explicit sync/async mode and eager execution.
+
+    Mirrors the logic in ``tests/e2e/offline_inference/test_glm_tts.py``
+    (``_get_deploy_config``) so that example runs match CI behavior.
+    """
+    with open(base_path) as f:
+        cfg = yaml.safe_load(f)
+    cfg["async_chunk"] = async_chunk
+    for stage in cfg.get("stages", []):
+        stage["enforce_eager"] = True
+        if stage.get("stage_id") == 0:
+            stage["async_scheduling"] = bool(async_chunk)
+    tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False, prefix="glm_tts_")
+    yaml.dump(cfg, tmp)
+    tmp.close()
+    return tmp.name
+
+
+def main(args):
+    """Run offline GLM-TTS inference."""
+    os.makedirs(args.output_dir, exist_ok=True)
+    base_deploy_config = args.deploy_config or DEFAULT_DEPLOY_CONFIG
+    deploy_config_path = _modify_deploy_config(base_deploy_config, args.async_chunk)
+
+    ref_audio_wav, ref_audio_sr = _load_ref_audio(args.ref_audio)
+    if not args.ref_text:
+        raise ValueError("GLM-TTS requires --ref-audio and --ref-text.")
+
+    inputs = [
+        {
+            "prompt": args.text,
+            "multi_modal_data": {
+                "audio": (ref_audio_wav.float().cpu().numpy(), ref_audio_sr),
+            },
+            "modalities": ["audio"],
+            "mm_processor_kwargs": {"prompt_text": args.ref_text},
+            "additional_information": build_glm_tts_prefill_metadata(
+                args.model,
+                args.text,
+                args.ref_text,
+            ),
+        }
+    ]
+
+    omni = Omni(
+        model=args.model,
+        stage_configs_path=deploy_config_path,
+        log_stats=args.log_stats,
+        stage_init_timeout=args.stage_init_timeout,
+    )
+
+    t_start = time.perf_counter()
+    outputs = omni.generate(inputs)
+    elapsed = (time.perf_counter() - t_start) * 1000
+
+    assert outputs, "No outputs returned"
+    audio_mm = outputs[0].multimodal_output
+    assert "audio" in audio_mm, "No audio output found"
+
+    audio = _concat_audio(audio_mm["audio"])
+    sr = _extract_sample_rate(audio_mm)
+    out_path = os.path.join(args.output_dir, "output.wav")
+    sf.write(out_path, audio, samplerate=sr, format="WAV")
+
+    logger.info("Saved %s (%.2fs @ %dHz)", out_path, len(audio) / sr, sr)
+    logger.info("Total inference: %.1f ms", elapsed)
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="GLM-TTS Text-to-Speech Example")
+    parser.add_argument("--model", type=str, required=True, help="Model path")
+    parser.add_argument("--text", type=str, default="你好，这是一个语音合成测试。")
+    parser.add_argument("--output-dir", type=str, default="./output")
+    parser.add_argument("--ref-audio", type=str, required=True, help="Reference WAV path/URL")
+    parser.add_argument("--ref-text", type=str, required=True, help="Transcript of ref audio")
+    parser.add_argument("--deploy-config", type=str, default=None)
+    parser.add_argument(
+        "--async-chunk",
+        action="store_true",
+        default=False,
+        help="Enable async_chunk mode (streaming DiT). Default: sync two-stage.",
+    )
+    parser.add_argument("--log-stats", action="store_true")
+    parser.add_argument("--stage-init-timeout", type=int, default=600)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    )
+    main(parse_args())
diff --git a/examples/offline_inference/text_to_speech/higgs_audio_v2/README.md b/examples/offline_inference/text_to_speech/higgs_audio_v2/README.md
new file mode 100644
index 000000000..95a200d88
--- /dev/null
+++ b/examples/offline_inference/text_to_speech/higgs_audio_v2/README.md
@@ -0,0 +1,38 @@
+# higgs-audio v2 — offline example
+
+Drives Stage 0 (DualFFN talker) + Stage 1 (HiggsAudio codec) for `bosonai/higgs-audio-v2-generation-3B-base` end-to-end through the vLLM-Omni engine and writes a 24 kHz mono WAV per prompt.
+
+## Prerequisites
+
+Voice clone needs `transformers>=5.3.0` — vllm-omni loads the audio codec via HF's `HiggsAudioV2TokenizerModel`, instantiated from the `k2-fsa/OmniVoice/audio_tokenizer/` subdirectory (only that ~806 MB subdir is downloaded). The boson-ai standalone tokenizer repo's `model.safetensors` is actually a copy of the 3B talker LM, so HF can't load it directly; the k2 bundle ships the same codec weights repackaged with HF-compatible key naming.
+
+```bash
+pip install -U "transformers>=5.3.0"
+```
+
+## Quick start
+
+Plain TTS:
+```bash
+python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
+    --texts "Hello world." "The quick brown fox jumps over the lazy dog." \
+    --output-dir results/wavs
+```
+
+## Voice cloning
+
+Pass both `--ref-audio` and `--ref-text` together:
+```bash
+python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
+    --texts "Hello, this is a cloned voice." \
+    --ref-audio /path/to/reference.wav \
+    --ref-text  "Exact transcript spoken in reference.wav." \
+    --output-dir results/wavs
+```
+
+## Notes
+
+- Output: 24 kHz mono WAV.
+- Deploy config: `vllm_omni/deploy/higgs_audio_v2.yaml` (auto-loaded by HF `model_type`).
+- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
+- For online serving, see [`examples/online_serving/text_to_speech/higgs_audio_v2/`](../../../online_serving/text_to_speech/higgs_audio_v2/).
diff --git a/examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py b/examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py
new file mode 100644
index 000000000..62a346621
--- /dev/null
+++ b/examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py
@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline higgs-audio v2 inference example.
+
+Runs Stage 0 (DualFFN talker) + Stage 1 (HiggsAudio codec) end-to-end
+through the vLLM-Omni engine without going through the HTTP server, and
+saves a 24 kHz mono WAV per prompt.
+
+Example:
+
+    python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \\
+        --texts "Hello world." \\
+                "The quick brown fox jumps over the lazy dog." \\
+        --output-dir results/wavs
+"""
+
+from __future__ import annotations
+
+import os
+
+# DeepGEMM FP8 kernels require an optional backend that may not be installed.
+# Disable the warmup before importing vLLM so engine startup falls back to the
+# regular gemm path. Users with deep_gemm installed can override these.
+os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
+os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
+
+import time
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+import torch
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni import Omni
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
+
+SAMPLE_RATE = 24_000
+DEFAULT_TEXTS = (
+    "Hello world.",
+    "The quick brown fox jumps over the lazy dog.",
+)
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="Offline higgs-audio v2 inference")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="bosonai/higgs-audio-v2-generation-3B-base",
+        help="Stage-0 talker model id or path.",
+    )
+    parser.add_argument(
+        "--texts",
+        type=str,
+        nargs="+",
+        default=list(DEFAULT_TEXTS),
+        help="One or more plain-text prompts to synthesize.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="results/wavs",
+        help="Directory to write per-prompt WAV files.",
+    )
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        default=None,
+        help="Override the deploy config path. Auto-loads "
+        "vllm_omni/deploy/higgs_audio_v2.yaml from the HF model_type by default.",
+    )
+    parser.add_argument(
+        "--max-new-tokens",
+        type=int,
+        default=500,
+        help="Cap on Stage-0 codec frames per request.",
+    )
+    parser.add_argument(
+        "--ref-audio",
+        type=str,
+        default=None,
+        help="Reference clip for voice clone (path to a WAV file). Paired with --ref-text.",
+    )
+    parser.add_argument(
+        "--ref-text",
+        type=str,
+        default=None,
+        help="Transcript of the reference clip. Required when --ref-audio is set.",
+    )
+    nullify_stage_engine_defaults(parser)
+    return parser.parse_args()
+
+
+def _slugify(text: str) -> str:
+    import re
+
+    slug = re.sub(r"\s+", "_", text.strip().lower())
+    slug = re.sub(r"[^a-z0-9_]+", "", slug)
+    return slug[:48] or "prompt"
+
+
+def _extract_pcm(multimodal_output: dict) -> torch.Tensor:
+    """Pull the final concatenated PCM tensor out of a request's multimodal_output."""
+    audio = multimodal_output.get("model_outputs")
+    if audio is None:
+        audio = multimodal_output.get("audio")
+    if audio is None:
+        raise ValueError(f"no audio key in multimodal_output: {list(multimodal_output.keys())}")
+    if isinstance(audio, list):
+        valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
+        if not valid:
+            raise ValueError("audio list is empty")
+        return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
+    return torch.as_tensor(audio).float().cpu().reshape(-1)
+
+
+def _pcm_to_int16(pcm: torch.Tensor) -> np.ndarray:
+    arr = pcm.numpy()
+    if arr.dtype.kind == "f":
+        arr = np.clip(arr, -1.0, 1.0)
+        arr = (arr * 32767.0).astype(np.int16)
+    else:
+        arr = arr.astype(np.int16)
+    return arr
+
+
+def main():
+    args = parse_args()
+    if (args.ref_audio is None) != (args.ref_text is None):
+        raise SystemExit("--ref-audio and --ref-text must be supplied together")
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    engine = Omni(model=args.model, deploy_config=args.deploy_config)
+
+    # Build prompt_token_ids using the same path serving_speech.py takes online.
+    from transformers import AutoProcessor
+
+    from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
+        build_plain_text_prompt,
+        build_voice_clone_prompt,
+        input_ids_to_python_list,
+    )
+
+    processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True)
+
+    # Voice-clone path: load the reference clip once. The HF processor will
+    # encode it via the bundled HiggsAudioV2TokenizerModel each time we build
+    # a prompt below. This is cheap on CPU for a few-second clip.
+    ref_wav: np.ndarray | None = None
+    ref_sr: int | None = None
+    if args.ref_audio is not None:
+        ref_wav, ref_sr = sf.read(args.ref_audio, always_2d=False)
+        if ref_wav.ndim == 2:
+            ref_wav = ref_wav.mean(axis=1)
+
+    print(f"Model       : {args.model}")
+    print(f"Prompts     : {len(args.texts)}")
+    print(f"Output dir  : {output_dir}")
+    print(f"Voice clone : {'yes' if ref_wav is not None else 'no'}")
+
+    # Run one prompt at a time. The Stage-0 talker's per-slot audio state is
+    # request-scoped; submitting multiple prompts in the same engine.generate()
+    # call would batch them in the AR runner and exercise a code path that is
+    # not validated for this model yet.
+    total_elapsed = 0.0
+    total_dur = 0.0
+    for text in args.texts:
+        if ref_wav is not None:
+            out = build_voice_clone_prompt(processor, text, ref_wav, int(ref_sr), args.ref_text)
+            prompt = {
+                "prompt_token_ids": out["prompt_token_ids"],
+                # Bare tensors (NOT list-wrapped): the msgspec serializer in
+                # vllm_omni.data_entry_keys routes torch.Tensor → tensor_data
+                # and list[Tensor] → list_data (which silently strips tensors).
+                "additional_information": {
+                    "audio_input_ids": out["audio_input_ids"],
+                    "audio_input_ids_mask": out["audio_input_ids_mask"],
+                },
+            }
+        else:
+            inputs = build_plain_text_prompt(processor, text)
+            prompt = {"prompt_token_ids": input_ids_to_python_list(inputs)}
+        t_start = time.perf_counter()
+        outputs = engine.generate([prompt])
+        elapsed = time.perf_counter() - t_start
+        total_elapsed += elapsed
+
+        mm = outputs[0].outputs[0].multimodal_output
+        pcm = _extract_pcm(mm)
+        slug = _slugify(text)
+        out_path = output_dir / f"{slug}.wav"
+        sf.write(str(out_path), _pcm_to_int16(pcm), SAMPLE_RATE, format="WAV", subtype="PCM_16")
+        dur = pcm.numel() / SAMPLE_RATE
+        total_dur += dur
+        print(f"  {slug:<50} dur={dur:5.2f}s  -> {out_path}")
+
+    rtf = total_elapsed / total_dur if total_dur > 0 else float("inf")
+    print(f"Total infer : {total_elapsed:.2f}s  total audio: {total_dur:.2f}s  RTF: {rtf:.3f}")
+
+
+if __name__ == "__main__":
+    os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    main()
diff --git a/examples/offline_inference/text_to_speech/omnivoice/end2end.py b/examples/offline_inference/text_to_speech/omnivoice/end2end.py
index 16d44fe7e..f3c23e60b 100644
--- a/examples/offline_inference/text_to_speech/omnivoice/end2end.py
+++ b/examples/offline_inference/text_to_speech/omnivoice/end2end.py
@@ -80,6 +80,12 @@ def run_e2e():
         default=600,
         help="Stage initialization timeout in seconds",
     )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=None,
+        help="Random seed for generation",
+    )
     nullify_stage_engine_defaults(parser)
     args = parser.parse_args()
 
@@ -122,7 +128,7 @@ def run_e2e():
     if mm_processor_kwargs:
         prompts["mm_processor_kwargs"] = mm_processor_kwargs
 
-    sampling_params_list = [OmniDiffusionSamplingParams()]
+    sampling_params_list = [OmniDiffusionSamplingParams(extra_args={"seed": args.seed})]
 
     print(f"Generating speech for: {args.text}")
 
diff --git a/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py b/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py
index 6d026fb86..d9f96838a 100644
--- a/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py
+++ b/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py
@@ -43,8 +43,8 @@ def _estimate_prompt_len(
     """
     try:
         from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
-        from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
-            Qwen3TTSTalkerForConditionalGeneration,
+        from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
+            Qwen3TTSPromptEmbedsBuilder,
         )
 
         if model_name not in _cache:
@@ -124,7 +124,7 @@ def _estimate_prompt_len(
             except Exception:
                 return None
 
-        return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
+        return Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
             additional_information=additional_information,
             task_type=task_type,
             tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
diff --git a/examples/offline_inference/vace/vace_video_generation.md b/examples/offline_inference/vace/vace_video_generation.md
index bbaf99452..8d8b00e48 100644
--- a/examples/offline_inference/vace/vace_video_generation.md
+++ b/examples/offline_inference/vace/vace_video_generation.md
@@ -1,24 +1,105 @@
 # VACE Video Generation
 
-[VACE](https://github.com/ali-vilab/VACE) (Video All-in-one Creation Engine) supports multiple video tasks through a single model.
+Generate videos from text prompts, images, or video conditions using vLLM-Omni's VACE diffusion pipeline.
 
-| Model | Architecture | Model Weights (bf16) | HuggingFace |
-|-------|-------------|----------------------|-------------|
-| Wan2.1-VACE (1.3B) | Wan2.1 | ~10 GB | [Wan-AI/Wan2.1-VACE-1.3B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers) |
-| Wan2.1-VACE (14B) | Wan2.1 | ~38 GB | [Wan-AI/Wan2.1-VACE-14B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) |
+- `vace_video_generation.py`: command-line script for multi-mode video generation with advanced options.
 
-## Text-to-Video (T2V)
+## Table of Contents
+
+- [Overview](#overview)
+- [Quick Start](#quick-start)
+- [Key Arguments](#key-arguments)
+- [More CLI Examples](#more-cli-examples)
+
+## Overview
+
+[VACE](https://github.com/ali-vilab/VACE) (Video All-in-one Creation Engine) supports multiple video tasks through a single unified model, including text-to-video, image-to-video, first-last-frame interpolation, inpainting, and reference image-guided generation.
+
+### Supported Models
+
+| Model | Architecture | Peak VRAM (GiB) * | Model Weights (GiB) | HuggingFace |
+|-------|-------------|-------------------|---------------------|-------------|
+| Wan2.1-VACE (1.3B) | Wan2.1 | TBD | ~10 | [Wan-AI/Wan2.1-VACE-1.3B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers) |
+| Wan2.1-VACE (14B) | Wan2.1 | TBD | ~38 | [Wan-AI/Wan2.1-VACE-14B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) |
+
+!!! info
+*Peak VRAM: based on basic single-card usage, 480x832 resolution, 81 frames, 30 inference steps, without any acceleration/optimization features.
+
+Default model: `Wan-AI/Wan2.1-VACE-14B-diffusers`
+
+## Quick Start
+
+### Python API
+
+Text-to-video generation:
+
+```python
+import torch
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+if __name__ == "__main__":
+    omni = Omni(model="Wan-AI/Wan2.1-VACE-1.3B-diffusers")
+    prompt = "A sleek robot stands in a vast warehouse filled with boxes"
+    outputs = omni.generate(
+        prompt,
+        OmniDiffusionSamplingParams(
+            height=480,
+            width=832,
+            num_frames=81,
+            num_inference_steps=30,
+            guidance_scale=5.0,
+        ),
+    )
+    video = outputs[0].images
+
+    from diffusers.utils import export_to_video
+    export_to_video(list(video[0]), "t2v_output.mp4", fps=16)
+    omni.close()
+```
+
+### Local CLI Usage
 
 ```bash
 python vace_video_generation.py \
   --mode t2v \
   --prompt "A sleek robot stands in a vast warehouse filled with boxes" \
   --height 480 --width 832 --num-frames 81 \
-  --num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
   --output t2v_output.mp4
 ```
 
-## Image-to-Video (I2V)
+## Key Arguments
+
+| Argument | Type | Default | Description |
+| -------- | ---- | ------- | ----------- |
+| `--mode` | str | `"t2v"` | VACE task mode: `t2v`, `i2v`, `v2lf`, `flf2v`, `inpaint`, `r2v` |
+| `--model` | str | `"Wan-AI/Wan2.1-VACE-14B-diffusers"` | Model ID or local path |
+| `--prompt` | str | `"A cat walking in a garden"` | Text description of desired video |
+| `--negative-prompt` | str | `""` | Negative prompt for classifier-free guidance |
+| `--image` | str | `None` | Input image path (for I2V, R2V, FLF2V, inpaint modes) |
+| `--last-image` | str | `None` | Last frame image path (for FLF2V mode) |
+| `--height` | int | `480` | Output video height in pixels (should be a multiple of 16) |
+| `--width` | int | `832` | Output video width in pixels (should be a multiple of 16) |
+| `--num-frames` | int | `81` | Number of video frames to generate |
+| `--num-inference-steps` | int | `30` | Number of denoising steps (more steps = higher quality, slower) |
+| `--guidance-scale` | float | `5.0` | Classifier-free guidance scale |
+| `--flow-shift` | float | `5.0` | Scheduler flow shift parameter |
+| `--seed` | int | `42` | Random seed for deterministic sampling |
+| `--fps` | int | `16` | Frames per second for the saved MP4 |
+| `--output` | str | `"vace_output.mp4"` | Path to save the generated video |
+| `--vae-use-tiling` | flag | on | Enable VAE tiling for memory optimization |
+| `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
+| `--ulysses-degree` | int | `1` | Ulysses sequence parallel degree for multi-GPU inference |
+| `--ring-degree` | int | `1` | Ring sequence parallel degree for hybrid Ulysses + Ring inference |
+| `--tensor-parallel-size` | int | — | Tensor parallel size |
+| `--enforce-eager` | flag | off | Disable torch.compile |
+
+> If you encounter OOM errors, try `--vae-use-tiling` or multi-GPU parallelism options (`--ulysses-degree`, `--cfg-parallel-size`).
+
+## More CLI Examples
+
+### Image-to-Video (I2V)
 
 First frame is kept, remaining frames are generated:
 
@@ -31,7 +112,7 @@ python vace_video_generation.py \
   --output i2v_output.mp4
 ```
 
-## First-Last-Frame Interpolation (FLF2V)
+### First-Last-Frame Interpolation (FLF2V)
 
 ```bash
 python vace_video_generation.py \
@@ -42,7 +123,7 @@ python vace_video_generation.py \
   --output flf2v_output.mp4
 ```
 
-## Inpainting
+### Inpainting
 
 Center vertical stripe is masked and regenerated:
 
@@ -55,7 +136,7 @@ python vace_video_generation.py \
   --output inpaint_output.mp4
 ```
 
-## Reference Image-guided (R2V)
+### Reference Image-guided (R2V)
 
 ```bash
 python vace_video_generation.py \
@@ -65,24 +146,3 @@ python vace_video_generation.py \
   --height 480 --width 832 --num-frames 81 \
   --output r2v_output.mp4
 ```
-
-## Key Arguments
-
-- `--mode`: VACE task mode (`t2v`, `i2v`, `flf2v`, `inpaint`, `r2v`).
-- `--model`: Model ID (default: `Wan-AI/Wan2.1-VACE-1.3B-diffusers`).
-- `--image`: Input image for I2V, inpainting, and R2V modes.
-- `--last-image`: Last frame image for FLF2V mode.
-- `--prompt`: Text description of desired video.
-- `--height/--width`: Output resolution (default 480x832). Dimensions should be multiples of 16.
-- `--num-frames`: Number of frames (default 81).
-- `--guidance-scale`: CFG scale (default 5.0).
-- `--flow-shift`: Scheduler flow shift (default 5.0).
-- `--num-inference-steps`: Number of denoising steps (default 30).
-- `--fps`: Frames per second for the saved MP4 (default 16).
-- `--output`: Path to save the generated video.
-- `--vae-use-tiling`: Enable VAE tiling for memory optimization.
-- `--ulysses-degree`: Ulysses sequence parallelism degree for multi-GPU.
-- `--cfg-parallel-size`: CFG parallel size for multi-GPU.
-- `--tensor-parallel-size`: Tensor parallel size.
-
-> If you encounter OOM errors, try `--vae-use-tiling` or multi-GPU parallelism options.
diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md b/examples/offline_inference/x_to_video_audio/x_to_video_audio.md
index 7621a6058..f5071ce6d 100644
--- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md
+++ b/examples/offline_inference/x_to_video_audio/x_to_video_audio.md
@@ -101,3 +101,6 @@ Key arguments:
 - `--video-negative-prompt`: negative prompt for video generation.
 - `--audio-negative-prompt`: negative prompt for audio generation.
 - `--enable-cpu-offload`: enable CPU offload (defaults False).
+- `--cache-backend`: enable `cache_dit` for acceleration.
+- `--quantization`: online (dynamic) quantization method — `fp8` or `int8`.
+  (VAEs, the T5 text encoder, norms and modulation stay bf16.)
diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py
index 544c285ab..30de8aa45 100644
--- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py
+++ b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py
@@ -4,6 +4,7 @@
 import argparse
 import re
 import time
+from typing import Any
 
 import numpy as np
 from PIL import Image
@@ -41,6 +42,12 @@ def parse_args() -> argparse.Namespace:
         choices=[1, 2, 3, 4],
         help="Number of GPUs used for classifier free guidance parallel size (max 4 branches).",
     )
+    parser.add_argument(
+        "--tensor-parallel-size",
+        type=int,
+        default=1,
+        help=("Number of GPUs used for tensor parallelism (TP) inside the DiT."),
+    )
     parser.add_argument(
         "--video-negative-prompt",
         default="jitter, bad hands, blur, distortion",
@@ -52,6 +59,13 @@ def parse_args() -> argparse.Namespace:
         help="Negative prompt for audio.",
     )
     parser.add_argument("--output", default="dreamid_output.mp4", help="Output video path.")
+    parser.add_argument(
+        "--quantization",
+        type=str,
+        default=None,
+        choices=["fp8", "int8"],
+        help=("Online (dynamic) quantization method for the DreamID-Omni transformer. "),
+    )
     parser.add_argument(
         "--enable-cpu-offload",
         action="store_true",
@@ -63,6 +77,7 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
+    parser.add_argument("--cache-backend", type=str, default=None, choices=["cache_dit"], help="Cache backend.")
     parser.add_argument(
         "--use-hsdp",
         action="store_true",
@@ -99,6 +114,27 @@ def load_image_and_audio(image_paths, audio_paths):
     return image, audio
 
 
+def _extract_peak_memory_mb(result: Any) -> float:
+    """Pull worker-reported peak VRAM (MiB) generation result.
+
+    Mirrors vllm_omni/entrypoints/openai/serving_video.py:_extract_peak_memory_mb.
+    """
+    if isinstance(result, list):
+        result = result[0] if result else None
+    if result is None:
+        return 0.0
+    val = getattr(result, "peak_memory_mb", 0.0)
+    if not val:
+        inner = getattr(result, "request_output", None)
+        if isinstance(inner, list):
+            inner = inner[0] if inner else None
+        val = getattr(inner, "peak_memory_mb", 0.0)
+    try:
+        return float(val or 0.0)
+    except (TypeError, ValueError):
+        return 0.0
+
+
 def main() -> None:
     args = parse_args()
     if args.prompt is None and args.prompt_file is None:
@@ -141,18 +177,39 @@ def main() -> None:
 
     parallel_config = DiffusionParallelConfig(
         cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
         use_hsdp=args.use_hsdp,
         hsdp_shard_size=args.hsdp_shard_size,
         hsdp_replicate_size=args.hsdp_replicate_size,
     )
 
-    omni = Omni(
+    cache_config = None
+    if args.cache_backend == "cache_dit":
+        cache_config = {
+            "Fn_compute_blocks": 1,
+            "Bn_compute_blocks": 0,
+            "max_warmup_steps": 4,
+            "max_cached_steps": 20,
+            "residual_diff_threshold": 0.24,
+            "max_continuous_cached_steps": 3,
+            "enable_taylorseer": False,
+            "taylorseer_order": 1,
+            "scm_steps_mask_policy": None,
+            "scm_steps_policy": "dynamic",
+        }
+
+    omni_kwargs = dict(
         model=args.model,
         parallel_config=parallel_config,
         model_type=args.model_type,
         enable_cpu_offload=args.enable_cpu_offload,
         enable_layerwise_offload=args.enable_layerwise_offload,
+        cache_backend=args.cache_backend,
+        cache_config=cache_config,
     )
+    if args.quantization is not None:
+        omni_kwargs["quantization"] = args.quantization
+    omni = Omni(**omni_kwargs)
     start = time.perf_counter()
     outputs = omni.generate(prompt, sampling_params)
     elapsed = time.perf_counter() - start
@@ -190,6 +247,9 @@ def main() -> None:
         f.write(video_bytes)
     print(f"Saved generated video to {output_path}")
     print(f"Total time: {elapsed:.2f}s")
+    peak_mb = _extract_peak_memory_mb(outputs)
+    if peak_mb:
+        print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
 
 
 if __name__ == "__main__":
diff --git a/examples/online_serving/ming_flash_omni/README.md b/examples/online_serving/ming_flash_omni/README.md
index 533a30e1c..76b0bef9d 100644
--- a/examples/online_serving/ming_flash_omni/README.md
+++ b/examples/online_serving/ming_flash_omni/README.md
@@ -10,6 +10,7 @@ Please refer to [README.md](../../../README.md)
 |------|---------------|--------|
 | Thinker + Talker (omni-speech, default) | `vllm serve ... --omni` | Text + Audio |
 | Thinker only (multimodal understanding) | `vllm serve ... --omni --deploy-config vllm_omni/deploy/ming_flash_omni_thinker_only.yaml` | Text |
+| Thinker + Imagegen (text-to-image / img2img) | `vllm serve ... --omni --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml` | Image |
 
 For standalone TTS (talker only), see the [Ming-flash-omni-TTS section in the Text-To-Speech hub](../text_to_speech/README.md#ming-flash-omni-tts).
 
@@ -56,8 +57,91 @@ bash run_curl_multimodal_generation.sh use_image
 bash run_curl_multimodal_generation.sh use_audio
 bash run_curl_multimodal_generation.sh use_video
 bash run_curl_multimodal_generation.sh use_mixed_modalities
+bash run_curl_multimodal_generation.sh use_image_gen
 ```
 
+## Image generation (text-to-image)
+
+Ming-flash-omni-2.0 also exposes an image-generation (diffusion) stage.
+Launch with the image deploy YAML, which adds an image-gen stage behind
+the thinker:
+
+```bash
+vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \
+    --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml \
+    --stage-init-timeout 1800 \
+    --init-timeout 1800 \
+    --port 8091
+```
+
+Then request image output by passing `"modalities": ["image"]`:
+
+```bash
+curl -s http://127.0.0.1:8091/v1/chat/completions \
+    -H "Content-Type: application/json" \
+    -d '{
+      "model": "Jonathan1909/Ming-flash-omni-2.0",
+      "messages": [{"role": "user", "content": "Please draw a cute cat."}],
+      "modalities": ["image"]
+    }' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen.png
+```
+
+### Optional knobs
+
+Pass image-gen overrides via the diffusion-stage `sampling_params_list[1].extra_args.image_gen`:
+
+| Key | Default | Description |
+| --- | --- | --- |
+| `height` / `width` | from config (1024) | Output resolution (multiples of `vae_scale_factor * 2`, currently 16). |
+| `steps` | 30 | Number of FlowMatchEuler denoise steps. |
+| `cfg` | 2.0 | Classifier-free guidance scale. |
+| `seed` | 42 | Per-request RNG seed (deterministic when ByT5 is also seed-stable). |
+| `byte5_text` | (auto) | Override the glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "...". ` format. |
+| `negative_prompt` | empty | Real CFG negative conditioning (must be set on **stage-0 thinker** `sampling_params` so `expand_cfg_prompts` spawns the companion). |
+
+Example with all the knobs:
+
+```bash
+curl http://127.0.0.1:8091/v1/chat/completions \
+    -H "Content-Type: application/json" \
+    -d '{
+      "model": "Jonathan1909/Ming-flash-omni-2.0",
+      "modalities": ["image"],
+      "sampling_params_list": [
+        {"temperature":0.4,"top_p":0.9,"top_k":1,"max_tokens":1,"seed":42,
+         "extra_args":{"image_gen":{"negative_prompt":"ugly, blurry, distorted"}}},
+        {"seed":42,"extra_args":{"image_gen":{
+            "steps":6,"cfg":1.5,"height":512,"width":512,"seed":123,
+            "byte5_text":["理解与生成统一"]}}}
+      ],
+      "messages": [{"role": "user", "content": "Draw a poster."}]
+    }' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen_knobs.png
+```
+
+### img2img (reference image)
+
+Add an `image_url` content part to the user message; the parser routes it
+into the diffusion stage as `extra[reference_image]`:
+
+```jsonc
+"messages": [{
+  "role": "user",
+  "content": [
+    {"type": "text", "text": "Change the background to a sandy beach at sunset."},
+    {"type": "image_url", "image_url": {"url": "data:image/png;base64,<base64>"}}
+  ]
+}]
+```
+
+### GPU layout
+
+The shipped `ming_flash_omni_image.yaml` allocates the thinker on GPUs 0–3
+(TP=4) and the diffusion stage on GPU 4 (TP=1). Copy the YAML and edit
+`devices` per stage to relocate; with fewer GPUs available, drop
+the thinker TP to 2 and run the diffusion stage on a free card. Image-gen
+warmup takes roughly an extra 30–60 s on top of the thinker — set
+`--stage-init-timeout 1800` if the default 300 s is too tight.
+
 ## Modality control
 
 | `modalities` | Server config | Output |
@@ -65,6 +149,7 @@ bash run_curl_multimodal_generation.sh use_mixed_modalities
 | `["text"]` or omitted | Thinker only | Text |
 | `["audio"]` | Thinker + Talker | Audio (speech) |
 | `["text", "audio"]` | Thinker + Talker | Text + Audio |
+| `["image"]` | Thinker + Imagegen (image deploy YAML) | Image (PNG, base64 in `choices[0].message.content`) |
 
 For ready-to-copy curl examples (text / audio / multimodal input, SSE
 streaming, reasoning mode), see the recipe at
diff --git a/examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh b/examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh
index 715f23edd..92bb32c47 100755
--- a/examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh
+++ b/examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh
@@ -7,14 +7,15 @@ PORT="${PORT:-8091}"
 QUERY_TYPE="${1:-text}"
 
 # Validate query type
-if [[ ! "$QUERY_TYPE" =~ ^(text|use_audio|use_image|use_video|use_mixed_modalities)$ ]]; then
+if [[ ! "$QUERY_TYPE" =~ ^(text|use_audio|use_image|use_video|use_mixed_modalities|use_image_gen)$ ]]; then
     echo "Error: Invalid query type '$QUERY_TYPE'"
-    echo "Usage: $0 [text|use_audio|use_image|use_video|use_mixed_modalities]"
+    echo "Usage: $0 [text|use_audio|use_image|use_video|use_mixed_modalities|use_image_gen]"
     echo "  text: Text-only query"
     echo "  use_audio: Audio + Text query"
     echo "  use_image: Image + Text query"
     echo "  use_video: Video + Text query"
     echo "  use_mixed_modalities: Audio + Image + Video + Text query"
+    echo "  use_image_gen: Text-to-image (diffusion stage) — saves the generated PNG locally"
     exit 1
 fi
 
@@ -100,6 +101,43 @@ esac
 echo "Running query type: $QUERY_TYPE"
 echo ""
 
+if [[ "$QUERY_TYPE" == "use_image_gen" ]]; then
+    # Image-gen branch: dual-stage diffusion endpoint, modalities=["image"].
+    # The shipped ming_flash_omni_image.yaml runs the thinker on cards 0-3
+    # (TP=4) and the diffusion stage on card 4 (TP=1). See the image-gen
+    # section in README.md for knobs and img2img.
+    out_path="${OUT_PATH:-/tmp/ming_imagegen.png}"
+    request_body=$(cat <<'EOF'
+{
+  "model": "Jonathan1909/Ming-flash-omni-2.0",
+  "modalities": ["image"],
+  "messages": [
+    {"role": "user", "content": "Draw a beautiful girl with short black hair and red dress."}
+  ]
+}
+EOF
+)
+    output=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \
+        -X POST "http://localhost:${PORT}/v1/chat/completions" \
+        -H "Content-Type: application/json" \
+        -d "$request_body")
+    # Server returns base64-encoded PNG either as a data URL in
+    # ``choices[0].message.content[0].image_url.url`` or as a raw string.
+    b64=$(echo "$output" | jq -r '.choices[0].message.content
+        | if type == "array"
+            then (map(select(.image_url? != null)) | .[0].image_url.url // "")
+            elif type == "string" then . else "" end')
+    if [[ -z "$b64" ]]; then
+        echo "Error: no image returned. Raw response:"
+        echo "$output" | jq '.'
+        exit 1
+    fi
+    if [[ "$b64" == data:image* ]]; then b64="${b64#*,}"; fi
+    printf '%s' "$b64" | base64 --decode > "$out_path"
+    echo "Saved generated image: $out_path ($(stat -c%s "$out_path" 2>/dev/null || stat -f%z "$out_path") bytes)"
+    exit 0
+fi
+
 request_body=$(cat <<EOF
 {
   "model": "Jonathan1909/Ming-flash-omni-2.0",
diff --git a/examples/online_serving/qwen3_omni/README.md b/examples/online_serving/qwen3_omni/README.md
index a012cf21f..a13ab52df 100644
--- a/examples/online_serving/qwen3_omni/README.md
+++ b/examples/online_serving/qwen3_omni/README.md
@@ -18,10 +18,6 @@ Asynchronous chunk streaming operates as **enabled by default** within this bund
 Additionally, NPU, ROCm, and XPU per-platform configuration deltas are deterministically merged from the
 `platforms`: section of the corresponding YAML.
 
-**Note:** The OpenAI-style **`/v1/realtime`** WebSocket interface (facilitating streaming PCM audio input alongside audio and transcription output)
-is currently **unsupported** while the `async_chunk` configuration attribute is enabled.
-It is requisite to instantiate the default omni architecture or utilize a deployment configuration specifying `async_chunk: false` to facilitate real-time streaming sessions.
-
 To explicitly utilize a custom deployment YAML, mandate the configuration path accordingly:
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
@@ -289,7 +285,7 @@ python openai_realtime_client.py \
 | `--num-requests` | `1` | Number of sequential sessions (see `--concurrency`) |
 | `--concurrency` | `1` | Max concurrent WebSocket sessions when `--num-requests` > 1 |
 
-Ensure the server is running **without** `async_chunk` if you use `/v1/realtime`, for example:
+Ensure the server is running, for example:
 
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
diff --git a/examples/online_serving/sensenova_u1/README.md b/examples/online_serving/sensenova_u1/README.md
index 161e5de3b..86eadb1d6 100644
--- a/examples/online_serving/sensenova_u1/README.md
+++ b/examples/online_serving/sensenova_u1/README.md
@@ -6,11 +6,21 @@
 vllm serve SenseNova/SenseNova-U1-8B-MoT --omni --port 8091
 ```
 
+With Cache-DiT acceleration:
+
+```bash
+vllm serve SenseNova/SenseNova-U1-8B-MoT --omni --port 8091 \
+    --cache-backend cache_dit
+```
+
 Or use the convenience script:
 
 ```bash
 cd examples/online_serving/sensenova_u1
 bash run_server.sh
+
+# Cache-DiT
+CACHE_BACKEND=cache_dit bash run_server.sh
 ```
 
 ### Tensor Parallelism (TP)
diff --git a/examples/online_serving/sensenova_u1/run_server.sh b/examples/online_serving/sensenova_u1/run_server.sh
index 8edf54cc3..9dd8a9662 100644
--- a/examples/online_serving/sensenova_u1/run_server.sh
+++ b/examples/online_serving/sensenova_u1/run_server.sh
@@ -3,10 +3,18 @@
 
 MODEL="${MODEL:-SenseNova/SenseNova-U1-8B-MoT}"
 PORT="${PORT:-8091}"
+CACHE_BACKEND="${CACHE_BACKEND:-none}"
 
 echo "Starting SenseNova-U1 server..."
 echo "Model: $MODEL"
 echo "Port: $PORT"
+echo "Cache backend: $CACHE_BACKEND"
+
+EXTRA_ARGS=()
+if [ "$CACHE_BACKEND" != "none" ]; then
+    EXTRA_ARGS+=(--cache-backend "$CACHE_BACKEND")
+fi
 
 vllm serve "$MODEL" --omni \
-    --port "$PORT"
+    --port "$PORT" \
+    "${EXTRA_ARGS[@]}"
diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md
index 75b85580d..12dafde41 100644
--- a/examples/online_serving/text_to_speech/README.md
+++ b/examples/online_serving/text_to_speech/README.md
@@ -15,9 +15,10 @@ For the full list of supported architectures across all modalities, see
 | Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo |
 |---|---|---|---|---|---|
 | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
+| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
 | Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | — (caption-controlled) | — | caption fields (`instructions`) | — |
 | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✓ (`ref_audio` required) | ✓ (PCM stream) | — | ✓ |
-| OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — |
+| OmniVoice | `k2-fsa/OmniVoice` | ✓ | — | — | — |
 | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
 | VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ |
 | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ |
@@ -94,6 +95,45 @@ For full request-shape documentation (all parameters, response formats, error co
 
 ---
 
+## GLM-TTS
+
+2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`.
+
+### Launch
+```bash
+vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
+# or:
+bash examples/online_serving/text_to_speech/glm_tts/run_server.sh /path/to/GLM-TTS
+```
+
+### Sending requests
+```bash
+# Voice cloning (required)
+python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
+    --text "你好，这是语音克隆测试。" \
+    --ref-audio file:///path/to/ref.wav \
+    --ref-text "这是参考音频的文本内容。"
+
+# Custom format
+python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
+    --text "Hello, this is a voice cloning test." \
+    --ref-audio file:///path/to/ref.wav \
+    --ref-text "Transcript of the reference audio." \
+    --response-format mp3 -o output.mp3
+```
+
+### Gradio demo
+```bash
+bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
+```
+
+### Notes
+- Output: 24 kHz mono WAV via HiFT vocoder.
+- `ref_audio` + `ref_text` are **required** together on every request. Reference audio should be 3-10 seconds.
+- Voice cloning feature extraction (WhisperVQ, CampPlus, mel) runs on the model side — no external dependency on the serving layer.
+
+---
+
 ## Fish Speech S2 Pro
 
 4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.
@@ -103,6 +143,26 @@ For full request-shape documentation (all parameters, response formats, error co
 pip install fish-speech
 ```
 
+### Kvcache attention fast path
+
+Fish Speech S2 Pro uses a Triton decode-only kvcache attention fast path by
+default on CUDA builds. Set `VLLM_OMNI_FISH_KVCACHE_ATTN=0` to disable it, or
+`VLLM_OMNI_FISH_KVCACHE_ATTN=required` to fail fast if the fast path cannot be
+installed.
+
+```bash
+# Verify fast path availability.
+python - <<'PY'
+from vllm_omni.attention import fish_kvcache_attn
+
+print(fish_kvcache_attn.is_available())
+print(fish_kvcache_attn.load_error())
+PY
+
+# Optional: disable the runtime fast path.
+export VLLM_OMNI_FISH_KVCACHE_ATTN=0
+```
+
 ### Launch
 ```bash
 vllm serve fishaudio/s2-pro --omni --port 8091
@@ -261,16 +321,29 @@ vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
 ### CLI client
 ```bash
 cd examples/online_serving/text_to_speech/omnivoice
+# Text-only (auto voice)
 python speech_client.py --text "Hello, how are you?"
+
+# Language hint
 python speech_client.py --text "Bonjour, comment allez-vous?" --language French
+# Voice cloning (reference audio + optional ref_text)
+python speech_client.py \
+--text "Bonjour, comment allez-vous?" \
+--ref-audio /path/to/ref_audio.wav \
+--ref-text "Bonjour, comment allez-vous?"
+
+# Style instruction (voice design-style control)
+python speech_client.py \
+--text "Bonjour, comment allez-vous?" \
+--language French \
+--instructions "loud voice"
+
+# Deterministic output with seed parameter
+python speech_client.py --text "Hello, how are you?" --seed 42
 ```
 
-The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--language`, `--output`.
+The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--language`, `--voice`, `--ref-audio`, `--ref-text`, `--instructions`, `--seed`, and `--output`.
 
-### Notes
-- Voice cloning and voice design require offline inference; see the [offline OmniVoice section](../../offline_inference/text_to_speech/README.md#omnivoice).
-
----
 
 ## Qwen3-TTS
 
diff --git a/examples/online_serving/text_to_speech/glm_tts/gradio_demo.py b/examples/online_serving/text_to_speech/glm_tts/gradio_demo.py
new file mode 100644
index 000000000..12e5d4b11
--- /dev/null
+++ b/examples/online_serving/text_to_speech/glm_tts/gradio_demo.py
@@ -0,0 +1,291 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Gradio demo for GLM-TTS online serving via /v1/audio/speech API.
+
+Supports:
+  - Voice-cloned text-to-speech synthesis from reference audio (upload or URL)
+  - Streaming (progressive PCM chunks) and non-streaming modes
+
+Usage:
+    # Start the server first (see run_server.sh), then:
+    python gradio_demo.py --api-base http://localhost:8091
+
+    # Or use run_gradio_demo.sh to start both server and demo together.
+"""
+
+import argparse
+import base64
+import io
+
+try:
+    import gradio as gr
+except ImportError:
+    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
+import httpx
+import numpy as np
+import soundfile as sf
+
+PCM_SAMPLE_RATE = 24000
+DEFAULT_API_BASE = "http://localhost:8091"
+
+
+def encode_audio_to_base64(audio_data: tuple) -> str:
+    """Encode Gradio audio input (sample_rate, numpy_array) to base64 data URL."""
+    sample_rate, audio_np = audio_data
+    if audio_np.dtype != np.int16:
+        if audio_np.dtype in (np.float32, np.float64):
+            audio_np = np.clip(audio_np, -1.0, 1.0)
+            audio_np = (audio_np * 32767).astype(np.int16)
+        else:
+            audio_np = audio_np.astype(np.int16)
+    buf = io.BytesIO()
+    sf.write(buf, audio_np, sample_rate, format="WAV")
+    wav_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+    return f"data:audio/wav;base64,{wav_b64}"
+
+
+def build_payload(
+    text: str,
+    ref_audio: tuple | None,
+    ref_audio_url: str,
+    ref_text: str,
+    response_format: str = "wav",
+    stream: bool = False,
+    max_new_tokens: int = 4096,
+) -> dict:
+    """Build the /v1/audio/speech request payload."""
+    if not text or not text.strip():
+        raise gr.Error("Please enter text to synthesize.")
+
+    payload: dict = {
+        "input": text.strip(),
+        "voice": "default",
+        "response_format": "pcm" if stream else response_format,
+        "stream": stream,
+        "max_new_tokens": max_new_tokens,
+    }
+
+    # Voice cloning: ref_audio takes priority over URL.
+    ref_url_stripped = ref_audio_url.strip() if ref_audio_url else ""
+    if ref_audio is not None:
+        payload["ref_audio"] = encode_audio_to_base64(ref_audio)
+    elif ref_url_stripped:
+        payload["ref_audio"] = ref_url_stripped
+    else:
+        raise gr.Error("GLM-TTS requires a reference audio file or URL.")
+
+    if not ref_text or not ref_text.strip():
+        raise gr.Error("GLM-TTS requires a transcript of the reference audio.")
+    payload["ref_text"] = ref_text.strip()
+
+    return payload
+
+
+def generate_speech(api_base: str, text: str, ref_audio, ref_audio_url, ref_text, response_format, max_new_tokens):
+    """Non-streaming: call /v1/audio/speech and return full audio."""
+    payload = build_payload(
+        text, ref_audio, ref_audio_url, ref_text, response_format, stream=False, max_new_tokens=max_new_tokens
+    )
+
+    try:
+        with httpx.Client(timeout=300.0) as client:
+            resp = client.post(
+                f"{api_base}/v1/audio/speech",
+                json=payload,
+                headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"},
+            )
+    except httpx.TimeoutException:
+        raise gr.Error("Request timed out. The server may be busy.")
+    except httpx.ConnectError:
+        raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?")
+
+    if resp.status_code != 200:
+        raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
+
+    content_type = resp.headers.get("content-type", "")
+    if "application/json" in content_type:
+        try:
+            raise gr.Error(f"Server error: {resp.json()}")
+        except ValueError:
+            pass
+
+    try:
+        if response_format == "pcm":
+            audio_np = np.frombuffer(resp.content, dtype=np.int16).astype(np.float32) / 32767.0
+            return (PCM_SAMPLE_RATE, audio_np)
+        audio_np, sample_rate = sf.read(io.BytesIO(resp.content))
+        if audio_np.ndim > 1:
+            audio_np = audio_np[:, 0]
+        return (sample_rate, audio_np.astype(np.float32))
+    except Exception as e:
+        raise gr.Error(f"Failed to decode audio: {e}")
+
+
+def generate_speech_stream(
+    api_base, text, ref_audio, ref_audio_url, ref_text, response_format, max_new_tokens, chunk_seconds=0.5
+):
+    """Streaming: yield PCM chunks for progressive playback."""
+    payload = build_payload(
+        text, ref_audio, ref_audio_url, ref_text, response_format, stream=True, max_new_tokens=max_new_tokens
+    )
+
+    min_chunk_samples = int(PCM_SAMPLE_RATE * chunk_seconds)
+    prebuffer_chunks = 2
+    pending: list[np.ndarray] = []
+    pending_len = 0
+    chunks_yielded = 0
+    prebuffer: list[np.ndarray] = []
+
+    try:
+        leftover = b""
+        with httpx.Client(timeout=300.0) as client:
+            with client.stream(
+                "POST",
+                f"{api_base}/v1/audio/speech",
+                json=payload,
+                headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"},
+            ) as resp:
+                if resp.status_code != 200:
+                    resp.read()
+                    raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
+                for chunk in resp.iter_bytes():
+                    if not chunk:
+                        continue
+                    raw = leftover + chunk
+                    usable = len(raw) - (len(raw) % 2)
+                    leftover = raw[usable:]
+                    if usable == 0:
+                        continue
+                    float_samples = np.frombuffer(raw[:usable], dtype=np.int16).copy().astype(np.float32) / 32767.0
+                    pending.append(float_samples)
+                    pending_len += len(float_samples)
+                    if pending_len >= min_chunk_samples:
+                        audio_chunk = np.concatenate(pending)
+                        pending.clear()
+                        pending_len = 0
+                        if chunks_yielded < prebuffer_chunks:
+                            prebuffer.append(audio_chunk)
+                            chunks_yielded += 1
+                            if chunks_yielded == prebuffer_chunks:
+                                yield (PCM_SAMPLE_RATE, np.concatenate(prebuffer))
+                                prebuffer.clear()
+                        else:
+                            yield (PCM_SAMPLE_RATE, audio_chunk)
+        remaining = prebuffer + pending
+        if remaining:
+            yield (PCM_SAMPLE_RATE, np.concatenate(remaining))
+    except httpx.TimeoutException:
+        raise gr.Error("Request timed out. The server may be busy.")
+    except httpx.ConnectError:
+        raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?")
+
+
+def on_stream_change(stream: bool):
+    """Lock format to PCM when streaming."""
+    if stream:
+        return gr.update(value="pcm", interactive=False)
+    return gr.update(value="wav", interactive=True)
+
+
+def build_interface(api_base: str, stream_chunk_seconds: float = 0.5):
+    """Build the Gradio interface."""
+    with gr.Blocks(title="GLM-TTS Demo") as demo:
+        gr.Markdown("# GLM-TTS - Voice Cloning")
+        gr.Markdown(f"**Server:** `{api_base}` | **Model:** GLM-TTS | **Output:** 24kHz")
+
+        with gr.Row():
+            with gr.Column(scale=3):
+                text_input = gr.Textbox(
+                    label="Text to Synthesize",
+                    placeholder="Enter text here (supports Chinese and English)...",
+                    lines=4,
+                )
+
+                with gr.Accordion("Reference Voice (required)", open=True):
+                    gr.Markdown("Upload or link a short reference audio (3-10s) and provide its transcript.")
+                    ref_audio = gr.Audio(
+                        label="Reference Audio",
+                        type="numpy",
+                        sources=["upload", "microphone"],
+                    )
+                    ref_audio_url = gr.Textbox(
+                        label="Reference Audio URL (alternative to upload)",
+                        placeholder="https://example.com/reference.wav",
+                        lines=1,
+                    )
+                    ref_text = gr.Textbox(
+                        label="Reference Audio Transcript (required for cloning)",
+                        placeholder="Exact transcript of the reference audio...",
+                        lines=2,
+                    )
+
+                with gr.Row():
+                    response_format = gr.Dropdown(
+                        choices=["wav", "mp3", "flac", "pcm"],
+                        value="wav",
+                        label="Audio Format",
+                        scale=1,
+                    )
+                    stream_checkbox = gr.Checkbox(
+                        label="Stream output",
+                        value=False,
+                        info="Progressive PCM streaming",
+                        scale=1,
+                    )
+                    max_new_tokens = gr.Slider(
+                        minimum=256,
+                        maximum=8192,
+                        value=4096,
+                        step=256,
+                        label="Max New Tokens",
+                        scale=1,
+                    )
+
+                generate_btn = gr.Button("Generate Speech", variant="primary", size="lg")
+
+            with gr.Column(scale=2):
+                audio_output = gr.Audio(
+                    label="Generated Audio",
+                    interactive=False,
+                    streaming=True,
+                    autoplay=True,
+                )
+                gr.Markdown(
+                    "### About\n"
+                    "- **GLM-TTS** by Zhipu AI: Llama-based AR + Flow Matching DiT\n"
+                    "- **Voice cloning**: upload 3-10s reference + transcript\n"
+                    "- **Streaming**: real-time PCM output\n"
+                    "- **24kHz** output via Vocos vocoder"
+                )
+
+        stream_checkbox.change(fn=on_stream_change, inputs=[stream_checkbox], outputs=[response_format])
+
+        all_inputs = [text_input, ref_audio, ref_audio_url, ref_text, response_format, max_new_tokens]
+
+        def dispatch(stream_enabled, *args):
+            if stream_enabled:
+                yield from generate_speech_stream(api_base, *args, chunk_seconds=stream_chunk_seconds)
+            else:
+                yield generate_speech(api_base, *args)
+
+        generate_btn.click(fn=dispatch, inputs=[stream_checkbox] + all_inputs, outputs=[audio_output])
+        demo.queue()
+    return demo
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Gradio demo for GLM-TTS")
+    parser.add_argument("--api-base", default=DEFAULT_API_BASE, help=f"API base URL (default: {DEFAULT_API_BASE})")
+    parser.add_argument("--host", default="0.0.0.0", help="Gradio host (default: 0.0.0.0)")
+    parser.add_argument("--port", type=int, default=7860, help="Gradio port (default: 7860)")
+    parser.add_argument("--share", action="store_true", help="Share publicly")
+    parser.add_argument("--stream-chunk-seconds", type=float, default=0.5, help="Seconds per streaming chunk")
+    args = parser.parse_args()
+
+    print(f"Connecting to vLLM server at: {args.api_base}")
+    demo = build_interface(args.api_base, stream_chunk_seconds=args.stream_chunk_seconds)
+    demo.launch(server_name=args.host, server_port=args.port, share=args.share)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py b/examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py
new file mode 100644
index 000000000..a0ff181c3
--- /dev/null
+++ b/examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py
@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""OpenAI-compatible client for GLM-TTS via /v1/audio/speech endpoint.
+
+GLM-TTS is a two-stage TTS system (AR + DiT) that generates audio from text
+conditioned on reference speech. Each request requires ref_audio + ref_text.
+
+Usage:
+    # Voice cloning
+    python openai_speech_client.py --text "你好" --ref-audio file:///path/to/ref.wav --ref-text "参考文本"
+
+    # Streaming response, for async_chunk server mode
+    python openai_speech_client.py --text "你好" --stream --ref-audio file:///path/to/ref.wav --ref-text "参考文本"
+
+    # Specify output format
+    python openai_speech_client.py --text "你好" --ref-audio file:///path/to/ref.wav \
+        --ref-text "参考文本" --response-format mp3 -o output.mp3
+"""
+
+import argparse
+
+import httpx
+
+# Default server configuration
+DEFAULT_API_BASE = "http://localhost:8091"
+DEFAULT_API_KEY = "EMPTY"
+
+
+def run_tts_generation(args) -> None:
+    """Run TTS generation via OpenAI-compatible /v1/audio/speech API."""
+    if not args.ref_audio or not args.ref_text:
+        raise ValueError("GLM-TTS requires --ref-audio and --ref-text for voice cloning.")
+
+    payload = {
+        "model": args.model,
+        "voice": "default",
+        "input": args.text,
+        "response_format": args.response_format,
+        "stream": bool(args.stream),
+        "ref_audio": args.ref_audio,
+        "ref_text": args.ref_text,
+    }
+    if args.max_new_tokens:
+        payload["max_new_tokens"] = args.max_new_tokens
+
+    print(f"Model: {args.model}")
+    print(f"Text: {args.text}")
+    print(f"Voice cloning: ref_audio={args.ref_audio}, ref_text={args.ref_text}")
+    print(f"Stream: {args.stream}")
+    print("Generating audio...")
+
+    api_url = f"{args.api_base}/v1/audio/speech"
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {args.api_key}",
+    }
+
+    if args.stream:
+        output_path = args.output or "tts_output.pcm"
+        with httpx.Client(timeout=300.0) as client, open(output_path, "wb") as f:
+            with client.stream("POST", api_url, json=payload, headers=headers) as response:
+                if response.status_code != 200:
+                    print(f"Error: {response.status_code}")
+                    response.read()
+                    print(response.text)
+                    return
+                for chunk in response.iter_bytes():
+                    f.write(chunk)
+        print(f"Streaming audio saved to: {output_path}")
+    else:
+        with httpx.Client(timeout=300.0) as client:
+            response = client.post(api_url, json=payload, headers=headers)
+        if response.status_code != 200:
+            print(f"Error: {response.status_code}")
+            print(response.text)
+            return
+        try:
+            text = response.content.decode("utf-8")
+            if text.startswith('{"error"'):
+                print(f"Error: {text}")
+                return
+        except UnicodeDecodeError:
+            pass
+        output_path = args.output or f"tts_output.{args.response_format}"
+        with open(output_path, "wb") as f:
+            f.write(response.content)
+        print(f"Audio saved to: {output_path}")
+
+
+def parse_args():
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="OpenAI-compatible client for GLM-TTS via /v1/audio/speech",
+    )
+
+    # Server configuration
+    parser.add_argument(
+        "--api-base",
+        type=str,
+        default=DEFAULT_API_BASE,
+        help=f"API base URL (default: {DEFAULT_API_BASE})",
+    )
+    parser.add_argument(
+        "--api-key",
+        type=str,
+        default=DEFAULT_API_KEY,
+        help="API key (default: EMPTY)",
+    )
+    parser.add_argument(
+        "--model",
+        "-m",
+        type=str,
+        default="glm-tts",
+        help="Model name/path",
+    )
+
+    # Input text
+    parser.add_argument(
+        "--text",
+        type=str,
+        required=True,
+        help="Text to synthesize",
+    )
+
+    # Generation parameters
+    parser.add_argument(
+        "--max-new-tokens",
+        type=int,
+        default=None,
+        help="Maximum new tokens to generate (default: model default)",
+    )
+
+    # Output
+    parser.add_argument(
+        "--response-format",
+        type=str,
+        default="wav",
+        choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
+        help="Audio output format (default: wav)",
+    )
+    parser.add_argument(
+        "--stream",
+        action="store_true",
+        help="Request a streaming audio response (use with async_chunk server mode).",
+    )
+    parser.add_argument(
+        "--output",
+        "-o",
+        type=str,
+        default=None,
+        help="Output audio file path (default: tts_output.<format>)",
+    )
+
+    # Voice cloning parameters
+    parser.add_argument(
+        "--ref-audio",
+        type=str,
+        default=None,
+        help="Reference audio URL, file:// URI, or base64 data URL for voice cloning",
+    )
+    parser.add_argument(
+        "--ref-text",
+        type=str,
+        default=None,
+        help="Transcript of the reference audio (required with --ref-audio)",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run_tts_generation(args)
diff --git a/examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh b/examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
new file mode 100644
index 000000000..75cecdf9e
--- /dev/null
+++ b/examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
@@ -0,0 +1,48 @@
+#!/bin/bash
+# Launch GLM-TTS server + Gradio demo together.
+#
+# Usage:
+#   ./run_gradio_demo.sh
+#   CUDA_VISIBLE_DEVICES=0 PORT=8091 GRADIO_PORT=7860 ./run_gradio_demo.sh
+
+set -e
+
+MODEL="${MODEL:-zai-org/GLM-TTS}"
+PORT="${PORT:-8091}"
+GRADIO_PORT="${GRADIO_PORT:-7860}"
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+REPO_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)"
+
+echo "Starting GLM-TTS server (port $PORT)..."
+FLASHINFER_DISABLE_VERSION_CHECK=1 \
+vllm-omni serve "$MODEL" \
+    --deploy-config "$REPO_ROOT/vllm_omni/deploy/glm_tts.yaml" \
+    --host 0.0.0.0 \
+    --port "$PORT" \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni &
+SERVER_PID=$!
+
+cleanup() {
+    echo "Stopping server (PID $SERVER_PID)..."
+    kill $SERVER_PID 2>/dev/null
+    wait $SERVER_PID 2>/dev/null
+}
+trap cleanup EXIT
+
+# Wait for server to be ready.
+echo "Waiting for server to start..."
+for i in $(seq 1 120); do
+    if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then
+        echo "Server ready."
+        break
+    fi
+    sleep 2
+done
+
+echo "Starting Gradio demo (port $GRADIO_PORT)..."
+python "$SCRIPT_DIR/gradio_demo.py" \
+    --api-base "http://localhost:$PORT" \
+    --port "$GRADIO_PORT"
diff --git a/examples/online_serving/text_to_speech/glm_tts/run_server.sh b/examples/online_serving/text_to_speech/glm_tts/run_server.sh
new file mode 100644
index 000000000..f5cc3a1f8
--- /dev/null
+++ b/examples/online_serving/text_to_speech/glm_tts/run_server.sh
@@ -0,0 +1,38 @@
+#!/bin/bash
+# Launch vLLM-Omni server for GLM-TTS models
+#
+# Usage:
+#   ./run_server.sh                           # Default model path, async_chunk mode
+#   ./run_server.sh /path/to/GLM-TTS          # Custom model path, async_chunk mode
+#   ./run_server.sh /path/to/GLM-TTS sync     # Sync two-stage mode
+#
+# NOTE: The model path should point to the repo ROOT (not llm/ subdirectory).
+# model_subdir/tokenizer_subdir in the pipeline config resolve subdirectories.
+
+set -e
+
+MODEL="${1:-zai-org/GLM-TTS}"
+MODE="${2:-async}"
+
+EXTRA_ARGS=()
+case "$MODE" in
+    async|async_chunk)
+        ;;
+    sync|no_async_chunk)
+        EXTRA_ARGS+=("--no-async-chunk")
+        ;;
+    *)
+        echo "Unknown mode: $MODE (expected async or sync)" >&2
+        exit 1
+        ;;
+esac
+
+echo "Starting GLM-TTS server with model: $MODEL (mode: $MODE)"
+
+vllm-omni serve "$MODEL" \
+    --deploy-config vllm_omni/deploy/glm_tts.yaml \
+    --host 0.0.0.0 \
+    --port 8091 \
+    --trust-remote-code \
+    --omni \
+    "${EXTRA_ARGS[@]}"
diff --git a/examples/online_serving/text_to_speech/higgs_audio_v2/README.md b/examples/online_serving/text_to_speech/higgs_audio_v2/README.md
new file mode 100644
index 000000000..7b5e272df
--- /dev/null
+++ b/examples/online_serving/text_to_speech/higgs_audio_v2/README.md
@@ -0,0 +1,63 @@
+# higgs-audio v2 online example
+
+This directory contains the online-serving entry points for boson-ai's higgs-audio v2 as integrated by vllm-omni: a 2-stage TTS pipeline (Llama-3.2-3B talker with DualFFN audio expert + HiggsAudio codec decoder) emitting 24 kHz mono speech.
+
+## Prerequisites
+
+Voice clone uses HF's `HiggsAudioV2TokenizerModel` loaded from `k2-fsa/OmniVoice/audio_tokenizer/` (the boson-ai standalone tokenizer Hub repo's `model.safetensors` is the 3B talker LM, not the codec). Only that ~806 MB subdir is downloaded.
+
+```bash
+pip install -U "transformers>=5.3.0"
+```
+
+## Files
+
+- `run_server.sh` — launch the vllm-omni server with the bundled `vllm_omni/deploy/higgs_audio_v2.yaml` deploy config.
+- `batch_speech_client.py` — send a list of prompts to `/v1/audio/speech` and save the returned WAV / PCM bytes to a directory; optionally passes `--ref-audio` + `--ref-text` for shallow voice clone.
+
+## Launching the server
+
+```bash
+GPUS=6,7 PORT=8094 bash examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh
+```
+
+Environment overrides:
+
+- `MODEL` — HF id of the talker (default `bosonai/higgs-audio-v2-generation-3B-base`).
+- `PORT` — server port (default `8094`).
+- `GPUS` — `CUDA_VISIBLE_DEVICES` value (default `6,7`).
+- `GPU_UTIL` — `--gpu-memory-utilization` (default `0.4`).
+
+The script also exports `VLLM_USE_DEEP_GEMM=0` / `VLLM_MOE_USE_DEEP_GEMM=0` so the example works on images without the optional `deep_gemm` backend.
+
+The deploy YAML ships with `async_chunk: false` and `codec_streaming: true`, i.e. Stage 0 finishes its codec frames before Stage 1 starts decoding, and Stage 1 streams WAV/PCM bytes to the client chunk-by-chunk.
+
+## Driving the server
+
+Plain TTS:
+
+```bash
+python examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py \
+    --base-url http://localhost:8094 \
+    --model bosonai/higgs-audio-v2-generation-3B-base \
+    --output-dir /tmp/higgs_audio_v2_batch \
+    --prompts "Hello world." \
+              "The quick brown fox jumps over the lazy dog."
+```
+
+Voice clone — pass a reference clip and its transcript (both required together):
+
+```bash
+python examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py \
+    --base-url http://localhost:8094 \
+    --model bosonai/higgs-audio-v2-generation-3B-base \
+    --output-dir /tmp/higgs_audio_v2_clone \
+    --ref-audio /path/to/reference.wav \
+    --ref-text  "Exact transcript spoken in reference.wav." \
+    --prompts "Hello, this is a cloned voice."
+```
+
+## Notes
+
+- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
+- Out of scope (rejected with explicit 4xx by the request validator): multi-speaker `[SPEAKERn]` tags inside `input`, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.
diff --git a/examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py b/examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py
new file mode 100644
index 000000000..1e227b14d
--- /dev/null
+++ b/examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py
@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Batch client for the higgs-audio v2 online server.
+
+Sends a fixed list of prompts to ``/v1/audio/speech`` and saves the returned
+WAV files (or raw PCM bytes when ``--format pcm``) into ``--output-dir``.
+
+Usage (plain text -> speech):
+
+  python examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py \
+      --base-url http://localhost:8094 \
+      --output-dir /tmp/higgs_audio_v2_batch \
+      --prompts "Hello world." "The quick brown fox jumps over the lazy dog."
+
+Usage (shallow voice clone — pass a reference clip + its transcript):
+
+  python examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py \
+      --base-url http://localhost:8094 \
+      --output-dir /tmp/higgs_audio_v2_clone \
+      --ref-audio path/to/reference.wav \
+      --ref-text "the transcript of the reference clip" \
+      --prompts "Hello world."
+"""
+
+from __future__ import annotations
+
+import argparse
+import base64
+import sys
+from pathlib import Path
+
+DEFAULT_PROMPTS = (
+    "Hello world.",
+    "The quick brown fox jumps over the lazy dog.",
+    "It was the night before my birthday.",
+    "Innovation distinguishes between a leader and a follower.",
+)
+
+
+def _slug(text: str) -> str:
+    import re
+
+    s = re.sub(r"\s+", "_", text.strip().lower())
+    return re.sub(r"[^a-z0-9_]+", "", s)[:32] or "prompt"
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--base-url", default="http://localhost:8094")
+    parser.add_argument("--model", default="higgs_audio_v2")
+    parser.add_argument("--prompts", nargs="+", default=list(DEFAULT_PROMPTS))
+    parser.add_argument("--output-dir", type=Path, default=Path("/tmp/higgs_audio_v2_batch"))
+    parser.add_argument("--format", choices=("wav", "pcm"), default="wav")
+    parser.add_argument("--max-new-tokens", type=int, default=300)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--timeout-s", type=float, default=120.0)
+    parser.add_argument(
+        "--ref-audio",
+        type=Path,
+        default=None,
+        help="Reference clip for voice clone (path to a WAV file). Must be paired with --ref-text.",
+    )
+    parser.add_argument(
+        "--ref-text",
+        type=str,
+        default=None,
+        help="Transcript of the reference clip. Required when --ref-audio is set.",
+    )
+    args = parser.parse_args()
+
+    if (args.ref_audio is None) != (args.ref_text is None):
+        print("--ref-audio and --ref-text must be supplied together", file=sys.stderr)
+        return 2
+
+    ref_audio_data_url: str | None = None
+    if args.ref_audio is not None:
+        if not args.ref_audio.exists():
+            print(f"ref-audio file not found: {args.ref_audio}", file=sys.stderr)
+            return 2
+        mime = "audio/wav" if args.ref_audio.suffix.lower() == ".wav" else "audio/mpeg"
+        ref_b64 = base64.b64encode(args.ref_audio.read_bytes()).decode("ascii")
+        ref_audio_data_url = f"data:{mime};base64,{ref_b64}"
+
+    try:
+        import httpx
+    except ImportError:
+        print(
+            "this client needs `httpx`. Install with `pip install httpx`.",
+            file=sys.stderr,
+        )
+        return 2
+
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    url = args.base_url.rstrip("/") + "/v1/audio/speech"
+    failures = 0
+    with httpx.Client(timeout=args.timeout_s) as client:
+        for prompt in args.prompts:
+            payload = {
+                "model": args.model,
+                "input": prompt,
+                "response_format": args.format,
+                "max_new_tokens": args.max_new_tokens,
+                "seed": args.seed,
+            }
+            if ref_audio_data_url is not None:
+                payload["ref_audio"] = ref_audio_data_url
+                payload["ref_text"] = args.ref_text
+            resp = client.post(url, json=payload)
+            if resp.status_code != 200:
+                print(f"[FAIL] {prompt!r} -> {resp.status_code}: {resp.text[:200]}", file=sys.stderr)
+                failures += 1
+                continue
+            suffix = ".wav" if args.format == "wav" else ".pcm"
+            out = args.output_dir / f"{_slug(prompt)}{suffix}"
+            out.write_bytes(resp.content)
+            print(f"[ ok ] {prompt!r} -> {out} ({len(resp.content)} bytes)")
+
+    return 1 if failures else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh b/examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh
new file mode 100755
index 000000000..12db89558
--- /dev/null
+++ b/examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh
@@ -0,0 +1,37 @@
+#!/bin/bash
+# Launch vLLM-Omni server for higgs-audio v2.
+#
+# v1 scope: plain text -> 24 kHz speech only. Voice cloning, multi-speaker,
+# ChatML rich content, and language overrides are rejected by the validator
+# with explicit 4xx (see vllm_omni/entrypoints/openai/serving_speech.py).
+#
+# Usage:
+#   ./run_server.sh                 # default port 8094, GPUs 6 and 7
+#   PORT=8095 GPUS=6,7 ./run_server.sh
+#   MODEL=bosonai/higgs-audio-v2-generation-3B-base ./run_server.sh
+
+set -e
+
+MODEL="${MODEL:-bosonai/higgs-audio-v2-generation-3B-base}"
+PORT="${PORT:-8094}"
+GPUS="${GPUS:-6,7}"
+GPU_UTIL="${GPU_UTIL:-0.4}"
+
+echo "Starting higgs-audio v2 server"
+echo "  MODEL=$MODEL"
+echo "  PORT=$PORT"
+echo "  CUDA_VISIBLE_DEVICES=$GPUS"
+
+# DeepGEMM FP8 kernels are optional and trip warmup on builds without
+# the deep_gemm backend; disable them so the example works out of the box.
+# Users with deep_gemm installed can re-enable via the same env vars.
+CUDA_VISIBLE_DEVICES="$GPUS" \
+VLLM_USE_DEEP_GEMM=0 \
+VLLM_MOE_USE_DEEP_GEMM=0 \
+vllm-omni serve "$MODEL" \
+    --deploy-config vllm_omni/deploy/higgs_audio_v2.yaml \
+    --host 0.0.0.0 \
+    --port "$PORT" \
+    --gpu-memory-utilization "$GPU_UTIL" \
+    --trust-remote-code \
+    --omni
diff --git a/examples/online_serving/text_to_speech/omnivoice/speech_client.py b/examples/online_serving/text_to_speech/omnivoice/speech_client.py
index b8e6f3889..9c9b8128e 100644
--- a/examples/online_serving/text_to_speech/omnivoice/speech_client.py
+++ b/examples/online_serving/text_to_speech/omnivoice/speech_client.py
@@ -6,9 +6,14 @@ Examples:
 
     # Specify language
     python speech_client.py --text "Bonjour, comment allez-vous?" --language French
+
+    # Use a specific uploaded/supported voice
+    python speech_client.py --text "Hello" --voice my_uploaded_voice
 """
 
 import argparse
+import base64
+import os
 
 import httpx
 
@@ -16,20 +21,61 @@ DEFAULT_API_BASE = "http://localhost:8091"
 DEFAULT_API_KEY = "EMPTY"
 
 
+def encode_audio_to_base64(audio_path: str) -> str:
+    """Encode a local audio file to a base64 data URL."""
+    if not os.path.exists(audio_path):
+        raise FileNotFoundError(f"Audio file not found: {audio_path}")
+
+    ext = audio_path.lower().rsplit(".", 1)[-1]
+    mime = {
+        "wav": "audio/wav",
+        "mp3": "audio/mpeg",
+        "flac": "audio/flac",
+        "ogg": "audio/ogg",
+    }.get(ext, "audio/wav")
+
+    with open(audio_path, "rb") as f:
+        b64 = base64.b64encode(f.read()).decode("utf-8")
+    return f"data:{mime};base64,{b64}"
+
+
 def run_tts(args) -> None:
     """Generate speech via /v1/audio/speech API."""
     payload = {
         "model": args.model,
         "input": args.text,
-        "voice": "default",
         "response_format": args.response_format,
     }
+    if args.seed is not None:
+        payload["extra_params"] = {}
+        payload["extra_params"]["seed"] = args.seed
 
+    if args.voice:
+        payload["voice"] = args.voice
     if args.language:
         payload["language"] = args.language
 
+    if args.ref_audio:
+        ref = args.ref_audio
+        if ref.startswith(("http://", "https://", "data:")):
+            payload["ref_audio"] = ref
+        else:
+            payload["ref_audio"] = encode_audio_to_base64(ref)
+
+    if args.ref_text:
+        payload["ref_text"] = args.ref_text
+
+    if args.instructions:
+        payload["instructions"] = args.instructions
+
     print(f"Model: {args.model}")
     print(f"Text: {args.text}")
+    if args.seed:
+        print(f"Seed: {args.seed}")
+
+    if args.voice:
+        print(f"Voice: {args.voice}")
+
     if args.language:
         print(f"Language: {args.language}")
     print("Generating audio...")
@@ -39,7 +85,6 @@ def run_tts(args) -> None:
         "Content-Type": "application/json",
         "Authorization": f"Bearer {args.api_key}",
     }
-
     with httpx.Client(timeout=300.0) as client:
         response = client.post(api_url, json=payload, headers=headers)
 
@@ -68,6 +113,11 @@ def main():
     parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
     parser.add_argument("--model", "-m", default="k2-fsa/OmniVoice", help="Model name")
     parser.add_argument("--text", required=True, help="Text to synthesize")
+    parser.add_argument(
+        "--voice",
+        default=None,
+        help="Voice name (omit for auto voice; must match a supported or uploaded speaker if set)",
+    )
     parser.add_argument("--language", default=None, help="Language hint (e.g., English, Chinese, French)")
     parser.add_argument(
         "--response-format",
@@ -75,6 +125,30 @@ def main():
         choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
         help="Audio format (default: wav)",
     )
+    parser.add_argument(
+        "--ref-audio",
+        type=str,
+        default=None,
+        help="Reference audio for voice cloning (local path, URL, or data: URI)",
+    )
+    parser.add_argument(
+        "--ref-text",
+        type=str,
+        default=None,
+        help="Reference text for voice cloning",
+    )
+    parser.add_argument(
+        "--instructions",
+        type=str,
+        default=None,
+        help="Voice style/emotion instructions",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=None,
+        help="Random seed for generation, default: None for stochastic output)",
+    )
     parser.add_argument("--output", "-o", default=None, help="Output file path")
     args = parser.parse_args()
     run_tts(args)
diff --git a/mkdocs.yml b/mkdocs.yml
index 6461c65f2..74f2145fe 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -38,13 +38,12 @@ theme:
   features:
     - content.action.edit
     - content.code.copy
-    - navigation.instant
-    - navigation.instant.progress
     - navigation.tracking
     - navigation.tabs
     - navigation.tabs.sticky
     - navigation.sections
     - navigation.indexes
+    - navigation.prune
     - navigation.top
     - search.suggest
     - search.highlight
@@ -56,6 +55,7 @@ theme:
   custom_dir: docs/mkdocs/overrides
 
 hooks:
+  - docs/mkdocs/hooks/search_index.py
   - docs/mkdocs/hooks/generate_api_readme.py
   - docs/mkdocs/hooks/url_schemes.py
   - docs/mkdocs/hooks/generate_examples.py
@@ -117,10 +117,13 @@ plugins:
             separate_signature: true
             show_overloads: true
             signature_crossrefs: true
+            show_source: false
           inventories:
           - https://docs.python.org/3/objects.inv
           - https://typing-extensions.readthedocs.io/en/latest/objects.inv
-          - https://docs.aiohttp.org/en/stable/objects.inv
+          # Disabled because Read the Docs strict builds fail if this
+          # external inventory intermittently resets the connection.
+          # - https://docs.aiohttp.org/en/stable/objects.inv
           - https://pillow.readthedocs.io/en/stable/objects.inv
           - https://numpy.org/doc/stable/objects.inv
           # Temporarily disabled due to decompression errors
@@ -165,9 +168,8 @@ extra_css:
   - mkdocs/stylesheets/extra.css
 
 extra_javascript:
+  - mkdocs/javascript/memory.js
   - mkdocs/javascript/mathjax.js
-  - https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js
-  - https://unpkg.com/mermaid@10/dist/mermaid.min.js
   - mkdocs/javascript/mermaid.js
   - mkdocs/javascript/edit_and_feedback.js
   - mkdocs/javascript/slack_and_forum.js
diff --git a/pyproject.toml b/pyproject.toml
index bb9cdc688..306f2e7cc 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -241,12 +241,15 @@ extend-ignore-identifiers-re = [
     ".*MoBA",
     ".*temperal_downsample",
     ".*_nd",
+    ".*nd$",
+    ".*EmbedND$",
     ".*nothink.*",
     ".*NOTHINK.*",
     ".*nin.*",
     ".*[Oo]no_[Aa]nna.*",
     ".*cann.*",
     ".*[Nn][Dd]3.*",
+    "byt5.*",
 ]
 extend-ignore-re = [
     "\\.\\.\\.n,d->\\.\\.\\.nd",
@@ -259,3 +262,4 @@ Vai = "Vai"
 tockens = "tockens"
 CANN = "CANN"
 ND = "ND"
+byt5 = "byt5"
diff --git a/recipes/GLM/GLM-Image.md b/recipes/GLM/GLM-Image.md
new file mode 100644
index 000000000..bd11558c9
--- /dev/null
+++ b/recipes/GLM/GLM-Image.md
@@ -0,0 +1,112 @@
+# GLM-Image for text-to-image and image editing
+
+## Summary
+
+- Vendor: Z.ai
+- Model: `GLM/GLM-Image`
+- Task: Text-to-image (T2I) and image-to-image
+- Mode: Online serving with the OpenAI-compatible API
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want a known-good starting point for serving
+`GLM/GLM-Image` with vLLM-Omni on **two 80 GB NVIDIA A800** GPUs (Ampere-class,
+same default layout as the upstream **2×A100 80GB** example: Stage 0 AR on GPU 0,
+Stage 1 diffusion on GPU 1) and validate the deployment with the existing
+`examples/online_serving/glm_image` clients.
+
+## References
+
+- Upstream or canonical docs:
+  [`docs/user_guide/examples/online_serving/glm_image.md`](../../docs/user_guide/examples/online_serving/glm_image.md)
+- Related issue or discussion:
+  [#2888](https://github.com/vllm-project/vllm-omni/pull/2888)
+
+## Hardware Support
+
+This recipe documents **dual-GPU** CUDA layouts on A800 80 GB
+for the same software stack. Add more platforms (for example ROCm / NPU) as
+community validation lands.
+
+## GPU
+
+### 2× A800 80GB
+
+#### Environment
+
+These versions were taken from a working **editable** install: activate `vllm-omni/.venv` (or your equivalent), then align `pip` / Git with the rows below when reproducing this recipe.
+
+- OS: Linux
+- Python: 3.12
+- Driver / runtime: NVIDIA CUDA stack with **two** A800 80 GB GPUs visible (set `CUDA_VISIBLE_DEVICES` on your host if needed)
+- vLLM: **0.19.0**
+- vLLM-Omni: **0.19.0rc2.dev138+g38d5f2d53** (editable install from this repo; Git **`38d5f2d5`**, `git describe` ≈ **`v0.19.0rc1-138-g38d5f2d5`**)
+- Transformers: **5.5.4** (same `.venv` as above; required so `glm_image` configs load for Stage 0)
+
+#### Command
+
+Start the server from the repository root:
+
+```bash
+vllm serve zai-org/GLM-Image --omni --port 8091
+```
+
+To use the bundled stage config explicitly (same default as above):
+
+```bash
+vllm serve zai-org/GLM-Image \
+  --omni \
+  --port 8091 \
+  --deploy-config vllm_omni/deploy/glm_image.yaml
+```
+
+#### Verification
+
+Run one of the existing example clients after the server is ready:
+
+```bash
+curl -s http://172.18.69.133:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [
+      {"role": "user", "content": "A beautiful landscape painting"}
+    ],
+    "extra_body": {
+      "height": 1920,
+      "width": 1920,
+      "num_inference_steps": 50,
+      "true_cfg_scale": 1.5,
+      "seed": 42
+    }
+  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > land.png
+```
+After the command finishes, check for the output files:
+
+```bash
+ls output.png
+```
+
+#### Sample end-to-end metrics
+
+One representative **offline** GLM-Image E2E run on this recipe’s **2× A800 80GB**.
+Overall summary from the run’s metrics. Rough wall-time split: **Stage 0 (AR)** ~**25 s**,
+**Stage 1 (diffusion)** ~**34 s** (see `e2e_stage_*_wall_time_ms` below).
+
+| Field | Value |
+| --- | ---: |
+| e2e_requests | 1 |
+| e2e_wall_time_ms | 61,148.679 |
+| e2e_total_tokens | 1,300 |
+| e2e_avg_time_per_request_ms | 61,148.679 |
+| e2e_avg_tokens_per_s | 21.260 |
+| e2e_stage_0_wall_time_ms | 24,708.760 |
+| e2e_stage_1_wall_time_ms | 33,787.442 |
+
+#### Notes
+
+- Memory usage: Roughly **~38 GiB + KV** on Stage 0 (AR) and **~20 GiB** on Stage 1 (DiT+VAE) per the user guide; two 80 GB cards match the default split.
+- Key flags: `--omni` is required; `--stage-configs-path` is optional unless you use a custom YAML (for example single-GPU).
+- Keep **Transformers ≥ 5.5.1** (this recipe used **5.5.4**) so `glm_image` configs resolve; otherwise Stage 0 can fail at `ModelConfig` validation.
+- Known limitations: This starter recipe follows the dual-GPU online path documented under `examples/online_serving/glm_image`. The first request may be slower due to warmup.
+- Generation time: about **61 s** wall time end-to-end for the sample above (50 inference steps, 1024×1024).
diff --git a/recipes/Helios/Helios.md b/recipes/Helios/Helios.md
new file mode 100644
index 000000000..ebc446b74
--- /dev/null
+++ b/recipes/Helios/Helios.md
@@ -0,0 +1,109 @@
+# Helios for video generation
+
+## Summary
+
+- Vendor: Helios
+- Model: `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled`
+- Task: Text-to-video, image-to-video, and video-to-video generation
+- Mode: Offline inference with the bundled Helios example
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want a known-good starting point for running Helios
+video generation with vLLM-Omni. The concrete command below focuses on
+`BestWishYsh/Helios-Base` text-to-video generation on one NVIDIA H20 GPU. The
+same example script also supports Helios-Mid, Helios-Distilled, image-to-video,
+and video-to-video; see the linked example docs for those model-specific flags.
+
+## References
+
+- Upstream repository: <https://github.com/PKU-YuanGroup/Helios>
+- Model weights:
+  <https://huggingface.co/BestWishYsh/Helios-Base>
+- Upstream or canonical docs:
+  [`docs/user_guide/examples/offline_inference/helios.md`](../../docs/user_guide/examples/offline_inference/helios.md)
+- Related example under `examples/`:
+  [`examples/offline_inference/helios/README.md`](../../examples/offline_inference/helios/README.md)
+
+## Hardware Support
+
+## GPU
+
+### 1x NVIDIA H20
+
+#### Environment
+
+- OS: Ubuntu Linux x86_64
+- Python: 3.12.12
+- Driver / runtime: NVIDIA driver `580.126.20`, CUDA `13.0`
+- Hardware: 1x NVIDIA H20 GPU from an 8x H20 host
+- vLLM version: `0.19.0`
+- vLLM-Omni version or commit: `a3903810`
+
+#### Command
+
+Run the baseline Helios-Base text-to-video example from the repository root:
+
+```bash
+cd examples/offline_inference/helios
+
+python end2end.py \
+  --model BestWishYsh/Helios-Base \
+  --sample-type t2v \
+  --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
+  --guidance-scale 5.0 \
+  --output helios_t2v_base.mp4
+```
+
+To use cache-dit acceleration, run on a vLLM-Omni checkout that includes Helios cache-dit support and enable the cache backend:
+
+```bash
+cd examples/offline_inference/helios
+
+python end2end.py \
+  --cache-backend cache_dit \
+  --enable-cache-dit-summary \
+  --model BestWishYsh/Helios-Base \
+  --sample-type t2v \
+  --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
+  --guidance-scale 5.0 \
+  --output helios_t2v_base.mp4
+```
+
+#### Verification
+
+The script should print the resolved generation configuration, total generation
+time, and output path. A successful run ends with output similar to:
+
+```text
+Total generation time: <seconds> seconds (<milliseconds> ms)
+Saved generated video to helios_t2v_base.mp4
+```
+
+#### Important flags and stage configs
+
+- `--model BestWishYsh/Helios-Base` selects the base Helios checkpoint used by
+  this recipe.
+- `--sample-type t2v` selects text-to-video generation. Use `i2v` with
+  `--image-path` for image-to-video, or `v2v` with `--video-path` for
+  video-to-video.
+- `--guidance-scale 5.0` matches the Helios-Base recommendation. Use
+  `--guidance-scale 1.0` for Helios-Distilled.
+- `--cache-backend cache_dit` enables the cache-dit acceleration path.
+- `--enable-cache-dit-summary` prints cache-dit summary information after
+  diffusion forward passes.
+- No separate stage config is required for this offline recipe; the bundled
+  Helios example configures the pipeline through `end2end.py` arguments.
+- For Helios-Mid, add `--is-enable-stage2`,
+  `--pyramid-num-inference-steps-list 20 20 20`, `--use-cfg-zero-star`,
+  `--use-zero-init`, and `--zero-steps 1`.
+- For Helios-Distilled, add `--is-enable-stage2`,
+  `--pyramid-num-inference-steps-list 2 2 2`, and
+  `--is-amplify-first-chunk`.
+
+#### Known limitations
+
+- Helios generates video in 33-frame chunks. For best performance, set
+  `--num-frames` to a multiple of `33`; non-multiple values are rounded up to
+  the nearest multiple of `33`.
diff --git a/recipes/HiDream/HiDream-I1-Image.md b/recipes/HiDream/HiDream-I1-Image.md
new file mode 100644
index 000000000..958d03d6d
--- /dev/null
+++ b/recipes/HiDream/HiDream-I1-Image.md
@@ -0,0 +1,106 @@
+# HiDream-I1-Full
+
+> Text-to-image serving and benchmark
+
+## Summary
+
+- Vendor: HiDream.ai
+- Model: `HiDream-ai/HiDream-I1-Full`
+- Task: Text-to-image generation
+- Mode: Online serving with the OpenAI-compatible API
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want a known-good starting point for serving
+`HiDream-ai/HiDream-I1-Full` on a single 80 GB A800, validate the normal online-serving path.
+
+## References
+
+- Model: <https://huggingface.co/HiDream-ai/HiDream-I1-Full>
+- Upstream or canonical docs:
+  [`docs/user_guide/examples/online_serving/text_to_image.md`](../../docs/user_guide/examples/online_serving/text_to_image.md)
+- Related example under `examples/`:
+  [`examples/online_serving/text_to_image/README.md`](../../examples/online_serving/text_to_image/README.md)
+
+
+## Hardware Support
+
+This recipe currently documents one CUDA GPU serving configuration.
+
+## GPU
+
+### 1x A800 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- Driver / runtime: NVIDIA CUDA environment with an A100 80 GB GPU
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+#### Command
+
+Start the baseline server:
+
+```bash
+export MODEL_NAME_OR_PATH=HiDream-ai/HiDream-I1-Full
+vllm serve ${MODEL_NAME_OR_PATH} \
+   --omni \
+   --port 8092 \
+   --auxiliary-text-encoder meta-llama/Llama-3.1-8B-Instruct \
+   --tensor-parallel-size 1 \
+   --vae_use_slicing \
+   --vae_use_tiling
+```
+
+
+You can also use the example launcher and pass the extra flags through:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model HiDream-ai/HiDream-I1-Full \
+  --prompt "The setting sun of late autumn dyes the riverside with a warm orange hue" \
+  --seed 42 \
+  --guidance-scale 5.0 \
+  --tensor-parallel-size 1 \
+  --num-images-per-prompt 1 \
+  --num-inference-steps 50 \
+  --auxiliary-text-encoder meta-llama/Llama-3.1-8B-Instruct \
+  --output /workspace/mnt/cmss-wmd/result/output.png
+
+```
+
+#### Verification
+
+Run the existing client example after the server is ready:
+
+```bash
+python examples/online_serving/text_to_image/openai_chat_client.py \
+  --server http://localhost:8092 \
+  --prompt "The setting sun of late autumn dyes the riverside with a warm orange hue" \
+  --output /tmp/hidream_i1_recipe.png
+```
+
+For a direct API smoke test:
+
+```bash
+curl -s http://localhost:8092/v1/chat/completions   -H "Content-Type: application/json"   -d '{
+    "messages": [
+      {"role": "user", "content": "The setting sun of late autumn dyes the riverside with a warm orange hue"}
+    ],
+    "extra_body": {
+      "num_inference_steps": 50,
+      "seed": 42,
+      "height": 1024,
+      "width": 1024
+    }
+  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output_online.png
+```
+
+
+
+#### Notes
+
+- Key flags: `--auxiliary-text-encoder` designates the path for the auxiliary text model meta-llama/Llama-3.1-8B-Instruct. For HiDream-I1-Full, unspecified use defaults meta-llama/Llama-3.1-8B-Instruct (downloaded from official Hugging Face), and custom paths are supported.
diff --git a/recipes/LTX/LTX-2.3.md b/recipes/LTX/LTX-2.3.md
index 36a9eac09..269ea6425 100644
--- a/recipes/LTX/LTX-2.3.md
+++ b/recipes/LTX/LTX-2.3.md
@@ -86,7 +86,6 @@ curl -X POST http://localhost:8000/v1/videos \
 - Duration: Controlled by `num_frames` (frames = duration_seconds * 24 + 1)
 - Known limitations:
   - No image-to-video support yet (LTX23ImageToVideoPipeline is a placeholder)
-  - No CFG-parallel support (single-GPU only)
   - Requires `diffusers >= 0.38.0` (not yet on PyPI, install from git)
 
 ## Hardware Support
diff --git a/recipes/Qwen/Qwen-Image.md b/recipes/Qwen/Qwen-Image.md
index 327160973..c40d71c1e 100644
--- a/recipes/Qwen/Qwen-Image.md
+++ b/recipes/Qwen/Qwen-Image.md
@@ -130,3 +130,108 @@ and compare the output JSON or terminal metrics.
 - Key flags: `--step-execution` enables the step-wise runtime; `--max-num-seqs` controls how many compatible requests may stay active together.
 - Keep `--max-num-seqs 1` when you want the more conservative path, when traffic is mostly single-request, or when you are debugging correctness before measuring throughput.
 - Current batching is still shape-sensitive: different step progress can co-batch, but different resolutions do not yet co-batch.
+
+### 2x B200 ModelOpt mixed FP8/NVFP4
+
+#### Environment
+
+- OS: Linux
+- Python: 3.12
+- Driver / runtime: NVIDIA CUDA environment with two B200 GPUs
+- vLLM version: Match the vLLM-Omni checkout used for deployment
+- vLLM-Omni version or commit: Use the commit that contains ModelOpt mixed
+  FP8/NVFP4 checkpoint loading for diffusion models
+
+#### Command
+
+This configuration uses a pre-quantized ModelOpt mixed checkpoint for
+Qwen-Image-2512. The validated checkpoint keeps boundary, modulation,
+normalization, and output layers in BF16, uses FP8 for sensitive attention
+projections, and uses NVFP4 for heavier MLP / attention-output linear layers.
+
+Validated checkpoint:
+
+```text
+feizhai123/qwen-image-2512-modelopt-mixed-fp8-sensitive-nvfp4-heavy
+```
+
+Start the server with the CUTLASS quantized linear backend:
+
+```bash
+CUDA_VISIBLE_DEVICES=2,3 \
+vllm serve feizhai123/qwen-image-2512-modelopt-mixed-fp8-sensitive-nvfp4-heavy \
+  --omni \
+  --host 0.0.0.0 \
+  --port 8091 \
+  --trust-remote-code \
+  --tensor-parallel-size 2 \
+  --linear-backend cutlass \
+  --force-cutlass-fp8
+```
+
+#### Verification
+
+Run a direct image generation request after the server is ready:
+
+```bash
+curl http://localhost:8091/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "feizhai123/qwen-image-2512-modelopt-mixed-fp8-sensitive-nvfp4-heavy",
+    "prompt": "A cinematic Warhammer 40,000 Ultramarine warrior in blue power armor, armor stained with alien ichor, swinging a chainsword at a Tyranid creature, gothic battlefield, dramatic lighting, high detail, no text, no watermark",
+    "negative_prompt": "blurry, low quality, text, watermark, logo, malformed armor, bad hands",
+    "size": "1024x1024",
+    "num_inference_steps": 20,
+    "true_cfg_scale": 4.0,
+    "guidance_scale": 4.0,
+    "seed": 48
+  }'
+```
+
+Check that:
+
+- The server responds on `http://localhost:8091/health`.
+- The response contains one generated image.
+- Logs show CUTLASS selection for ModelOpt FP8 / mixed-precision linear
+  layers instead of falling back to the default automatic backend choice.
+
+#### Benchmark
+
+The numbers below are loaded-once request-level measurements, not an HTTP
+concurrency benchmark. Each case loads the model once, runs warmup requests,
+then measures 100 sequential requests.
+
+```text
+Hardware: 2x B200
+Tensor parallel size: 2
+Requests: 100
+Concurrency: 1
+Resolution: 1024x1024
+Denoising steps: 20
+```
+
+| Config | Mean | P90 | P99 | Peak VRAM |
+| --- | ---: | ---: | ---: | ---: |
+| BF16 | 2.696s | 2.700s | 2.707s | 99,000 MiB |
+| ModelOpt FP8 CUTLASS | 2.663s | 2.671s | 2.676s | 85,760 MiB |
+| Mixed FP8/NVFP4 CUTLASS | 2.584s | 2.590s | 2.594s | 84,412 MiB |
+
+The same 2-GPU setup was also used to sweep NVFP4 linear backends. Direct
+CUTLASS was the fastest validated backend for this checkpoint.
+
+| Backend | Mean | P90 | P99 | Peak VRAM |
+| --- | ---: | ---: | ---: | ---: |
+| auto | 2.927s | 2.948s | 3.535s | 85,340 MiB |
+| flashinfer-cudnn | 3.523s | 3.528s | 3.766s | 84,500 MiB |
+| flashinfer-trtllm | 2.711s | 2.727s | 3.010s | 84,996 MiB |
+| flashinfer-cutlass | 2.935s | 2.950s | 3.011s | 85,340 MiB |
+| cutlass | 2.578s | 2.584s | 2.592s | 84,412 MiB |
+
+#### Notes
+
+- Use `--linear-backend cutlass` to select the validated quantized linear
+  backend from the CLI.
+- `--force-cutlass-fp8` keeps ModelOpt FP8 layers on the CUTLASS FP8 path.
+- The benchmark table intentionally reports Qwen-Image-2512 as
+  concurrency-1 request-level data; do not compare it directly with an online
+  HTTP concurrency benchmark.
diff --git a/recipes/README.md b/recipes/README.md
index b16599da0..2c92e8999 100644
--- a/recipes/README.md
+++ b/recipes/README.md
@@ -31,18 +31,23 @@ recipes/
 | [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB |
 | [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB |
 | [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB |
+| [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
 | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
 | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
+| [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
+| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
 | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay | 1x A100 80GB |
 | [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
 | [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB |
 | [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB |
 | [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x H200 (144GB) |
 | [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB |
-| [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark | 4x H100/H800 80GB |
+| [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
 | [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
 | [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
+| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
+| [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
 
 Within a single recipe file, include different hardware support sections such
 as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like
diff --git a/recipes/Tencent/HunyuanImage-3.0-Instruct.md b/recipes/Tencent/HunyuanImage-3.0-Instruct.md
index 9ccb37bed..3e6adf130 100644
--- a/recipes/Tencent/HunyuanImage-3.0-Instruct.md
+++ b/recipes/Tencent/HunyuanImage-3.0-Instruct.md
@@ -1,6 +1,8 @@
 # HunyuanImage-3.0-Instruct
 
-> DiT-only text-to-image serving and benchmark
+> DiT-only text-to-image serving and benchmark with FP8, tensor parallelism,
+> sequence parallelism, CFG parallelism, and ModelOpt mixed FP8/NVFP4
+> checkpoints.
 
 ## Summary
 
@@ -16,13 +18,15 @@ Use this recipe when you want to run or benchmark the HunyuanImage-3.0 DiT
 stage directly. This is the recommended first setup when validating DiT
 throughput, memory, FP8 kernels, sequence parallelism, or CFG parallelism.
 
-The recipe covers three 4-GPU configurations:
+The recipe covers three 4-GPU FP8 configurations and one 2-GPU ModelOpt mixed
+FP8/NVFP4 configuration:
 
 | Configuration | Parallelism | Notes |
 | --- | --- | --- |
 | `tp4_fp8` | TP=4 | Lowest per-GPU memory, higher communication overhead |
 | `tp2_fp8_sp2` | TP=2, SP=2, Ulysses=2 | Splits sequence work across two GPUs per TP group |
 | `tp2_fp8_cfgp2` | TP=2, CFG=2 | Runs CFG branches in parallel; fastest validated DiT setup |
+| `tp2_mixed_fp8_nvfp4` | TP=2, EP enabled | Uses FP8 dense layers and NVFP4 routed experts on B200 |
 
 ## References
 
@@ -222,3 +226,127 @@ of:
 - If you see OOM on 80GB GPUs, reduce image size, request concurrency, or use
   the TP=4 configuration before increasing batch size. GPU memory utilization
   is not a useful primary tuning knob for this DiT-only recipe.
+
+### 2x B200 ModelOpt mixed FP8/NVFP4
+
+#### Environment
+
+- OS: Linux
+- Python: 3.12
+- Driver / runtime: NVIDIA CUDA environment with two B200 GPUs
+- vLLM version: Match the vLLM-Omni checkout used for deployment
+- vLLM-Omni version or commit: Use the commit that contains ModelOpt mixed
+  FP8/NVFP4 checkpoint loading for diffusion models
+
+#### Command
+
+This configuration uses a pre-quantized ModelOpt mixed checkpoint for the
+HunyuanImage-3.0 DiT stage. Dense attention and shared dense projections use
+FP8, while routed MoE expert projections use NVFP4. Embeddings, norms,
+routers, output layers, and non-DiT components stay in BF16.
+
+Validated checkpoint:
+
+```text
+feizhai123/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8
+```
+
+Start the server with CUTLASS selected from the CLI for quantized linear and
+MoE kernels:
+
+```bash
+CUDA_VISIBLE_DEVICES=2,3 \
+vllm serve feizhai123/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8 \
+  --omni \
+  --host 0.0.0.0 \
+  --port 8091 \
+  --trust-remote-code \
+  --tensor-parallel-size 2 \
+  --enable-expert-parallel \
+  --linear-backend cutlass \
+  --moe-backend cutlass \
+  --force-cutlass-fp8
+```
+
+#### Verification
+
+Run a direct image generation request after the server is ready:
+
+```bash
+curl -s http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "feizhai123/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8",
+    "messages": [
+      {"role": "user", "content": "A cinematic portrait of an Adeptus Custodes warrior standing on Terra before the Imperial Palace, ornate golden armor, guardian spear, crimson cloak, marble steps, gothic architecture, holy golden light, highly detailed, no text, no watermark"}
+    ],
+    "extra_body": {
+      "height": 1024,
+      "width": 1024,
+      "num_inference_steps": 20,
+      "guidance_scale": 4.0,
+      "true_cfg_scale": 4.0,
+      "seed": 48
+    }
+  }' | jq -r '.choices[0].message.content[0].image_url.url' \
+     | cut -d',' -f2- \
+     | base64 -d > hunyuan_image3_mixed_nvfp4.png
+```
+
+Check that:
+
+- The server responds on `http://localhost:8091/health`.
+- The response contains one generated image.
+- Logs show CUTLASS selection for ModelOpt FP8 dense layers and ModelOpt
+  NVFP4 routed expert layers.
+
+#### Benchmark
+
+The numbers below are online HTTP measurements with 16 concurrent requests.
+Each case serves 100 image-generation requests on the same 2-GPU B200 setup.
+
+```text
+Hardware: 2x B200
+Tensor parallel size: 2
+Requests: 100
+Concurrency: 16
+Resolution: 1024x1024
+Denoising steps: 20
+```
+
+| Config | Mean | P90 | P99 | Peak VRAM |
+| --- | ---: | ---: | ---: | ---: |
+| BF16 | 70.729s | 77.716s | 77.758s | 191,920 MiB |
+| ModelOpt FP8 | 58.603s | 64.435s | 64.469s | 115,816 MiB |
+| Mixed FP8/NVFP4 | 54.739s | 60.244s | 60.317s | 83,560 MiB |
+
+#### Accuracy
+
+Run the HunyuanImage-3.0 quantized DiT accuracy check with the BF16, FP8, and
+mixed FP8/NVFP4 checkpoints:
+
+```bash
+CUDA_VISIBLE_DEVICES=2,3 \
+HUNYUAN_IMAGE3_QUANT_DEVICES=0,1 \
+HUNYUAN_IMAGE3_QUANT_TP=2 \
+HUNYUAN_IMAGE3_BF16_MODEL=/path/to/hunyuan-image3-bf16 \
+HUNYUAN_IMAGE3_FP8_MODEL=/path/to/hunyuan-image3-modelopt-fp8 \
+HUNYUAN_IMAGE3_NVFP4_MODEL=/path/to/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8 \
+PYTHONPATH=/path/to/vllm-omni:/path/to/vllm:${PYTHONPATH:-} \
+python -m pytest -s -v \
+  tests/e2e/accuracy/test_hunyuan_image3.py \
+  -k quantized_dit_matches_bf16_accuracy
+```
+
+The test compares quantized images with the BF16 image using CLIP-based image
+and text scores, plus structural SSIM / PSNR checks. The current structural
+thresholds are `SSIM >= 0.20` and `PSNR >= 10.0`.
+
+#### Notes
+
+- Use `--linear-backend cutlass` to select the validated quantized linear
+  backend from the CLI.
+- Use `--moe-backend cutlass` for the routed expert path.
+- `--force-cutlass-fp8` keeps the FP8 dense layers on the CUTLASS FP8 path.
+- The benchmark table is an online HTTP concurrency-16 result; do not compare
+  it directly with loaded-once request-level benchmarks.
diff --git a/recipes/fishaudio/Fish-Speech-S2-Pro.md b/recipes/fishaudio/Fish-Speech-S2-Pro.md
index 91ac3fe03..dca9f537e 100644
--- a/recipes/fishaudio/Fish-Speech-S2-Pro.md
+++ b/recipes/fishaudio/Fish-Speech-S2-Pro.md
@@ -50,6 +50,26 @@ fi
 pip install fish-speech
 ```
 
+#### Kvcache attention fast path
+
+Fish Speech S2 Pro uses a Triton decode-only kvcache attention fast path by
+default on CUDA builds. Set `VLLM_OMNI_FISH_KVCACHE_ATTN=0` to disable it, or
+`VLLM_OMNI_FISH_KVCACHE_ATTN=required` to fail fast if the fast path cannot be
+installed.
+
+```bash
+# Verify fast path availability.
+python - <<'PY'
+from vllm_omni.attention import fish_kvcache_attn
+
+print(fish_kvcache_attn.is_available())
+print(fish_kvcache_attn.load_error())
+PY
+
+# Optional: disable the runtime fast path.
+export VLLM_OMNI_FISH_KVCACHE_ATTN=0
+```
+
 #### Command
 
 ```bash
@@ -102,3 +122,30 @@ curl -X POST http://localhost:8091/v1/audio/speech \
   both `ref_audio` and `ref_text` parameters.
 - Memory usage: Model loads at ~48.3 GiB, peaks at ~48.9 GiB during inference
   headroom for video frames and audio caches.
+
+### 2x H100 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- Required package: `fish-speech` (for DAC codec)
+- CUDA 12.8
+- vLLM version: 0.19.0+
+
+#### Command
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1 vllm serve fishaudio/s2-pro --omni --port 8091 \
+    --deploy-config vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml
+```
+
+This profile pins Stage0 (Slow/Fast AR) to GPU 0 and Stage1 (DAC decoder)
+to GPU 1 to remove AR/DAC contention, sets `max_num_seqs=64` on both
+stages so concurrencies above 8 don't queue, and reduces
+`connector_get_sleep_s` from 10 ms to 1 ms so DAC chunks aren't held in
+the connector once they're ready.
+
+For benchmark numbers see the discussion in
+vllm-project/vllm-omni#2515 (initial H100x2 sweep by @Sy0307) and
+vllm-project/vllm-omni#3323 (H20x2 sweep with the shipped config).
diff --git a/recipes/mistralai/Voxtral-TTS.md b/recipes/mistralai/Voxtral-TTS.md
new file mode 100644
index 000000000..595d5c29e
--- /dev/null
+++ b/recipes/mistralai/Voxtral-TTS.md
@@ -0,0 +1,107 @@
+# Voxtral TTS for text-to-speech
+
+## Summary
+
+- Vendor: Mistral AI
+- Model: `mistralai/Voxtral-4B-TTS-2603`
+- Task: Text-to-speech with model-provided voice presets or voice cloning
+- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want a known-good starting point for serving Voxtral TTS
+models with vLLM-Omni and validate the deployment with the existing TTS client
+examples in this repository.
+
+## References
+
+- Upstream model card: <https://huggingface.co/mistralai/Voxtral-4B-TTS-2603>
+- Related examples under `examples/`:
+  [`examples/offline_inference/text_to_speech/voxtral_tts/end2end.py`](../../examples/offline_inference/text_to_speech/voxtral_tts/end2end.py)
+  [`examples/online_serving/text_to_speech/voxtral_tts/gradio_demo.py`](../../examples/online_serving/text_to_speech/voxtral_tts/gradio_demo.py)
+- Related issue or discussion:
+  [RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645)
+
+## Hardware Support
+
+## GPU
+
+### 1x RTX 4090 24GB
+
+#### Environment
+
+- OS: Ubuntu 22.04.5 LTS
+- Python: 3.12.11
+- PyTorch: 2.11.0+cu130
+- Driver / runtime: CUDA 13.0
+- GPU: NVIDIA GeForce RTX 4090, 24 GB
+- vLLM version: 0.20.0
+- vLLM-Omni version or commit: 0.20.0rc2.dev99+g857356d5b
+- Required package: mistral_common >= 1.10.0
+
+#### Command
+
+Online:
+
+```bash
+vllm serve mistralai/Voxtral-4B-TTS-2603 \
+    --omni \
+    --port 8091
+```
+
+The deploy config is auto-loaded from `vllm_omni/deploy/voxtral_tts.yaml`
+
+#### Verification
+
+**Quick smoke test with curl:**
+
+```bash
+curl -X POST http://127.0.0.1:8091/v1/audio/speech \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "mistralai/Voxtral-4B-TTS-2603",
+    "input": "Hello, this is Voxtral TTS running with vLLM-Omni.",
+    "voice": "casual_female",
+    "language": "English",
+    "response_format": "wav"
+  }' \
+  --output voxtral.wav
+```
+
+**Streaming audio:**
+
+```bash
+curl -X POST http://127.0.0.1:8091/v1/audio/speech \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "mistralai/Voxtral-4B-TTS-2603",
+    "input": "Hello, this is Voxtral TTS streaming PCM.",
+    "voice": "casual_female",
+    "language": "English",
+    "stream": true,
+    "response_format": "pcm"
+  }' \
+  --output voxtral_stream.pcm
+```
+Convert the raw PCM output to WAV
+```bash
+ffmpeg -f s16le -ar 24000 -ac 1 -i voxtral_stream.pcm voxtral_stream.wav -y
+ffprobe -hide_banner voxtral_stream.wav
+```
+
+#### Notes
+
+- Memory usage: Stage 0 (`audio_generation`) ~18.95 GiB, Stage 1
+  (`audio_tokenizer`) ~1.55 GiB. The observed server-startup peak was
+  `21006 MiB / 24564 MiB` (`20.51 GiB / 23.99 GiB`, about 85.5%).
+  Both stages share GPU 0 via the deploy config (`gpu_memory_utilization: 0.8`
+  for Stage 0, `0.1` for Stage 1).
+- Key flags: `--omni` is required. `--port 8091` is optional and only selects
+  the HTTP port.
+- Known limitation: Voxtral TTS outputs mono audio at 24 kHz. And with the public
+mistralai/Voxtral-4B-TTS-2603 checkpoint. running the
+voice-cloning e2e test ([`Voice cloning (capability gated upstream)`](../../examples/offline_inference/text_to_speech/README.md#voice-cloning-capability-gated-upstream)
+) or adding a `ref_audio` request parameter will fail
+because the public checkpoint does not provide the encoder weights needed to
+convert reference audio into conditioning features. The expected failure is: `RuntimeError: encode_waveforms requires encoder weights which are not available in the open-source checkpoint.`
diff --git a/recipes/zai-org/GLM-TTS.md b/recipes/zai-org/GLM-TTS.md
new file mode 100644
index 000000000..2bf1a202a
--- /dev/null
+++ b/recipes/zai-org/GLM-TTS.md
@@ -0,0 +1,105 @@
+# GLM-TTS for Chinese/English TTS on 1x GPU
+
+## Summary
+
+- Vendor: zai-org
+- Model: `zai-org/GLM-TTS`
+- Task: Zero-shot voice-cloned text-to-speech synthesis
+- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to serve GLM-TTS as a two-stage TTS system (AR + DiT
+flow-matching) for Chinese and English speech synthesis. Every request is
+conditioned on reference audio and its transcript.
+
+## References
+
+- Upstream or canonical docs:
+  [zai-org/GLM-TTS on HuggingFace](https://huggingface.co/zai-org/GLM-TTS)
+- GitHub repository:
+  [zai-org/GLM-TTS](https://github.com/zai-org/GLM-TTS)
+- Related example under `examples/`:
+  [`examples/online_serving/text_to_speech/README.md#glm-tts`](../../examples/online_serving/text_to_speech/README.md#glm-tts)
+- Offline inference example:
+  [`examples/offline_inference/text_to_speech/README.md#glm-tts`](../../examples/offline_inference/text_to_speech/README.md#glm-tts)
+
+## Hardware Support
+
+### GPU
+
+### 1x A40 48GB
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- Driver / runtime: NVIDIA CUDA environment with A40 48GB GPU
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+#### Command
+
+Start the server from the repository root:
+
+```bash
+vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
+```
+
+Async chunking is enabled by default in the bundled deployment config. For
+the sync (non-streaming) path:
+
+```bash
+vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk
+```
+
+Use a custom deploy config for advanced cases:
+
+```bash
+vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 \
+  --deploy-config /path/to/your_glm_tts_overrides.yaml
+```
+
+#### Verification
+
+Run the bundled OpenAI-compatible client with reference audio:
+
+```bash
+python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
+  --text "你好，这是一个语音合成测试。" \
+  --ref-audio file:///path/to/ref.wav \
+  --ref-text "这是参考音频的文本内容。"
+```
+
+For a quick API smoke test:
+
+```bash
+curl http://localhost:8091/v1/audio/speech \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "zai-org/GLM-TTS",
+    "input": "你好，这是一个语音合成测试。",
+    "response_format": "wav",
+    "ref_audio": "file:///path/to/ref.wav",
+    "ref_text": "这是参考音频的文本内容。"
+  }' --output test.wav
+```
+
+Voice cloning with reference audio:
+
+```bash
+python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
+  --text "你好，这是语音克隆测试。" \
+  --ref-audio file:///path/to/ref.wav \
+  --ref-text "这是参考音频的文本内容。"
+```
+
+#### Notes
+
+- Hardware scope: the default bundled config is CUDA-only and verified on 1x A40 48GB (~16.6 GiB peak); fits 24GB cards. Split stages across GPUs for higher concurrency.
+- Memory usage: ~18-20GB total (AR ~10GB, DiT ~8GB); both stages share GPU 0 by default.
+- Audio output: 24kHz mono WAV via HiFT vocoder (Vocos2D 32kHz fallback with resampling).
+- Key flags: `--omni` is required; `--trust-remote-code` is needed for the GLM-TTS phoneme tokenizer; the DiT stage enables bucketed model-internal CUDA graphs with eager fallback.
+- Voice cloning: requires `ref_audio` + `ref_text` together. Reference audio should be 3-10 seconds. Feature extraction (WhisperVQ tokenizer, CampPlus ONNX, mel) runs on the model side.
+- Known limitations: First request may be slow due to lazy model loading (WhisperVQ, CampPlus ONNX). Warm-cache RTF is approximately 0.6-0.7x on A40.
diff --git a/setup.py b/setup.py
index 057212d67..403d8a416 100644
--- a/setup.py
+++ b/setup.py
@@ -62,6 +62,12 @@ def detect_target_device() -> str:
         else:
             print(f"Warning: Invalid VLLM_OMNI_TARGET_DEVICE '{target_device}', falling back to auto-detection")
 
+    # Priority 1.5: ReadTheDocs builds use CPU requirements (no GPU, 1 GB RAM cap).
+    # Avoids pulling ~2 GB of CUDA libs that push the mkdocs build into swap.
+    if os.environ.get("READTHEDOCS"):
+        print("Detected READTHEDOCS=1; using cpu requirements for docs build")
+        return "cpu"
+
     # Priority 2: Torch backend detection
     # This is a code path for when user is using
     # --no-build-isolation flag
diff --git a/tests/assets/glm_tts/jiayan_zh.wav b/tests/assets/glm_tts/jiayan_zh.wav
new file mode 100644
index 000000000..1dd528747
Binary files /dev/null and b/tests/assets/glm_tts/jiayan_zh.wav differ
diff --git a/tests/attention/test_fish_kvcache_attn.py b/tests/attention/test_fish_kvcache_attn.py
new file mode 100644
index 000000000..89b43e627
--- /dev/null
+++ b/tests/attention/test_fish_kvcache_attn.py
@@ -0,0 +1,1137 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm_omni.attention import fish_kvcache_attn, fish_kvcache_backend
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _FakeImpl:
+    dcp_world_size = 1
+    alibi_slopes = None
+    sliding_window = None
+    scale = 1.0
+
+    def __init__(self):
+        self.original_calls = 0
+
+    def forward(
+        self,
+        layer,
+        query,
+        key,
+        value,
+        kv_cache,
+        attn_metadata,
+        output,
+        output_scale=None,
+        output_block_scale=None,
+    ):
+        del layer, query, key, value, kv_cache, attn_metadata, output_scale, output_block_scale
+        self.original_calls += 1
+        output.fill_(2)
+        return output
+
+
+class _FakeAttentionLayer:
+    def __init__(self):
+        self.impl = _FakeImpl()
+
+
+class _FakeSelfAttn:
+    def __init__(self):
+        self.attn = _FakeAttentionLayer()
+
+
+class _FakeLayer:
+    def __init__(self):
+        self.self_attn = _FakeSelfAttn()
+
+
+class _FakeModel:
+    def __init__(self):
+        self.layers = [_FakeLayer()]
+
+
+_DEFAULT_UPPER_BOUND = object()
+
+
+def _metadata(
+    *,
+    batch_size: int = 2,
+    num_actual_tokens: int | None = None,
+    max_query_len: int = 1,
+    max_seq_len: int = 128,
+    use_cascade: bool = False,
+    seq_lens_cpu_upper_bound: torch.Tensor | None | object = _DEFAULT_UPPER_BOUND,
+    seq_lens: torch.Tensor | None = None,
+):
+    if num_actual_tokens is None:
+        num_actual_tokens = batch_size
+    if seq_lens is None:
+        seq_lens = torch.ones((batch_size,), dtype=torch.int32)
+    if seq_lens_cpu_upper_bound is _DEFAULT_UPPER_BOUND and max_query_len == 1:
+        seq_lens_cpu_upper_bound = seq_lens.detach().cpu()
+    elif seq_lens_cpu_upper_bound is _DEFAULT_UPPER_BOUND:
+        seq_lens_cpu_upper_bound = None
+    return type(
+        "Metadata",
+        (),
+        {
+            "use_cascade": use_cascade,
+            "num_actual_tokens": num_actual_tokens,
+            "block_table": torch.zeros((batch_size, 8), dtype=torch.int32),
+            "seq_lens": seq_lens,
+            "seq_lens_cpu_upper_bound": seq_lens_cpu_upper_bound,
+            "max_query_len": max_query_len,
+            "max_seq_len": max_seq_len,
+        },
+    )()
+
+
+def test_fish_kvcache_enabled_by_default(monkeypatch):
+    monkeypatch.delenv("VLLM_OMNI_FISH_KVCACHE_ATTN", raising=False)
+
+    assert fish_kvcache_attn.is_fish_kvcache_attn_enabled()
+    assert fish_kvcache_backend._fish_kvcache_enabled()
+    assert not fish_kvcache_attn.is_fish_kvcache_attn_required()
+
+
+@pytest.mark.parametrize("value", ["", "1", "true", "yes", "on", "required"])
+def test_fish_kvcache_enabled_values_are_consistent(monkeypatch, value):
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", value)
+    assert fish_kvcache_attn.is_fish_kvcache_attn_enabled()
+    assert fish_kvcache_backend._fish_kvcache_enabled()
+    assert fish_kvcache_attn.is_fish_kvcache_attn_required() is (value == "required")
+
+
+@pytest.mark.parametrize("value", ["0", "false", "no", "off", "disabled", "disable"])
+def test_fish_kvcache_disabled_values_are_consistent(monkeypatch, value):
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", value)
+
+    assert not fish_kvcache_attn.is_fish_kvcache_attn_enabled()
+    assert not fish_kvcache_backend._fish_kvcache_enabled()
+    assert not fish_kvcache_attn.is_fish_kvcache_attn_required()
+
+
+@pytest.mark.parametrize("sliding_window", [None, (-1, -1)])
+def test_fish_kvcache_attn_guard_accepts_decode_only_shape(monkeypatch, sliding_window):
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
+    monkeypatch.setattr(fish_kvcache_attn, "is_available", lambda: True)
+
+    assert fish_kvcache_attn.can_use_fish_kvcache_attn(
+        query=torch.empty((4, 32, 128), dtype=torch.float16),
+        key_cache=torch.empty((8, 16, 8, 128), dtype=torch.float16),
+        value_cache=torch.empty((8, 16, 8, 128), dtype=torch.float16),
+        block_table=torch.zeros((4, 8), dtype=torch.int32),
+        seq_lens=torch.ones((4,), dtype=torch.int32),
+        max_query_len=1,
+        max_seq_len=128,
+        dcp_world_size=1,
+        use_cascade=False,
+        alibi_slopes=None,
+        sliding_window=sliding_window,
+    )
+
+
+def test_fish_kvcache_attn_guard_rejects_unsupported_block_size(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
+    monkeypatch.setattr(fish_kvcache_attn, "is_available", lambda: True)
+
+    assert not fish_kvcache_attn.can_use_fish_kvcache_attn(
+        query=torch.empty((4, 32, 128), dtype=torch.float16),
+        key_cache=torch.empty((8, 32, 8, 128), dtype=torch.float16),
+        value_cache=torch.empty((8, 32, 8, 128), dtype=torch.float16),
+        block_table=torch.zeros((4, 8), dtype=torch.int32),
+        seq_lens=torch.ones((4,), dtype=torch.int32),
+        max_query_len=1,
+        max_seq_len=128,
+        dcp_world_size=1,
+        use_cascade=False,
+        alibi_slopes=None,
+        sliding_window=None,
+    )
+
+
+def test_fish_kvcache_backend_wraps_only_model_instance(monkeypatch):
+    fish_kvcache_backend.reset_fish_kvcache_attn_stats()
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
+    monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
+    monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None)
+    monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", lambda **_: True)
+
+    hits = []
+
+    def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, scale, max_seq_len):
+        del key_cache, value_cache, block_table, seq_lens, scale, max_seq_len
+        hits.append(tuple(query.shape))
+        out.fill_(1)
+        return out
+
+    monkeypatch.setattr(fish_kvcache_backend, "fish_decode_kvcache_attn", fake_decode)
+
+    model = _FakeModel()
+    other_impl = _FakeImpl()
+
+    assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
+    assert not hasattr(other_impl, "_fish_kvcache_attn_installed")
+
+    metadata = _metadata()
+    query = torch.zeros((2, 32, 128), dtype=torch.float16)
+    output = torch.zeros_like(query)
+    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+
+    result = model.layers[0].self_attn.attn.impl.forward(
+        None,
+        query,
+        None,
+        None,
+        kv_cache,
+        metadata,
+        output,
+    )
+
+    assert result is output
+    assert hits == [(2, 32, 128)]
+    assert model.layers[0].self_attn.attn.impl.original_calls == 0
+    assert torch.equal(output, torch.ones_like(output))
+    assert fish_kvcache_backend.get_fish_kvcache_attn_stats()["small_hit_count"] == 1
+
+
+def test_fish_kvcache_backend_install_is_idempotent(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
+    monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
+    monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None)
+    monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", lambda **_: False)
+
+    model = _FakeModel()
+    assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
+    wrapped_forward = model.layers[0].self_attn.attn.impl.forward
+    assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 0
+    assert model.layers[0].self_attn.attn.impl.forward is wrapped_forward
+
+    query = torch.zeros((2, 32, 128), dtype=torch.float16)
+    output = torch.zeros_like(query)
+    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+
+    model.layers[0].self_attn.attn.impl.forward(
+        None,
+        query,
+        None,
+        None,
+        kv_cache,
+        _metadata(),
+        output,
+    )
+
+    assert model.layers[0].self_attn.attn.impl.original_calls == 1
+
+
+def test_fish_kvcache_backend_falls_back_to_original_forward_on_guard_miss(monkeypatch):
+    fish_kvcache_backend.reset_fish_kvcache_attn_stats()
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
+    monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
+    monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None)
+    monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", lambda **_: False)
+
+    decode_calls = []
+
+    def fake_decode(*args, **kwargs):
+        del args, kwargs
+        decode_calls.append(True)
+
+    monkeypatch.setattr(fish_kvcache_backend, "fish_decode_kvcache_attn", fake_decode)
+
+    model = _FakeModel()
+    assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
+
+    query = torch.zeros((2, 32, 128), dtype=torch.float16)
+    output = torch.zeros_like(query)
+    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+
+    result = model.layers[0].self_attn.attn.impl.forward(
+        None,
+        query,
+        None,
+        None,
+        kv_cache,
+        _metadata(),
+        output,
+    )
+
+    assert result is output
+    assert model.layers[0].self_attn.attn.impl.original_calls == 1
+    assert decode_calls == []
+    assert torch.equal(output, torch.full_like(output, 2))
+    stats = fish_kvcache_backend.get_fish_kvcache_attn_stats()
+    assert stats["fallback_count_by_reason"] == {"guard_miss": 1}
+
+
+def test_fish_kvcache_backend_uses_decode_seq_lens_upper_bound(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
+    monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
+    monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None)
+
+    guard_max_seq_lens = []
+
+    def fake_can_use(**kwargs):
+        guard_max_seq_lens.append(kwargs["max_seq_len"])
+        return True
+
+    decode_max_seq_lens = []
+
+    def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, scale, max_seq_len):
+        del query, key_cache, value_cache, block_table, seq_lens, scale
+        decode_max_seq_lens.append(max_seq_len)
+        out.fill_(1)
+        return out
+
+    monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", fake_can_use)
+    monkeypatch.setattr(fish_kvcache_backend, "fish_decode_kvcache_attn", fake_decode)
+
+    model = _FakeModel()
+    assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
+
+    query = torch.zeros((2, 32, 128), dtype=torch.float16)
+    output = torch.zeros_like(query)
+    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+
+    model.layers[0].self_attn.attn.impl.forward(
+        None,
+        query,
+        None,
+        None,
+        kv_cache,
+        _metadata(
+            max_query_len=1,
+            max_seq_len=16384,
+            seq_lens_cpu_upper_bound=torch.tensor([512, 768], dtype=torch.int32),
+        ),
+        output,
+    )
+
+    assert guard_max_seq_lens == [1024]
+    assert decode_max_seq_lens == [1024]
+
+
+def test_fish_kvcache_backend_uses_cpu_upper_bound_above_metadata_max(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
+    monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
+    monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None)
+
+    guard_max_seq_lens = []
+
+    def fake_can_use(**kwargs):
+        guard_max_seq_lens.append(kwargs["max_seq_len"])
+        return True
+
+    monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", fake_can_use)
+    monkeypatch.setattr(
+        fish_kvcache_backend,
+        "fish_decode_kvcache_attn",
+        lambda query, key_cache, value_cache, block_table, seq_lens, out, *, scale, max_seq_len: out.fill_(1),
+    )
+
+    model = _FakeModel()
+    assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
+
+    query = torch.zeros((2, 32, 128), dtype=torch.float16)
+    output = torch.zeros_like(query)
+    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+
+    model.layers[0].self_attn.attn.impl.forward(
+        None,
+        query,
+        None,
+        None,
+        kv_cache,
+        _metadata(
+            max_query_len=1,
+            max_seq_len=64,
+            seq_lens_cpu_upper_bound=torch.tensor([1024, 1024], dtype=torch.int32),
+        ),
+        output,
+    )
+
+    assert guard_max_seq_lens == [1024]
+
+
+def test_fish_kvcache_backend_slices_upper_bound_by_active_batch(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
+    monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
+    monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None)
+
+    guard_max_seq_lens = []
+
+    def fake_can_use(**kwargs):
+        guard_max_seq_lens.append(kwargs["max_seq_len"])
+        return True
+
+    def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, scale, max_seq_len):
+        del query, key_cache, value_cache, block_table, seq_lens, scale, max_seq_len
+        out.fill_(1)
+        return out
+
+    monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", fake_can_use)
+    monkeypatch.setattr(fish_kvcache_backend, "fish_decode_kvcache_attn", fake_decode)
+
+    model = _FakeModel()
+    assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
+
+    query = torch.zeros((2, 32, 128), dtype=torch.float16)
+    output = torch.zeros_like(query)
+    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+
+    model.layers[0].self_attn.attn.impl.forward(
+        None,
+        query,
+        None,
+        None,
+        kv_cache,
+        _metadata(
+            batch_size=4,
+            num_actual_tokens=2,
+            max_query_len=1,
+            max_seq_len=16384,
+            seq_lens_cpu_upper_bound=torch.tensor([512, 768, 4096, 4096], dtype=torch.int32),
+        ),
+        output,
+    )
+
+    assert guard_max_seq_lens == [1024]
+
+
+def test_fish_kvcache_backend_falls_back_on_short_upper_bound_shape(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
+    monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
+    monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None)
+
+    guard_max_seq_lens = []
+
+    def fake_can_use(**kwargs):
+        guard_max_seq_lens.append(kwargs["max_seq_len"])
+        return True
+
+    def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, scale, max_seq_len):
+        del query, key_cache, value_cache, block_table, seq_lens, scale, max_seq_len
+        out.fill_(1)
+        return out
+
+    monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", fake_can_use)
+    monkeypatch.setattr(fish_kvcache_backend, "fish_decode_kvcache_attn", fake_decode)
+
+    model = _FakeModel()
+    assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
+
+    query = torch.zeros((2, 32, 128), dtype=torch.float16)
+    output = torch.zeros_like(query)
+    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+
+    result = model.layers[0].self_attn.attn.impl.forward(
+        None,
+        query,
+        None,
+        None,
+        kv_cache,
+        _metadata(
+            batch_size=2,
+            max_query_len=1,
+            max_seq_len=16384,
+            seq_lens_cpu_upper_bound=torch.tensor([512], dtype=torch.int32),
+        ),
+        output,
+    )
+
+    assert result is output
+    assert guard_max_seq_lens == []
+    assert model.layers[0].self_attn.attn.impl.original_calls == 1
+    assert torch.equal(output, torch.full_like(output, 2))
+
+
+def test_fish_kvcache_backend_falls_back_on_non_cpu_upper_bound(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
+    monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
+    monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None)
+
+    guard_max_seq_lens = []
+
+    def fake_can_use(**kwargs):
+        guard_max_seq_lens.append(kwargs["max_seq_len"])
+        return True
+
+    def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, scale, max_seq_len):
+        del query, key_cache, value_cache, block_table, seq_lens, scale, max_seq_len
+        out.fill_(1)
+        return out
+
+    monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", fake_can_use)
+    monkeypatch.setattr(fish_kvcache_backend, "fish_decode_kvcache_attn", fake_decode)
+
+    model = _FakeModel()
+    assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
+
+    query = torch.zeros((2, 32, 128), dtype=torch.float16)
+    output = torch.zeros_like(query)
+    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+
+    result = model.layers[0].self_attn.attn.impl.forward(
+        None,
+        query,
+        None,
+        None,
+        kv_cache,
+        _metadata(
+            batch_size=2,
+            max_query_len=1,
+            max_seq_len=16384,
+            seq_lens_cpu_upper_bound=torch.empty((2,), device="meta", dtype=torch.int32),
+        ),
+        output,
+    )
+
+    assert result is output
+    assert guard_max_seq_lens == []
+    assert model.layers[0].self_attn.attn.impl.original_calls == 1
+    assert torch.equal(output, torch.full_like(output, 2))
+
+
+def test_fish_kvcache_backend_required_rejects_non_cpu_upper_bound(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "required")
+    monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
+    monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None)
+    monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", lambda **_: True)
+
+    model = _FakeModel()
+    assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
+
+    query = torch.zeros((2, 32, 128), dtype=torch.float16)
+    output = torch.zeros_like(query)
+    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+
+    with pytest.raises(RuntimeError, match="must be a CPU tensor"):
+        model.layers[0].self_attn.attn.impl.forward(
+            None,
+            query,
+            None,
+            None,
+            kv_cache,
+            _metadata(
+                batch_size=2,
+                max_query_len=1,
+                max_seq_len=16384,
+                seq_lens_cpu_upper_bound=torch.empty((2,), device="meta", dtype=torch.int32),
+            ),
+            output,
+        )
+
+
+def test_fish_kvcache_backend_required_rejects_underestimated_upper_bound(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "required")
+    monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
+    monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None)
+    monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", lambda **_: True)
+
+    model = _FakeModel()
+    assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
+
+    query = torch.zeros((2, 32, 128), dtype=torch.float16)
+    output = torch.zeros_like(query)
+    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+
+    with pytest.raises(RuntimeError, match="underestimates"):
+        model.layers[0].self_attn.attn.impl.forward(
+            None,
+            query,
+            None,
+            None,
+            kv_cache,
+            _metadata(
+                max_query_len=1,
+                max_seq_len=16384,
+                seq_lens=torch.tensor([512, 2048], dtype=torch.int32),
+                seq_lens_cpu_upper_bound=torch.tensor([512, 768], dtype=torch.int32),
+            ),
+            output,
+        )
+
+
+def test_fish_kvcache_backend_required_rejects_missing_upper_bound(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "required")
+    monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
+    monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None)
+    monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", lambda **_: True)
+
+    model = _FakeModel()
+    assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
+
+    query = torch.zeros((2, 32, 128), dtype=torch.float16)
+    output = torch.zeros_like(query)
+    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+
+    with pytest.raises(RuntimeError, match="requires seq_lens_cpu_upper_bound"):
+        model.layers[0].self_attn.attn.impl.forward(
+            None,
+            query,
+            None,
+            None,
+            kv_cache,
+            _metadata(max_query_len=1, max_seq_len=16384, seq_lens_cpu_upper_bound=None),
+            output,
+        )
+
+
+def test_fish_kvcache_backend_required_rejects_zero_installed_layers(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "required")
+    monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
+    monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None)
+
+    with pytest.raises(RuntimeError, match="installed 0 attention layers"):
+        fish_kvcache_backend.install_fish_kvcache_attn_backend(type("EmptyModel", (), {"layers": []})())
+
+
+def test_fish_kvcache_backend_required_raises_on_guard_miss(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "required")
+    monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
+    monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None)
+    monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", lambda **_: False)
+
+    model = _FakeModel()
+    assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
+
+    query = torch.zeros((2, 32, 128), dtype=torch.float16)
+    output = torch.zeros_like(query)
+    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+
+    with pytest.raises(RuntimeError, match="required"):
+        model.layers[0].self_attn.attn.impl.forward(
+            None,
+            query,
+            None,
+            None,
+            kv_cache,
+            _metadata(),
+            output,
+        )
+
+
+def test_fish_kvcache_backend_required_allows_prefill_fallback(monkeypatch):
+    fish_kvcache_backend.reset_fish_kvcache_attn_stats()
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "required")
+    monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
+    monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None)
+    monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", lambda **_: False)
+
+    model = _FakeModel()
+    assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
+
+    query = torch.zeros((2, 32, 128), dtype=torch.float16)
+    output = torch.zeros_like(query)
+    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+
+    result = model.layers[0].self_attn.attn.impl.forward(
+        None,
+        query,
+        None,
+        None,
+        kv_cache,
+        _metadata(max_query_len=2, max_seq_len=2),
+        output,
+    )
+
+    assert result is output
+    assert model.layers[0].self_attn.attn.impl.original_calls == 1
+    assert torch.equal(output, torch.full_like(output, 2))
+    stats = fish_kvcache_backend.get_fish_kvcache_attn_stats()
+    assert stats["fallback_count_by_reason"] == {"non_decode": 1}
+
+
+def test_fish_kvcache_backend_falls_back_without_cpu_upper_bound(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
+    monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
+    monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None)
+
+    guard_max_seq_lens = []
+
+    def fake_can_use(**kwargs):
+        guard_max_seq_lens.append(kwargs["max_seq_len"])
+        return True
+
+    decode_max_seq_lens = []
+
+    def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, scale, max_seq_len):
+        del query, key_cache, value_cache, block_table, seq_lens, scale
+        decode_max_seq_lens.append(max_seq_len)
+        out.fill_(1)
+        return out
+
+    monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", fake_can_use)
+    monkeypatch.setattr(fish_kvcache_backend, "fish_decode_kvcache_attn", fake_decode)
+
+    model = _FakeModel()
+    assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
+
+    query = torch.zeros((2, 32, 128), dtype=torch.float16)
+    output = torch.zeros_like(query)
+    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+
+    result = model.layers[0].self_attn.attn.impl.forward(
+        None,
+        query,
+        None,
+        None,
+        kv_cache,
+        _metadata(max_query_len=1, max_seq_len=16384, seq_lens_cpu_upper_bound=None),
+        output,
+    )
+
+    assert result is output
+    assert guard_max_seq_lens == []
+    assert decode_max_seq_lens == []
+    assert model.layers[0].self_attn.attn.impl.original_calls == 1
+    assert torch.equal(output, torch.full_like(output, 2))
+
+
+def test_fish_kvcache_attn_guard_accepts_decode_metadata_above_small_cap(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
+    monkeypatch.setattr(fish_kvcache_attn, "is_available", lambda: True)
+
+    assert fish_kvcache_attn.can_use_fish_kvcache_attn(
+        query=torch.empty((4, 32, 128), dtype=torch.float16),
+        key_cache=torch.empty((1024, 16, 8, 128), dtype=torch.float16),
+        value_cache=torch.empty((1024, 16, 8, 128), dtype=torch.float16),
+        block_table=torch.zeros((4, 1024), dtype=torch.int32),
+        seq_lens=torch.ones((4,), dtype=torch.int32),
+        max_query_len=1,
+        max_seq_len=16384,
+        dcp_world_size=1,
+        use_cascade=False,
+        alibi_slopes=None,
+        sliding_window=None,
+    )
+
+
+def test_fish_kvcache_attn_guard_rejects_seq_len_above_block_table_capacity(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
+    monkeypatch.setattr(fish_kvcache_attn, "is_available", lambda: True)
+
+    assert not fish_kvcache_attn.can_use_fish_kvcache_attn(
+        query=torch.empty((4, 32, 128), dtype=torch.float16),
+        key_cache=torch.empty((8, 16, 8, 128), dtype=torch.float16),
+        value_cache=torch.empty((8, 16, 8, 128), dtype=torch.float16),
+        block_table=torch.zeros((4, 8), dtype=torch.int32),
+        seq_lens=torch.ones((4,), dtype=torch.int32),
+        max_query_len=1,
+        max_seq_len=129,
+        dcp_world_size=1,
+        use_cascade=False,
+        alibi_slopes=None,
+        sliding_window=None,
+    )
+
+
+def test_fish_kvcache_decode_reuses_long_workspace(monkeypatch):
+    fish_kvcache_attn._WORKSPACE_CACHE.clear()
+    calls = []
+
+    def fake_decode(
+        query,
+        key_cache,
+        value_cache,
+        block_table,
+        seq_lens,
+        out,
+        scale,
+        max_seq_len,
+        small_path_max_seq_len,
+        long_split_tokens,
+        partial_m,
+        partial_l,
+        partial_acc,
+    ):
+        del query, key_cache, value_cache, block_table, seq_lens, scale, max_seq_len
+        del small_path_max_seq_len, long_split_tokens
+        calls.append((partial_m.data_ptr(), partial_l.data_ptr(), partial_acc.data_ptr(), tuple(partial_acc.shape)))
+        return out
+
+    monkeypatch.setattr(
+        fish_kvcache_attn._triton_backend(),
+        "fish_decode_kvcache_attn_triton",
+        fake_decode,
+    )
+
+    q = torch.zeros((2, 4, 128), dtype=torch.float16)
+    k_cache = torch.zeros((128, 16, 2, 128), dtype=torch.float16)
+    v_cache = torch.zeros_like(k_cache)
+    block_table = torch.zeros((2, 128), dtype=torch.int32)
+    seq_lens = torch.full((2,), 2048, dtype=torch.int32)
+    out = torch.empty_like(q)
+
+    for _ in range(2):
+        fish_kvcache_attn.fish_decode_kvcache_attn(
+            q,
+            k_cache,
+            v_cache,
+            block_table,
+            seq_lens,
+            out,
+            scale=128**-0.5,
+            max_seq_len=2048,
+        )
+
+    assert len(calls) == 2
+    assert calls[0] == calls[1]
+    assert calls[0][3] == (2, 2, 4, 128)
+
+
+def test_fish_kvcache_decode_workspace_miss_rejects_cuda_graph_capture(monkeypatch):
+    fish_kvcache_attn._WORKSPACE_CACHE.clear()
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
+
+    q = torch.zeros((2, 4, 128), dtype=torch.float16)
+
+    with pytest.raises(RuntimeError, match="workspace was not prewarmed"):
+        fish_kvcache_attn._get_decode_workspace(q, 2048)
+
+
+def test_fish_kvcache_backend_prewarms_capture_workspaces(monkeypatch):
+    fish_kvcache_attn._WORKSPACE_CACHE.clear()
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
+    monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
+
+    text_config = type(
+        "TextConfig",
+        (),
+        {"num_attention_heads": 4, "head_dim": 128},
+    )()
+    hf_config = type("HFConfig", (), {"text_config": text_config})()
+    model_config = type(
+        "ModelConfig",
+        (),
+        {
+            "hf_config": hf_config,
+            "max_model_len": 2048,
+            "model_arch": "FishSpeechSlowARForConditionalGeneration",
+        },
+    )()
+
+    assert (
+        fish_kvcache_backend.prewarm_fish_kvcache_attn_capture_workspaces(
+            model_config=model_config,
+            device=torch.device("cpu"),
+            dtype=torch.float16,
+            capture_sizes=[1, 4],
+        )
+        == 3
+    )
+
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
+    q = torch.zeros((4, 4, 128), dtype=torch.float16)
+    partial_m, partial_l, partial_acc = fish_kvcache_attn._get_decode_workspace(q, 2048)
+
+    assert tuple(partial_m.shape) == (2, 4, 4)
+    assert tuple(partial_l.shape) == (2, 4, 4)
+    assert tuple(partial_acc.shape) == (2, 4, 4, 128)
+
+
+def _reference_decode_attention(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    seq_lens: torch.Tensor,
+    scale: float,
+) -> torch.Tensor:
+    block_size = int(k_cache.shape[1])
+    num_q_heads = int(q.shape[1])
+    num_kv_heads = int(k_cache.shape[2])
+    outputs = []
+    for batch_idx in range(q.shape[0]):
+        seq_len = int(seq_lens[batch_idx].item())
+        if seq_len <= 0:
+            outputs.append(torch.zeros_like(q[batch_idx], dtype=torch.float32))
+            continue
+        k_rows = []
+        v_rows = []
+        for logical_block in range((seq_len + block_size - 1) // block_size):
+            physical_block = int(block_table[batch_idx, logical_block].item())
+            start = logical_block * block_size
+            take = min(block_size, seq_len - start)
+            k_rows.append(k_cache[physical_block, :take])
+            v_rows.append(v_cache[physical_block, :take])
+        k = torch.cat(k_rows, dim=0).to(torch.float32)
+        v = torch.cat(v_rows, dim=0).to(torch.float32)
+        per_head = []
+        for q_head in range(num_q_heads):
+            kv_head = q_head // (num_q_heads // num_kv_heads)
+            scores = torch.matmul(k[:, kv_head, :], q[batch_idx, q_head].to(torch.float32)) * scale
+            weights = torch.softmax(scores, dim=0)
+            per_head.append(torch.matmul(weights, v[:, kv_head, :]))
+        outputs.append(torch.stack(per_head, dim=0))
+    return torch.stack(outputs, dim=0).to(q.dtype)
+
+
+@pytest.mark.parametrize("seq_len", [64, 128, 256, 512, 1024, 1025, 2048])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("batch_size", [1, 3])
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or not fish_kvcache_attn.is_available(),
+    reason="Fish kvcache Triton attention is not available",
+)
+def test_fish_kvcache_native_op_matches_reference(seq_len, dtype, batch_size):
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    block_size = 16
+    num_q_heads = 4
+    num_kv_heads = 2
+    head_dim = 128
+    max_blocks = (seq_len + block_size - 1) // block_size
+    num_blocks = batch_size * max_blocks + 3
+    scale = head_dim**-0.5
+
+    q = torch.randn((batch_size, num_q_heads, head_dim), device=device, dtype=dtype)
+    k_cache = torch.randn((num_blocks, block_size, num_kv_heads, head_dim), device=device, dtype=dtype)
+    v_cache = torch.randn((num_blocks, block_size, num_kv_heads, head_dim), device=device, dtype=dtype)
+    block_table = torch.empty((batch_size, max_blocks), device=device, dtype=torch.int32)
+    for batch_idx in range(batch_size):
+        pages = torch.arange(batch_idx * max_blocks, (batch_idx + 1) * max_blocks, device=device, dtype=torch.int32)
+        block_table[batch_idx] = torch.flip(pages, dims=[0])
+    seq_lens = torch.full((batch_size,), seq_len, device=device, dtype=torch.int32)
+
+    out = torch.empty_like(q)
+    fish_kvcache_attn.fish_decode_kvcache_attn(
+        q,
+        k_cache,
+        v_cache,
+        block_table,
+        seq_lens,
+        out,
+        scale=scale,
+        max_seq_len=seq_len,
+    )
+
+    expected = _reference_decode_attention(q, k_cache, v_cache, block_table, seq_lens, scale)
+    torch.testing.assert_close(out, expected, atol=3e-2, rtol=3e-2)
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or not fish_kvcache_attn.is_available(),
+    reason="Fish kvcache Triton attention is not available",
+)
+def test_fish_kvcache_native_op_zero_seq_len_writes_zero():
+    device = torch.device("cuda")
+    dtype = torch.float16
+    block_size = 16
+    num_q_heads = 4
+    num_kv_heads = 2
+    head_dim = 128
+    max_seq_len = 512
+    max_blocks = max_seq_len // block_size
+    scale = head_dim**-0.5
+
+    q = torch.randn((1, num_q_heads, head_dim), device=device, dtype=dtype)
+    k_cache = torch.randn((max_blocks, block_size, num_kv_heads, head_dim), device=device, dtype=dtype)
+    v_cache = torch.randn_like(k_cache)
+    block_table = torch.arange(max_blocks, device=device, dtype=torch.int32).reshape(1, max_blocks)
+    seq_lens = torch.zeros((1,), device=device, dtype=torch.int32)
+    out = torch.full_like(q, float("nan"))
+
+    fish_kvcache_attn.fish_decode_kvcache_attn(
+        q,
+        k_cache,
+        v_cache,
+        block_table,
+        seq_lens,
+        out,
+        scale=scale,
+        max_seq_len=max_seq_len,
+    )
+
+    assert torch.isfinite(out).all()
+    assert torch.equal(out, torch.zeros_like(out))
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or not fish_kvcache_attn.is_available(),
+    reason="Fish kvcache Triton attention is not available",
+)
+def test_fish_kvcache_native_op_handles_mixed_short_and_long_seq_lens():
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    dtype = torch.float16
+    block_size = 16
+    num_q_heads = 4
+    num_kv_heads = 2
+    head_dim = 128
+    seq_lens = torch.tensor([512, 2048, 1025], device=device, dtype=torch.int32)
+    max_seq_len = int(seq_lens.max().item())
+    batch_size = int(seq_lens.numel())
+    max_blocks = (max_seq_len + block_size - 1) // block_size
+    num_blocks = batch_size * max_blocks
+    scale = head_dim**-0.5
+
+    q = torch.randn((batch_size, num_q_heads, head_dim), device=device, dtype=dtype)
+    k_cache = torch.randn((num_blocks, block_size, num_kv_heads, head_dim), device=device, dtype=dtype)
+    v_cache = torch.randn((num_blocks, block_size, num_kv_heads, head_dim), device=device, dtype=dtype)
+    block_table = torch.empty((batch_size, max_blocks), device=device, dtype=torch.int32)
+    for batch_idx in range(batch_size):
+        pages = torch.arange(batch_idx * max_blocks, (batch_idx + 1) * max_blocks, device=device, dtype=torch.int32)
+        block_table[batch_idx] = pages
+
+    out = torch.empty_like(q)
+    fish_kvcache_attn.fish_decode_kvcache_attn(
+        q,
+        k_cache,
+        v_cache,
+        block_table,
+        seq_lens,
+        out,
+        scale=scale,
+        max_seq_len=max_seq_len,
+    )
+
+    expected = _reference_decode_attention(q, k_cache, v_cache, block_table, seq_lens, scale)
+    torch.testing.assert_close(out, expected, atol=3e-2, rtol=3e-2)
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or not fish_kvcache_attn.is_available(),
+    reason="Fish kvcache Triton attention is not available",
+)
+def test_fish_kvcache_native_op_handles_zero_short_and_long_seq_lens():
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    dtype = torch.float16
+    block_size = 16
+    num_q_heads = 4
+    num_kv_heads = 2
+    head_dim = 128
+    seq_lens = torch.tensor([0, 512, 2048], device=device, dtype=torch.int32)
+    max_seq_len = int(seq_lens.max().item())
+    batch_size = int(seq_lens.numel())
+    max_blocks = (max_seq_len + block_size - 1) // block_size
+    num_blocks = batch_size * max_blocks
+    scale = head_dim**-0.5
+
+    q = torch.randn((batch_size, num_q_heads, head_dim), device=device, dtype=dtype)
+    k_cache = torch.randn((num_blocks, block_size, num_kv_heads, head_dim), device=device, dtype=dtype)
+    v_cache = torch.randn((num_blocks, block_size, num_kv_heads, head_dim), device=device, dtype=dtype)
+    block_table = torch.empty((batch_size, max_blocks), device=device, dtype=torch.int32)
+    for batch_idx in range(batch_size):
+        pages = torch.arange(batch_idx * max_blocks, (batch_idx + 1) * max_blocks, device=device, dtype=torch.int32)
+        block_table[batch_idx] = pages
+    out = torch.full_like(q, float("nan"))
+
+    fish_kvcache_attn.fish_decode_kvcache_attn(
+        q,
+        k_cache,
+        v_cache,
+        block_table,
+        seq_lens,
+        out,
+        scale=scale,
+        max_seq_len=max_seq_len,
+    )
+
+    expected = _reference_decode_attention(q, k_cache, v_cache, block_table, seq_lens, scale)
+    assert torch.isfinite(out).all()
+    torch.testing.assert_close(out, expected, atol=3e-2, rtol=3e-2)
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or not fish_kvcache_attn.is_available(),
+    reason="Fish kvcache Triton attention is not available",
+)
+def test_fish_kvcache_native_op_can_be_captured_by_cuda_graph():
+    device = torch.device("cuda")
+    dtype = torch.float16
+    q = torch.randn((2, 4, 128), device=device, dtype=dtype)
+    k_cache = torch.randn((8, 16, 2, 128), device=device, dtype=dtype)
+    v_cache = torch.randn((8, 16, 2, 128), device=device, dtype=dtype)
+    block_table = torch.arange(8, device=device, dtype=torch.int32).reshape(2, 4)
+    seq_lens = torch.full((2,), 64, device=device, dtype=torch.int32)
+    out = torch.empty_like(q)
+    expected = torch.empty_like(q)
+
+    fish_kvcache_attn.fish_decode_kvcache_attn(
+        q,
+        k_cache,
+        v_cache,
+        block_table,
+        seq_lens,
+        expected,
+        scale=128**-0.5,
+        max_seq_len=64,
+    )
+    torch.accelerator.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        fish_kvcache_attn.fish_decode_kvcache_attn(
+            q,
+            k_cache,
+            v_cache,
+            block_table,
+            seq_lens,
+            out,
+            scale=128**-0.5,
+            max_seq_len=64,
+        )
+    graph.replay()
+    torch.accelerator.synchronize()
+
+    torch.testing.assert_close(out, expected, atol=3e-2, rtol=3e-2)
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or not fish_kvcache_attn.is_available(),
+    reason="Fish kvcache Triton attention is not available",
+)
+def test_fish_kvcache_native_long_op_can_be_captured_by_cuda_graph():
+    device = torch.device("cuda")
+    dtype = torch.float16
+    seq_len = 2048
+    block_size = 16
+    num_q_heads = 4
+    num_kv_heads = 2
+    head_dim = 128
+    max_blocks = (seq_len + block_size - 1) // block_size
+
+    q = torch.randn((2, num_q_heads, head_dim), device=device, dtype=dtype)
+    k_cache = torch.randn((2 * max_blocks, block_size, num_kv_heads, head_dim), device=device, dtype=dtype)
+    v_cache = torch.randn((2 * max_blocks, block_size, num_kv_heads, head_dim), device=device, dtype=dtype)
+    block_table = torch.arange(2 * max_blocks, device=device, dtype=torch.int32).reshape(2, max_blocks)
+    seq_lens = torch.full((2,), seq_len, device=device, dtype=torch.int32)
+    out = torch.empty_like(q)
+    expected = torch.empty_like(q)
+
+    fish_kvcache_attn.fish_decode_kvcache_attn(
+        q,
+        k_cache,
+        v_cache,
+        block_table,
+        seq_lens,
+        expected,
+        scale=head_dim**-0.5,
+        max_seq_len=seq_len,
+    )
+    torch.accelerator.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        fish_kvcache_attn.fish_decode_kvcache_attn(
+            q,
+            k_cache,
+            v_cache,
+            block_table,
+            seq_lens,
+            out,
+            scale=head_dim**-0.5,
+            max_seq_len=seq_len,
+        )
+    graph.replay()
+    torch.accelerator.synchronize()
+
+    torch.testing.assert_close(out, expected, atol=3e-2, rtol=3e-2)
diff --git a/tests/core/sched/test_generation_scheduler_restore.py b/tests/core/sched/test_generation_scheduler_restore.py
index 5cc1cab70..096471e5b 100644
--- a/tests/core/sched/test_generation_scheduler_restore.py
+++ b/tests/core/sched/test_generation_scheduler_restore.py
@@ -28,7 +28,7 @@ class FakeAdapter:
             req = running.pop()
             self.waiting_for_chunk_running_requests.append(req)
 
-    def restore_queues(self, waiting, running):
+    def restore_queues(self, waiting, running, scheduler_requests=None):
         """Put requests back."""
         self.restore_called = True
         running.extend(self.waiting_for_chunk_running_requests)
diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py
new file mode 100644
index 000000000..f6539a40a
--- /dev/null
+++ b/tests/core/sched/test_omni_ar_scheduler_streaming.py
@@ -0,0 +1,92 @@
+"""Unit tests for Omni AR streaming-session async placeholder handling."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+# Imports must run in this order: vllm_omni applies patches to vllm.v1.request before
+# Request / StreamingUpdate are bound in this module. Ruff isort would reorder them.
+# isort: off
+import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
+from vllm.sampling_params import SamplingParams
+from vllm.v1.request import Request, RequestStatus, StreamingUpdate
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+
+# isort: on
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_scheduler(*, stage_id: int = 0) -> OmniARScheduler:
+    sched = OmniARScheduler.__new__(OmniARScheduler)
+    sched._new_prompt_len_snapshot = {}
+    sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(stage_id=stage_id))
+    sched.num_waiting_for_streaming_input = 0
+    sched.log_stats = False
+    sched.chunk_transfer_adapter = None
+    return sched
+
+
+def _make_request() -> Request:
+    return Request(
+        request_id="req-ar-streaming-test",
+        prompt_token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=8),
+        pooling_params=None,
+        arrival_time=100.0,
+        block_hasher=None,
+    )
+
+
+def _make_update(prompt_token_ids: list[int] | None = None) -> StreamingUpdate:
+    return StreamingUpdate(
+        mm_features=None,
+        prompt_token_ids=[10, 20] if prompt_token_ids is None else prompt_token_ids,
+        max_tokens=32,
+        arrival_time=200.0,
+        sampling_params=SamplingParams(max_tokens=16),
+    )
+
+
+def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() -> None:
+    sched = _make_scheduler(stage_id=0)
+    session = _make_request()
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    session.append_output_token_ids([7, 8, 9])
+    session.num_computed_tokens = 6
+    session.num_output_placeholders = 1
+    session.spec_token_ids = [-1]
+
+    sched._update_request_as_session(session, _make_update([10, 20]))
+
+    assert session.discard_latest_async_tokens is True
+    assert session.num_output_placeholders == 0
+    assert session.spec_token_ids == []
+    # The async placeholder makes token 9 unconfirmed, so only 7 and 8 are
+    # carried into the next streaming prompt before the new chunk tokens.
+    assert session.prompt_token_ids == [1, 2, 3, 7, 8, 10, 20]
+    assert list(session._all_token_ids) == [1, 2, 3, 7, 8, 10, 20]
+    assert session._output_token_ids == []
+    assert session.num_prompt_tokens == 7
+    assert sched._new_prompt_len_snapshot[session.request_id] == 2
+
+
+def test_stage0_streaming_update_keeps_all_computed_tokens_without_placeholder() -> None:
+    sched = _make_scheduler(stage_id=0)
+    session = _make_request()
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    session.append_output_token_ids([7, 8, 9])
+    session.num_computed_tokens = 6
+    session.num_output_placeholders = 0
+
+    sched._update_request_as_session(session, _make_update([10, 20]))
+
+    assert session.discard_latest_async_tokens is False
+    assert session.num_output_placeholders == 0
+    assert session.prompt_token_ids == [1, 2, 3, 7, 8, 9, 10, 20]
+    assert list(session._all_token_ids) == [1, 2, 3, 7, 8, 9, 10, 20]
+    assert session._output_token_ids == []
+    assert session.num_prompt_tokens == 8
+    assert sched._new_prompt_len_snapshot[session.request_id] == 2
diff --git a/tests/core/sched/test_omni_scheduler_mixin.py b/tests/core/sched/test_omni_generation_scheduler_update_session.py
similarity index 79%
rename from tests/core/sched/test_omni_scheduler_mixin.py
rename to tests/core/sched/test_omni_generation_scheduler_update_session.py
index e04a9c39f..cf492e7cc 100644
--- a/tests/core/sched/test_omni_scheduler_mixin.py
+++ b/tests/core/sched/test_omni_generation_scheduler_update_session.py
@@ -1,6 +1,6 @@
-"""Unit tests for OmniSchedulerMixin streaming session replacement.
+"""Unit tests for generation streaming session replacement.
 
-These tests pin the behavior of `_replace_session_with_streaming_update` against
+These tests pin the behavior of `_update_request_as_session` against
 current vLLM `Request` / `StreamingUpdate` (and Omni patches). When upgrading
 vLLM, failures here should highlight incompatible changes to request state or
 update payloads early.
@@ -19,19 +19,34 @@ import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine import EngineCoreEventType
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
-from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
 
 # isort: on
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-class _SchedulerStub(OmniSchedulerMixin):
-    """Minimal scheduler surface required by OmniSchedulerMixin."""
+class _SkippedWaitingStub:
+    def __contains__(self, request: Request) -> bool:
+        return False
+
+
+class _ChunkTransferAdapterStub:
+    def __init__(self) -> None:
+        self.segment_finished_requests: set[str] = set()
+
+
+class _SchedulerStub(OmniGenerationScheduler):
+    """Minimal scheduler surface required by OmniGenerationScheduler."""
 
     def __init__(self, *, log_stats: bool = False) -> None:
         self.num_waiting_for_streaming_input = 0
         self.log_stats = log_stats
+        self.chunk_transfer_adapter = _ChunkTransferAdapterStub()
+        self.skipped_waiting = _SkippedWaitingStub()
+
+    def _enqueue_waiting_request(self, session: Request) -> None:
+        raise AssertionError("unexpected enqueue for skipped_waiting miss")
 
 
 def _make_request(**kwargs) -> Request:
@@ -71,7 +86,7 @@ class TestReplaceSessionWithStreamingUpdate:
 
         update = _make_update(prompt_token_ids=[40, 41, 42])
         sched.num_waiting_for_streaming_input = 3
-        sched._replace_session_with_streaming_update(session, update)
+        sched._update_request_as_session(session, update)
 
         assert session._output_token_ids == []
         assert list(session._all_token_ids) == [40, 41, 42]
@@ -88,7 +103,7 @@ class TestReplaceSessionWithStreamingUpdate:
         session = _make_request()
         session.status = RequestStatus.RUNNING
         update = _make_update(prompt_token_ids=None)
-        sched._replace_session_with_streaming_update(session, update)
+        sched._update_request_as_session(session, update)
 
         assert session.prompt_token_ids == ()
         assert list(session._all_token_ids) == []
@@ -108,7 +123,7 @@ class TestReplaceSessionWithStreamingUpdate:
             pytest.skip("StreamingUpdate has no additional_information (Omni patch inactive?)")
         update = replace(base, additional_information=None)
 
-        sched._replace_session_with_streaming_update(session, update)
+        sched._update_request_as_session(session, update)
         assert session.additional_information is None
 
     def test_does_not_decrement_waiting_when_not_streaming_status(self) -> None:
@@ -116,14 +131,14 @@ class TestReplaceSessionWithStreamingUpdate:
         session = _make_request()
         session.status = RequestStatus.RUNNING
         sched.num_waiting_for_streaming_input = 5
-        sched._replace_session_with_streaming_update(session, _make_update())
+        sched._update_request_as_session(session, _make_update())
         assert sched.num_waiting_for_streaming_input == 5
 
     def test_records_queued_event_when_log_stats_enabled(self) -> None:
         sched = _SchedulerStub(log_stats=True)
         session = _make_request()
         session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
-        sched._replace_session_with_streaming_update(session, _make_update())
+        sched._update_request_as_session(session, _make_update())
 
         assert session.events
         assert session.events[-1].type == EngineCoreEventType.QUEUED
diff --git a/tests/core/sched/test_omni_scheduler_mixin_timeouts.py b/tests/core/sched/test_omni_scheduler_mixin_timeouts.py
new file mode 100644
index 000000000..053c926cc
--- /dev/null
+++ b/tests/core/sched/test_omni_scheduler_mixin_timeouts.py
@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit coverage for _process_pending_input_timeouts.
+
+Verifies that the mixin correctly *delegates* timed-out requests to the
+base scheduler's ``finish_requests`` API with ``RequestStatus.FINISHED_ERROR``.
+The end-to-end effect (queue removal + status set + per-request cleanup +
+client-facing FINISHED_ERROR emission) is the responsibility of upstream
+vLLM's ``finish_requests`` implementation and is covered by upstream tests;
+this file only asserts the wiring from the mixin to that API.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+
+
+class _FakeCoordinator:
+    def __init__(self, timed_out_ids):
+        self._timed_out_ids = set(timed_out_ids)
+        self.calls = []
+
+    def collect_timed_out_request_ids(self, timeout_s):
+        self.calls.append(timeout_s)
+        return set(self._timed_out_ids)
+
+
+class _FakeScheduler(OmniSchedulerMixin):
+    def __init__(self, requests, coordinator):
+        self.requests = requests
+        self.input_coordinator = coordinator
+        self.finish_calls = []
+
+    def finish_requests(self, req_ids, status):
+        self.finish_calls.append((set(req_ids), status))
+
+
+def test_process_pending_input_timeouts_delegates_to_finish_requests():
+    """Timed-out request present in self.requests is forwarded to finish_requests."""
+    req_id = "stuck-req"
+    requests = {req_id: SimpleNamespace(request_id=req_id)}
+    coord = _FakeCoordinator(timed_out_ids={req_id})
+    scheduler = _FakeScheduler(requests, coord)
+
+    scheduler._process_pending_input_timeouts()
+
+    assert len(coord.calls) == 1, "coordinator should be polled once"
+    assert coord.calls[0] > 0, "timeout must be positive when enabled"
+
+    assert len(scheduler.finish_calls) == 1
+    finished_ids, status = scheduler.finish_calls[0]
+    assert finished_ids == {req_id}
+    # RequestStatus is the upstream enum; the mixin imports it as
+    # RequestStatus.FINISHED_ERROR.  Check by name to avoid hard import here.
+    assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR")
+
+
+def test_process_pending_input_timeouts_skips_already_freed_request():
+    """Timed-out id no longer in self.requests must not be forwarded."""
+    coord = _FakeCoordinator(timed_out_ids={"already-freed"})
+    scheduler = _FakeScheduler(requests={}, coordinator=coord)
+
+    scheduler._process_pending_input_timeouts()
+
+    assert coord.calls == [coord.calls[0]] and coord.calls[0] > 0
+    assert scheduler.finish_calls == []
+
+
+def test_process_pending_input_timeouts_noop_without_coordinator():
+    """No coordinator => no finish_requests call, no crash."""
+
+    class _NoCoord(OmniSchedulerMixin):
+        def __init__(self):
+            self.requests = {}
+            self.input_coordinator = None
+            self.finish_calls = []
+
+        def finish_requests(self, req_ids, status):
+            self.finish_calls.append((set(req_ids), status))
+
+    scheduler = _NoCoord()
+    scheduler._process_pending_input_timeouts()
+    assert scheduler.finish_calls == []
+
+
+def test_process_pending_input_timeouts_disabled_when_timeout_zero(monkeypatch):
+    """Setting DEFAULT_INPUT_WAIT_TIMEOUT_S <= 0 disables the safety net."""
+    from vllm_omni.core.sched import omni_scheduler_mixin
+
+    monkeypatch.setattr(omni_scheduler_mixin, "DEFAULT_INPUT_WAIT_TIMEOUT_S", 0.0)
+
+    coord = _FakeCoordinator(timed_out_ids={"r1"})
+    scheduler = _FakeScheduler(requests={"r1": SimpleNamespace(request_id="r1")}, coordinator=coord)
+    scheduler._process_pending_input_timeouts()
+    assert coord.calls == [], "coordinator must not be polled when timeout is disabled"
+    assert scheduler.finish_calls == []
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/core/sched/test_omni_scheduling_coordinator.py b/tests/core/sched/test_omni_scheduling_coordinator.py
index 8c7cc20b0..1b36cd784 100644
--- a/tests/core/sched/test_omni_scheduling_coordinator.py
+++ b/tests/core/sched/test_omni_scheduling_coordinator.py
@@ -16,7 +16,7 @@ import torch
 import vllm_omni.core.sched.omni_scheduling_coordinator as coord_mod
 from vllm_omni.core.sched.omni_scheduling_coordinator import (
     OmniSchedulingCoordinator,
-    uses_qwen3_omni_full_payload_input_coordinator,
+    uses_full_payload_input_coordinator,
 )
 
 # ------------------------------------------------------------------ #
@@ -92,47 +92,96 @@ class MockQueue:
 
 
 class TestFullPayloadCoordinatorSelection(unittest.TestCase):
-    def test_qwen3_omni_talker_and_code2wav_use_full_payload_input_coordinator(self):
-        for model_stage in ("talker", "code2wav"):
+    """Tests for the (model_arch, model_stage) whitelist gate.
+
+    The init_omni_connectors arch allowlist is keyed by ``model_arch`` and
+    is a superset of the stages registered here -- consumer-wait stages
+    must be registered explicitly in ``_FULL_PAYLOAD_INPUT_STAGES``, while
+    the init allowlist covers both producer- and consumer-side runners.
+    These tests pin which ``(arch, stage)`` pairs the gate fires for today.
+    """
+
+    # Expected whitelist (model_arch, model_stage).  Hardcoded to avoid the
+    # tautology of importing _FULL_PAYLOAD_INPUT_STAGES and asserting it
+    # against itself; any drift between this matrix and the whitelist will
+    # fail loudly here.
+    EXPECTED_FULL_PAYLOAD_INPUT_STAGES: frozenset[tuple[str, str]] = frozenset(
+        {
+            ("Qwen3OmniMoeForConditionalGeneration", "talker"),
+            ("Qwen3OmniMoeForConditionalGeneration", "code2wav"),
+            ("Qwen2_5OmniForConditionalGeneration", "talker"),
+            ("Qwen2_5OmniForConditionalGeneration", "code2wav"),
+            ("CovoAudioForConditionalGeneration", "code2wav"),
+            ("MiMoAudioModel", "code2wav"),
+            ("Qwen3TTSCode2Wav", "code2wav"),
+            ("CosyVoice3Model", "cosyvoice3_code2wav"),
+            ("DyninOmniForConditionalGeneration", "token2image"),
+            ("DyninOmniForConditionalGeneration", "token2audio"),
+        }
+    )
+
+    def test_whitelist_matches_expected_matrix(self):
+        """_FULL_PAYLOAD_INPUT_STAGES must equal the hardcoded expected matrix.
+
+        Catches both accidental additions (which would silently enable the
+        consumer-wait gate for a new arch) and accidental removals (which
+        would silently disable an enabled arch).
+        """
+        from vllm_omni.core.sched.omni_scheduling_coordinator import _FULL_PAYLOAD_INPUT_STAGES
+
+        self.assertEqual(
+            frozenset(_FULL_PAYLOAD_INPUT_STAGES),
+            self.EXPECTED_FULL_PAYLOAD_INPUT_STAGES,
+            msg="_FULL_PAYLOAD_INPUT_STAGES drifted from the expected matrix; "
+            "update EXPECTED_FULL_PAYLOAD_INPUT_STAGES if intentional.",
+        )
+
+    def test_all_whitelisted_arch_stage_pairs_fire_gate(self):
+        """Every (arch, stage) pair in the expected matrix must fire
+        the gate when stage_id > 0 and async_chunk=False.
+        """
+        for arch, stage in self.EXPECTED_FULL_PAYLOAD_INPUT_STAGES:
             model_config = SimpleNamespace(
                 stage_id=1,
                 async_chunk=False,
-                model_arch="Qwen3OmniMoeForConditionalGeneration",
-                model_stage=model_stage,
+                model_arch=arch,
+                model_stage=stage,
+            )
+            self.assertTrue(
+                uses_full_payload_input_coordinator(model_config),
+                msg=f"expected gate to fire for {arch}/{stage}",
             )
 
-            self.assertTrue(uses_qwen3_omni_full_payload_input_coordinator(model_config))
-
-    def test_async_chunk_and_non_qwen3_omni_do_not_use_full_payload_input_coordinator(self):
+    def test_other_arch_or_stage_or_mode_does_not_fire(self):
         cases = [
             SimpleNamespace(
-                stage_id=1,
-                async_chunk=True,
-                model_arch="Qwen3OmniMoeForConditionalGeneration",
-                model_stage="talker",
+                stage_id=1, async_chunk=True, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage="talker"
+            ),
+            SimpleNamespace(
+                stage_id=0, async_chunk=False, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage="thinker"
             ),
             SimpleNamespace(
                 stage_id=1,
                 async_chunk=False,
-                model_arch="Qwen3TTSForConditionalGeneration",
-                model_stage="code2wav",
-            ),
-            SimpleNamespace(
-                stage_id=1,
-                async_chunk=False,
-                model_arch="Qwen2_5OmniForConditionalGeneration",
-                model_stage="talker",
-            ),
-            SimpleNamespace(
-                stage_id=0,
-                async_chunk=False,
                 model_arch="Qwen3OmniMoeForConditionalGeneration",
-                model_stage="thinker",
+                model_stage="some_future_stage",
+            ),
+            SimpleNamespace(
+                stage_id=1, async_chunk=False, model_arch="Qwen3TTSForConditionalGeneration", model_stage="code2wav"
+            ),
+            SimpleNamespace(
+                stage_id=1, async_chunk=False, model_arch="MingFlashOmniForConditionalGeneration", model_stage="talker"
+            ),
+            SimpleNamespace(stage_id=1, async_chunk=False, model_arch=None, model_stage="talker"),
+            SimpleNamespace(
+                stage_id=1, async_chunk=False, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage=None
             ),
         ]
-
         for model_config in cases:
-            self.assertFalse(uses_qwen3_omni_full_payload_input_coordinator(model_config))
+            self.assertFalse(
+                uses_full_payload_input_coordinator(model_config),
+                msg=f"expected gate OFF for {model_config}",
+            )
 
 
 class TestChunkCoordinatorStateTransition(unittest.TestCase):
diff --git a/tests/core/test_prefix_cache.py b/tests/core/test_prefix_cache.py
index b5d0e96d3..326eb6b34 100644
--- a/tests/core/test_prefix_cache.py
+++ b/tests/core/test_prefix_cache.py
@@ -13,7 +13,7 @@ DEFAULT_SHAPE = torch.Size([NUM_BLOCKS, BLOCK_SIZE, HIDDEN_SIZE])
 
 
 class MockInputBatch:
-    def __init__(self, num_computed_tokens_cpu):
+    def __init__(self, num_computed_tokens_cpu, block_table=None):
         self.req_ids = ["req1", "req2"]
         self.req_id_to_index = {req_id: i for i, req_id in enumerate(self.req_ids)}
         self.num_computed_tokens_cpu = num_computed_tokens_cpu
@@ -21,11 +21,33 @@ class MockInputBatch:
         # Block table is only mocked for validation of length;
         # we don't actually need to add valid values here since
         # we patch the table when testing.
-        class _DummyBlockTable:
-            pass
+        if block_table is None:
 
-        self.block_table = _DummyBlockTable()
-        self.block_table.block_tables = [None]
+            class _DummyBlockTable:
+                pass
+
+            self.block_table = _DummyBlockTable()
+            self.block_table.block_tables = [None]
+        else:
+
+            class _TensorWrapper:
+                def __init__(self, tensor):
+                    self.cpu = tensor
+
+            class _BlockGroup:
+                def __init__(self, tensor):
+                    self.block_table = _TensorWrapper(tensor)
+
+            class _BlockTable:
+                def __init__(self, tensor):
+                    self._group = _BlockGroup(tensor)
+                    self.block_tables = [self._group.block_table]
+
+                def __getitem__(self, idx):
+                    assert idx == 0
+                    return self._group
+
+            self.block_table = _BlockTable(block_table)
 
 
 def get_omni_pcache_with_mm_tensors(feat_dims, seq_len) -> OmniTensorPrefixCache:
@@ -171,6 +193,94 @@ def test_update_with_multimodal_outputs(feat_dims):
             assert torch.all(slot_states == new_output)
 
 
+def test_deferred_multimodal_cache_write_commits_on_completion():
+    cache = get_omni_pcache()
+    mm_key = "codes.audio"
+    block_table = torch.tensor(
+        [
+            [2, 3, 4, 5],
+            [6, 7, 8, 9],
+        ],
+        dtype=torch.long,
+    )
+    input_batch = MockInputBatch(
+        num_computed_tokens_cpu=torch.tensor([4, 0], dtype=torch.long),
+        block_table=block_table,
+    )
+
+    first_chunk = torch.arange(8, dtype=torch.long).reshape(4, 2)
+    cache.update_omni_tensor_prefix_cache(
+        hidden_states=None,
+        multimodal_outputs={mm_key: first_chunk},
+        num_tokens_unpadded=4,
+        slot_mapping=torch.arange(8, 12),
+        skip_mm_cache_keys={mm_key},
+    )
+
+    rows = cache.mm_outputs_cache[mm_key].view(-1, 2)
+    assert torch.all(rows[8:12] == 0)
+
+    cache.stage_deferred_mm_outputs(
+        query_start_loc=torch.tensor([0, 4], dtype=torch.long),
+        input_batch=input_batch,
+        multimodal_outputs={mm_key: first_chunk},
+        num_scheduled_tokens={"req1": 4, "req2": 0},
+        deferred_mm_cache_keys={mm_key},
+    )
+
+    second_chunk = torch.arange(8, 16, dtype=torch.long).reshape(4, 2)
+    input_batch.num_computed_tokens_cpu = torch.tensor([8, 0], dtype=torch.long)
+    cache.stage_deferred_mm_outputs(
+        query_start_loc=torch.tensor([0, 4], dtype=torch.long),
+        input_batch=input_batch,
+        multimodal_outputs={mm_key: second_chunk},
+        num_scheduled_tokens={"req1": 4, "req2": 0},
+        deferred_mm_cache_keys={mm_key},
+    )
+
+    cache.commit_deferred_mm_outputs({"req1"}, input_batch)
+
+    expected = torch.cat([first_chunk, second_chunk], dim=0)
+    rows = cache.mm_outputs_cache[mm_key].view(-1, 2)
+    assert torch.equal(rows[8:16], expected)
+
+
+def test_deferred_multimodal_cache_can_be_merged_on_full_block_hit():
+    cache = get_omni_pcache()
+    mm_key = "codes.audio"
+    block_table = torch.tensor(
+        [
+            [2, 3, 4, 5],
+            [6, 7, 8, 9],
+        ],
+        dtype=torch.long,
+    )
+    input_batch = MockInputBatch(
+        num_computed_tokens_cpu=torch.tensor([8, 0], dtype=torch.long),
+        block_table=block_table,
+    )
+    cached_codes = torch.arange(16, dtype=torch.long).reshape(8, 2)
+    cache.stage_deferred_mm_outputs(
+        query_start_loc=torch.tensor([0, 8], dtype=torch.long),
+        input_batch=input_batch,
+        multimodal_outputs={mm_key: cached_codes},
+        num_scheduled_tokens={"req1": 8, "req2": 0},
+        deferred_mm_cache_keys={mm_key},
+    )
+    cache.commit_deferred_mm_outputs({"req1"}, input_batch)
+
+    new_codes = torch.tensor([[99, 100]], dtype=torch.long)
+    cache.add_prefix_cached_new_req_id("req1")
+    merged = cache.get_merged_multimodal_states(
+        query_start_loc=torch.tensor([0, 1], dtype=torch.long),
+        input_batch=input_batch,
+        multimodal_outputs={mm_key: new_codes},
+        num_scheduled_tokens={"req1": 1, "req2": 0},
+    )
+
+    assert torch.equal(merged[mm_key]["req1"], torch.cat([cached_codes, new_codes], dim=0))
+
+
 ### Tests for Merging
 def fake_get_cached_block_ids(self, req_idx, *args, **kwargs):
     """Fake block table lookup.
diff --git a/tests/dfx/perf/scripts/run_benchmark.py b/tests/dfx/perf/scripts/run_benchmark.py
index 55ad2f8b2..7ca376323 100644
--- a/tests/dfx/perf/scripts/run_benchmark.py
+++ b/tests/dfx/perf/scripts/run_benchmark.py
@@ -191,7 +191,16 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
         raise ValueError("The number of prompts does not match the QPS or max_concurrency")
 
     args = ["--host", host, "--port", str(port)]
-    exclude_keys = {"request_rate", "baseline", "num_prompts", "max_concurrency", "task", "enabled", "eval_phase"}
+    exclude_keys = {
+        "request_rate",
+        "baseline",
+        "num_prompts",
+        "max_concurrency",
+        "task",
+        "enabled",
+        "eval_phase",
+        "trust_remote_code",
+    }
 
     for key, value in params.items():
         if key in exclude_keys or value is None:
@@ -207,6 +216,14 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
         elif not isinstance(value, bool):
             args.extend([arg_name, str(value)])
 
+    for config in BENCHMARK_CONFIGS:
+        if config.get("test_name") != test_name:
+            continue
+        server_params = config.get("server_params") or {}
+        if server_params.get("trust_remote_code") or params.get("trust_remote_code"):
+            args.append("--trust-remote-code")
+        break
+
     # QPS / request-rate sweep
     for i, (qps, num_prompt) in enumerate(zip(qps_list, num_prompt_list)):
         args = args + ["--request-rate", str(qps), "--num-prompts", str(num_prompt)]
diff --git a/tests/dfx/perf/scripts/run_diffusion_benchmark.py b/tests/dfx/perf/scripts/run_diffusion_benchmark.py
index 22e844427..92ce2740d 100644
--- a/tests/dfx/perf/scripts/run_diffusion_benchmark.py
+++ b/tests/dfx/perf/scripts/run_diffusion_benchmark.py
@@ -41,6 +41,72 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
 
+
+# ---------------------------------------------------------------------------
+# Inline field processing
+# ---------------------------------------------------------------------------
+def _process_inline_fields(obj: Any, parent_key: str = "") -> None:
+    """Recursively process '*-inline' fields into temp files."""
+
+    if isinstance(obj, list):
+        for item in obj:
+            _process_inline_fields(item, parent_key)
+        return
+
+    if not isinstance(obj, dict):
+        return
+
+    import atexit
+
+    import yaml
+
+    for key in list(obj.keys()):
+        value = obj[key]
+
+        if not key.endswith("-inline"):
+            _process_inline_fields(value, key)
+            continue
+
+        base_key = key[:-7]
+        full_key = f"{parent_key}.{key}" if parent_key else key
+
+        try:
+            if not isinstance(value, dict):
+                raise ValueError("must be a dict")
+
+            file_type = value.get("type")
+            content = value.get("content")
+
+            if file_type not in {"yaml", "jsonl"}:
+                raise ValueError(f"invalid type: {file_type}")
+
+            fd, path = tempfile.mkstemp(
+                suffix=f".{file_type}",
+                prefix=f"{base_key}_",
+            )
+
+            atexit.register(Path(path).unlink, missing_ok=True)
+
+            with os.fdopen(fd, "w", encoding="utf-8") as f:
+                if file_type == "jsonl":
+                    items = content if isinstance(content, list) else [content]
+                    f.writelines(json.dumps(x, ensure_ascii=False) + "\n" for x in items)
+                else:
+                    yaml.dump(
+                        content,
+                        f,
+                        allow_unicode=True,
+                        sort_keys=False,
+                        indent=2,
+                    )
+
+            obj[base_key] = path
+            del obj[key]
+
+        except Exception as e:
+            print(f"Warning: failed processing '{full_key}': {e}")
+
+
 # ---------------------------------------------------------------------------
 # Paths
 # ---------------------------------------------------------------------------
@@ -102,11 +168,14 @@ def _resolve_refs(configs: list[dict[str, Any]], config_dir: Path) -> list[dict[
 
 
 def load_configs(config_path: str) -> list[dict[str, Any]]:
+    """Load benchmark configs from JSON file and process inline fields."""
     try:
         abs_path = Path(config_path).resolve()
         with open(abs_path, encoding="utf-8") as f:
             configs = json.load(f)
-        return _resolve_refs(configs, abs_path.parent)
+        configs = _resolve_refs(configs, abs_path.parent)
+        _process_inline_fields(configs)
+        return configs
     except json.JSONDecodeError as e:
         raise ValueError(f"JSON parsing error: {str(e)}")
     except FileNotFoundError:
@@ -654,9 +723,47 @@ def run_benchmark(
 # ---------------------------------------------------------------------------
 
 
-def assert_result(result: dict[str, Any], params: dict[str, Any], *, assert_baseline: bool) -> None:
-    """Assert request completion; optionally assert metrics against ``baseline`` in *params*."""
-    num_prompts = params.get("num-prompts", 10)
+def _resolve_baseline_value(
+    baseline_raw: Any,
+    *,
+    sweep_index: int | None,
+    max_concurrency: Any = None,
+    request_rate: Any = None,
+) -> Any:
+    """Pick the baseline threshold for this sweep step."""
+    if baseline_raw is None:
+        return 100000
+    if isinstance(baseline_raw, dict):
+        if max_concurrency is not None:
+            for key in (max_concurrency, str(max_concurrency)):
+                if key in baseline_raw:
+                    return baseline_raw[key]
+        if request_rate is not None:
+            for key in (request_rate, str(request_rate)):
+                if key in baseline_raw:
+                    return baseline_raw[key]
+        raise KeyError(
+            f"baseline dict has no key for max_concurrency={max_concurrency!r} "
+            f"or request_rate={request_rate!r}; keys={list(baseline_raw.keys())!r}"
+        )
+    if isinstance(baseline_raw, (list, tuple)):
+        if sweep_index is None:
+            raise ValueError("sweep_index is required when baseline is a list or tuple")
+        return baseline_raw[sweep_index]
+    return baseline_raw
+
+
+def assert_result(
+    result: dict[str, Any],
+    params: dict[str, Any],
+    num_prompts: int,
+    *,
+    sweep_index: int | None = None,
+    max_concurrency: Any = None,
+    request_rate: Any = None,
+    assert_baseline: bool = True,
+) -> None:
+    """Assert that benchmark metrics satisfy the configured baselines."""
     completed = result.get("completed_requests", result.get("completed", 0))
     assert completed == num_prompts, f"Expected {num_prompts} completed requests, got {completed}"
 
@@ -666,9 +773,15 @@ def assert_result(result: dict[str, Any], params: dict[str, Any], *, assert_base
         print("Skipping performance assertions.")
         return
 
-    for metric, threshold in params.get("baseline", {}).items():
+    for metric, baseline_raw in params.get("baseline", {}).items():
         current = result.get(metric)
         assert current is not None, f"Metric '{metric}' not found in result: {list(result.keys())}"
+        threshold = _resolve_baseline_value(
+            baseline_raw,
+            sweep_index=sweep_index,
+            max_concurrency=max_concurrency,
+            request_rate=request_rate,
+        )
         if "throughput" in metric:
             assert current >= threshold, f"{metric}: {current:.4f} < baseline {threshold}"
         else:
@@ -692,6 +805,111 @@ def _resolve_benchmark_endpoint(server_cfg: dict[str, Any], params: dict[str, An
     return _default_benchmark_endpoint_for_task(cast(str, params.get("task", "t2i")))
 
 
+def _to_list(value: Any) -> list[Any]:
+    if value is None:
+        return []
+    return [value] if not isinstance(value, (list, tuple)) else list(value)
+
+
+def _build_run_params(
+    params: dict[str, Any],
+    *,
+    num_prompts: int,
+    sweep_index: int | None = None,
+    request_rate: Any | None = None,
+    max_concurrency: Any | None = None,
+) -> dict[str, Any]:
+    run_params = {
+        key: value for key, value in params.items() if key not in {"request-rate", "max-concurrency", "num-prompts"}
+    }
+    run_params["num-prompts"] = num_prompts
+    if request_rate is not None:
+        run_params["request-rate"] = request_rate
+    if max_concurrency is not None:
+        run_params["max-concurrency"] = max_concurrency
+    if "baseline" in params:
+        run_params["baseline"] = {
+            metric: _resolve_baseline_value(
+                baseline_raw,
+                sweep_index=sweep_index,
+                max_concurrency=max_concurrency,
+                request_rate=request_rate,
+            )
+            for metric, baseline_raw in params["baseline"].items()
+        }
+    return run_params
+
+
+def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
+    request_rate_list = _to_list(params.get("request-rate"))
+    num_prompt_list = _to_list(params.get("num-prompts", 10))
+    max_concurrency_list = _to_list(params.get("max-concurrency"))
+
+    max_len = max(len(request_rate_list), len(max_concurrency_list))
+    if len(num_prompt_list) == 1 and max_len > 1:
+        num_prompt_list = num_prompt_list * max_len
+    elif max_len == 1 and len(num_prompt_list) > 1:
+        if len(request_rate_list) == 1:
+            request_rate_list = request_rate_list * len(num_prompt_list)
+        if len(max_concurrency_list) == 1:
+            max_concurrency_list = max_concurrency_list * len(num_prompt_list)
+        max_len = max(len(request_rate_list), len(max_concurrency_list))
+    elif len(num_prompt_list) != max_len and max_len > 0:
+        raise ValueError("The number of prompts does not match the request-rate or max-concurrency")
+
+    sweep_runs: list[dict[str, Any]] = []
+
+    for i, (request_rate, num_prompts) in enumerate(zip(request_rate_list, num_prompt_list)):
+        sweep_runs.append(
+            {
+                "params": _build_run_params(
+                    params,
+                    request_rate=request_rate,
+                    num_prompts=num_prompts,
+                    sweep_index=i,
+                ),
+                "num_prompts": num_prompts,
+                "sweep_index": i,
+                "request_rate": request_rate,
+                "max_concurrency": None,
+            }
+        )
+
+    for i, (max_concurrency, num_prompts) in enumerate(zip(max_concurrency_list, num_prompt_list)):
+        sweep_runs.append(
+            {
+                "params": _build_run_params(
+                    params,
+                    max_concurrency=max_concurrency,
+                    num_prompts=num_prompts,
+                    sweep_index=i,
+                    request_rate="inf",
+                ),
+                "num_prompts": num_prompts,
+                "sweep_index": i,
+                "request_rate": None,
+                "max_concurrency": max_concurrency,
+            }
+        )
+
+    if not sweep_runs:
+        default_num_prompts = num_prompt_list[0]
+        sweep_runs.append(
+            {
+                "params": _build_run_params(
+                    params,
+                    num_prompts=default_num_prompts,
+                ),
+                "num_prompts": default_num_prompts,
+                "sweep_index": None,
+                "request_rate": None,
+                "max_concurrency": None,
+            }
+        )
+
+    return sweep_runs
+
+
 # ---------------------------------------------------------------------------
 # Test entry point
 # ---------------------------------------------------------------------------
@@ -715,36 +933,47 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, req
     test_name = benchmark_params["test_name"]
     params = benchmark_params["params"]
     server_cfg = getattr(diffusion_server, "server_cfg", {})
-    endpoint = _resolve_benchmark_endpoint(server_cfg, params)
+    sweep_runs = _iter_sweep_runs(params)
 
-    result = run_benchmark(
-        host=diffusion_server.host,
-        port=diffusion_server.port,
-        model=diffusion_server.model,
-        params=params,
-        test_name=test_name,
-        endpoint=endpoint,
-        server_cfg=server_cfg,
-        source_file=cast(str, CONFIG_FILE_PATH),
-    )
+    for sweep_run in sweep_runs:
+        endpoint = _resolve_benchmark_endpoint(server_cfg, sweep_run["params"])
+        result = run_benchmark(
+            host=diffusion_server.host,
+            port=diffusion_server.port,
+            model=diffusion_server.model,
+            params=sweep_run["params"],
+            test_name=test_name,
+            endpoint=endpoint,
+            server_cfg=server_cfg,
+            source_file=cast(str, CONFIG_FILE_PATH),
+        )
 
-    print(f"\n{'=' * 60}")
-    print(f"Results for {test_name} (server={diffusion_server.server_type}, endpoint={endpoint}):")
-    for key in (
-        "throughput_qps",
-        "latency_mean",
-        "latency_median",
-        "latency_p50",
-        "latency_p99",
-        "peak_memory_mb_max",
-        "peak_memory_mb_mean",
-        "peak_memory_mb_median",
-    ):
-        if key in result:
-            print(f"  {key}: {result[key]:.4f}")
+        print(f"\n{'=' * 60}")
+        print(f"Results for {test_name} (server={diffusion_server.server_type}, endpoint={endpoint}):")
+        for key in (
+            "throughput_qps",
+            "latency_mean",
+            "latency_median",
+            "latency_p50",
+            "latency_p99",
+            "peak_memory_mb_max",
+            "peak_memory_mb_mean",
+            "peak_memory_mb_median",
+        ):
+            if key in result:
+                print(f"  {key}: {result[key]:.4f}")
 
-    print(f"\n  Aggregated results: {AGGREGATED_RESULT_FILE}")
-    print("=" * 60)
+        print(f"\n  Aggregated results: {AGGREGATED_RESULT_FILE}")
+        print("=" * 60)
 
-    assert_baseline = request.config.getoption("--assert-baseline", default=False)
-    assert_result(result, params, assert_baseline=assert_baseline)
+        assert_baseline = request.config.getoption("--assert-baseline", default=False)
+
+        assert_result(
+            result,
+            params,
+            sweep_run["num_prompts"],
+            sweep_index=sweep_run["sweep_index"],
+            max_concurrency=sweep_run["max_concurrency"],
+            request_rate=sweep_run["request_rate"],
+            assert_baseline=assert_baseline,
+        )
diff --git a/tests/dfx/perf/tests/test_bagel_vllm_omni.json b/tests/dfx/perf/tests/test_bagel_vllm_omni.json
new file mode 100644
index 000000000..09fc96030
--- /dev/null
+++ b/tests/dfx/perf/tests/test_bagel_vllm_omni.json
@@ -0,0 +1,152 @@
+[
+  {
+    "test_name": "test_bagel_single_device_single_stage_t2i",
+    "description": "Single-stage BAGEL (TP=1, CFG=1), t2i Task",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "ByteDance-Seed/BAGEL-7B-MoT",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true,
+        "deploy-config": "../vllm_omni/deploy/bagel_single_stage.yaml",
+        "no-async-chunk": true,
+        "tensor-parallel-size": 1,
+        "cfg-parallel-size": 1,
+        "gpu-memory-utilization": 0.90,
+        "max-model-len": 8192,
+        "max-num-batched-tokens": 9216
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "512x512_t2i_steps20",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 20,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.3678,
+          "latency_mean": 2.691,
+          "peak_memory_mb_max": 33083.6,
+          "peak_memory_mb_mean": 33082.5
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_bagel_single_device_single_stage_i2i",
+    "description": "Single-stage BAGEL (TP=1, CFG=1), i2i Task",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "ByteDance-Seed/BAGEL-7B-MoT",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true,
+        "deploy-config": "../vllm_omni/deploy/bagel_single_stage.yaml",
+        "no-async-chunk": true,
+        "tensor-parallel-size": 1,
+        "cfg-parallel-size": 1,
+        "gpu-memory-utilization": 0.90,
+        "max-model-len": 8192,
+        "max-num-batched-tokens": 9216
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "512x512_i2i_steps20",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 20,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "num-input-images": 1,
+        "baseline": {
+          "throughput_qps": 0.3124,
+          "latency_mean": 3.1686,
+          "peak_memory_mb_max": 33448.8,
+          "peak_memory_mb_mean": 33447.7
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_bagel_single_device_multi_stage_t2i",
+    "description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, t2i Task",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "ByteDance-Seed/BAGEL-7B-MoT",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true,
+        "deploy-config": "../vllm_omni/deploy/bagel.yaml",
+        "no-async-chunk": true,
+        "tensor-parallel-size": 1,
+        "cfg-parallel-size": 1,
+        "gpu-memory-utilization": 0.45,
+        "max-model-len": 8192,
+        "max-num-batched-tokens": 9216
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "512x512_t2i_steps20",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 20,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.3746,
+          "latency_mean": 2.6425,
+          "peak_memory_mb_max": 33077.0,
+          "peak_memory_mb_mean": 33077.0
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_bagel_single_device_multi_stage_i2i",
+    "description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, i2i Task",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "ByteDance-Seed/BAGEL-7B-MoT",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true,
+        "deploy-config": "../vllm_omni/deploy/bagel.yaml",
+        "no-async-chunk": true,
+        "tensor-parallel-size": 1,
+        "cfg-parallel-size": 1,
+        "gpu-memory-utilization": 0.45,
+        "max-model-len": 8192,
+        "max-num-batched-tokens": 9216
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "512x512_i2i_steps20",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 20,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "num-input-images": 1,
+        "baseline": {
+          "throughput_qps": 0.1349,
+          "latency_mean": 7.339,
+          "peak_memory_mb_max": 33321.2,
+          "peak_memory_mb_mean": 33321.2
+        }
+      }
+    ]
+  }
+]
diff --git a/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json b/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json
new file mode 100644
index 000000000..661f6d263
--- /dev/null
+++ b/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json
@@ -0,0 +1,131 @@
+[
+    {
+        "test_name": "test_hunyuan_image3_ti2i_tp4_tp4_inline",
+        "description": "AR (TP=4) + DiT (TP=4) - using inline dataset and deploy-config",
+        "server_type": "vllm-omni",
+        "benchmark_endpoint": "/v1/images/edits",
+        "server_params": {
+            "model": "tencent/HunyuanImage-3.0-Instruct",
+            "serve_args": {
+                "deploy-config-inline": {
+                    "type": "yaml",
+                    "content": {
+                        "pipeline": "hunyuan_image_3_moe",
+                        "async_chunk": false,
+                        "trust_remote_code": true,
+                        "connectors": {
+                            "rdma_connector": {
+                                "name": "MooncakeTransferEngineConnector",
+                                "extra": {
+                                    "host": "auto",
+                                    "zmq_port": 50051,
+                                    "protocol": "rdma",
+                                    "device_name": "",
+                                    "memory_pool_size": 4294967296,
+                                    "memory_pool_device": "cpu"
+                                }
+                            }
+                        },
+                        "stages": [
+                            {
+                                "stage_id": 0,
+                                "is_comprehension": true,
+                                "final_output": true,
+                                "final_output_type": "text",
+                                "max_num_seqs": 4,
+                                "gpu_memory_utilization": 0.95,
+                                "enforce_eager": false,
+                                "enable_prefix_caching": false,
+                                "max_num_batched_tokens": 32768,
+                                "devices": "0,1,2,3",
+                                "tensor_parallel_size": 4,
+                                "hf_overrides": {
+                                    "rope_parameters": {
+                                        "mrope_section": [0, 32, 32],
+                                        "rope_type": "default"
+                                    }
+                                },
+                                "omni_kv_config": {
+                                    "need_send_cache": true
+                                },
+                                "output_connectors": {
+                                    "to_stage_1": "rdma_connector"
+                                },
+                                "default_sampling_params": {
+                                    "temperature": 0.0,
+                                    "top_p": 1,
+                                    "top_k": -1,
+                                    "max_tokens": 8192,
+                                    "detokenize": true,
+                                    "skip_special_tokens": false,
+                                    "stop_token_ids": [128025]
+                                }
+                            },
+                            {
+                                "stage_id": 1,
+                                "max_num_seqs": 1,
+                                "enforce_eager": true,
+                                "devices": "4,5,6,7",
+                                "distributed_executor_backend": "mp",
+                                "omni_kv_config": {
+                                    "need_recv_cache": true
+                                },
+                                "parallel_config": {
+                                    "tensor_parallel_size": 4,
+                                    "enable_expert_parallel": true
+                                },
+                                "input_connectors": {
+                                    "from_stage_0": "rdma_connector"
+                                },
+                                "default_sampling_params": {
+                                    "num_inference_steps": 50,
+                                    "guidance_scale": 0
+                                }
+                            }
+                        ],
+                        "edges": [
+                            {
+                                "from": 0,
+                                "to": 1,
+                                "window_size": -1,
+                                "max_inflight": 1
+                            }
+                        ]
+                    }
+                },
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "inline_dataset_5_requests",
+                "dataset": "custom",
+                "dataset-path-inline": {
+                    "type": "jsonl",
+                    "content": [
+                        {
+                            "prompt": "新年宠物海报，Q版圆润的可爱标题“新年快乐汪”，副标题“HAPPY NEW YEAR”。 鱼眼镜头，背景是房间门口，近景，上传的主体歪头笑，围着红色围巾，戴着红色毛线帽，高清，绒毛细节，面部特写。 宝丽莱相纸，超现实主义，写实主义，胶片摄影，打印颗粒感肌理。肌理，超写实，复古感。",
+                            "image_urls": ["https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_0_0.png"],
+                            "num_inference_steps": 8,
+                            "seed": 42,
+                            "guidance_scale": 1.0
+                        }
+                    ]
+                },
+                "task": "ti2i",
+                "bot-task": "think_recaption",
+                "max-concurrency": 4,
+                "num-prompts": 16,
+                "skip-performance-assertion": true,
+                "baseline": {
+                    "throughput_qps": 0.1,
+                    "latency_p99": 44.46,
+                    "latency_mean": 35.57,
+                    "latency_std_dev": 0.5,
+                    "peak_memory_mb_max": 70000,
+                    "peak_memory_mb_mean": 70000
+                }
+            }
+        ]
+    }
+]
diff --git a/tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_cfgp2.json b/tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json
similarity index 57%
rename from tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_cfgp2.json
rename to tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json
index e6271545f..e6d7ce659 100644
--- a/tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_cfgp2.json
+++ b/tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json
@@ -1,35 +1,33 @@
 [
     {
-        "test_name": "test_hunyuan_image_tp2_fp8_cfgp2",
-        "description": "TP=2 Quantization=fp8 CfgP=2 baseline",
+        "test_name": "test_hunyuan_image_tp2_cfgp2",
+        "description": "TP=2 CfgP=2 baseline",
         "server_type": "vllm-omni",
         "server_params": {
             "model": "tencent/HunyuanImage-3.0-Instruct",
             "serve_args": {
+                "deploy_config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
                 "tensor-parallel-size": 2,
                 "cfg-parallel-size": 2,
-                "quantization": "fp8",
-                "distributed-executor-backend": "mp",
-                "enforce-eager": true,
                 "enable-diffusion-pipeline-profiler": true
             }
         },
         "benchmark_params": [
             {
-                "name": "1024x1024_steps50",
+                "name": "1024x1024_steps8",
                 "dataset": "random",
                 "task": "t2i",
                 "width": 1024,
                 "height": 1024,
-                "num-inference-steps": 50,
+                "num-inference-steps": 8,
                 "num-prompts": 10,
                 "max-concurrency": 1,
                 "skip-performance-assertion": true,
                 "baseline": {
-                    "throughput_qps": 0.1035,
-                    "latency_p99": 9.9057,
-                    "peak_memory_mb_max": 66470,
-                    "peak_memory_mb_mean": 66470
+                    "throughput_qps": 0.21,
+                    "latency_p99": 4.7469,
+                    "peak_memory_mb_max": 101100,
+                    "peak_memory_mb_mean": 101100
                 }
             }
         ]
diff --git a/tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_sp2.json b/tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json
similarity index 54%
rename from tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_sp2.json
rename to tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json
index 0112fc51a..fbbe1643e 100644
--- a/tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_sp2.json
+++ b/tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json
@@ -1,35 +1,33 @@
 [
     {
-        "test_name": "test_hunyuan_image_tp2_fp8_sp2",
-        "description": "TP=2 Quantization=fp8 SP=2 baseline",
+        "test_name": "test_hunyuan_image_tp2_sp2",
+        "description": "TP=2 SP=2 baseline",
         "server_type": "vllm-omni",
         "server_params": {
             "model": "tencent/HunyuanImage-3.0-Instruct",
             "serve_args": {
+                "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
                 "tensor-parallel-size": 2,
-                "usp": 2,
-                "quantization": "fp8",
-                "distributed-executor-backend": "mp",
-                "enforce-eager": true,
+                "ulysses-degree": 2,
                 "enable-diffusion-pipeline-profiler": true
             }
         },
         "benchmark_params": [
             {
-                "name": "1024x1024_steps50",
+                "name": "1024x1024_steps8",
                 "dataset": "random",
                 "task": "t2i",
                 "width": 1024,
                 "height": 1024,
-                "num-inference-steps": 50,
+                "num-inference-steps": 8,
                 "num-prompts": 10,
                 "max-concurrency": 1,
                 "skip-performance-assertion": true,
                 "baseline": {
-                    "throughput_qps": 0.08,
-                    "latency_p99": 12.0731,
-                    "peak_memory_mb_max": 66314,
-                    "peak_memory_mb_mean": 66314
+                    "throughput_qps": 0.20,
+                    "latency_p99": 5.1025,
+                    "peak_memory_mb_max": 97402,
+                    "peak_memory_mb_mean": 97402
                 }
             }
         ]
diff --git a/tests/dfx/perf/tests/test_hunyuan_image_tp4_fp8.json b/tests/dfx/perf/tests/test_hunyuan_image_tp4.json
similarity index 56%
rename from tests/dfx/perf/tests/test_hunyuan_image_tp4_fp8.json
rename to tests/dfx/perf/tests/test_hunyuan_image_tp4.json
index 8f6bad2f8..8f4c0ec50 100644
--- a/tests/dfx/perf/tests/test_hunyuan_image_tp4_fp8.json
+++ b/tests/dfx/perf/tests/test_hunyuan_image_tp4.json
@@ -1,34 +1,32 @@
 [
     {
-        "test_name": "test_hunyuan_image_tp4_fp8",
-        "description": "TP=4 Quantization=fp8 baseline",
+        "test_name": "test_hunyuan_image_tp4",
+        "description": "TP=4 baseline",
         "server_type": "vllm-omni",
         "server_params": {
             "model": "tencent/HunyuanImage-3.0-Instruct",
             "serve_args": {
+                "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
                 "tensor-parallel-size": 4,
-                "quantization": "fp8",
-                "distributed-executor-backend": "mp",
-                "enforce-eager": true,
                 "enable-diffusion-pipeline-profiler": true
             }
         },
         "benchmark_params": [
             {
-                "name": "1024x1024_steps50",
+                "name": "1024x1024_steps8",
                 "dataset": "random",
                 "task": "t2i",
                 "width": 1024,
                 "height": 1024,
-                "num-inference-steps": 50,
+                "num-inference-steps": 8,
                 "num-prompts": 10,
                 "max-concurrency": 1,
                 "skip-performance-assertion": true,
                 "baseline": {
-                    "throughput_qps": 0.08,
-                    "latency_p99": 13.1227,
-                    "peak_memory_mb_max": 46838,
-                    "peak_memory_mb_mean": 46838
+                    "throughput_qps": 0.21,
+                    "latency_p99": 4.9235,
+                    "peak_memory_mb_max": 57576,
+                    "peak_memory_mb_mean": 57576
                 }
             }
         ]
diff --git a/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json b/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json
new file mode 100644
index 000000000..0e7cb63f7
--- /dev/null
+++ b/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json
@@ -0,0 +1,115 @@
+[
+    {
+        "test_name": "test_qwen3_omni_3gpu_replica2",
+        "server_params": {
+            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+            "stage_overrides": {
+                "1": {"num_replicas": 2, "devices": "1,2"},
+                "2": {"num_replicas": 2, "devices": "1,2"}
+            },
+            "extra_cli_args": [
+                "--async-chunk"
+            ]
+        },
+        "benchmark_params": [
+            {
+                "dataset_name": "random",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [128, 128, 128, 128],
+                "max_concurrency": [8, 16, 24, 32],
+                "random_input_len": 2500,
+                "random_output_len": 900,
+                "ignore_eos": true,
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [1000, 2000, 3000, 4000],
+                    "mean_tpot_ms": [200, 200, 200, 200],
+                    "mean_audio_ttfp_ms": [10000, 15000, 20000, 25000],
+                    "mean_audio_rtf": [1.5, 1.5, 1.5, 1.5]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [10],
+                "request_rate": [0.1],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 1,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(0, 60, 3)": 1.0
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [2000],
+                    "mean_audio_ttfp_ms": [2000],
+                    "mean_audio_rtf": [0.25]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [40],
+                "request_rate": [0.5],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 2,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.5,
+                    "(720, 1280, 2)": 0.5
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [6000],
+                    "mean_audio_ttfp_ms": [6000],
+                    "mean_audio_rtf": [0.7]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [100],
+                "request_rate": [1.0],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 3,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1,
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.34,
+                    "(720, 1280, 2)": 0.33,
+                    "(0, 60, 3)": 0.33
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [12000],
+                    "mean_audio_ttfp_ms": [12000],
+                    "mean_audio_rtf": [1.0]
+                }
+            }
+        ]
+    }
+]
diff --git a/tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
similarity index 83%
rename from tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json
rename to tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
index d13f37053..1903bc670 100644
--- a/tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json
+++ b/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
@@ -1,10 +1,10 @@
 [
     {
-        "test_name": "test_qwen_image_edit_2509_single_device",
+        "test_name": "test_qwen_image_edit_2511_single_device",
         "description": "Single-device baseline (two input images)",
         "server_type": "vllm-omni",
         "server_params": {
-            "model": "Qwen/Qwen-Image-Edit-2509",
+            "model": "Qwen/Qwen-Image-Edit-2511",
             "serve_args": {
                 "enable-diffusion-pipeline-profiler": true
             }
@@ -49,11 +49,11 @@
         ]
     },
     {
-        "test_name": "test_qwen_image_edit_2509_ulysses2_cfg2_vae_patch4",
+        "test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_vae_patch4",
         "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
         "server_type": "vllm-omni",
         "server_params": {
-            "model": "Qwen/Qwen-Image-Edit-2509",
+            "model": "Qwen/Qwen-Image-Edit-2511",
             "serve_args": {
                 "ulysses-degree": 2,
                 "cfg-parallel-size": 2,
@@ -63,24 +63,6 @@
             }
         },
         "benchmark_params": [
-            {
-                "name": "512x512_steps20_i2i_2img",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 512,
-                "height": 512,
-                "num-inference-steps": 20,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "num-input-images": 2,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.15,
-                    "latency_mean": 10,
-                    "peak_memory_mb_max": 55000,
-                    "peak_memory_mb_mean": 55000
-                }
-            },
             {
                 "name": "1536x1536_steps35_i2i_2img",
                 "dataset": "random",
@@ -102,11 +84,11 @@
         ]
     },
     {
-        "test_name": "test_qwen_image_edit_2509_ulysses2_cfg2_cache_dit",
+        "test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_cache_dit",
         "description": "Ulysses SP=2 + CFG=2 + CacheDiT",
         "server_type": "vllm-omni",
         "server_params": {
-            "model": "Qwen/Qwen-Image-Edit-2509",
+            "model": "Qwen/Qwen-Image-Edit-2511",
             "serve_args": {
                 "ulysses-degree": 2,
                 "cfg-parallel-size": 2,
diff --git a/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
index 8ae370df6..21695b752 100644
--- a/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
+++ b/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
@@ -61,23 +61,6 @@
             }
         },
         "benchmark_params": [
-            {
-                "name": "512x512_steps20_i2i",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 512,
-                "height": 512,
-                "num-inference-steps": 20,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.21,
-                    "latency_mean": 10,
-                    "peak_memory_mb_max": 56000,
-                    "peak_memory_mb_mean": 56000
-                }
-            },
             {
                 "name": "1536x1536_steps35_i2i",
                 "dataset": "random",
diff --git a/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
index e05f1b945..a6eacf88e 100644
--- a/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
+++ b/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
@@ -181,5 +181,39 @@
                 }
             }
         ]
+    },
+    {
+        "test_name": "test_qwen_image_single_device_step_execution_high_concurrency",
+        "description": "Single-device Qwen-Image step-level execution benchmark at 512x512, 20 steps, concurrency sweep",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Qwen/Qwen-Image",
+            "serve_args": {
+                "enable-diffusion-pipeline-profiler": true,
+                "step-execution": true,
+                "max-num-seqs": 8
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "512x512_steps20_high_concurrency",
+                "dataset": "random",
+                "task": "t2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 20,
+                "num-prompts": [20, 40, 80, 160],
+                "max-concurrency": [1, 2, 4, 8],
+                "warmup-requests": 8,
+                "warmup-concurrency": 8,
+                "warmup-num-inference-steps": 20,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": [0.30, 0.50, 0.60, 0.70],
+                    "latency_mean": [3.50, 4.00, 7.00, 12.00],
+                    "peak_memory_mb_mean": [67000, 67000, 67000, 67000]
+                }
+            }
+        ]
     }
 ]
diff --git a/tests/dfx/perf/tests/test_runner_metadata.py b/tests/dfx/perf/tests/test_runner_metadata.py
index 1276a8470..a195c4546 100644
--- a/tests/dfx/perf/tests/test_runner_metadata.py
+++ b/tests/dfx/perf/tests/test_runner_metadata.py
@@ -17,7 +17,16 @@ def test_task_excluded_from_cli_args():
         "percentile-metrics": "audio_rtf,audio_ttfp",
         "baseline": {"mean_audio_rtf": [0.5]},
     }
-    exclude_keys = {"request_rate", "baseline", "num_prompts", "max_concurrency", "task", "enabled", "eval_phase"}
+    exclude_keys = {
+        "request_rate",
+        "baseline",
+        "num_prompts",
+        "max_concurrency",
+        "task",
+        "enabled",
+        "eval_phase",
+        "trust_remote_code",
+    }
     args = []
     for key, value in params.items():
         if key in exclude_keys or value is None:
diff --git a/tests/dfx/perf/tests/test_voxcpm2.json b/tests/dfx/perf/tests/test_voxcpm2.json
new file mode 100644
index 000000000..f19f95960
--- /dev/null
+++ b/tests/dfx/perf/tests/test_voxcpm2.json
@@ -0,0 +1,198 @@
+[
+  {
+    "test_name": "test_voxcpm2",
+    "server_params": {
+      "model": "openbmb/VoxCPM2",
+      "trust_remote_code": true
+    },
+    "benchmark_params": [
+      {
+        "task": "voice_clone",
+        "eval_phase": "latency",
+        "dataset_name": "seed-tts",
+        "backend": "openai-audio-speech",
+        "endpoint": "/v1/audio/speech",
+        "dataset_path": "linyueqian/seed-tts-eval-subset",
+        "num_prompts": [
+          20
+        ],
+        "max_concurrency": [
+          1
+        ],
+        "seed_tts_locale": "en",
+        "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "median_audio_ttfp_ms": [
+            600
+          ],
+          "median_audio_rtf": [
+            0.2
+          ]
+        }
+      },
+      {
+        "task": "voice_clone",
+        "eval_phase": "throughput",
+        "dataset_name": "seed-tts",
+        "backend": "openai-audio-speech",
+        "endpoint": "/v1/audio/speech",
+        "dataset_path": "linyueqian/seed-tts-eval-subset",
+        "num_prompts": [
+          80
+        ],
+        "max_concurrency": [
+          8
+        ],
+        "seed_tts_locale": "en",
+        "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "median_audio_ttfp_ms": [
+            3500
+          ],
+          "median_audio_rtf": [
+            0.85
+          ],
+          "audio_throughput": [
+            8
+          ]
+        }
+      },
+      {
+        "task": "voice_clone",
+        "eval_phase": "quality",
+        "enabled": false,
+        "dataset_name": "seed-tts",
+        "backend": "openai-audio-speech",
+        "endpoint": "/v1/audio/speech",
+        "dataset_path": "linyueqian/seed-tts-eval-subset",
+        "num_prompts": [
+          200
+        ],
+        "max_concurrency": [
+          4
+        ],
+        "seed_tts_locale": "en",
+        "seed_tts_wer_eval": true,
+        "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "mean_audio_rtf": [
+            0.45
+          ]
+        }
+      },
+      {
+        "task": "default_voice",
+        "eval_phase": "latency",
+        "dataset_name": "seed-tts-text",
+        "backend": "openai-audio-speech",
+        "endpoint": "/v1/audio/speech",
+        "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
+        "num_prompts": [
+          20
+        ],
+        "max_concurrency": [
+          1
+        ],
+        "seed_tts_locale": "en",
+        "extra_body": {
+          "voice": "default"
+        },
+        "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "median_audio_ttfp_ms": [
+            350
+          ],
+          "median_audio_rtf": [
+            0.15
+          ]
+        }
+      },
+      {
+        "task": "default_voice",
+        "eval_phase": "throughput",
+        "dataset_name": "seed-tts-text",
+        "backend": "openai-audio-speech",
+        "endpoint": "/v1/audio/speech",
+        "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
+        "num_prompts": [
+          80
+        ],
+        "max_concurrency": [
+          8
+        ],
+        "seed_tts_locale": "en",
+        "extra_body": {
+          "voice": "default"
+        },
+        "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "median_audio_ttfp_ms": [
+            3000
+          ],
+          "median_audio_rtf": [
+            0.5
+          ],
+          "audio_throughput": [
+            10
+          ]
+        }
+      },
+      {
+        "task": "default_voice",
+        "eval_phase": "stress",
+        "dataset_name": "seed-tts-text",
+        "backend": "openai-audio-speech",
+        "endpoint": "/v1/audio/speech",
+        "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
+        "num_prompts": [
+          100
+        ],
+        "request_rate": [
+          2.0
+        ],
+        "seed_tts_locale": "en",
+        "extra_body": {
+          "voice": "default"
+        },
+        "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "median_audio_ttfp_ms": [
+            500
+          ],
+          "median_audio_rtf": [
+            0.3
+          ],
+          "audio_throughput": [
+            5
+          ]
+        }
+      },
+      {
+        "task": "default_voice",
+        "eval_phase": "quality",
+        "enabled": false,
+        "dataset_name": "seed-tts-text",
+        "backend": "openai-audio-speech",
+        "endpoint": "/v1/audio/speech",
+        "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
+        "num_prompts": [
+          200
+        ],
+        "max_concurrency": [
+          4
+        ],
+        "seed_tts_locale": "en",
+        "extra_body": {
+          "voice": "default"
+        },
+        "seed_tts_wer_eval": true,
+        "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "mean_audio_rtf": [
+            0.35
+          ]
+        }
+      }
+    ]
+  }
+]
diff --git a/tests/dfx/reliability/helpers.py b/tests/dfx/reliability/helpers.py
index 9ebd0ce97..fd01e8a8e 100644
--- a/tests/dfx/reliability/helpers.py
+++ b/tests/dfx/reliability/helpers.py
@@ -4,10 +4,16 @@ This module keeps fault injection callable from tests directly:
 - GPU OOM (CUDA sidecar memory hog)
 - process kill by pattern and signal
 - post-ready hooks via ``fault_injector`` / ``omni_server_after_fault`` fixtures
+
+Worker PID classification (``worker_pids`` in snapshots) is still used for logging /
+markers; post-fault **process** cleanup assertions use the full captured
+``tree_pids`` (serve root + descendants at fault time) — see
+:func:`assert_no_server_tree_process_residual_and_gpu_release`.
 """
 
 from __future__ import annotations
 
+import concurrent.futures
 import http.client
 import json
 import logging
@@ -27,6 +33,16 @@ import pytest
 
 logger = logging.getLogger(__name__)
 
+# Substrings matched against ``psutil.Process.name`` + argv text for classifying GPU
+# workers under the test server tree. Extend per-suite via
+# ``server.reliability_worker_markers_extra`` or replace via
+# ``server.reliability_worker_markers`` (see ``_resolve_runtime_worker_markers``).
+DEFAULT_RUNTIME_WORKER_MARKERS: tuple[str, ...] = (
+    "multiprocessing.spawn",
+    # Covers ``VLLM::Worker``, ``VLLM::StageEngineCoreProc_*`` (Omni stage engines), etc.
+    "VLLM::",
+)
+
 
 @dataclass
 class OomHandle:
@@ -452,6 +468,61 @@ def _safe_proc_info(pid: int) -> tuple[str, str]:
         return "<unknown>", "<unavailable>"
 
 
+def _normalize_worker_marker_substrings(seq: Sequence[Any], *, context: str) -> tuple[str, ...]:
+    """Return deduped non-empty marker substrings; raises if any entry is blank."""
+    out: list[str] = []
+    for item in seq:
+        s = str(item).strip()
+        if not s:
+            raise ValueError(f"{context}: worker marker must be a non-empty string, got {item!r}")
+        out.append(s)
+    return tuple(dict.fromkeys(out))
+
+
+def _resolve_runtime_worker_markers(
+    server: Any,
+    injector_markers_extra: Sequence[str] | None,
+) -> tuple[str, ...]:
+    """Resolve which substrings classify worker PIDs for fault snapshots.
+
+    If ``server.reliability_worker_markers`` is set to a sequence, **only** those
+    markers are used (no defaults). Otherwise defaults from
+    :data:`DEFAULT_RUNTIME_WORKER_MARKERS` are merged with:
+
+    - ``server.reliability_worker_markers_extra`` when present, and
+    - ``injector_markers_extra`` from :func:`make_server_root_kill_fault_injector` /
+      :func:`make_server_tree_kill_fault_injector`.
+    """
+    replace = getattr(server, "reliability_worker_markers", None)
+    if replace is not None:
+        return _normalize_worker_marker_substrings(replace, context="server.reliability_worker_markers")
+
+    merged: list[str] = []
+    merged.extend(DEFAULT_RUNTIME_WORKER_MARKERS)
+    extra_attr = getattr(server, "reliability_worker_markers_extra", None)
+    if extra_attr is not None:
+        merged.extend(
+            _normalize_worker_marker_substrings(extra_attr, context="server.reliability_worker_markers_extra")
+        )
+    if injector_markers_extra is not None:
+        merged.extend(
+            _normalize_worker_marker_substrings(
+                injector_markers_extra,
+                context="injector worker_markers_extra",
+            )
+        )
+    return tuple(dict.fromkeys(merged))
+
+
+def _pid_looks_like_runtime_worker(pid: int, markers: Sequence[str]) -> bool:
+    """True when process name or argv contains any resolved worker marker substring."""
+    if not markers:
+        return False
+    name, cmdline = _safe_proc_info(pid)
+    text = f"{name} {cmdline}"
+    return any(marker in text for marker in markers)
+
+
 def _list_server_process_tree(server: Any) -> list[int]:
     """Return [root, descendants...] PIDs for the current test server instance."""
     root_proc = getattr(server, "proc", None)
@@ -468,6 +539,29 @@ def _list_server_process_tree(server: Any) -> list[int]:
     return [root_pid, *descendants]
 
 
+def _pids_in_server_tree_substring_match(server: Any, pattern: str, *, limit: int) -> list[int]:
+    """Return up to ``limit`` PIDs in the server tree whose name+cmdline contains ``pattern`` literally.
+
+    ``pgrep -f`` matches the process command line used by procps and does not consult the
+    short process name; vLLM-Omni often exposes titles like ``VLLM::Worker`` or
+    ``VLLM::StageEngineCoreProc_*`` via prctl/setproctitle in ``ps``/``name()`` only.
+    This helper keeps kill injection aligned with :func:`_safe_proc_info` / fault snapshots.
+    """
+    if limit <= 0:
+        return []
+    tree = _list_server_process_tree(server)
+    if not tree or not pattern:
+        return []
+    hits: list[int] = []
+    for pid in tree:
+        if len(hits) >= limit:
+            break
+        name, cmdline = _safe_proc_info(int(pid))
+        if pattern in f"{name} {cmdline}":
+            hits.append(int(pid))
+    return hits
+
+
 def _log_server_process_tree(server: Any) -> None:
     """Print server process tree for debugging fault injection targets."""
     pids = _list_server_process_tree(server)
@@ -486,6 +580,315 @@ FaultInjector = Callable[[Any], None]
 """Callable invoked with the live ``OmniServer`` after it is ready (see ``omni_server_after_fault``)."""
 
 
+def run_fault_injection_with_rate_load(
+    *,
+    submit_request: Callable[[], Any],
+    inject_fault: Callable[[], None],
+    num_requests: int = 10,
+    request_rate: float = 0.3,
+    submit_interval_sec: float | None = None,
+    completion_timeout_sec: float = 120.0,
+) -> dict[str, Any]:
+    """Submit requests with a fixed rate, then inject fault once an in-flight request is observed."""
+    if num_requests <= 0:
+        raise ValueError("num_requests must be > 0")
+    if request_rate <= 0:
+        raise ValueError("request_rate must be > 0")
+
+    interval_sec = submit_interval_sec if submit_interval_sec is not None else (1.0 / request_rate)
+    if interval_sec < 0:
+        raise ValueError("submit_interval_sec must be >= 0")
+
+    max_workers = min(max(2, num_requests), 8)
+    injected = False
+    futures: list[concurrent.futures.Future[Any]] = []
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+        for idx in range(num_requests):
+            futures.append(executor.submit(submit_request))
+            in_flight = any(not future.done() for future in futures)
+            if in_flight and not injected:
+                inject_fault()
+                injected = True
+            if idx != num_requests - 1 and interval_sec > 0:
+                time.sleep(interval_sec)
+
+        if not injected:
+            in_flight = any(not future.done() for future in futures)
+            if in_flight:
+                inject_fault()
+                injected = True
+
+        if not injected:
+            pytest.skip(
+                "no in-flight request observed before fault injection; "
+                "increase request cost or request rate for this environment"
+            )
+
+        done, pending = concurrent.futures.wait(
+            futures,
+            timeout=completion_timeout_sec,
+            return_when=concurrent.futures.ALL_COMPLETED,
+        )
+        assert not pending, (
+            "some in-flight-load requests did not converge after fault injection: "
+            f"pending={len(pending)} done={len(done)}"
+        )
+    failure_observed = False
+    completed = 0
+    success = 0
+    exceptions = 0
+    for future in futures:
+        completed += 1
+        try:
+            future.result()
+            success += 1
+        except Exception:  # noqa: BLE001
+            exceptions += 1
+            failure_observed = True
+    return {
+        "num_submitted": num_requests,
+        "completed": completed,
+        "success": success,
+        "exceptions": exceptions,
+        "failure_observed": failure_observed,
+        "inflight_observed": True,
+    }
+
+
+def list_alive_pids(pids: Sequence[int]) -> list[int]:
+    """Return PIDs from ``pids`` that still exist in the kernel and are **not** zombies.
+
+    After SIGTERM/SIGKILL the serve child may exit before the test harness calls
+    ``Popen.wait()``; until reaped, ``psutil.pid_exists`` stays true for the defunct
+    slot. Those PIDs must not count as "residual server processes" for leak checks.
+    """
+    out: list[int] = []
+    for pid in pids:
+        pid_i = int(pid)
+        if not psutil.pid_exists(pid_i):
+            continue
+        try:
+            if psutil.Process(pid_i).status() == psutil.STATUS_ZOMBIE:
+                continue
+        except psutil.Error:
+            continue
+        out.append(pid_i)
+    return out
+
+
+def query_gpu_compute_pid_used_memory_mb() -> dict[int, int] | None:
+    """Query NVIDIA compute-process memory map (pid -> used MB).
+
+    Returns ``None`` when ``nvidia-smi`` is unavailable/unreadable on current host.
+    """
+    out = subprocess.run(
+        ["nvidia-smi", "--query-compute-apps=pid,used_memory", "--format=csv,noheader,nounits"],
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+    if out.returncode != 0:
+        logger.warning("[reliability][gpu] nvidia-smi query failed: %s", out.stderr.strip())
+        return None
+
+    pid_to_mem_mb: dict[int, int] = {}
+    for line in out.stdout.splitlines():
+        row = line.strip()
+        if not row:
+            continue
+        pieces = [part.strip() for part in row.split(",", maxsplit=1)]
+        if len(pieces) != 2:
+            continue
+        pid_raw, mem_raw = pieces
+        if not pid_raw.isdigit():
+            continue
+        mem_digits = "".join(ch for ch in mem_raw if ch.isdigit())
+        if not mem_digits:
+            continue
+        pid_to_mem_mb[int(pid_raw)] = int(mem_digits)
+    return pid_to_mem_mb
+
+
+def worker_residual_timeout_after_kill_signal(signal_name: str) -> float:
+    """Wall-clock budget for fault-snapshot PIDs / GPU to clear after ``signal_name``."""
+    if signal_name == "SIGKILL":
+        return 30.0
+    if signal_name == "SIGINT":
+        return 120.0
+    if signal_name == "SIGTERM":
+        return 90.0
+    return 30.0
+
+
+def assert_no_server_tree_process_residual_and_gpu_release(
+    server: Any,
+    *,
+    scenario: str,
+    timeout_sec: float = 30.0,
+    poll_interval_sec: float = 0.5,
+) -> None:
+    """Assert no live PIDs remain from the fault-time server tree and no GPU use by those PIDs.
+
+    Uses ``reliability_fault_snapshot["tree_pids"]`` (root + all descendants captured at
+    injection time), not the marker-filtered ``worker_pids`` subset — so helpers like
+    ``multiprocessing.resource_tracker`` are included in the residual check.
+
+    **Zombie** PIDs (exited children not yet reaped by the test ``Popen``) are excluded
+    from the "alive" list; ``psutil.pid_exists`` alone would false-positive on them.
+
+    Note: if a child is reparented to PID 1 while staying alive, its PID is unchanged
+    and remains detected; if it **exits and a new unrelated process reuses the same
+    numeric PID**, this check may false-positive (rare on short windows).
+    """
+    snapshot = getattr(server, "reliability_fault_snapshot", None)
+    if not isinstance(snapshot, dict):
+        pytest.fail(f"[{scenario}] missing reliability fault snapshot on server")
+    tree_pids = [int(pid) for pid in snapshot.get("tree_pids", [])]
+    if not tree_pids:
+        pytest.skip(f"[{scenario}] no server process tree PIDs captured for this run")
+
+    tree_pid_set = set(tree_pids)
+    deadline = time.monotonic() + timeout_sec
+    last_alive: list[int] = []
+    last_gpu_leaks: dict[int, int] = {}
+    while time.monotonic() < deadline:
+        last_alive = list_alive_pids(tree_pids)
+        gpu_map = query_gpu_compute_pid_used_memory_mb()
+        if gpu_map is None:
+            pytest.skip(f"[{scenario}] nvidia-smi unavailable; skip GPU release assertion")
+        last_gpu_leaks = {pid: mem_mb for pid, mem_mb in gpu_map.items() if pid in tree_pid_set and mem_mb > 0}
+        if not last_alive and not last_gpu_leaks:
+            return
+        time.sleep(poll_interval_sec)
+
+    assert not last_alive, f"[{scenario}] residual server-tree processes remain alive: {last_alive}"
+    assert not last_gpu_leaks, f"[{scenario}] server-tree PID GPU memory not released: {last_gpu_leaks}"
+
+
+def assert_no_worker_residual_and_gpu_release(
+    server: Any,
+    *,
+    scenario: str,
+    timeout_sec: float = 30.0,
+    poll_interval_sec: float = 0.5,
+) -> None:
+    """Deprecated: use :func:`assert_no_server_tree_process_residual_and_gpu_release`."""
+    assert_no_server_tree_process_residual_and_gpu_release(
+        server,
+        scenario=scenario,
+        timeout_sec=timeout_sec,
+        poll_interval_sec=poll_interval_sec,
+    )
+
+
+def _capture_server_fault_snapshot(
+    server: Any,
+    *,
+    worker_markers_extra: Sequence[str] | None = None,
+) -> dict[str, list[int]]:
+    """Capture current server process snapshot for post-fault assertions.
+
+    ``tree_pids`` is ``[root, ...descendants]`` and is used by
+    :func:`assert_no_server_tree_process_residual_and_gpu_release`.
+    ``worker_pids`` is the marker-filtered subset for debugging only.
+    """
+    markers = _resolve_runtime_worker_markers(server, worker_markers_extra)
+    tree_pids = _list_server_process_tree(server)
+    root_pid = tree_pids[0] if tree_pids else None
+    worker_pids: list[int] = []
+    for pid in tree_pids:
+        if root_pid is not None and int(pid) == int(root_pid):
+            continue
+        if _pid_looks_like_runtime_worker(pid, markers):
+            worker_pids.append(pid)
+    snapshot = {
+        "root_pids": [root_pid] if root_pid is not None else [],
+        "tree_pids": tree_pids,
+        "worker_pids": worker_pids,
+        "worker_markers": list(markers),
+    }
+    setattr(server, "reliability_fault_snapshot", snapshot)
+    return snapshot
+
+
+def make_server_root_kill_fault_injector(
+    *,
+    signal_name: str = "SIGKILL",
+    post_kill_wait_seconds: float = 0.0,
+    worker_markers_extra: Sequence[str] | None = None,
+) -> FaultInjector:
+    """Kill only the current server root process (simulate first strike on PID1-like process).
+
+    ``worker_markers_extra`` is merged into worker PID detection for the fault snapshot
+    (after :data:`DEFAULT_RUNTIME_WORKER_MARKERS` unless
+    ``server.reliability_worker_markers`` replaces the set entirely).
+    """
+
+    def _inject(server: Any) -> None:
+        _log_server_process_tree(server)
+        snapshot = _capture_server_fault_snapshot(server, worker_markers_extra=worker_markers_extra)
+        root_pids = snapshot["root_pids"]
+        if not root_pids:
+            pytest.skip("no server root process found for root-kill injection")
+        root_pid = int(root_pids[0])
+        sig = getattr(signal, signal_name, None)
+        if sig is None:
+            raise ValueError(f"Unsupported signal_name: {signal_name}")
+        name, cmdline = _safe_proc_info(root_pid)
+        print(
+            f"[reliability][process-kill] root-kill pid={root_pid} name={name} signal={signal_name} cmdline={cmdline}",
+            flush=True,
+        )
+        os.kill(root_pid, sig)
+        if post_kill_wait_seconds > 0:
+            time.sleep(post_kill_wait_seconds)
+
+    return _inject
+
+
+def make_server_tree_kill_fault_injector(
+    *,
+    signal_name: str = "SIGKILL",
+    post_kill_wait_seconds: float = 0.0,
+    inter_kill_wait_seconds: float = 0.2,
+    worker_markers_extra: Sequence[str] | None = None,
+) -> FaultInjector:
+    """Kill current server root first, then kill all remaining descendants.
+
+    See :func:`make_server_root_kill_fault_injector` for ``worker_markers_extra``.
+    """
+
+    def _inject(server: Any) -> None:
+        _log_server_process_tree(server)
+        snapshot = _capture_server_fault_snapshot(server, worker_markers_extra=worker_markers_extra)
+        tree_pids = [int(pid) for pid in snapshot["tree_pids"]]
+        if not tree_pids:
+            pytest.skip("no server process tree found for tree-kill injection")
+
+        sig = getattr(signal, signal_name, None)
+        if sig is None:
+            raise ValueError(f"Unsupported signal_name: {signal_name}")
+
+        root_pid = tree_pids[0]
+        kill_order = [root_pid, *[pid for pid in tree_pids if pid != root_pid]]
+        for pid in kill_order:
+            if not psutil.pid_exists(pid):
+                continue
+            name, cmdline = _safe_proc_info(pid)
+            print(
+                f"[reliability][process-kill] tree-kill pid={pid} name={name} signal={signal_name} cmdline={cmdline}",
+                flush=True,
+            )
+            os.kill(pid, sig)
+            if inter_kill_wait_seconds > 0:
+                time.sleep(inter_kill_wait_seconds)
+
+        if post_kill_wait_seconds > 0:
+            time.sleep(post_kill_wait_seconds)
+
+    return _inject
+
+
 def make_process_kill_fault_injector(
     *,
     grep_patterns: str | Sequence[str],
@@ -493,11 +896,18 @@ def make_process_kill_fault_injector(
     limit: int = 1,
     post_kill_wait_seconds: float = 0.0,
 ) -> FaultInjector:
-    """Build a post-ready injector that kills processes matched by ``pgrep -f``.
+    """Build a post-ready injector that kills processes matched by ordered ``grep_patterns``.
 
-    Tries each pattern in order until at least one PID is killed. If none match,
-    the returned callable issues ``pytest.skip`` (same behavior as the previous
-    inline reliability test).
+    Matching is tried in two phases per pattern (first hit wins):
+
+    1. **Server process tree** (``psutil``): literal substring match against
+       ``Process.name()`` plus argv text. This matches how :func:`_safe_proc_info`
+       logs processes and catches vLLM-style titles (``VLLM::...``) that often do not
+       appear in ``pgrep -f``\'s command-line view.
+    2. **``pgrep -f``** (legacy): scoped to the server PID tree when it is known;
+       uses procps regular-expression rules for the pattern.
+
+    If neither phase finds a target, the returned callable issues ``pytest.skip``.
 
     Args:
         grep_patterns: One pattern or an ordered list of patterns.
@@ -514,9 +924,45 @@ def make_process_kill_fault_injector(
             logger.warning(
                 "[reliability][process-kill] no server process tree found; fallback to global pgrep matching"
             )
+
+        def _kill_and_wait(filtered: list[int], pattern: str, *, source: str) -> None:
+            sig = getattr(signal, signal_name, None)
+            if sig is None:
+                raise ValueError(f"Unsupported signal_name: {signal_name}")
+            for pid in filtered:
+                name, cmdline = _safe_proc_info(pid)
+                print(
+                    f"[reliability][process-kill] killing pid={pid} name={name} signal={signal_name} "
+                    f"source={source} cmdline={cmdline}",
+                    flush=True,
+                )
+                os.kill(pid, sig)
+            print(
+                f"[reliability][process-kill] matched pattern={pattern!r} source={source} "
+                f"killed_pids={filtered} killed_count={len(filtered)}",
+                flush=True,
+            )
+            if post_kill_wait_seconds > 0:
+                print(
+                    f"[reliability][process-kill] waiting {post_kill_wait_seconds:.2f}s after kill",
+                    flush=True,
+                )
+                time.sleep(post_kill_wait_seconds)
+
         for pattern in patterns:
             print(
-                f"[reliability][process-kill] trying pattern={pattern} signal={signal_name} limit={limit}",
+                f"[reliability][process-kill] trying server_tree substring pattern={pattern!r} "
+                f"signal={signal_name} limit={limit}",
+                flush=True,
+            )
+            tree_hits = _pids_in_server_tree_substring_match(server, pattern, limit=limit)
+            if tree_hits:
+                _kill_and_wait(tree_hits, pattern, source="server_tree")
+                return
+
+        for pattern in patterns:
+            print(
+                f"[reliability][process-kill] trying pgrep -f pattern={pattern!r} signal={signal_name} limit={limit}",
                 flush=True,
             )
             pids = inject_process_kill(
@@ -535,26 +981,7 @@ def make_process_kill_fault_injector(
                 )
                 continue
             if filtered:
-                sig = getattr(signal, signal_name, None)
-                if sig is None:
-                    raise ValueError(f"Unsupported signal_name: {signal_name}")
-                for pid in filtered:
-                    name, cmdline = _safe_proc_info(pid)
-                    print(
-                        f"[reliability][process-kill] killing pid={pid} name={name} signal={signal_name} cmdline={cmdline}",
-                        flush=True,
-                    )
-                    os.kill(pid, sig)
-                print(
-                    f"[reliability][process-kill] matched pattern={pattern} killed_pids={filtered} killed_count={len(filtered)}",
-                    flush=True,
-                )
-                if post_kill_wait_seconds > 0:
-                    print(
-                        f"[reliability][process-kill] waiting {post_kill_wait_seconds:.2f}s after kill",
-                        flush=True,
-                    )
-                    time.sleep(post_kill_wait_seconds)
+                _kill_and_wait(filtered, pattern, source="pgrep")
                 return
         logger.warning(
             "[reliability][process-kill] no process matched patterns=%s signal=%s limit=%s",
diff --git a/tests/dfx/reliability/test_reliability_qwen3_omni.py b/tests/dfx/reliability/test_reliability_qwen3_omni.py
index 37f63b942..49028a03f 100644
--- a/tests/dfx/reliability/test_reliability_qwen3_omni.py
+++ b/tests/dfx/reliability/test_reliability_qwen3_omni.py
@@ -21,13 +21,19 @@ from tests.dfx.conftest import (
     resolve_oom_device_spec,
 )
 from tests.dfx.reliability.helpers import (
+    FaultInjector,
+    assert_no_server_tree_process_residual_and_gpu_release,
     extract_openai_error_contract_from_bytes,
     get_health_raw,
     inject_gpu_oom,
     make_process_kill_fault_injector,
+    make_server_root_kill_fault_injector,
+    make_server_tree_kill_fault_injector,
     post_chat_completions_raw,
     post_json_raw,
+    run_fault_injection_with_rate_load,
     stop_gpu_oom_hogs,
+    worker_residual_timeout_after_kill_signal,
 )
 from tests.helpers.mark import hardware_test
 from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
@@ -97,7 +103,73 @@ FAULT_ERROR_KEYWORDS = (
     "500",
     "503",
 )
-RUNTIME_WORKER_PATTERN = "VLLM::Worker"
+PROCESS_KILL_ERROR_KEYWORDS = (
+    "timeout",
+    "did not complete within",
+    "connection",
+    "engine",
+    "orchestrator",
+    "dead",
+    "internal",
+    "500",
+    "503",
+)
+# Prefix match for vLLM-titled children (``VLLM::Worker``, ``VLLM::StageEngineCoreProc_*``, …).
+# Must be a literal substring of ``name``/``cmdline`` — no trailing whitespace (would break match).
+RUNTIME_WORKER_PATTERN = "VLLM::"
+# RFC#2366 signal x target matrix:
+# - worker: SIGTERM, SIGKILL
+# - serve-root no-load: SIGINT; SIGTERM/SIGKILL skipped (vllm issue#43060)
+# - serve-root with-load: SIGINT; SIGTERM skipped (issue#3683); SIGKILL skipped (issue#43060)
+# - serve-tree no-load: SIGTERM/SIGKILL
+# - serve-tree with-load: SIGKILL; SIGTERM skipped (issue#3683)
+_SERVE_ROOT_NON_SIGINT_SKIP = pytest.mark.skip(reason="vllm issue#43060")
+_WITH_LOAD_SIGTERM_SKIP = pytest.mark.skip(reason="issue#3683")
+SERVE_SIGNAL_PARAMS = [
+    pytest.param("SIGTERM", id="sigterm"),
+    pytest.param("SIGINT", id="sigint"),
+    pytest.param("SIGKILL", id="sigkill", marks=_SERVE_ROOT_NON_SIGINT_SKIP),
+]
+SERVE_WITH_LOAD_SIGNAL_PARAMS = [
+    pytest.param("SIGTERM", id="sigterm", marks=_WITH_LOAD_SIGTERM_SKIP),
+    pytest.param("SIGINT", id="sigint", marks=_WITH_LOAD_SIGTERM_SKIP),
+    pytest.param("SIGKILL", id="sigkill", marks=_WITH_LOAD_SIGTERM_SKIP),
+]
+TREE_SIGNAL_PARAMS = [
+    pytest.param("SIGTERM", id="sigterm"),
+    pytest.param("SIGKILL", id="sigkill"),
+]
+TREE_WITH_LOAD_SIGNAL_PARAMS = [
+    pytest.param("SIGTERM", id="sigterm", marks=_WITH_LOAD_SIGTERM_SKIP),
+    pytest.param("SIGKILL", id="sigkill"),
+]
+WORKER_SIGNAL_FAULT_PARAMS = [
+    pytest.param(
+        make_process_kill_fault_injector(
+            grep_patterns=RUNTIME_WORKER_PATTERN,
+            signal_name="SIGTERM",
+            limit=1,
+            post_kill_wait_seconds=2.0,
+        ),
+        id="runtime_process_chain_sigterm",
+    ),
+    pytest.param(
+        make_process_kill_fault_injector(
+            grep_patterns=RUNTIME_WORKER_PATTERN,
+            signal_name="SIGKILL",
+            limit=1,
+            post_kill_wait_seconds=2.0,
+        ),
+        id="runtime_process_chain_sigkill",
+    ),
+]
+INFLIGHT_INJECTION_REQUEST_RATE = 0.3
+INFLIGHT_INJECTION_REQUEST_COUNT = 10
+
+# Post-fault chat probe: ``MAX_ELAPSED`` aligns with HTTP read timeout; assert allows +1s jitter
+# when the client raises at timeout (e.g. ~30.03s).
+_PROCESS_KILL_WORKER_CHAT_FF_HTTP_TIMEOUT_SEC = 30
+_PROCESS_KILL_WORKER_CHAT_FF_MAX_ELAPSED_SEC = 30
 
 
 class _HasServeArgs(Protocol):
@@ -194,6 +266,55 @@ def _looks_like_server_unreachable(exc: BaseException) -> bool:
     return "connection refused" in msg or "actively refused" in msg
 
 
+def _chat_request_config(omni_server: _HasServeArgs) -> dict[str, Any]:
+    messages = dummy_messages_from_mix_data(
+        system_prompt=_get_system_prompt(),
+        content_text="What is the capital of China? Answer in one short sentence.",
+    )
+    return {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": False,
+        "modalities": ["text"],
+    }
+
+
+def _assert_post_fault_chat_fast_fail(host: str, port: int, *, model: str, scenario: str) -> None:
+    payload = {
+        "model": model,
+        "messages": [{"role": "user", "content": "post-fault fast-fail check"}],
+        "stream": False,
+        "modalities": ["text"],
+    }
+    start = time.monotonic()
+    try:
+        status, body = post_json_raw(host, port, "/v1/chat/completions", payload, timeout_sec=20)
+        elapsed = time.monotonic() - start
+        assert elapsed < 15, f"[{scenario} fast_fail] /v1/chat/completions did not fail fast: {elapsed:.2f}s"
+        assert status >= 500, (
+            f"[{scenario} fast_fail] expected server-side error after fault, got status={status}, body={body[:200]!r}"
+        )
+    except Exception:  # noqa: BLE001
+        elapsed = time.monotonic() - start
+        assert elapsed < 15, f"[{scenario} fast_fail] exception was too slow after fault: {elapsed:.2f}s"
+
+
+def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -> None:
+    deadline = time.monotonic() + 20.0
+    last_observation = ""
+    while time.monotonic() < deadline:
+        try:
+            status, body = get_health_raw(host, port, timeout_sec=5)
+            last_observation = f"http={status}, body={body[:200]!r}"
+            if status == 503:
+                return
+        except Exception as exc:  # noqa: BLE001
+            last_observation = f"exception={exc!r}"
+            return
+        time.sleep(0.5)
+    pytest.fail(f"[{scenario} health] no terminal post-fault health observed: {last_observation}")
+
+
 QWEN_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLOY_CONFIGS_DIR)
 
 
@@ -354,17 +475,7 @@ def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_funct
 @pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
 @pytest.mark.parametrize(
     "fault_injector",
-    [
-        pytest.param(
-            make_process_kill_fault_injector(
-                grep_patterns="VLLM::Worker",
-                signal_name="SIGKILL",
-                limit=1,
-                post_kill_wait_seconds=2.0,
-            ),
-            id="runtime_process_chain",
-        ),
-    ],
+    WORKER_SIGNAL_FAULT_PARAMS,
     indirect=True,
 )
 @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
@@ -385,7 +496,7 @@ def test_reliability_fault_process_kill_request_failure(
     try:
         openai_client_function.send_omni_request(request_config, request_num=1)
     except Exception as exc:
-        assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
+        assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
     else:
         pytest.fail("expected request failure after process-kill injection")
 
@@ -394,48 +505,22 @@ def test_reliability_fault_process_kill_request_failure(
 @pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
 @pytest.mark.parametrize(
     "fault_injector",
-    [
-        pytest.param(
-            make_process_kill_fault_injector(
-                grep_patterns="VLLM::Worker",
-                signal_name="SIGKILL",
-                limit=1,
-                post_kill_wait_seconds=2.0,
-            ),
-            id="runtime_process_chain",
-        ),
-    ],
+    WORKER_SIGNAL_FAULT_PARAMS,
     indirect=True,
 )
 @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
 def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
     omni_server_after_fault_function,
 ) -> None:
-    """Black-box: after worker SIGKILL, /health→503, chat fails fast, concurrent chat does not hang."""
+    """Black-box: after worker process kill, chat fails fast, concurrent chat does not hang; /health→503 last.
+
+    ``SIGTERM`` on a stage engine often exceeds a 15s client-visible path; ``SIGKILL`` is
+    typically quicker. See module constants ``_PROCESS_KILL_WORKER_CHAT_FF_*``.
+    """
     host = omni_server_after_fault_function.host
     port = omni_server_after_fault_function.port
     model = omni_server_after_fault_function.model
 
-    deadline = time.monotonic() + 20.0
-    last_observation = ""
-    saw_503 = False
-    health_final_status: int | None = None
-    health_final_body = b""
-    while time.monotonic() < deadline:
-        try:
-            status, body = get_health_raw(host, port, timeout_sec=5)
-            last_observation = f"http={status}, body={body[:200]!r}"
-            health_final_status, health_final_body = status, body
-            if status == 503:
-                saw_503 = True
-                break
-        except Exception as exc:  # noqa: BLE001
-            last_observation = f"exception={exc!r}"
-        time.sleep(0.5)
-    assert saw_503, (
-        f"[process_kill health] expected /health to become 503 after fault injection, got {last_observation}"
-    )
-
     payload = {
         "model": model,
         "messages": [{"role": "user", "content": "Say hello in one short sentence."}],
@@ -444,20 +529,28 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
     }
     ff_status: int | None = None
     ff_body = b""
-    ff_exc: BaseException | None = None
     start = time.monotonic()
     try:
-        ff_status, ff_body = post_json_raw(host, port, "/v1/chat/completions", payload, timeout_sec=20)
+        ff_status, ff_body = post_json_raw(
+            host,
+            port,
+            "/v1/chat/completions",
+            payload,
+            timeout_sec=_PROCESS_KILL_WORKER_CHAT_FF_HTTP_TIMEOUT_SEC,
+        )
         elapsed = time.monotonic() - start
-        assert elapsed < 15, f"[process_kill fast_fail] request did not fail fast after fault: {elapsed:.2f}s"
+        assert elapsed <= _PROCESS_KILL_WORKER_CHAT_FF_MAX_ELAPSED_SEC + 1.0, (
+            f"[process_kill fast_fail] request did not fail fast after fault: {elapsed:.2f}s"
+        )
         assert ff_status >= 500, (
             f"[process_kill fast_fail] expected server-side failure after fault, "
             f"got status={ff_status}, body={ff_body[:200]!r}"
         )
-    except Exception as exc:
-        ff_exc = exc
+    except Exception:
         elapsed = time.monotonic() - start
-        assert elapsed < 15, f"[process_kill fast_fail] request exception was too slow after fault: {elapsed:.2f}s"
+        assert elapsed <= _PROCESS_KILL_WORKER_CHAT_FF_MAX_ELAPSED_SEC + 1.0, (
+            f"[process_kill fast_fail] request exception was too slow after fault: {elapsed:.2f}s"
+        )
 
     payload_json = json.dumps(
         {
@@ -490,32 +583,195 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
     assert elapsed < 30, f"[process_kill concurrent] fault-time request convergence is too slow: {elapsed:.2f}s"
 
     fault_observed = False
-    conc_debug: list[Any] = []
     for future in done:
         try:
             status, body = future.result()
-            conc_debug.append((status, body[:200]))
             if status >= 500:
                 fault_observed = True
         except Exception as exc:
-            conc_debug.append(repr(exc))
+            assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
             fault_observed = True
-    # DEBUG: remove before merge
-    print(
-        health_final_status,
-        health_final_body[:200],
-        ff_status,
-        ff_body[:200],
-        ff_exc,
-        conc_debug,
-    )
     assert fault_observed, (
         "[process_kill concurrent] expected at least one request to fail after process-kill fault injection"
     )
 
+    deadline = time.monotonic() + 20.0
+    last_observation = ""
+    final_health_status: int | None = None
+    while time.monotonic() < deadline:
+        try:
+            status, body = get_health_raw(host, port, timeout_sec=5)
+            last_observation = f"http={status}, body={body[:200]!r}"
+            final_health_status = status
+            if status == 503:
+                break
+        except Exception as exc:  # noqa: BLE001
+            last_observation = f"exception={exc!r}"
+        time.sleep(0.5)
+    if final_health_status != 503:
+        pytest.skip("issue#3050")
+    assert final_health_status == 503, (
+        "[process_kill health] expected /health 503 after fatal fault, "
+        f"got status={final_health_status}, last_observation={last_observation}"
+    )
+
 
 @pytest.mark.slow
-@pytest.mark.skip(reason="issue#2327")
+@pytest.mark.skip(reason="issue#3683")
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize(
+    "fault_injector",
+    WORKER_SIGNAL_FAULT_PARAMS,
+    indirect=True,
+)
+@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_worker_with_load_request_failure(
+    omni_server_function,
+    openai_client_function,
+    fault_injector: FaultInjector,
+) -> None:
+    request_config = _chat_request_config(omni_server_function)
+    scenario = "kill_worker_with_load"
+    load_result = run_fault_injection_with_rate_load(
+        submit_request=lambda: openai_client_function.send_omni_request(request_config, request_num=1),
+        inject_fault=lambda: fault_injector(omni_server_function),
+        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
+        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
+        completion_timeout_sec=120.0,
+    )
+    assert load_result["failure_observed"], (
+        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
+    )
+    host = omni_server_function.host
+    port = omni_server_function.port
+    _assert_post_fault_chat_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
+    _assert_post_fault_health_terminal(host, port, scenario=scenario)
+
+
+@pytest.mark.slow
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize("signal_name", SERVE_SIGNAL_PARAMS)
+@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_serve_root_no_load_fast_fail_and_cleanup(
+    omni_server_function,
+    signal_name: str,
+) -> None:
+    """Black-box: kill serve root without load; verify fast-fail/health/cleanup."""
+    scenario = f"kill_serve_root_no_load_{signal_name.lower()}"
+    injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
+    injector(omni_server_function)
+    host = omni_server_function.host
+    port = omni_server_function.port
+    _assert_post_fault_chat_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
+    _assert_post_fault_health_terminal(host, port, scenario=scenario)
+    assert_no_server_tree_process_residual_and_gpu_release(
+        omni_server_function,
+        scenario=scenario,
+        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
+    )
+
+
+@pytest.mark.slow
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize("signal_name", SERVE_WITH_LOAD_SIGNAL_PARAMS)
+@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup(
+    omni_server_function,
+    openai_client_function,
+    signal_name: str,
+) -> None:
+    request_config = _chat_request_config(omni_server_function)
+    scenario = f"kill_serve_root_with_load_{signal_name.lower()}"
+    injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
+    load_result = run_fault_injection_with_rate_load(
+        submit_request=lambda: openai_client_function.send_omni_request(request_config, request_num=1),
+        inject_fault=lambda: injector(omni_server_function),
+        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
+        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
+        completion_timeout_sec=300.0,
+    )
+    assert load_result["failure_observed"], (
+        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
+    )
+    host = omni_server_function.host
+    port = omni_server_function.port
+    _assert_post_fault_chat_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
+    _assert_post_fault_health_terminal(host, port, scenario=scenario)
+    assert_no_server_tree_process_residual_and_gpu_release(
+        omni_server_function,
+        scenario=scenario,
+        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
+    )
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize("signal_name", TREE_SIGNAL_PARAMS)
+@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup(
+    omni_server_function,
+    signal_name: str,
+) -> None:
+    """Black-box: kill server tree without load; verify fast-fail/health/cleanup."""
+    scenario = f"kill_serve_tree_no_load_{signal_name.lower()}"
+    injector = make_server_tree_kill_fault_injector(
+        signal_name=signal_name,
+        post_kill_wait_seconds=2.0,
+        inter_kill_wait_seconds=0.1,
+    )
+    injector(omni_server_function)
+    host = omni_server_function.host
+    port = omni_server_function.port
+    _assert_post_fault_chat_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
+    _assert_post_fault_health_terminal(host, port, scenario=scenario)
+    assert_no_server_tree_process_residual_and_gpu_release(
+        omni_server_function,
+        scenario=scenario,
+        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
+    )
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize("signal_name", TREE_WITH_LOAD_SIGNAL_PARAMS)
+@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
+    omni_server_function,
+    openai_client_function,
+    signal_name: str,
+) -> None:
+    request_config = _chat_request_config(omni_server_function)
+    scenario = f"kill_serve_tree_with_load_{signal_name.lower()}"
+    injector = make_server_tree_kill_fault_injector(
+        signal_name=signal_name,
+        post_kill_wait_seconds=2.0,
+        inter_kill_wait_seconds=0.1,
+    )
+    load_result = run_fault_injection_with_rate_load(
+        submit_request=lambda: openai_client_function.send_omni_request(request_config, request_num=1),
+        inject_fault=lambda: injector(omni_server_function),
+        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
+        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
+        completion_timeout_sec=120.0,
+    )
+    assert load_result["failure_observed"], (
+        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
+    )
+    host = omni_server_function.host
+    port = omni_server_function.port
+    _assert_post_fault_chat_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
+    _assert_post_fault_health_terminal(host, port, scenario=scenario)
+    assert_no_server_tree_process_residual_and_gpu_release(
+        omni_server_function,
+        scenario=scenario,
+        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
+    )
+
+
+@pytest.mark.slow
+@pytest.mark.skip(reason="OOM injection causes the serving instance to exit")
 @hardware_test(res={"cuda": "H100"}, num_cards=2)
 @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
 def test_reliability_fault_gpu_oom_state_converges_after_fault_removed(
diff --git a/tests/dfx/reliability/test_reliability_wan22.py b/tests/dfx/reliability/test_reliability_wan22.py
index e27db616f..84175429d 100644
--- a/tests/dfx/reliability/test_reliability_wan22.py
+++ b/tests/dfx/reliability/test_reliability_wan22.py
@@ -20,28 +20,25 @@ from tests.dfx.conftest import (
     supports_video_generation,
 )
 from tests.dfx.reliability.helpers import (
+    FaultInjector,
+    assert_no_server_tree_process_residual_and_gpu_release,
     get_health_raw,
     inject_gpu_oom,
     make_process_kill_fault_injector,
+    make_server_root_kill_fault_injector,
+    make_server_tree_kill_fault_injector,
+    run_fault_injection_with_rate_load,
     stop_gpu_oom_hogs,
+    worker_residual_timeout_after_kill_signal,
 )
 from tests.helpers.mark import hardware_test
 from tests.helpers.media import generate_synthetic_image
 
 RELIABILITY_SCENARIOS: list[dict[str, Any]] = [
     {
-        "test_name": "wan22_t2v_reliability_default",
+        "test_name": "wan22_i2v_reliability_default",
         "server_params": {
-            "model": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-            "server_args": [
-                "--num-gpus",
-                "1",
-                "--boundary-ratio",
-                "0.875",
-                "--flow-shift",
-                "5.0",
-                "--disable-log-stats",
-            ],
+            "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
         },
     }
 ]
@@ -73,9 +70,114 @@ PROCESS_KILL_ERROR_KEYWORDS = (
     "500",
     "503",
 )
+# RFC#2366 signal x target matrix (aligned with test_reliability_qwen3_omni.py):
+# - worker: SIGTERM, SIGKILL
+# - serve-root: SIGTERM, SIGINT; SIGKILL skipped (issue#3725)
+# - serve-tree no-load: SIGTERM, SIGKILL
+# - serve-tree with-load: SIGTERM, SIGKILL
+_SERVE_ROOT_SIGKILL_SKIP = pytest.mark.skip(reason="issue#3725")
+SERVE_SIGNAL_PARAMS = [
+    pytest.param("SIGTERM", id="sigterm"),
+    pytest.param("SIGINT", id="sigint"),
+    pytest.param("SIGKILL", id="sigkill", marks=_SERVE_ROOT_SIGKILL_SKIP),
+]
+TREE_SIGNAL_PARAMS = [
+    pytest.param("SIGTERM", id="sigterm"),
+    pytest.param("SIGKILL", id="sigkill"),
+]
+TREE_WITH_LOAD_SIGNAL_PARAMS = [
+    pytest.param("SIGTERM", id="sigterm"),
+    pytest.param("SIGKILL", id="sigkill"),
+]
+WORKER_SIGNAL_FAULT_PARAMS = [
+    pytest.param(
+        make_process_kill_fault_injector(
+            grep_patterns="vLLM-Omni::",
+            signal_name="SIGTERM",
+            limit=1,
+            post_kill_wait_seconds=2.0,
+        ),
+        id="runtime_process_chain_sigterm",
+    ),
+    pytest.param(
+        make_process_kill_fault_injector(
+            grep_patterns="vLLM-Omni::",
+            signal_name="SIGKILL",
+            limit=1,
+            post_kill_wait_seconds=2.0,
+        ),
+        id="runtime_process_chain_sigkill",
+    ),
+]
 
 WAN_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, E2E_STAGE_CONFIGS_DIR)
 DIFFUSION_VIDEO_PARAMS = [param for param in WAN_PARAMS if supports_video_generation(param.model)]
+INFLIGHT_INJECTION_REQUEST_RATE = 0.3
+INFLIGHT_INJECTION_REQUEST_COUNT = 10
+
+
+def _video_request_config() -> dict[str, Any]:
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
+    return {
+        "form_data": {
+            "prompt": "Generate a realistic road-driving video with camera motion.",
+            "width": 512,
+            "height": 512,
+            "fps": 8,
+            "num_frames": 8,
+            "guidance_scale": 1.0,
+            "flow_shift": 5.0,
+            "num_inference_steps": 4,
+            "seed": 42,
+        },
+        "image_reference": image_data_url,
+        "stream": False,
+    }
+
+
+def _assert_post_fault_video_fast_fail(host: str, port: int, *, scenario: str) -> None:
+    url = f"http://{host}:{port}/v1/videos"
+    payload = {
+        "prompt": "post-fault fast-fail check",
+        "width": "512",
+        "height": "512",
+        "fps": "8",
+        "num_frames": "8",
+        "num_inference_steps": "4",
+    }
+    start = time.monotonic()
+    try:
+        response = requests.post(
+            url,
+            data=payload,
+            headers={"Accept": "application/json"},
+            timeout=20,
+        )
+        elapsed = time.monotonic() - start
+        assert elapsed < 15, f"[{scenario} fast_fail] /v1/videos did not fail fast: {elapsed:.2f}s"
+        assert response.status_code >= 500, (
+            f"[{scenario} fast_fail] expected server-side error after fault, "
+            f"got status={response.status_code}, body={response.text[:200]!r}"
+        )
+    except Exception:  # noqa: BLE001
+        elapsed = time.monotonic() - start
+        assert elapsed < 15, f"[{scenario} fast_fail] exception was too slow after fault: {elapsed:.2f}s"
+
+
+def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -> None:
+    deadline = time.monotonic() + 20.0
+    last_observation = ""
+    while time.monotonic() < deadline:
+        try:
+            status, body = get_health_raw(host, port, timeout_sec=5)
+            last_observation = f"http={status}, body={body[:200]!r}"
+            if status == 503:
+                return
+        except Exception as exc:  # noqa: BLE001
+            last_observation = f"exception={exc!r}"
+            return
+        time.sleep(0.5)
+    pytest.fail(f"[{scenario} health] no terminal post-fault health observed: {last_observation}")
 
 
 @pytest.mark.slow
@@ -122,40 +224,15 @@ def test_reliability_fault_gpu_oom_video_large_request_failure(omni_server_funct
 @pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
 @pytest.mark.parametrize(
     "fault_injector",
-    [
-        pytest.param(
-            make_process_kill_fault_injector(
-                grep_patterns="multiprocessing.spawn",
-                signal_name="SIGKILL",
-                limit=1,
-                post_kill_wait_seconds=2.0,
-            ),
-            id="runtime_process_chain",
-        ),
-    ],
+    WORKER_SIGNAL_FAULT_PARAMS,
     indirect=True,
 )
 @pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_video_request_failure(
+def test_reliability_fault_process_kill_request_failure(
     omni_server_after_fault_function,
     openai_client_function,
 ) -> None:
-    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
-    request_config = {
-        "form_data": {
-            "prompt": "Generate a realistic road-driving video with camera motion.",
-            "width": 512,
-            "height": 512,
-            "fps": 8,
-            "num_frames": 8,
-            "guidance_scale": 1.0,
-            "flow_shift": 5.0,
-            "num_inference_steps": 4,
-            "seed": 42,
-        },
-        "image_reference": image_data_url,
-        "stream": False,
-    }
+    request_config = _video_request_config()
     try:
         openai_client_function.send_video_diffusion_request(request_config, request_num=1)
     except Exception as exc:
@@ -168,21 +245,11 @@ def test_reliability_fault_process_kill_video_request_failure(
 @pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
 @pytest.mark.parametrize(
     "fault_injector",
-    [
-        pytest.param(
-            make_process_kill_fault_injector(
-                grep_patterns="multiprocessing.spawn",
-                signal_name="SIGKILL",
-                limit=1,
-                post_kill_wait_seconds=2.0,
-            ),
-            id="runtime_process_chain",
-        ),
-    ],
+    WORKER_SIGNAL_FAULT_PARAMS,
     indirect=True,
 )
 @pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_video_health_fast_fail_and_concurrent(
+def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
     omni_server_after_fault_function,
     openai_client_function,
 ) -> None:
@@ -290,6 +357,37 @@ def test_reliability_fault_process_kill_video_health_fast_fail_and_concurrent(
     )
 
 
+@pytest.mark.slow
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize(
+    "fault_injector",
+    WORKER_SIGNAL_FAULT_PARAMS,
+    indirect=True,
+)
+@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_worker_with_load_request_failure(
+    omni_server_function,
+    openai_client_function,
+    fault_injector: FaultInjector,
+) -> None:
+    request_config = _video_request_config()
+    scenario = "kill_worker_with_load"
+    load_result = run_fault_injection_with_rate_load(
+        submit_request=lambda: openai_client_function.send_video_diffusion_request(request_config, request_num=1),
+        inject_fault=lambda: fault_injector(omni_server_function),
+        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
+        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
+        completion_timeout_sec=120.0,
+    )
+    assert load_result["failure_observed"], (
+        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
+    )
+    host = omni_server_function.host
+    port = omni_server_function.port
+    _assert_post_fault_video_fast_fail(host, port, scenario=scenario)
+    _assert_post_fault_health_terminal(host, port, scenario=scenario)
+
+
 @pytest.mark.slow
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 @pytest.mark.skip(reason="issue#2327")
@@ -409,3 +507,127 @@ def test_reliability_video_oom_recovers_after_fault_removed(
             "unhealthy terminal state should fail fast on /v1/videos, "
             f"got health={terminal_health}, request_status={recovery_status}, body={recovery_text_prefix!r}"
         )
+
+
+@pytest.mark.slow
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize("signal_name", SERVE_SIGNAL_PARAMS)
+@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup(
+    omni_server_function,
+    openai_client_function,
+    signal_name: str,
+) -> None:
+    """Black-box: during in-flight video requests, kill serve root and verify post-fault behavior/cleanup."""
+    request_config = _video_request_config()
+    scenario = f"kill_serve_root_with_load_{signal_name.lower()}"
+    injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
+    load_result = run_fault_injection_with_rate_load(
+        submit_request=lambda: openai_client_function.send_video_diffusion_request(request_config, request_num=1),
+        inject_fault=lambda: injector(omni_server_function),
+        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
+        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
+        completion_timeout_sec=120.0,
+    )
+    assert load_result["failure_observed"], (
+        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
+    )
+    host = omni_server_function.host
+    port = omni_server_function.port
+    _assert_post_fault_video_fast_fail(host, port, scenario=scenario)
+    _assert_post_fault_health_terminal(host, port, scenario=scenario)
+    assert_no_server_tree_process_residual_and_gpu_release(
+        omni_server_function,
+        scenario=scenario,
+        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
+    )
+
+
+@pytest.mark.slow
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize("signal_name", SERVE_SIGNAL_PARAMS)
+@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_serve_root_no_load_fast_fail_and_cleanup(
+    omni_server_function,
+    signal_name: str,
+) -> None:
+    """Black-box: kill serve root without load; verify fast-fail/health/cleanup."""
+    scenario = f"kill_serve_root_no_load_{signal_name.lower()}"
+    injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
+    injector(omni_server_function)
+    host = omni_server_function.host
+    port = omni_server_function.port
+    _assert_post_fault_video_fast_fail(host, port, scenario=scenario)
+    _assert_post_fault_health_terminal(host, port, scenario=scenario)
+    assert_no_server_tree_process_residual_and_gpu_release(
+        omni_server_function,
+        scenario=scenario,
+        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
+    )
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize("signal_name", TREE_SIGNAL_PARAMS)
+@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup(
+    omni_server_function,
+    signal_name: str,
+) -> None:
+    """Black-box: kill server tree without load; verify fast-fail/health/cleanup."""
+    scenario = f"kill_serve_tree_no_load_{signal_name.lower()}"
+    injector = make_server_tree_kill_fault_injector(
+        signal_name=signal_name,
+        post_kill_wait_seconds=2.0,
+        inter_kill_wait_seconds=0.1,
+    )
+    injector(omni_server_function)
+    host = omni_server_function.host
+    port = omni_server_function.port
+    _assert_post_fault_video_fast_fail(host, port, scenario=scenario)
+    _assert_post_fault_health_terminal(host, port, scenario=scenario)
+    assert_no_server_tree_process_residual_and_gpu_release(
+        omni_server_function,
+        scenario=scenario,
+        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
+    )
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize("signal_name", TREE_WITH_LOAD_SIGNAL_PARAMS)
+@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
+    omni_server_function,
+    openai_client_function,
+    signal_name: str,
+) -> None:
+    """Black-box: during in-flight video requests, kill server process tree and verify post-fault behavior/cleanup."""
+    request_config = _video_request_config()
+    scenario = f"kill_serve_tree_with_load_{signal_name.lower()}"
+    injector = make_server_tree_kill_fault_injector(
+        signal_name=signal_name,
+        post_kill_wait_seconds=2.0,
+        inter_kill_wait_seconds=0.1,
+    )
+    load_result = run_fault_injection_with_rate_load(
+        submit_request=lambda: openai_client_function.send_video_diffusion_request(request_config, request_num=1),
+        inject_fault=lambda: injector(omni_server_function),
+        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
+        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
+        completion_timeout_sec=120.0,
+    )
+    assert load_result["failure_observed"], (
+        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
+    )
+    host = omni_server_function.host
+    port = omni_server_function.port
+    _assert_post_fault_video_fast_fail(host, port, scenario=scenario)
+    _assert_post_fault_health_terminal(host, port, scenario=scenario)
+    assert_no_server_tree_process_residual_and_gpu_release(
+        omni_server_function,
+        scenario=scenario,
+        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
+    )
diff --git a/tests/diffusion/attention/test_sage_attn3.py b/tests/diffusion/attention/test_sage_attn3.py
new file mode 100644
index 000000000..e4d913974
--- /dev/null
+++ b/tests/diffusion/attention/test_sage_attn3.py
@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import importlib
+import sys
+import types
+
+import pytest
+import torch
+from vllm.platforms import current_platform
+
+SAGE_ATTN3_MODULE = "vllm_omni.diffusion.attention.backends.sage_attn3"
+
+
+def load_sage_attn3_module(monkeypatch: pytest.MonkeyPatch, kernel_impl):
+    fake_module = types.ModuleType("sageattn3")
+    fake_module.sageattn3_blackwell = kernel_impl
+    monkeypatch.setitem(sys.modules, "sageattn3", fake_module)
+    sys.modules.pop(SAGE_ATTN3_MODULE, None)
+    return importlib.import_module(SAGE_ATTN3_MODULE)
+
+
+def test_sage_attn3_forward_uses_blackwell_layout(monkeypatch: pytest.MonkeyPatch):
+    calls = {}
+
+    def fake_kernel(query, key, value, is_causal=False):
+        calls["query_shape"] = query.shape
+        calls["is_causal"] = is_causal
+        return query + key + value
+
+    backend_module = load_sage_attn3_module(monkeypatch, fake_kernel)
+    impl = backend_module.SageAttention3Impl(
+        num_heads=4,
+        head_size=64,
+        softmax_scale=1.0 / 8.0,
+        causal=False,
+    )
+
+    query = torch.randn(2, 8, 4, 64)
+    key = torch.randn(2, 8, 4, 64)
+    value = torch.randn(2, 8, 4, 64)
+
+    output = impl.forward_cuda(query, key, value)
+
+    assert calls["query_shape"] == (2, 4, 8, 64)
+    assert calls["is_causal"] is False
+    expected = (query.transpose(1, 2) + key.transpose(1, 2) + value.transpose(1, 2)).transpose(1, 2)
+    assert torch.allclose(output, expected)
+
+
+def test_sage_attn3_falls_back_to_sdpa_for_gqa(monkeypatch: pytest.MonkeyPatch):
+    def fake_kernel(*args, **kwargs):
+        raise AssertionError("sageattn3_blackwell should not be used for GQA")
+
+    backend_module = load_sage_attn3_module(monkeypatch, fake_kernel)
+    sdpa_calls = {}
+
+    def fake_sdpa(query, key, value, **kwargs):
+        sdpa_calls["query_shape"] = query.shape
+        sdpa_calls["key_shape"] = key.shape
+        sdpa_calls["enable_gqa"] = kwargs["enable_gqa"]
+        return query + 1
+
+    monkeypatch.setattr(backend_module.F, "scaled_dot_product_attention", fake_sdpa)
+
+    impl = backend_module.SageAttention3Impl(
+        num_heads=4,
+        head_size=64,
+        softmax_scale=1.0 / 8.0,
+        causal=False,
+    )
+
+    query = torch.randn(2, 8, 4, 64)
+    key = torch.randn(2, 8, 2, 64)
+    value = torch.randn(2, 8, 2, 64)
+
+    output = impl.forward_cuda(query, key, value)
+
+    assert sdpa_calls["query_shape"] == (2, 4, 8, 64)
+    assert sdpa_calls["key_shape"] == (2, 2, 8, 64)
+    assert sdpa_calls["enable_gqa"] is True
+    expected = (query.permute(0, 2, 1, 3) + 1).permute(0, 2, 1, 3)
+    assert torch.allclose(output, expected)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="sage_attn3 tests require CUDA platform")
+def test_cuda_platform_selects_sage_attn3_alias(monkeypatch: pytest.MonkeyPatch):
+    from vllm.platforms.interface import DeviceCapability
+
+    from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+    from vllm_omni.diffusion.envs import PACKAGES_CHECKER
+    from vllm_omni.platforms.cuda import platform as cuda_platform_module
+    from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
+
+    original_import_module = importlib.import_module
+
+    monkeypatch.setattr(
+        CudaOmniPlatform,
+        "get_device_capability",
+        classmethod(lambda cls, device_id=0: DeviceCapability(10, 0)),
+    )
+    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})
+    monkeypatch.setattr(
+        cuda_platform_module.importlib,
+        "import_module",
+        lambda module_name: object() if module_name == "sageattn3" else original_import_module(module_name),
+    )
+
+    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN_3", head_size=64)
+
+    assert backend_path == DiffusionAttentionBackendEnum.SAGE_ATTN_3.get_path()
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="sage_attn3 tests require CUDA platform")
+def test_cuda_platform_falls_back_when_sage_attn3_gpu_is_unsupported(monkeypatch: pytest.MonkeyPatch):
+    from vllm.platforms.interface import DeviceCapability
+
+    from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+    from vllm_omni.diffusion.envs import PACKAGES_CHECKER
+    from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
+
+    monkeypatch.setattr(
+        CudaOmniPlatform,
+        "get_device_capability",
+        classmethod(lambda cls, device_id=0: DeviceCapability(9, 0)),
+    )
+    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})
+
+    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN_3", head_size=64)
+
+    assert backend_path == DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
diff --git a/tests/diffusion/cache/test_cache_backends.py b/tests/diffusion/cache/test_cache_backends.py
index ecb993945..021473ca5 100644
--- a/tests/diffusion/cache/test_cache_backends.py
+++ b/tests/diffusion/cache/test_cache_backends.py
@@ -19,6 +19,7 @@ from vllm_omni.diffusion.cache.cache_dit_backend import (
     CUSTOM_DIT_ENABLERS,
     CacheDiTBackend,
 )
+from vllm_omni.diffusion.cache.magcache.backend import MagCacheBackend
 from vllm_omni.diffusion.cache.selector import get_cache_backend
 from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend
 from vllm_omni.diffusion.data import DiffusionCacheConfig
@@ -127,6 +128,10 @@ class TestCacheDiTBackend:
         """Test HunyuanImage3 custom cache-dit enabler is registered."""
         assert "HunyuanImage3Pipeline" in CUSTOM_DIT_ENABLERS
 
+    def test_dreamid_custom_enabler_registered(self):
+        """Test DreamIDOmni custom cache-dit enabler is registered."""
+        assert "DreamIDOmniPipeline" in CUSTOM_DIT_ENABLERS
+
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
     def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mock_block_adapter):
@@ -169,6 +174,47 @@ class TestCacheDiTBackend:
         with pytest.raises(ValueError, match="HunyuanImage3Pipeline"):
             backend.enable(mock_pipeline)
 
+    @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
+    @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
+    def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter):
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline"
+        mock_pipeline.transformer = Mock()
+        mock_pipeline.transformer.fused_blocks = Mock()
+
+        mock_cache_dit.enable_cache = Mock()
+        mock_cache_dit.refresh_context = Mock()
+
+        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
+        backend.enable(mock_pipeline)
+
+        assert backend.enabled is True
+        assert backend._refresh_func is not None
+        mock_block_adapter.assert_called_once()
+        adapter_kwargs = mock_block_adapter.call_args.kwargs
+        assert adapter_kwargs["transformer"] is mock_pipeline.transformer
+        assert adapter_kwargs["blocks"] is mock_pipeline.transformer.fused_blocks
+        assert adapter_kwargs["forward_pattern"] == adapter_kwargs["forward_pattern"].__class__.Pattern_0
+        assert adapter_kwargs["has_separate_cfg"] is True
+        assert len(adapter_kwargs["params_modifiers"]) == 1
+        mock_cache_dit.enable_cache.assert_called_once()
+
+        backend.refresh(mock_pipeline, num_inference_steps=12)
+        mock_cache_dit.refresh_context.assert_called_once()
+        call_args = mock_cache_dit.refresh_context.call_args
+        assert call_args[0][0] is mock_pipeline.transformer
+        assert call_args[1]["num_inference_steps"] == 12
+
+    def test_enable_dreamid_pipeline_requires_fused_blocks(self):
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline"
+        mock_pipeline.transformer = Mock(spec=[])
+
+        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
+
+        with pytest.raises(ValueError, match="DreamIDOmniPipeline"):
+            backend.enable(mock_pipeline)
+
 
 class TestTeaCacheBackend:
     """Test TeaCacheBackend implementation."""
@@ -268,3 +314,150 @@ class TestCacheSelector:
         """Test getting invalid backend raises error."""
         with pytest.raises(ValueError, match="Unsupported cache backend"):
             get_cache_backend("invalid_backend", {})
+
+
+class TestMagCacheBackend:
+    """Test MagCacheBackend implementation."""
+
+    from vllm_omni.diffusion.cache.magcache.backend import MagCacheBackend
+
+    def test_init(self):
+        """Test initialization."""
+        config = DiffusionCacheConfig(mag_threshold=0.1, mag_max_skip_steps=2, mag_calibrate=True)
+        backend = MagCacheBackend(config)
+        assert backend.config.mag_threshold == 0.1
+        assert backend.config.mag_max_skip_steps == 2
+        assert backend.enabled is False
+
+    @patch("vllm_omni.diffusion.cache.magcache.backend.apply_mag_cache_hook")
+    def test_enable(self, mock_apply_hook):
+        """Test enabling MagCache on pipeline."""
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "FluxPipeline"
+        mock_transformer = Mock()
+        mock_transformer.__class__.__name__ = "FluxTransformer2DModel"
+        mock_pipeline.transformer = mock_transformer
+
+        mock_ratios = [1.0] * 28
+        config = DiffusionCacheConfig(
+            mag_ratios=mock_ratios,
+        )
+        backend = MagCacheBackend(config)
+        backend.enable(mock_pipeline)
+
+        assert backend.enabled is True
+        mock_apply_hook.assert_called_once()
+
+        call_args = mock_apply_hook.call_args
+        assert call_args[0][0] == mock_transformer
+
+    @patch("vllm_omni.diffusion.cache.magcache.backend.apply_mag_cache_hook")
+    def test_enable_with_calibration(self, mock_apply_hook):
+        """Test enabling MagCache in calibration mode."""
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "FluxPipeline"
+        mock_transformer = Mock()
+        mock_transformer.__class__.__name__ = "FluxTransformer2DModel"
+        mock_pipeline.transformer = mock_transformer
+
+        config = DiffusionCacheConfig(
+            mag_calibrate=True,
+        )
+        backend = MagCacheBackend(config)
+        backend.enable(mock_pipeline)
+
+        assert backend.enabled is True
+        mock_apply_hook.assert_called_once()
+
+    def test_refresh(self):
+        """Test refreshing MagCache state calls enable when not registered."""
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "FluxPipeline"
+        mock_transformer = Mock()
+        mock_transformer.__class__.__name__ = "FluxTransformer2DModel"
+        mock_pipeline.transformer = mock_transformer
+
+        mock_transformer.named_children = Mock(return_value=[])
+
+        mock_ratios = [1.0] * 28
+        config = DiffusionCacheConfig(
+            mag_ratios=mock_ratios,
+        )
+        backend = MagCacheBackend(config)
+
+        assert backend._registered is False
+        backend.refresh(mock_pipeline, num_inference_steps=50)
+        assert backend._registered is True
+
+    def test_is_enabled(self):
+        """Test is_enabled method."""
+        mock_ratios = [1.0] * 28
+        config = DiffusionCacheConfig(mag_ratios=mock_ratios)
+        backend = MagCacheBackend(config)
+        assert backend.is_enabled() is False
+
+    def test_get_mag_cache_backend(self):
+        """Test getting MagCache backend via selector."""
+        mock_ratios = [1.0] * 28
+        config_dict = {
+            "mag_ratios": mock_ratios,
+            "num_inference_steps": 28,
+            "threshold": 0.06,
+            "max_skip_steps": 3,
+            "retention_ratio": 0.2,
+        }
+        backend = get_cache_backend("mag_cache", config_dict)
+        assert backend is not None
+        assert isinstance(backend, MagCacheBackend)
+        assert backend.config.threshold == 0.06
+
+    @patch("vllm_omni.diffusion.cache.magcache.backend.apply_mag_cache_hook")
+    def test_enable_single_block(self, mock_apply_hook):
+        """Test enabling MagCache on single transformer block."""
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "FluxPipeline"
+
+        mock_block = Mock()
+        mock_block.__class__.__name__ = "FluxTransformer2DModel"
+        mock_blocks = [mock_block]
+
+        mock_transformer = Mock()
+        mock_transformer.__class__.__name__ = "FluxTransformer2DModel"
+        mock_transformer.blocks = mock_blocks
+        mock_pipeline.transformer = mock_transformer
+
+        mock_ratios = [1.0] * 28
+        config = DiffusionCacheConfig(
+            mag_ratios=mock_ratios,
+        )
+        backend = MagCacheBackend(config)
+        backend.enable(mock_pipeline)
+
+        assert backend.enabled is True
+        mock_apply_hook.assert_called_once()
+
+        call_args = mock_apply_hook.call_args
+        assert call_args[0][0] == mock_transformer
+
+    @patch("vllm_omni.diffusion.cache.magcache.backend.apply_mag_cache_hook")
+    def test_enable_multi_block(self, mock_apply_hook):
+        """Test enabling MagCache on multiple transformer blocks."""
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "FluxPipeline"
+
+        mock_blocks = [Mock() for _ in range(24)]
+
+        mock_transformer = Mock()
+        mock_transformer.__class__.__name__ = "FluxTransformer2DModel"
+        mock_transformer.blocks = mock_blocks
+        mock_pipeline.transformer = mock_transformer
+
+        mock_ratios = [1.0] * 28
+        config = DiffusionCacheConfig(
+            mag_ratios=mock_ratios,
+        )
+        backend = MagCacheBackend(config)
+        backend.enable(mock_pipeline)
+
+        assert backend.enabled is True
+        mock_apply_hook.assert_called_once()
diff --git a/tests/diffusion/cache/test_cache_dit.py b/tests/diffusion/cache/test_cache_dit.py
index 0b7ef7235..0fd07393d 100644
--- a/tests/diffusion/cache/test_cache_dit.py
+++ b/tests/diffusion/cache/test_cache_dit.py
@@ -8,14 +8,21 @@ Model specific tests for CacheDiT enablement.
 from unittest.mock import Mock, patch
 
 import pytest
+import torch
+from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2
 
 import vllm_omni.diffusion.cache.cache_dit_backend as cd_backend
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTBackend
 from vllm_omni.diffusion.data import DiffusionCacheConfig
+from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
+from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 SEPARATE_CFG_ENABLERS = [
+    cd_backend.enable_cache_for_dreamid_omni,
     cd_backend.enable_cache_for_ltx2,
+    cd_backend.enable_cache_for_helios,
     cd_backend.enable_cache_for_wan22,
     cd_backend.enable_cache_for_longcat_image,
 ]
@@ -38,3 +45,85 @@ def test_separate_cfg(mock_cache_dit, mock_block_adapter, enabler):
     mock_cache_dit.enable_cache.assert_called_once()
     adapter_kwargs = mock_block_adapter.call_args.kwargs
     assert adapter_kwargs["has_separate_cfg"] is True
+
+
+# This test is skipped on ROCm since rocm_unquantized_gemm doesn't support CPU backend
+@pytest.mark.skipif(
+    current_omni_platform.is_rocm(),
+    reason="vLLM ROCm custom ops lack CPU fallback",
+)
+def test_ltx2_cache_dit_receives_audio_as_encoder(init_fake_tp_group):
+    """CacheDiT Pattern_0 treats the second positional arg as encoder_hidden_states,
+    which is a collision for one of the kwargs in LTX2 since we treat the audio
+    hidden states as encoder_hidden_states.
+
+    This test ensures that a tiny LTX2 transformer can be initialized and run
+    through the cache DiT backend without hitting a collision on the kwargs.
+    """
+    seq_len = 4
+    video_in = torch.full((1, seq_len, 16), 1.0)
+    audio_in = torch.full((1, seq_len, 16), 2.0)
+    text_in = torch.full((1, seq_len, 16), 3.0)
+    audio_text_in = torch.full((1, seq_len, 16), 4.0)
+
+    model = LTX2VideoTransformer3DModel(
+        in_channels=16,
+        out_channels=16,
+        patch_size=1,
+        patch_size_t=1,
+        num_attention_heads=2,
+        attention_head_dim=8,
+        cross_attention_dim=16,
+        audio_in_channels=16,
+        audio_out_channels=16,
+        audio_num_attention_heads=2,
+        audio_attention_head_dim=8,
+        audio_cross_attention_dim=16,
+        num_layers=2,
+        caption_channels=16,
+    )
+
+    # NOTE: This is currently using the LTX2 custom enabler, but the custom
+    # enablers will be consolidated after
+    # https://github.com/vllm-project/vllm-omni/pull/2527 lands.
+    LTX2Pipeline = type("LTX2Pipeline", (), {})
+    pipeline = LTX2Pipeline()
+    pipeline.transformer = model
+    backend = CacheDiTBackend(DiffusionCacheConfig())
+    backend.enable(pipeline)
+    backend.refresh(pipeline, num_inference_steps=5)
+
+    # Wrap call_Fn_blocks in CacheDiT so that we can verify the
+    # hidden/encoder states are what we expect them to be
+    captured = {}
+    original = CachedBlocks_Pattern_0_1_2.call_Fn_blocks
+
+    def call_Fn_blocks_and_capture(self, hidden_states, encoder_hidden_states, *a, **kw):
+        captured["hidden_states"] = hidden_states
+        captured["encoder_hidden_states"] = encoder_hidden_states
+        return original(self, hidden_states, encoder_hidden_states, *a, **kw)
+
+    # Also, map projections to identity so that we can just check
+    # the captured tensors directly instead of having to reproject
+    identity = torch.nn.Identity()
+    with (
+        patch.object(model, "proj_in", identity),
+        patch.object(model, "audio_proj_in", identity),
+        patch.object(CachedBlocks_Pattern_0_1_2, "call_Fn_blocks", call_Fn_blocks_and_capture),
+        torch.no_grad(),
+    ):
+        model(
+            hidden_states=video_in,
+            audio_hidden_states=audio_in,
+            encoder_hidden_states=text_in,
+            audio_encoder_hidden_states=audio_text_in,
+            timestep=torch.tensor([[1000.0] * seq_len]),
+            num_frames=1,
+            height=2,
+            width=2,
+            audio_num_frames=seq_len,
+        )
+
+    # Pattern_0 maps (hidden_states, encoder_hidden_states) to (video, audio)
+    assert torch.equal(captured["hidden_states"], video_in)
+    assert torch.equal(captured["encoder_hidden_states"], audio_in)
diff --git a/tests/diffusion/cache/test_prompt_embed_cache.py b/tests/diffusion/cache/test_prompt_embed_cache.py
new file mode 100644
index 000000000..184279c99
--- /dev/null
+++ b/tests/diffusion/cache/test_prompt_embed_cache.py
@@ -0,0 +1,529 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Unit tests for prompt-embedding cache for diffusion pipelines.
+
+Covers the module :mod:`vllm_omni.diffusion.cache.prompt_embed_cache`:
+
+* :class:`PromptEmbedCache` LRU semantics, stats and thread-safety.
+* :func:`_hashable` safe-input classification (scalars, torch dtype/device,
+  numpy scalars, nested containers, tensors/PIL bypass).
+* :func:`install_prompt_embed_cache` / :func:`uninstall_prompt_embed_cache`
+  behaviour: caching, bypass on tensors, bypass on precomputed embeds,
+  idempotent install, restore on uninstall, detachment of cached tensors.
+* :func:`resolve_prompt_embed_cache_config` env-var overrides.
+"""
+
+from __future__ import annotations
+
+import threading
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.cache.prompt_embed_cache import (
+    _CACHE_MISS,
+    _NOT_HASHABLE,
+    PromptEmbedCache,
+    _build_key,
+    _detach_output,
+    _hashable,
+    install_prompt_embed_cache,
+    resolve_prompt_embed_cache_config,
+    uninstall_prompt_embed_cache,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+# ---------------------------------------------------------------------------
+# _hashable
+# ---------------------------------------------------------------------------
+
+
+class TestHashable:
+    """Tests for the ``_hashable`` helper."""
+
+    @pytest.mark.parametrize(
+        "value",
+        [None, "prompt", "", 0, 1, -5, 3.14, True, False, b"bytes"],
+    )
+    def test_scalars_are_returned_unchanged(self, value):
+        assert _hashable(value) == value
+
+    def test_torch_device_is_normalized(self):
+        out = _hashable(torch.device("cpu"))
+        assert out == ("__torch_device__", "cpu")
+        assert hash(out) is not None  # must actually be hashable
+
+    def test_torch_dtype_is_normalized(self):
+        out = _hashable(torch.float32)
+        assert out[0] == "__torch_dtype__"
+        assert "float32" in out[1]
+        assert hash(out) is not None
+
+    def test_numpy_scalar_is_unwrapped(self):
+        np = pytest.importorskip("numpy")
+        assert _hashable(np.int64(7)) == 7
+        assert _hashable(np.float32(1.5)) == pytest.approx(1.5)
+
+    def test_numpy_ndarray_is_not_hashable(self):
+        np = pytest.importorskip("numpy")
+        assert _hashable(np.array([1, 2, 3])) is _NOT_HASHABLE
+
+    def test_nested_list_and_tuple(self):
+        out = _hashable(["a", ("b", 1), [2, 3]])
+        assert out[0] == "__seq__"
+        # fully hashable
+        assert hash(out) is not None
+
+    def test_dict_is_normalized_by_sorted_keys(self):
+        a = _hashable({"b": 1, "a": 2})
+        b = _hashable({"a": 2, "b": 1})
+        assert a == b
+        assert hash(a) is not None
+
+    def test_tensor_is_not_hashable(self):
+        assert _hashable(torch.zeros(2)) is _NOT_HASHABLE
+
+    def test_list_containing_tensor_is_not_hashable(self):
+        assert _hashable(["a", torch.zeros(1)]) is _NOT_HASHABLE
+
+    def test_dict_containing_tensor_is_not_hashable(self):
+        assert _hashable({"x": torch.zeros(1)}) is _NOT_HASHABLE
+
+    def test_custom_object_is_not_hashable(self):
+        class Obj:
+            pass
+
+        assert _hashable(Obj()) is _NOT_HASHABLE
+
+
+# ---------------------------------------------------------------------------
+# _detach_output
+# ---------------------------------------------------------------------------
+
+
+class TestDetachOutput:
+    def test_detaches_tensor(self):
+        t = torch.zeros(2, requires_grad=True)
+        out = _detach_output(t)
+        assert isinstance(out, torch.Tensor)
+        assert not out.requires_grad
+
+    def test_detaches_inside_tuple(self):
+        t = torch.ones(1, requires_grad=True)
+        out = _detach_output((t, None))
+        assert isinstance(out, tuple)
+        assert not out[0].requires_grad
+        assert out[1] is None
+
+    def test_detaches_inside_list_and_dict(self):
+        t = torch.ones(1, requires_grad=True)
+        out = _detach_output({"a": [t]})
+        assert not out["a"][0].requires_grad
+
+    def test_preserves_namedtuple_type(self):
+        from collections import namedtuple
+
+        NT = namedtuple("NT", ["x", "y"])
+        t = torch.ones(1, requires_grad=True)
+        out = _detach_output(NT(x=t, y=1))
+        assert isinstance(out, NT)
+        assert not out.x.requires_grad
+        assert out.y == 1
+
+    def test_passthrough_for_non_tensor(self):
+        assert _detach_output("hi") == "hi"
+        assert _detach_output(None) is None
+        assert _detach_output(5) == 5
+
+
+# ---------------------------------------------------------------------------
+# PromptEmbedCache
+# ---------------------------------------------------------------------------
+
+
+class TestPromptEmbedCache:
+    def test_rejects_non_positive_max_size(self):
+        with pytest.raises(ValueError):
+            PromptEmbedCache(max_size=0)
+        with pytest.raises(ValueError):
+            PromptEmbedCache(max_size=-1)
+
+    def test_miss_returns_sentinel_and_bumps_miss_counter(self):
+        cache = PromptEmbedCache(max_size=2)
+        assert cache.get("k") is _CACHE_MISS
+        assert cache.stats()["misses"] == 1
+        assert cache.stats()["hits"] == 0
+
+    def test_put_then_get_returns_value_and_bumps_hit_counter(self):
+        cache = PromptEmbedCache(max_size=2)
+        cache.put("k", 42)
+        assert cache.get("k") == 42
+        assert cache.stats()["hits"] == 1
+
+    def test_none_is_a_cacheable_value(self):
+        cache = PromptEmbedCache(max_size=2)
+        cache.put("k", None)
+        # Distinguishable from a miss.
+        assert cache.get("k") is None
+        assert cache.stats()["hits"] == 1
+
+    def test_lru_eviction_keeps_recently_used(self):
+        cache = PromptEmbedCache(max_size=2)
+        cache.put("a", 1)
+        cache.put("b", 2)
+        # Access "a" to mark it MRU.
+        assert cache.get("a") == 1
+        cache.put("c", 3)  # should evict "b"
+        assert cache.get("b") is _CACHE_MISS
+        assert cache.get("a") == 1
+        assert cache.get("c") == 3
+        assert cache.stats()["size"] == 2
+
+    def test_get_refreshes_recency(self):
+        cache = PromptEmbedCache(max_size=2)
+        cache.put("a", 1)
+        cache.put("b", 2)
+        # Hit on "a" → "b" becomes LRU.
+        cache.get("a")
+        cache.put("c", 3)
+        assert cache.get("b") is _CACHE_MISS
+        assert cache.get("a") == 1
+
+    def test_put_overwrites_existing_key(self):
+        cache = PromptEmbedCache(max_size=2)
+        cache.put("k", 1)
+        cache.put("k", 2)
+        assert cache.get("k") == 2
+        assert cache.stats()["size"] == 1
+
+    def test_clear_empties_store_but_keeps_counters(self):
+        cache = PromptEmbedCache(max_size=2)
+        cache.put("k", 1)
+        cache.get("k")
+        cache.clear()
+        assert cache.stats()["size"] == 0
+        # clear only drops entries, counters reflect prior activity.
+        assert cache.stats()["hits"] == 1
+
+    def test_stats_reports_expected_fields(self):
+        cache = PromptEmbedCache(max_size=4)
+        stats = cache.stats()
+        assert set(stats.keys()) == {"size", "max_size", "hits", "misses", "bypassed"}
+        assert stats["max_size"] == 4
+        assert stats["size"] == 0
+
+    def test_thread_safety_under_concurrent_writes(self):
+        cache = PromptEmbedCache(max_size=1024)
+
+        def worker(start: int) -> None:
+            for i in range(start, start + 200):
+                cache.put(f"k{i}", i)
+                cache.get(f"k{i}")
+
+        threads = [threading.Thread(target=worker, args=(i * 200,)) for i in range(4)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        stats = cache.stats()
+        # 4 * 200 = 800 writes, each followed by a successful read.
+        assert stats["hits"] == 800
+        assert stats["size"] == 800
+
+
+# ---------------------------------------------------------------------------
+# _build_key
+# ---------------------------------------------------------------------------
+
+
+def _make_sig(fn):
+    import inspect
+
+    return inspect.signature(fn)
+
+
+class TestBuildKey:
+    def test_positional_and_keyword_collide(self):
+        def encode(prompt, device=None):
+            return None
+
+        sig = _make_sig(encode)
+        k1 = _build_key(sig, "m", ("hello",), {"device": "cpu"})
+        k2 = _build_key(sig, "m", (), {"prompt": "hello", "device": "cpu"})
+        assert k1 == k2
+        assert k1 is not None
+
+    def test_different_prompts_produce_different_keys(self):
+        def encode(prompt):
+            return None
+
+        sig = _make_sig(encode)
+        assert _build_key(sig, "m", ("a",), {}) != _build_key(sig, "m", ("b",), {})
+
+    def test_different_model_tags_produce_different_keys(self):
+        def encode(prompt):
+            return None
+
+        sig = _make_sig(encode)
+        assert _build_key(sig, "m1", ("a",), {}) != _build_key(sig, "m2", ("a",), {})
+
+    def test_tensor_argument_bypasses(self):
+        def encode(prompt, extra=None):
+            return None
+
+        sig = _make_sig(encode)
+        assert _build_key(sig, "m", ("hi",), {"extra": torch.zeros(1)}) is None
+
+    def test_precomputed_prompt_embeds_bypasses(self):
+        def encode(prompt=None, prompt_embeds=None):
+            return None
+
+        sig = _make_sig(encode)
+        # None → not a bypass.
+        assert _build_key(sig, "m", (), {"prompt": "hi"}) is not None
+        # A non-None precomputed embed triggers bypass even if it would
+        # otherwise be "hashable".
+        assert _build_key(sig, "m", (), {"prompt": "hi", "prompt_embeds": "x"}) is None
+
+    def test_invalid_binding_returns_none(self):
+        def encode(prompt):
+            return None
+
+        sig = _make_sig(encode)
+        # Unknown kwarg → TypeError on bind → bypass.
+        assert _build_key(sig, "m", (), {"nope": 1}) is None
+
+
+# ---------------------------------------------------------------------------
+# install_prompt_embed_cache / uninstall_prompt_embed_cache
+# ---------------------------------------------------------------------------
+
+
+class _FakePipeline:
+    """Minimal pipeline with a signature-bearing ``encode_prompt``."""
+
+    def __init__(self):
+        self.call_count = 0
+
+    def encode_prompt(
+        self,
+        prompt: str,
+        device=None,
+        num_images_per_prompt: int = 1,
+        negative_prompt: str | None = None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+    ):
+        self.call_count += 1
+        # Return a fresh tensor so identity changes between real calls.
+        return torch.tensor([float(self.call_count)])
+
+
+class TestInstallAndUninstall:
+    def test_returns_none_when_pipeline_has_no_encode_prompt(self):
+        class Empty:
+            pass
+
+        assert install_prompt_embed_cache(Empty()) is None
+
+    def test_cache_hits_on_identical_args(self):
+        pipe = _FakePipeline()
+        cache = install_prompt_embed_cache(pipe, max_size=4)
+        assert cache is not None
+
+        out1 = pipe.encode_prompt("cat", device="cpu")
+        out2 = pipe.encode_prompt("cat", device="cpu")
+        # Second call must be served from cache (wrapped fn not invoked).
+        assert pipe.call_count == 1
+        # Cache returns the same tensor object.
+        assert torch.equal(out1, out2)
+        stats = cache.stats()
+        assert stats["hits"] == 1
+        assert stats["misses"] == 1
+        assert stats["bypassed"] == 0
+
+    def test_positional_and_keyword_share_cache_slot(self):
+        pipe = _FakePipeline()
+        cache = install_prompt_embed_cache(pipe)
+        pipe.encode_prompt("cat")
+        pipe.encode_prompt(prompt="cat")
+        assert pipe.call_count == 1
+        assert cache.stats()["hits"] == 1
+
+    def test_different_prompts_miss(self):
+        pipe = _FakePipeline()
+        cache = install_prompt_embed_cache(pipe)
+        pipe.encode_prompt("cat")
+        pipe.encode_prompt("dog")
+        assert pipe.call_count == 2
+        assert cache.stats()["misses"] == 2
+
+    def test_tensor_argument_bypasses_and_increments_counter(self):
+        pipe = _FakePipeline()
+        cache = install_prompt_embed_cache(pipe)
+        # Put a tensor into a non-precomputed-embed slot to trigger bypass.
+        pipe.encode_prompt("cat", device=torch.device("cpu"), num_images_per_prompt=1)
+
+        # Now pass a fake tensor via a harmless positional-like call: use
+        # the negative_prompt slot with an unhashable object.
+        class Unhashable:
+            pass
+
+        pipe.encode_prompt("cat", negative_prompt=Unhashable())  # type: ignore[arg-type]
+        pipe.encode_prompt("cat", negative_prompt=Unhashable())  # type: ignore[arg-type]
+        assert cache.stats()["bypassed"] == 2
+        # All three calls actually executed the underlying function.
+        assert pipe.call_count == 3
+
+    def test_precomputed_embeds_bypass_cache(self):
+        pipe = _FakePipeline()
+        cache = install_prompt_embed_cache(pipe)
+        pipe.encode_prompt("cat", prompt_embeds=torch.zeros(1))
+        pipe.encode_prompt("cat", prompt_embeds=torch.zeros(1))
+        assert cache.stats()["bypassed"] == 2
+        assert cache.stats()["hits"] == 0
+        assert pipe.call_count == 2
+
+    def test_install_is_idempotent(self):
+        pipe = _FakePipeline()
+        c1 = install_prompt_embed_cache(pipe, max_size=4)
+        c2 = install_prompt_embed_cache(pipe, max_size=16)
+        assert c1 is c2
+        # The second call must not re-wrap (max_size stays from first install).
+        assert c1.max_size == 4
+
+    def test_disabled_cache_passes_through(self):
+        pipe = _FakePipeline()
+        cache = install_prompt_embed_cache(pipe, enabled=False)
+        pipe.encode_prompt("cat")
+        pipe.encode_prompt("cat")
+        assert pipe.call_count == 2
+        assert cache.stats()["hits"] == 0
+        assert cache.stats()["misses"] == 0
+
+    def test_toggle_enabled_flag_takes_effect(self):
+        pipe = _FakePipeline()
+        cache = install_prompt_embed_cache(pipe, enabled=True)
+        pipe.encode_prompt("cat")
+        pipe.encode_prompt("cat")
+        assert cache.stats()["hits"] == 1
+
+        cache.enabled = False
+        pipe.encode_prompt("cat")
+        pipe.encode_prompt("cat")
+        # Cache-bypassed calls go through; hit count unchanged.
+        assert cache.stats()["hits"] == 1
+
+    def test_cached_tensor_is_detached(self):
+        class GradPipeline:
+            def encode_prompt(self, prompt):
+                return torch.ones(1, requires_grad=True)
+
+        pipe = GradPipeline()
+        install_prompt_embed_cache(pipe)
+        # First call populates the cache.
+        pipe.encode_prompt("cat")
+        # Second call returns the cached tensor which must be detached.
+        cached = pipe.encode_prompt("cat")
+        assert isinstance(cached, torch.Tensor)
+        assert not cached.requires_grad
+
+    def test_lru_eviction_through_wrapper(self):
+        pipe = _FakePipeline()
+        cache = install_prompt_embed_cache(pipe, max_size=2)
+        pipe.encode_prompt("a")
+        pipe.encode_prompt("b")
+        pipe.encode_prompt("c")  # evicts "a"
+        pipe.encode_prompt("a")  # miss again → re-run
+        assert pipe.call_count == 4
+        assert cache.stats()["size"] == 2
+
+    def test_uninstall_restores_original(self):
+        pipe = _FakePipeline()
+        original = pipe.encode_prompt
+        install_prompt_embed_cache(pipe)
+        assert pipe.encode_prompt is not original
+        uninstall_prompt_embed_cache(pipe)
+        # The bound method reference will differ, but identity of underlying
+        # function matches and wrapper attributes are gone.
+        assert not hasattr(pipe, "_prompt_embed_cache") or pipe._prompt_embed_cache is None
+        # After uninstall calls go straight to the real function.
+        pipe.encode_prompt("cat")
+        pipe.encode_prompt("cat")
+        assert pipe.call_count == 2
+
+    def test_uninstall_on_unwrapped_pipeline_is_noop(self):
+        class Empty:
+            def encode_prompt(self, prompt):
+                return None
+
+        pipe = Empty()
+        uninstall_prompt_embed_cache(pipe)  # must not raise
+
+
+# ---------------------------------------------------------------------------
+# resolve_prompt_embed_cache_config
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def clean_env(monkeypatch):
+    monkeypatch.delenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE", raising=False)
+    monkeypatch.delenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", raising=False)
+
+
+class TestResolveConfig:
+    def test_defaults_when_nothing_set(self, clean_env):
+        enable, size = resolve_prompt_embed_cache_config()
+        assert enable is False
+        assert size == 32
+
+    def test_explicit_args_used(self, clean_env):
+        enable, size = resolve_prompt_embed_cache_config(enable=True, max_size=8)
+        assert enable is True
+        assert size == 8
+
+    @pytest.mark.parametrize("raw", ["1", "true", "TRUE", "yes", "on"])
+    def test_env_enable_truthy(self, clean_env, monkeypatch, raw):
+        monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE", raw)
+        enable, _ = resolve_prompt_embed_cache_config(enable=False)
+        assert enable is True
+
+    @pytest.mark.parametrize("raw", ["0", "false", "FALSE", "no", "off"])
+    def test_env_enable_falsy_overrides_true(self, clean_env, monkeypatch, raw):
+        monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE", raw)
+        enable, _ = resolve_prompt_embed_cache_config(enable=True)
+        assert enable is False
+
+    def test_env_enable_invalid_leaves_arg_untouched(self, clean_env, monkeypatch):
+        monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE", "garbage")
+        enable, _ = resolve_prompt_embed_cache_config(enable=True)
+        assert enable is True
+
+    def test_env_size_overrides_arg(self, clean_env, monkeypatch):
+        monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", "7")
+        _, size = resolve_prompt_embed_cache_config(max_size=99)
+        assert size == 7
+
+    def test_env_size_zero_or_negative_ignored(self, clean_env, monkeypatch):
+        monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", "0")
+        _, size = resolve_prompt_embed_cache_config(max_size=12)
+        assert size == 12
+
+        monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", "-5")
+        _, size = resolve_prompt_embed_cache_config(max_size=12)
+        assert size == 12
+
+    def test_env_size_non_integer_ignored(self, clean_env, monkeypatch):
+        monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", "notanint")
+        _, size = resolve_prompt_embed_cache_config(max_size=12)
+        assert size == 12
+
+    def test_max_size_none_falls_back_to_default(self, clean_env):
+        _, size = resolve_prompt_embed_cache_config(enable=True, max_size=None)
+        assert size == 32
diff --git a/tests/diffusion/diffusion_backend/test_diffusers_backend.py b/tests/diffusion/diffusion_backend/test_diffusers_backend.py
index 92dd73358..b64d6f738 100644
--- a/tests/diffusion/diffusion_backend/test_diffusers_backend.py
+++ b/tests/diffusion/diffusion_backend/test_diffusers_backend.py
@@ -60,6 +60,7 @@ def _make_request(**overrides) -> OmniDiffusionRequest:
             output_type="pil",
             generator_device="cpu",
         ),
+        "request_id": "diffusers-backend",
     }
     defaults.update(overrides)
     return OmniDiffusionRequest(**defaults)
diff --git a/tests/diffusion/distributed/test_autoencoder_kl_hunyuan.py b/tests/diffusion/distributed/test_autoencoder_kl_hunyuan.py
index 58ce14546..e65313d38 100644
--- a/tests/diffusion/distributed/test_autoencoder_kl_hunyuan.py
+++ b/tests/diffusion/distributed/test_autoencoder_kl_hunyuan.py
@@ -7,6 +7,9 @@ import torch
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan import (
     DistributedAutoencoderKLHunyuan,
 )
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan_video_15 import (
+    DistributedAutoencoderKLHunyuanVideo15,
+)
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -36,6 +39,33 @@ class _DummyDistributedAutoencoderKLHunyuan(DistributedAutoencoderKLHunyuan):
         return b
 
 
+class _DummyDistributedAutoencoderKLHunyuanVideo15(DistributedAutoencoderKLHunyuanVideo15):
+    def __init__(self):
+        torch.nn.Module.__init__(self)
+        self.tile_sample_min_height = 8
+        self.tile_sample_min_width = 12
+        self.tile_latent_min_height = 8
+        self.tile_latent_min_width = 12
+        self.tile_overlap_factor = 0.25
+        self.use_tiling = False
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return torch.float32
+
+    def decoder(self, x: torch.Tensor) -> torch.Tensor:
+        return x + 10
+
+    def encoder(self, x: torch.Tensor) -> torch.Tensor:
+        return x + 20
+
+    def blend_v(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor:
+        return b
+
+    def blend_h(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor:
+        return b
+
+
 def test_hunyuan_vae_use_tiling_aliases_spatial_tiling():
     # Verify use_tiling property maps to use_spatial_tiling.
     vae = _DummyDistributedAutoencoderKLHunyuan()
@@ -77,3 +107,43 @@ def test_hunyuan_vae_encode_tiles_round_trip():
     assert grid_spec.tile_spec == {"blend_extent": 2, "row_limit": 6}
     assert len(tile_tasks) == 4
     assert torch.equal(output, x + 20)
+
+
+def test_hunyuan_video15_vae_decode_tiles_round_trip():
+    vae = _DummyDistributedAutoencoderKLHunyuanVideo15()
+    z = torch.arange(192, dtype=torch.float32).reshape(1, 1, 1, 12, 16)
+
+    tile_tasks, grid_spec = vae.tile_split(z)
+    decoded_tiles = {task.grid_coord: vae.tile_exec(task) for task in tile_tasks}
+    output = vae.tile_merge(decoded_tiles, grid_spec)
+
+    assert grid_spec.split_dims == (3, 4)
+    assert grid_spec.grid_shape == (2, 2)
+    assert grid_spec.tile_spec == {
+        "blend_height": 2,
+        "blend_width": 3,
+        "row_limit_height": 6,
+        "row_limit_width": 9,
+    }
+    assert len(tile_tasks) == 4
+    assert torch.equal(output, z + 10)
+
+
+def test_hunyuan_video15_vae_encode_tiles_round_trip():
+    vae = _DummyDistributedAutoencoderKLHunyuanVideo15()
+    x = torch.arange(192, dtype=torch.float32).reshape(1, 1, 1, 12, 16)
+
+    tile_tasks, grid_spec = vae.encode_tile_split(x)
+    encoded_tiles = {task.grid_coord: vae.encode_tile_exec(task) for task in tile_tasks}
+    output = vae.tile_merge(encoded_tiles, grid_spec)
+
+    assert grid_spec.split_dims == (3, 4)
+    assert grid_spec.grid_shape == (2, 2)
+    assert grid_spec.tile_spec == {
+        "blend_height": 2,
+        "blend_width": 3,
+        "row_limit_height": 6,
+        "row_limit_width": 9,
+    }
+    assert len(tile_tasks) == 4
+    assert torch.equal(output, x + 20)
diff --git a/tests/diffusion/models/bagel/test_forward_cache_update_vae.py b/tests/diffusion/models/bagel/test_forward_cache_update_vae.py
new file mode 100644
index 000000000..e49b408f4
--- /dev/null
+++ b/tests/diffusion/models/bagel/test_forward_cache_update_vae.py
@@ -0,0 +1,169 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for Bagel.forward_cache_update_vae packed-sequence assembly."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from transformers import PretrainedConfig
+from vllm.transformers_utils.configs.bagel import BagelConfig
+
+from vllm_omni.diffusion.models.bagel.bagel_transformer import (
+    Bagel,
+    BaseNavitOutputWithPast,
+    NaiveCache,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+HIDDEN_SIZE = 32
+NUM_LAYERS = 2
+
+
+def _make_bagel_config() -> BagelConfig:
+    llm_config = PretrainedConfig()
+    llm_config.hidden_size = HIDDEN_SIZE
+    llm_config.layer_module = "Qwen2MoTDecoderLayer"
+    llm_config.num_hidden_layers = NUM_LAYERS
+    llm_config.num_attention_heads = 4
+
+    vae_config = PretrainedConfig()
+    vae_config.downsample = 8
+    vae_config.z_channels = 4
+
+    return BagelConfig(
+        llm_config=llm_config,
+        vae_config=vae_config,
+        visual_gen=True,
+        visual_und=False,
+        latent_patch_size=2,
+        max_latent_size=32,
+    )
+
+
+def _make_tracking_language_model() -> MagicMock:
+    """Return a mock LM that records forward() kwargs."""
+    calls: list[dict] = []
+    hidden_size = HIDDEN_SIZE
+
+    def forward(*_args, **kwargs):
+        calls.append(dict(kwargs))
+        if kwargs.get("return_embeddings_only"):
+            text_ids = kwargs["packed_text_ids"]
+            emb = torch.randn(text_ids.shape[0], hidden_size)
+            return BaseNavitOutputWithPast(packed_query_sequence=emb)
+
+        packed_query_sequence = kwargs.get("packed_query_sequence")
+        if packed_query_sequence is None:
+            text_ids = kwargs["packed_text_ids"]
+            packed_query_sequence = torch.randn(text_ids.shape[0], hidden_size)
+        return BaseNavitOutputWithPast(
+            packed_query_sequence=packed_query_sequence,
+            past_key_values=kwargs.get("past_key_values"),
+        )
+
+    lm = MagicMock()
+    lm.forward = forward
+    lm._calls = calls
+    return lm
+
+
+def _vae_transform(image: Image.Image) -> torch.Tensor:
+    arr = torch.from_numpy(np.array(image.convert("RGB"), dtype=np.float32))
+    arr = arr / 127.5 - 1.0
+    return arr.permute(2, 0, 1)
+
+
+@pytest.fixture
+def bagel_and_vae_input():
+    language_model = _make_tracking_language_model()
+    bagel = Bagel(language_model=language_model, vit_model=None, config=_make_bagel_config())
+
+    image = Image.new("RGB", (64, 64), color=(128, 64, 32))
+    new_token_ids = {"start_of_image": 100, "end_of_image": 101}
+
+    gen_input, _newlens, _new_rope = bagel.prepare_vae_images(
+        curr_kvlens=[0],
+        curr_rope=[0],
+        images=[image],
+        transforms=_vae_transform,
+        new_token_ids=new_token_ids,
+    )
+
+    vae_model = MagicMock()
+    patch_shapes = gen_input["patchified_vae_latent_shapes"]
+    p = bagel.latent_patch_size
+    c = bagel.latent_channel
+
+    def encode(padded_images: torch.Tensor) -> torch.Tensor:
+        # Match AutoEncoder.encode layout: [B, z_channels, h*p, w*p] per image.
+        batch = padded_images.shape[0]
+        assert batch == len(patch_shapes)
+        latents = [torch.randn(c, h * p, w * p) for h, w in patch_shapes]
+        return torch.stack(latents, dim=0)
+
+    vae_model.encode = encode
+
+    expected_seq_len = int(gen_input["packed_seqlens"].sum().item())
+    num_text_tokens = int(gen_input["packed_text_ids"].numel())
+
+    return SimpleNamespace(
+        bagel=bagel,
+        language_model=language_model,
+        vae_model=vae_model,
+        gen_input=gen_input,
+        past_key_values=NaiveCache(NUM_LAYERS),
+        expected_seq_len=expected_seq_len,
+        num_text_tokens=num_text_tokens,
+    )
+
+
+class TestForwardCacheUpdateVae:
+    """Regression tests for single-stage img2img VAE cache prefill."""
+
+    def test_forwards_full_packed_query_sequence(self, bagel_and_vae_input):
+        """Update forward must use packed_query_sequence spanning text + VAE tokens."""
+        ctx = bagel_and_vae_input
+        assert ctx.expected_seq_len > ctx.num_text_tokens, "test setup should include many VAE tokens"
+
+        ctx.bagel.forward_cache_update_vae(
+            ctx.vae_model,
+            ctx.past_key_values,
+            **ctx.gen_input,
+        )
+
+        update_calls = [c for c in ctx.language_model._calls if not c.get("return_embeddings_only")]
+        assert len(update_calls) == 1
+
+        call = update_calls[0]
+        assert "packed_query_sequence" in call
+        assert "packed_text_ids" not in call
+
+        packed = call["packed_query_sequence"]
+        assert packed.shape[0] == ctx.expected_seq_len
+        assert call.get("mode") == "gen"
+
+    def test_packed_sequence_covers_text_and_vae_index_ranges(self, bagel_and_vae_input):
+        """Text and VAE index ranges must lie within packed_query_sequence."""
+        ctx = bagel_and_vae_input
+        ctx.bagel.forward_cache_update_vae(
+            ctx.vae_model,
+            ctx.past_key_values,
+            **ctx.gen_input,
+        )
+
+        update_call = next(c for c in ctx.language_model._calls if not c.get("return_embeddings_only"))
+        packed = update_call["packed_query_sequence"]
+
+        text_idx = ctx.gen_input["packed_text_indexes"]
+        vae_idx = ctx.gen_input["packed_vae_token_indexes"]
+        assert int(text_idx.max()) < packed.shape[0]
+        assert int(vae_idx.max()) < packed.shape[0]
+        assert int(text_idx.min()) >= 0
+        assert int(vae_idx.min()) >= 0
diff --git a/tests/diffusion/models/bagel/test_trajectory_recording.py b/tests/diffusion/models/bagel/test_trajectory_recording.py
index a5ceb0cc2..730b918e4 100644
--- a/tests/diffusion/models/bagel/test_trajectory_recording.py
+++ b/tests/diffusion/models/bagel/test_trajectory_recording.py
@@ -56,10 +56,7 @@ def _make_generate_args(num_tokens=NUM_TOKENS, hidden_dim=HIDDEN_DIM, cfg=False)
         packed_vae_token_indexes=torch.arange(2, seq_len, dtype=torch.long),
         packed_seqlens=torch.tensor([seq_len], dtype=torch.int),
         packed_position_ids=torch.arange(seq_len, dtype=torch.long),
-        packed_indexes=torch.arange(seq_len, dtype=torch.long),
         past_key_values=NaiveCache(1),
-        key_values_lens=torch.tensor([0], dtype=torch.int),
-        packed_key_value_indexes=torch.zeros(0, dtype=torch.long),
         num_timesteps=NUM_TIMESTEPS,
         timestep_shift=1.0,
         cfg_text_scale=1.0,
@@ -68,11 +65,8 @@ def _make_generate_args(num_tokens=NUM_TOKENS, hidden_dim=HIDDEN_DIM, cfg=False)
     if cfg:
         base |= dict(
             cfg_text_scale=4.0,
-            cfg_text_packed_query_indexes=torch.arange(seq_len, dtype=torch.long),
             cfg_text_packed_position_ids=torch.arange(seq_len, dtype=torch.long),
             cfg_text_past_key_values=NaiveCache(1),
-            cfg_text_key_values_lens=torch.tensor([0], dtype=torch.int),
-            cfg_text_packed_key_value_indexes=torch.zeros(0, dtype=torch.long),
         )
     return base
 
diff --git a/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py b/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py
index f7ada6ebf..1365a3477 100644
--- a/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py
+++ b/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py
@@ -48,6 +48,7 @@ def _make_request(prompts=None, **sp_kwargs) -> OmniDiffusionRequest:
     return OmniDiffusionRequest(
         prompts=prompts or [{"prompt": "a cat dancing"}],
         sampling_params=sp,
+        request_id="dmd2-sanitize",
     )
 
 
diff --git a/tests/diffusion/models/dmd2/test_dmd2_scheduler.py b/tests/diffusion/models/dmd2/test_dmd2_scheduler.py
index 7aac36970..e5a1b6d83 100644
--- a/tests/diffusion/models/dmd2/test_dmd2_scheduler.py
+++ b/tests/diffusion/models/dmd2/test_dmd2_scheduler.py
@@ -48,7 +48,7 @@ def _make_pipeline(cls):
 
 def _make_request(**sp_kwargs) -> OmniDiffusionRequest:
     sp = OmniDiffusionSamplingParams(**sp_kwargs)
-    return OmniDiffusionRequest(prompts=[{"prompt": "a cat"}], sampling_params=sp)
+    return OmniDiffusionRequest(prompts=[{"prompt": "a cat"}], sampling_params=sp, request_id="dmd2-scheduler")
 
 
 @pytest.fixture(
diff --git a/tests/diffusion/models/ernie_image/test_ernie_image_pe.py b/tests/diffusion/models/ernie_image/test_ernie_image_pe.py
index d1dc90715..a632ab3be 100644
--- a/tests/diffusion/models/ernie_image/test_ernie_image_pe.py
+++ b/tests/diffusion/models/ernie_image/test_ernie_image_pe.py
@@ -73,7 +73,7 @@ def test_should_apply_pe_respects_sampling_extra_args():
 
 
 def test_should_apply_pe_disables_dummy_warmup_request():
-    req = SimpleNamespace(request_ids=["dummy_req_id"], sampling_params=SimpleNamespace(extra_args={}))
+    req = SimpleNamespace(is_dummy_run=lambda: True, sampling_params=SimpleNamespace(extra_args={}))
 
     assert ErnieImagePipeline._should_apply_pe(req) is False
 
diff --git a/tests/diffusion/models/glm_image/test_glm_image_quantization.py b/tests/diffusion/models/glm_image/test_glm_image_quantization.py
new file mode 100644
index 000000000..8ab046f0c
--- /dev/null
+++ b/tests/diffusion/models/glm_image/test_glm_image_quantization.py
@@ -0,0 +1,741 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for GLM-Image quantization support (W4A16/AutoRound).
+
+These tests verify that the GLM-Image DiT transformer correctly accepts and uses
+quantization configs for W4A16/AutoRound quantization support.
+"""
+
+import os
+
+import pytest
+import torch
+from pytest_mock import MockerFixture
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
+    ColumnParallelGELU,
+    ColumnParallelSiLU,
+    GlmImageAdaLayerNormContinuous,
+    GlmImageAdaLayerNormZero,
+    GlmImageAttention,
+    GlmImageFeedForward,
+    GlmImageImageProjector,
+    GlmImagePrepare,
+    GlmImageRotaryPosEmbed,
+    GlmImageTransformer2DModel,
+    GlmImageTransformerBlock,
+    _positive_divisors,
+    validate_glm_image_tp_constraints,
+)
+from vllm_omni.model_executor.models.glm_image.pipeline import GLM_IMAGE_PIPELINE
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture(scope="function", autouse=True)
+def _init_distributed():
+    """Initialize the minimal distributed environment required by
+    ReplicatedLinear (tensor-parallel group must exist)."""
+    from vllm.distributed.parallel_state import (
+        cleanup_dist_env_and_memory,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29502")
+    init_distributed_environment(
+        world_size=1,
+        rank=0,
+        local_rank=0,
+        distributed_init_method="env://",
+    )
+    initialize_model_parallel()
+    yield
+    cleanup_dist_env_and_memory()
+
+
+@pytest.fixture(scope="function", autouse=True)
+def _force_default_gemm(monkeypatch):
+    """Force CPU-compatible GEMM dispatch for tests using CPU tensors.
+
+    vLLM's dispatch_unquantized_gemm() selects the backend by platform
+    (e.g. rocm_unquantized_gemm on AMD machines), not by tensor device.
+    CPU test tensors crash with NotImplementedError on ROCm.  Monkeypatch
+    the dispatcher to always return the default (torch.nn.functional.linear)
+    implementation which works on any device."""
+    from vllm.model_executor.layers.utils import default_unquantized_gemm
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
+        lambda: default_unquantized_gemm,
+    )
+
+
+@pytest.fixture(scope="function", autouse=True)
+def setup_mocks(mocker: MockerFixture):
+    """Set up common mocks for all tests."""
+    mocker.patch(
+        "vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size",
+        return_value=1,
+    )
+    mock_get_tp_group = mocker.patch("vllm.distributed.parallel_state.get_tp_group")
+    mock_tp_group = mocker.MagicMock()
+    mock_tp_group.world_size = 1
+    mock_get_tp_group.return_value = mock_tp_group
+    yield
+
+
+class TestPositiveDivisors:
+    """Test _positive_divisors helper function."""
+
+    def test_divisors_of_1(self):
+        assert _positive_divisors(1) == {1}
+
+    def test_divisors_of_12(self):
+        assert _positive_divisors(12) == {1, 2, 3, 4, 6, 12}
+
+    def test_divisors_of_prime(self):
+        assert _positive_divisors(7) == {1, 7}
+
+    def test_divisors_of_0_returns_empty(self):
+        assert _positive_divisors(0) == set()
+
+    def test_divisors_of_negative_returns_empty(self):
+        assert _positive_divisors(-5) == set()
+
+
+class TestValidateGlmImageTpConstraints:
+    """Test TP constraint validation for GLM-Image."""
+
+    def test_valid_tp_size_1(self):
+        """TP=1 should always be valid."""
+        result = validate_glm_image_tp_constraints(
+            dim=2560,
+            num_heads=64,
+            ffn_hidden_dim=10240,
+            tensor_parallel_size=1,
+        )
+        assert 1 in result
+
+    def test_valid_tp_size_2_for_divisible_dim(self):
+        """TP=2 is valid when all dims are divisible by 2."""
+        result = validate_glm_image_tp_constraints(
+            dim=2560,
+            num_heads=64,
+            ffn_hidden_dim=10240,
+            tensor_parallel_size=2,
+        )
+        assert 1 in result
+        assert 2 in result
+
+    def test_valid_tp_size_4_for_divisible_dim(self):
+        """TP=4 is valid when all dims are divisible by 4."""
+        result = validate_glm_image_tp_constraints(
+            dim=2560,
+            num_heads=64,
+            ffn_hidden_dim=10240,
+            tensor_parallel_size=4,
+        )
+        assert 1 in result
+        assert 2 in result
+        assert 4 in result
+
+    def test_invalid_tp_size_3_for_divisible_dim(self):
+        """TP=3 is invalid when dim is not divisible by 3."""
+        with pytest.raises(ValueError, match="dim % tensor_parallel_size == 0"):
+            validate_glm_image_tp_constraints(
+                dim=2560,  # 2560 % 3 != 0
+                num_heads=64,
+                ffn_hidden_dim=10240,
+                tensor_parallel_size=3,
+            )
+
+    def test_invalid_tp_size_zero(self):
+        """TP=0 should raise error."""
+        with pytest.raises(ValueError, match="tensor_parallel_size must be > 0"):
+            validate_glm_image_tp_constraints(
+                dim=2560,
+                num_heads=64,
+                ffn_hidden_dim=10240,
+                tensor_parallel_size=0,
+            )
+
+    def test_invalid_tp_size_negative(self):
+        """Negative TP size should raise error."""
+        with pytest.raises(ValueError, match="tensor_parallel_size must be > 0"):
+            validate_glm_image_tp_constraints(
+                dim=2560,
+                num_heads=64,
+                ffn_hidden_dim=10240,
+                tensor_parallel_size=-1,
+            )
+
+
+class TestGlmImageAdaLayerNormZeroQuantization:
+    """Test GlmImageAdaLayerNormZero with quantization config."""
+
+    def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
+        """Verify the class accepts quant_config parameter."""
+        mock_quant_config = mocker.MagicMock()
+        layer = GlmImageAdaLayerNormZero(
+            embedding_dim=512,
+            dim=2560,
+            quant_config=mock_quant_config,
+            prefix="test.norm1",
+        )
+        assert layer.linear.quant_config is mock_quant_config
+
+    def test_accepts_none_quant_config(self):
+        """Verify quant_config=None is accepted."""
+        layer = GlmImageAdaLayerNormZero(
+            embedding_dim=512,
+            dim=2560,
+            quant_config=None,
+        )
+        assert layer.linear.quant_config is None
+
+    def test_forward_handles_tuple_return_from_linear(self):
+        """Verify forward handles tuple returns from ReplicatedLinear."""
+        layer = GlmImageAdaLayerNormZero(
+            embedding_dim=512,
+            dim=2560,
+            quant_config=None,
+        )
+
+        batch_size = 2
+        seq_len = 10
+        hidden_states = torch.randn(batch_size, seq_len, 2560)
+        encoder_hidden_states = torch.randn(batch_size, seq_len, 2560)
+        temb = torch.randn(batch_size, 512)
+
+        # This should work regardless of whether linear returns tuple or tensor
+        result = layer(hidden_states, encoder_hidden_states, temb)
+        assert len(result) == 10  # Should return 10 chunks
+
+
+class TestGlmImageAdaLayerNormContinuousQuantization:
+    """Test GlmImageAdaLayerNormContinuous with quantization config."""
+
+    def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
+        """Verify the class accepts quant_config parameter."""
+        mock_quant_config = mocker.MagicMock()
+        layer = GlmImageAdaLayerNormContinuous(
+            embedding_dim=2560,
+            conditioning_embedding_dim=512,
+            quant_config=mock_quant_config,
+            prefix="test.norm_out",
+        )
+        assert layer.linear.quant_config is mock_quant_config
+
+    def test_accepts_none_quant_config(self):
+        """Verify quant_config=None is accepted."""
+        layer = GlmImageAdaLayerNormContinuous(
+            embedding_dim=2560,
+            conditioning_embedding_dim=512,
+            quant_config=None,
+        )
+        assert layer.linear.quant_config is None
+
+    def test_forward_handles_tuple_return_from_linear(self):
+        """Verify forward handles tuple returns from ReplicatedLinear."""
+        layer = GlmImageAdaLayerNormContinuous(
+            embedding_dim=2560,
+            conditioning_embedding_dim=512,
+            quant_config=None,
+        )
+
+        batch_size = 2
+        seq_len = 10
+        x = torch.randn(batch_size, seq_len, 2560)
+        conditioning_embedding = torch.randn(batch_size, 512)
+
+        result = layer(x, conditioning_embedding)
+        assert result.shape == (batch_size, seq_len, 2560)
+
+
+class TestGlmImageAttentionQuantization:
+    """Test GlmImageAttention with quantization config."""
+
+    def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
+        """Verify GlmImageAttention accepts quant_config parameter."""
+        mock_quant_config = mocker.MagicMock()
+        attn = GlmImageAttention(
+            dim=2560,
+            num_heads=64,
+            head_dim=40,
+            quant_config=mock_quant_config,
+            prefix="test.attn1",
+        )
+        assert attn.to_qkv.quant_config is mock_quant_config
+
+    def test_accepts_none_quant_config(self):
+        """Verify quant_config=None is accepted."""
+        attn = GlmImageAttention(
+            dim=2560,
+            num_heads=64,
+            head_dim=40,
+            quant_config=None,
+        )
+        assert attn.to_qkv.quant_config is None
+
+
+class TestColumnParallelModulesQuantization:
+    """Test ColumnParallelGELU and ColumnParallelSiLU with quantization config."""
+
+    def test_column_parallel_gelu_accepts_quant_config(self, mocker: MockerFixture):
+        """Verify ColumnParallelGELU accepts quant_config."""
+        mock_quant_config = mocker.MagicMock()
+        layer = ColumnParallelGELU(
+            dim_in=2560,
+            dim_out=10240,
+            quant_config=mock_quant_config,
+            prefix="test.gelu",
+        )
+        assert layer.proj.quant_config is mock_quant_config
+
+    def test_column_parallel_silu_accepts_quant_config(self, mocker: MockerFixture):
+        """Verify ColumnParallelSiLU accepts quant_config."""
+        mock_quant_config = mocker.MagicMock()
+        layer = ColumnParallelSiLU(
+            dim_in=2560,
+            dim_out=10240,
+            quant_config=mock_quant_config,
+            prefix="test.silu",
+        )
+        assert layer.proj.quant_config is mock_quant_config
+
+
+class TestGlmImageFeedForwardQuantization:
+    """Test GlmImageFeedForward with quantization config."""
+
+    def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
+        """Verify GlmImageFeedForward accepts quant_config parameter."""
+        mock_quant_config = mocker.MagicMock()
+        ff = GlmImageFeedForward(
+            dim=2560,
+            dim_out=2560,
+            inner_dim=10240,
+            activation_fn="gelu-approximate",
+            quant_config=mock_quant_config,
+            prefix="test.ff",
+        )
+        # Check that the first layer (ColumnParallelGELU) has quant_config
+        gelu_layer = ff.net[0]
+        assert gelu_layer.proj.quant_config is mock_quant_config
+
+    def test_accepts_none_quant_config(self):
+        """Verify quant_config=None is accepted."""
+        ff = GlmImageFeedForward(
+            dim=2560,
+            dim_out=2560,
+            inner_dim=10240,
+            activation_fn="gelu-approximate",
+            quant_config=None,
+        )
+        gelu_layer = ff.net[0]
+        assert gelu_layer.proj.quant_config is None
+
+    def test_linear_silu_activation(self):
+        """Test linear-silu activation function initialization works."""
+        # This test verifies the module can be instantiated with linear-silu activation.
+        # Full forward testing requires proper TP group setup, so we just verify construction.
+        ff = GlmImageFeedForward(
+            dim=2560,
+            dim_out=2560,
+            inner_dim=10240,
+            activation_fn="linear-silu",
+            quant_config=None,
+        )
+        # Verify the FFN has the correct structure
+        assert len(ff.net) == 3  # ColumnParallelSiLU, Identity, RowParallelLinear
+        assert isinstance(ff.net[0], ColumnParallelSiLU)
+
+
+class TestGlmImageTransformerBlockQuantization:
+    """Test GlmImageTransformerBlock with quantization config."""
+
+    def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
+        """Verify GlmImageTransformerBlock accepts quant_config parameter."""
+        mock_quant_config = mocker.MagicMock()
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+        block = GlmImageTransformerBlock(
+            dim=2560,
+            num_attention_heads=64,
+            attention_head_dim=40,
+            time_embed_dim=512,
+            parallel_config=parallel_config,
+            quant_config=mock_quant_config,
+            prefix="test.block",
+        )
+        # Check that inner modules have quant_config
+        assert block.norm1.linear.quant_config is mock_quant_config
+        assert block.attn1.to_qkv.quant_config is mock_quant_config
+        assert block.ff.net[0].proj.quant_config is mock_quant_config
+
+    def test_accepts_none_quant_config(self):
+        """Verify quant_config=None is accepted."""
+        block = GlmImageTransformerBlock(
+            dim=2560,
+            num_attention_heads=64,
+            attention_head_dim=40,
+            time_embed_dim=512,
+            quant_config=None,
+        )
+        assert block.norm1.linear.quant_config is None
+
+
+class TestGlmImagePrepareModule:
+    """Test GlmImagePrepare module."""
+
+    def test_prepare_module_exists(self):
+        """Verify GlmImagePrepare module exists and works."""
+        projector = GlmImageImageProjector(in_channels=16, hidden_size=2560, patch_size=2)
+        rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2)
+        prepare = GlmImagePrepare(
+            image_projector=projector,
+            rope=rope,
+            patch_size=2,
+        )
+
+        # Test forward pass
+        hidden_states = torch.randn(1, 16, 64, 64)  # [B, C, H, W]
+        result = prepare(hidden_states)
+
+        assert len(result) == 5
+        hidden_out, rope_cos, rope_sin, height, width = result
+        assert hidden_out.shape[0] == 1  # batch size
+        assert hidden_out.shape[1] == 1024  # seq_len (32 * 32)
+        assert hidden_out.shape[2] == 2560  # hidden_dim
+
+
+class TestGlmImagePipelineConfig:
+    """Test GLM_IMAGE_PIPELINE configuration."""
+
+    def test_glm_image_pipeline_config_exists(self):
+        """Verify GLM_IMAGE_PIPELINE config exists."""
+        assert GLM_IMAGE_PIPELINE is not None
+
+    def test_pipeline_has_correct_model_type(self):
+        """Verify pipeline has correct model_type."""
+        assert GLM_IMAGE_PIPELINE.model_type == "glm_image"
+
+    def test_pipeline_has_correct_model_arch(self):
+        """Verify pipeline has correct model_arch."""
+        assert GLM_IMAGE_PIPELINE.model_arch == "GlmImageForConditionalGeneration"
+
+    def test_pipeline_has_two_stages(self):
+        """Verify pipeline has two stages (AR + DiT)."""
+        assert len(GLM_IMAGE_PIPELINE.stages) == 2
+
+    def test_stage_0_is_ar_llm(self):
+        """Verify stage 0 is the AR LLM stage."""
+        stage_0 = GLM_IMAGE_PIPELINE.stages[0]
+        assert stage_0.stage_id == 0
+        assert stage_0.model_stage == "ar"
+        assert stage_0.execution_type.value == "llm_ar"
+        assert stage_0.owns_tokenizer is True
+
+    def test_stage_1_is_diffusion(self):
+        """Verify stage 1 is the DiT diffusion stage."""
+        stage_1 = GLM_IMAGE_PIPELINE.stages[1]
+        assert stage_1.stage_id == 1
+        assert stage_1.model_stage == "dit"
+        assert stage_1.execution_type.value == "diffusion"
+        assert stage_1.final_output is True
+        assert stage_1.final_output_type == "image"
+        assert stage_1.input_sources == (0,)  # Takes input from stage 0
+
+    def test_pipeline_has_diffusers_class_name(self):
+        """Verify pipeline has diffusers_class_name."""
+        assert GLM_IMAGE_PIPELINE.diffusers_class_name == "GlmImagePipeline"
+
+
+class TestGlmImageImageProjector:
+    """Test GlmImageImageProjector module."""
+
+    def test_projector_output_shape(self):
+        """Verify projector produces correct output shape."""
+        projector = GlmImageImageProjector(
+            in_channels=16,
+            hidden_size=2560,
+            patch_size=2,
+        )
+
+        # Input: [B, C, H, W] = [1, 16, 64, 64]
+        hidden_states = torch.randn(1, 16, 64, 64)
+        output = projector(hidden_states)
+
+        # After patchify: [B, H/2*W/2, D] = [1, 32*32, 2560] = [1, 1024, 2560]
+        assert output.shape == (1, 1024, 2560)
+
+    def test_projector_with_different_patch_sizes(self):
+        """Verify projector works with different patch sizes."""
+        for patch_size in [2, 4]:
+            projector = GlmImageImageProjector(
+                in_channels=16,
+                hidden_size=2560,
+                patch_size=patch_size,
+            )
+            h, w = 64, 64
+            hidden_states = torch.randn(1, 16, h, w)
+            output = projector(hidden_states)
+            expected_seq_len = (h // patch_size) * (w // patch_size)
+            assert output.shape[1] == expected_seq_len
+
+
+class TestGlmImageRotaryPosEmbed:
+    """Test GlmImageRotaryPosEmbed module."""
+
+    def test_rope_output_shape(self):
+        """Verify RoPE produces correct output shape."""
+        rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2, theta=10000.0)
+        hidden_states = torch.randn(1, 16, 64, 64)  # [B, C, H, W]
+
+        cos, sin = rope(hidden_states)
+
+        # After patchify: height=32, width=32
+        # Output: [32*32, dim] = [1024, 40]
+        assert cos.shape[0] == 1024
+        assert cos.shape[1] == 40
+        assert sin.shape == cos.shape
+
+    def test_rope_value_range(self):
+        """Verify RoPE produces values in valid range."""
+        rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2, theta=10000.0)
+        hidden_states = torch.randn(2, 16, 32, 32)
+
+        cos, sin = rope(hidden_states)
+
+        # cos and sin should be in [-1, 1]
+        assert cos.min() >= -1.0 and cos.max() <= 1.0
+        assert sin.min() >= -1.0 and sin.max() <= 1.0
+
+    def test_rope_consistency(self):
+        """Verify RoPE is consistent for same input."""
+        rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2, theta=10000.0)
+        hidden_states = torch.randn(1, 16, 32, 32)
+
+        cos1, sin1 = rope(hidden_states)
+        cos2, sin2 = rope(hidden_states)
+
+        # Same input should produce same output
+        assert torch.allclose(cos1, cos2)
+        assert torch.allclose(sin1, sin2)
+
+
+class TestGlmImageTransformer2DModelQuantization:
+    """Test GlmImageTransformer2DModel with quantization config."""
+
+    def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
+        """Verify the model accepts quant_config parameter."""
+        from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+        mock_quant_config = mocker.MagicMock()
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+
+        # Create a minimal mock od_config
+        mock_tf_config = mocker.MagicMock()
+        mock_tf_config.patch_size = 2
+        mock_tf_config.in_channels = 16
+        mock_tf_config.out_channels = 16
+        mock_tf_config.num_attention_heads = 64
+        mock_tf_config.attention_head_dim = 40
+        mock_tf_config.time_embed_dim = 512
+        mock_tf_config.condition_dim = 256
+        mock_tf_config.prior_vq_quantizer_codebook_size = 16384
+        mock_tf_config.text_embed_dim = 1024
+        mock_tf_config.num_layers = 2  # Small number for testing
+
+        mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
+        mock_od_config.tf_model_config = mock_tf_config
+        mock_od_config.parallel_config = parallel_config
+
+        model = GlmImageTransformer2DModel(
+            od_config=mock_od_config,
+            quant_config=mock_quant_config,
+        )
+
+        # Check that quantization config was passed to transformer blocks
+        for block in model.transformer_blocks:
+            assert block.norm1.linear.quant_config is mock_quant_config
+            assert block.attn1.to_qkv.quant_config is mock_quant_config
+            assert block.ff.net[0].proj.quant_config is mock_quant_config
+
+    def test_accepts_none_quant_config(self, mocker: MockerFixture):
+        """Verify quant_config=None is accepted."""
+        from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+
+        mock_tf_config = mocker.MagicMock()
+        mock_tf_config.patch_size = 2
+        mock_tf_config.in_channels = 16
+        mock_tf_config.out_channels = 16
+        mock_tf_config.num_attention_heads = 64
+        mock_tf_config.attention_head_dim = 40
+        mock_tf_config.time_embed_dim = 512
+        mock_tf_config.condition_dim = 256
+        mock_tf_config.prior_vq_quantizer_codebook_size = 16384
+        mock_tf_config.text_embed_dim = 1024
+        mock_tf_config.num_layers = 2
+
+        mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
+        mock_od_config.tf_model_config = mock_tf_config
+        mock_od_config.parallel_config = parallel_config
+
+        model = GlmImageTransformer2DModel(
+            od_config=mock_od_config,
+            quant_config=None,
+        )
+
+        # Check that quantization config is None
+        for block in model.transformer_blocks:
+            assert block.norm1.linear.quant_config is None
+            assert block.attn1.to_qkv.quant_config is None
+
+    def test_norm_out_has_no_quantization(self, mocker: MockerFixture):
+        """Verify norm_out (output layer) does NOT use quantization to preserve precision."""
+        from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+        mock_quant_config = mocker.MagicMock()
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+
+        mock_tf_config = mocker.MagicMock()
+        mock_tf_config.patch_size = 2
+        mock_tf_config.in_channels = 16
+        mock_tf_config.out_channels = 16
+        mock_tf_config.num_attention_heads = 64
+        mock_tf_config.attention_head_dim = 40
+        mock_tf_config.time_embed_dim = 512
+        mock_tf_config.condition_dim = 256
+        mock_tf_config.prior_vq_quantizer_codebook_size = 16384
+        mock_tf_config.text_embed_dim = 1024
+        mock_tf_config.num_layers = 2
+
+        mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
+        mock_od_config.tf_model_config = mock_tf_config
+        mock_od_config.parallel_config = parallel_config
+
+        model = GlmImageTransformer2DModel(
+            od_config=mock_od_config,
+            quant_config=mock_quant_config,
+        )
+
+        # norm_out.linear should NOT have quant_config to preserve output precision
+        assert model.norm_out.linear.quant_config is None
+
+    def test_model_has_sp_plan(self, mocker: MockerFixture):
+        """Verify model has _sp_plan defined for sequence parallelism."""
+        from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+
+        mock_tf_config = mocker.MagicMock()
+        mock_tf_config.patch_size = 2
+        mock_tf_config.in_channels = 16
+        mock_tf_config.out_channels = 16
+        mock_tf_config.num_attention_heads = 64
+        mock_tf_config.attention_head_dim = 40
+        mock_tf_config.time_embed_dim = 512
+        mock_tf_config.condition_dim = 256
+        mock_tf_config.prior_vq_quantizer_codebook_size = 16384
+        mock_tf_config.text_embed_dim = 1024
+        mock_tf_config.num_layers = 2
+
+        mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
+        mock_od_config.tf_model_config = mock_tf_config
+        mock_od_config.parallel_config = parallel_config
+
+        model = GlmImageTransformer2DModel(
+            od_config=mock_od_config,
+            quant_config=None,
+        )
+
+        # Verify _sp_plan exists
+        assert hasattr(model, "_sp_plan")
+        assert "prepare" in model._sp_plan
+        assert "proj_out" in model._sp_plan
+
+    def test_model_has_hsdp_shard_conditions(self, mocker: MockerFixture):
+        """Verify model has _hsdp_shard_conditions for HSDP parallelism."""
+        from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+
+        mock_tf_config = mocker.MagicMock()
+        mock_tf_config.patch_size = 2
+        mock_tf_config.in_channels = 16
+        mock_tf_config.out_channels = 16
+        mock_tf_config.num_attention_heads = 64
+        mock_tf_config.attention_head_dim = 40
+        mock_tf_config.time_embed_dim = 512
+        mock_tf_config.condition_dim = 256
+        mock_tf_config.prior_vq_quantizer_codebook_size = 16384
+        mock_tf_config.text_embed_dim = 1024
+        mock_tf_config.num_layers = 2
+
+        mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
+        mock_od_config.tf_model_config = mock_tf_config
+        mock_od_config.parallel_config = parallel_config
+
+        model = GlmImageTransformer2DModel(
+            od_config=mock_od_config,
+            quant_config=None,
+        )
+
+        assert hasattr(model, "_hsdp_shard_conditions")
+        assert len(model._hsdp_shard_conditions) > 0
+
+    def test_model_creates_kv_cache(self, mocker: MockerFixture):
+        """Verify model can create KV cache for image editing."""
+        from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+
+        mock_tf_config = mocker.MagicMock()
+        mock_tf_config.patch_size = 2
+        mock_tf_config.in_channels = 16
+        mock_tf_config.out_channels = 16
+        mock_tf_config.num_attention_heads = 64
+        mock_tf_config.attention_head_dim = 40
+        mock_tf_config.time_embed_dim = 512
+        mock_tf_config.condition_dim = 256
+        mock_tf_config.prior_vq_quantizer_codebook_size = 16384
+        mock_tf_config.text_embed_dim = 1024
+        mock_tf_config.num_layers = 2
+
+        mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
+        mock_od_config.tf_model_config = mock_tf_config
+        mock_od_config.parallel_config = parallel_config
+
+        model = GlmImageTransformer2DModel(
+            od_config=mock_od_config,
+            quant_config=None,
+        )
+
+        kv_cache = model.create_kv_cache()
+        assert kv_cache is not None
+        assert len(kv_cache) == 2  # num_layers
diff --git a/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py b/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py
index 342103350..8a35d8aaa 100644
--- a/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py
+++ b/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py
@@ -23,6 +23,7 @@ NUM_HEADS = 4
 NUM_KV_HEADS = 2
 HEAD_DIM = 16
 IMAGE_TOKEN_LEN = 8
+SUFFIX_TOKEN_LEN = 3
 SCALING = 1.0 / math.sqrt(HEAD_DIM)
 
 
@@ -79,6 +80,10 @@ def _make_known_kv(num_tokens, base=0.0):
     return k, v
 
 
+def _gen_timestep_index(bs: int, current_start: int) -> torch.Tensor:
+    return torch.full((bs, 1), current_start, dtype=torch.long)
+
+
 def _call_mgr(
     mgr,
     bs,
@@ -90,6 +95,8 @@ def _call_mgr(
     uncond_cfg_prefill=False,
     num_image_tokens=IMAGE_TOKEN_LEN,
     shard_image_size=None,
+    gen_timestep_scatter_index=None,
+    position_ids=None,
 ):
     query = torch.randn(bs * q_len, NUM_HEADS, HEAD_DIM)
     attn_mask = torch.zeros(bs, 1, seq_len, seq_len)
@@ -104,6 +111,8 @@ def _call_mgr(
         uncond_cfg_prefill=uncond_cfg_prefill,
         num_image_tokens=num_image_tokens,
         shard_image_size=shard_image_size,
+        gen_timestep_scatter_index=gen_timestep_scatter_index,
+        position_ids=position_ids,
     )
 
 
@@ -117,54 +126,65 @@ def test_no_ar_kv(bs):
     """
     No AR KV injected. Tests the basic first_step cache → update reuse path.
 
-    Sequence layout per batch on first_step (q_len=12, IMAGE_TOKEN_LEN=8):
-        [prompt(3) | image(8) | eoi(1)]
-        image_size = IMAGE_TOKEN_LEN + 1 = 9
-        cached_prompt_len = seq_len - image_size = 12 - 9 = 3
+    Sequence layout per batch on first_step (q_len=14, IMAGE_TOKEN_LEN=8):
+        [prompt(3) | current timestep/image(8) | suffix(3)]
+        gen_timestep_scatter_index = 3
+        cached_prompt_len = 3
 
     After first_step:
-        image_kv_cache_map stores prompt tokens [0:3] for each batch (total 3*bs tokens flat).
+        image_kv_cache_map stores prompt tokens [0:3] for each batch.
 
-    Update step (q_len=IMAGE_TOKEN_LEN=8, seq_len = cached_prompt(3) + eoi(1) + image(8) = 12):
-        _reuse_prompt_kv produces [cached_prompt(3) | new_image(8) | zero_eoi(1)] per batch.
+    Update step (q_len=IMAGE_TOKEN_LEN=8, seq_len = cached_prompt(3) + current(8) = 11):
+        _reuse_prompt_kv produces [cached_prompt(3) | new_current(8)] per batch.
     """
     mgr = _make_cache_mgr()
     assert mgr.image_kv_cache_map is None
 
     # --- first_step ---
-    q_len = 12
+    prompt_len = 3
+    q_len = prompt_len + IMAGE_TOKEN_LEN + SUFFIX_TOKEN_LEN
     k_flat, v_flat = _make_known_kv(bs * q_len, base=1.0)
 
-    _call_mgr(mgr, bs, q_len=q_len, seq_len=q_len, key_flat=k_flat, value_flat=v_flat, first_step=True)
+    _call_mgr(
+        mgr,
+        bs,
+        q_len=q_len,
+        seq_len=q_len,
+        key_flat=k_flat,
+        value_flat=v_flat,
+        first_step=True,
+        gen_timestep_scatter_index=_gen_timestep_index(bs, prompt_len),
+    )
 
     cached_key, cached_value = mgr.image_kv_cache_map
     # 3 prompt tokens cached per batch
-    assert cached_key.shape[0] == 3 * bs
+    assert cached_key.shape == (bs, prompt_len, NUM_KV_HEADS, HEAD_DIM)
+    assert torch.equal(mgr.image_kv_cache_lens, torch.full((bs,), prompt_len))
     for b in range(bs):
         flat_offset = b * q_len
-        cache_offset = b * 3
-        assert torch.allclose(cached_key[cache_offset : cache_offset + 3], k_flat[flat_offset : flat_offset + 3])
-        assert torch.allclose(cached_value[cache_offset : cache_offset + 3], v_flat[flat_offset : flat_offset + 3])
+        assert torch.allclose(cached_key[b, :prompt_len], k_flat[flat_offset : flat_offset + prompt_len])
+        assert torch.allclose(cached_value[b, :prompt_len], v_flat[flat_offset : flat_offset + prompt_len])
 
     # --- update step ---
     img_q_len = IMAGE_TOKEN_LEN
-    update_seq_len = 3 + 1 + img_q_len  # cached_prompt(3) + eoi(1) + image(8) = 12
+    update_seq_len = prompt_len + img_q_len  # cached_prompt(3) + current(8) = 11
     new_img_k, new_img_v = _make_known_kv(bs * img_q_len, base=50.0)
 
     key_input = new_img_k.reshape(bs, img_q_len, NUM_KV_HEADS, HEAD_DIM)
     val_input = new_img_v.reshape(bs, img_q_len, NUM_KV_HEADS, HEAD_DIM)
-    result_k, result_v = mgr._reuse_prompt_kv(key_input, val_input, update_seq_len, bs=bs)
+    position_ids = torch.arange(prompt_len, prompt_len + img_q_len).repeat(bs, 1)
+    result_k, result_v = mgr._reuse_prompt_kv(key_input, val_input, update_seq_len, bs=bs, position_ids=position_ids)
 
     assert result_k.shape == (bs, update_seq_len, NUM_KV_HEADS, HEAD_DIM)
     for b in range(bs):
-        cache_offset = b * 3
         img_offset = b * img_q_len
         # Cached prompt preserved
-        assert torch.allclose(result_k[b, :3], cached_key[cache_offset : cache_offset + 3])
+        assert torch.allclose(result_k[b, :prompt_len], cached_key[b, :prompt_len])
         # New image tokens
-        assert torch.allclose(result_k[b, 3 : 3 + img_q_len], new_img_k[img_offset : img_offset + img_q_len])
-        # Zero eoi at the end
-        assert torch.allclose(result_k[b, -1], torch.zeros(NUM_KV_HEADS, HEAD_DIM))
+        assert torch.allclose(
+            result_k[b, prompt_len : prompt_len + img_q_len],
+            new_img_k[img_offset : img_offset + img_q_len],
+        )
 
 
 # ============================================================
@@ -178,16 +198,17 @@ def test_ar_kv_no_cfg(sp_size):
     AR KV injected, bs=1, no CFG. Tests AR prefix prepend + cache + reuse.
 
     sp_size=1:
-        first_step input: q_len=12, ar_len=5, seq_len=17
-        Sequence layout: [ar(5) | prompt(3) | image(8) | eoi(1)]
-        image_size = IMAGE_TOKEN_LEN + 1 = 9
-        cached_prompt_len = seq_len - image_size = 17 - 9 = 8 (= ar(5) + prompt(3))
+        first_step input: q_len=14, ar_len=5, seq_len=19
+        Sequence layout: [ar(5) | prompt(3) | current(8) | suffix(3)]
+        gen_timestep_scatter_index = 3
+        cached_prompt_len = ar(5) + prompt(3) = 8
 
-        Update step: q_len=8, seq_len = cached_prompt(8) + eoi(1) + image(8) = 17
-        Result: [cached(8) | new_image(8) | zero_eoi(1)]
+        Update step: q_len=8, seq_len = cached_prompt(8) + current(8) = 16
+        Result: [cached(8) | new_current(8)]
 
     sp_size=2:
         shard_image_size = 4 (passed externally, simulating SP sharding)
+        gen_timestep_scatter_index = 8
         cached_prompt_len = seq_len - shard_image_size = 17 - 4 = 13 (= ar(5) + prompt(8))
         Note: with SP, more of the sequence is treated as "prompt" because
         image tokens are sharded across ranks.
@@ -204,15 +225,23 @@ def test_ar_kv_no_cfg(sp_size):
     mgr._injected_ar_kv = [(ar_k.clone(), ar_v.clone())]
 
     # --- first_step: call _cache_prompt_kv directly ---
-    bs, q_len = 1, 12
-    seq_len = q_len + ar_len  # 17
+    bs = 1
+    q_len = 12 if sp_size > 1 else 14
+    seq_len = q_len + ar_len
     k_raw, v_raw = _make_known_kv(q_len, base=1.0)
     key_4d = k_raw.reshape(1, q_len, NUM_KV_HEADS, HEAD_DIM)
     val_4d = v_raw.reshape(1, q_len, NUM_KV_HEADS, HEAD_DIM)
 
     shard_image_size = 4 if sp_size > 1 else None
+    current_start = 8 if sp_size > 1 else 3
     mgr.image_kv_cache_map = None
-    mgr._cache_prompt_kv(key_4d, val_4d, seq_len, shard_image_size)
+    mgr._cache_prompt_kv(
+        key_4d,
+        val_4d,
+        seq_len,
+        shard_image_size,
+        gen_timestep_scatter_index=_gen_timestep_index(bs, current_start),
+    )
 
     cached_key, cached_value = mgr.image_kv_cache_map
     if sp_size == 1:
@@ -222,25 +251,32 @@ def test_ar_kv_no_cfg(sp_size):
         # cached = seq_len - shard_image_size = 17 - 4 = 13
         expected_cached_len = seq_len - shard_image_size
 
-    assert cached_key.shape[0] == expected_cached_len
+    assert cached_key.shape == (bs, expected_cached_len, NUM_KV_HEADS, HEAD_DIM)
     # AR KV always at the front
-    assert torch.allclose(cached_key[:ar_len], ar_k)
-    assert torch.allclose(cached_value[:ar_len], ar_v)
+    assert torch.allclose(cached_key[0, :ar_len], ar_k)
+    assert torch.allclose(cached_value[0, :ar_len], ar_v)
     # Prompt tokens follow AR
     prompt_cached = expected_cached_len - ar_len
-    assert torch.allclose(cached_key[ar_len : ar_len + prompt_cached], k_raw[:prompt_cached])
+    assert torch.allclose(cached_key[0, ar_len : ar_len + prompt_cached], k_raw[:prompt_cached])
     # AR KV consumed
     assert mgr._injected_ar_kv is None
 
     # --- update step: call _reuse_prompt_kv directly ---
     if sp_size == 1:
         img_q_len = IMAGE_TOKEN_LEN
-        update_seq_len = expected_cached_len + 1 + img_q_len  # 8 + 1 + 8 = 17
+        update_seq_len = expected_cached_len + img_q_len  # 8 + 8 = 16
         new_img_k, new_img_v = _make_known_kv(img_q_len, base=50.0)
 
         key_input = new_img_k.reshape(1, img_q_len, NUM_KV_HEADS, HEAD_DIM)
         val_input = new_img_v.reshape(1, img_q_len, NUM_KV_HEADS, HEAD_DIM)
-        result_k, result_v = mgr._reuse_prompt_kv(key_input, val_input, update_seq_len, bs=1)
+        position_ids = torch.arange(expected_cached_len, expected_cached_len + img_q_len).reshape(1, img_q_len)
+        result_k, result_v = mgr._reuse_prompt_kv(
+            key_input,
+            val_input,
+            update_seq_len,
+            bs=1,
+            position_ids=position_ids,
+        )
 
         assert result_k.shape == (1, update_seq_len, NUM_KV_HEADS, HEAD_DIM)
         # AR + prompt preserved
@@ -248,8 +284,6 @@ def test_ar_kv_no_cfg(sp_size):
         assert torch.allclose(result_k[0, ar_len : ar_len + prompt_cached], k_raw[:prompt_cached])
         # New image tokens
         assert torch.allclose(result_k[0, expected_cached_len : expected_cached_len + img_q_len], new_img_k)
-        # Zero eoi
-        assert torch.allclose(result_k[0, -1], torch.zeros(NUM_KV_HEADS, HEAD_DIM))
     else:
         # SP path: _reuse_prompt_kv returns only cached prompt (no image concat)
         img_q_len = shard_image_size  # 4
@@ -292,13 +326,12 @@ def test_ar_kv_with_cfg(cfg_parallel, bs):
             → _injected_ar_kv becomes [(pos_ar(10), pos_av(10)), (neg_k(10), neg_v(10))]
 
         first_step (bs=2, q_len=12, seq_len=22):
-            Batch 0 (pos): [pos_ar(10) | prompt(3) | image(8) | eoi(1)]
-            Batch 1 (neg): [neg_ar(10) | prompt(3) | image(8) | eoi(1)]
-            cached_prompt_len per batch = 22 - 9 = 13
-            Total cached = 13 * 2 = 26
+            Batch 0 (pos): [pos_ar(10) | prompt(3) | current(8) | suffix(3)]
+            Batch 1 (neg): [neg_ar(10) | prompt(3) | current(8) | suffix(3)]
+            cached_prompt_len per batch = 10 + 3 = 13
 
-        Update (bs=2, seq_len = 13 + 1 + 8 = 22):
-            Result per batch: [cached(13) | new_image(8) | zero_eoi(1)]
+        Update (bs=2, seq_len = 13 + 8 = 21):
+            Result per batch: [cached(13) | new_current(8)]
 
     CFG Parallel (cfg_parallel=True, bs=1):
         This rank handles only the negative branch.
@@ -308,11 +341,11 @@ def test_ar_kv_with_cfg(cfg_parallel, bs):
             → _injected_ar_kv = [(neg_k(10), neg_v(10))]
 
         first_step (bs=1, q_len=12, seq_len=22):
-            [neg_ar(10) | prompt(3) | image(8) | eoi(1)]
-            cached_prompt_len = 22 - 9 = 13
+            [neg_ar(10) | prompt(3) | current(8) | suffix(3)]
+            cached_prompt_len = 13
 
-        Update (bs=1, seq_len = 13 + 1 + 8 = 22):
-            Result: [cached(13) | new_image(8) | zero_eoi(1)]
+        Update (bs=1, seq_len = 13 + 8 = 21):
+            Result: [cached(13) | new_current(8)]
     """
     positive_reuse_len = 10
     negative_reuse_len = 6
@@ -351,46 +384,62 @@ def test_ar_kv_with_cfg(cfg_parallel, bs):
         mgr._injected_ar_kv = [mgr._injected_ar_kv[1]]
 
     # --- first_step ---
-    q_len = 12
+    prompt_len = 3
+    q_len = prompt_len + IMAGE_TOKEN_LEN + SUFFIX_TOKEN_LEN
     seq_len = q_len + positive_reuse_len  # 22
     k_flat, v_flat = _make_known_kv(bs * q_len, base=1.0)
 
-    _call_mgr(mgr, bs, q_len=q_len, seq_len=seq_len, key_flat=k_flat, value_flat=v_flat, first_step=True)
+    _call_mgr(
+        mgr,
+        bs,
+        q_len=q_len,
+        seq_len=seq_len,
+        key_flat=k_flat,
+        value_flat=v_flat,
+        first_step=True,
+        gen_timestep_scatter_index=_gen_timestep_index(bs, prompt_len),
+    )
 
     cached_key, cached_value = mgr.image_kv_cache_map
-    cached_prompt_len_per_batch = seq_len - (IMAGE_TOKEN_LEN + 1)  # 22 - 9 = 13
-    assert cached_key.shape[0] == cached_prompt_len_per_batch * bs
+    cached_prompt_len_per_batch = positive_reuse_len + prompt_len
+    assert cached_key.shape == (bs, cached_prompt_len_per_batch, NUM_KV_HEADS, HEAD_DIM)
     assert mgr._injected_ar_kv is None
 
     if not cfg_parallel:
         # bs=2: batch 0 = pos, batch 1 = neg
         # Batch 0: pos_ar(10) + prompt(3)
-        assert torch.allclose(cached_key[:positive_reuse_len], pos_ar_k)
-        assert torch.allclose(cached_key[positive_reuse_len:13], k_flat[:3])
+        assert torch.allclose(cached_key[0, :positive_reuse_len], pos_ar_k)
+        assert torch.allclose(cached_key[0, positive_reuse_len:13], k_flat[:3])
         # Batch 1: neg_ar(10) + prompt(3)
-        assert torch.allclose(cached_key[13:23], neg_ar_k)
-        assert torch.allclose(cached_key[23:26], k_flat[q_len : q_len + 3])
+        assert torch.allclose(cached_key[1, :positive_reuse_len], neg_ar_k)
+        assert torch.allclose(cached_key[1, positive_reuse_len:13], k_flat[q_len : q_len + 3])
     else:
         # bs=1: only neg branch
-        assert torch.allclose(cached_key[:positive_reuse_len], neg_ar_k)
-        assert torch.allclose(cached_key[positive_reuse_len:13], k_flat[:3])
+        assert torch.allclose(cached_key[0, :positive_reuse_len], neg_ar_k)
+        assert torch.allclose(cached_key[0, positive_reuse_len:13], k_flat[:3])
 
     # --- update step ---
     img_q_len = IMAGE_TOKEN_LEN
-    update_seq_len = cached_prompt_len_per_batch + 1 + img_q_len  # 13 + 1 + 8 = 22
+    update_seq_len = cached_prompt_len_per_batch + img_q_len  # 13 + 8 = 21
     new_img_k, new_img_v = _make_known_kv(bs * img_q_len, base=50.0)
 
     key_input = new_img_k.reshape(bs, img_q_len, NUM_KV_HEADS, HEAD_DIM)
     val_input = new_img_v.reshape(bs, img_q_len, NUM_KV_HEADS, HEAD_DIM)
-    result_k, result_v = mgr._reuse_prompt_kv(key_input, val_input, update_seq_len, bs=bs)
+    position_ids = torch.arange(cached_prompt_len_per_batch, cached_prompt_len_per_batch + img_q_len).repeat(bs, 1)
+    result_k, result_v = mgr._reuse_prompt_kv(
+        key_input,
+        val_input,
+        update_seq_len,
+        bs=bs,
+        position_ids=position_ids,
+    )
 
     assert result_k.shape == (bs, update_seq_len, NUM_KV_HEADS, HEAD_DIM)
     for b in range(bs):
-        cache_offset = b * cached_prompt_len_per_batch
         # Cached prompt preserved
         assert torch.allclose(
             result_k[b, :cached_prompt_len_per_batch],
-            cached_key[cache_offset : cache_offset + cached_prompt_len_per_batch],
+            cached_key[b, :cached_prompt_len_per_batch],
         )
         # New image tokens
         img_offset = b * img_q_len
@@ -398,8 +447,6 @@ def test_ar_kv_with_cfg(cfg_parallel, bs):
             result_k[b, cached_prompt_len_per_batch : cached_prompt_len_per_batch + img_q_len],
             new_img_k[img_offset : img_offset + img_q_len],
         )
-        # Zero eoi
-        assert torch.allclose(result_k[b, -1], torch.zeros(NUM_KV_HEADS, HEAD_DIM))
 
 
 # ============================================================
@@ -412,7 +459,7 @@ def test_cross_request_isolation():
     Verify leftover image_kv_cache_map from a previous request is NOT treated as AR KV.
 
     Setup: mgr has stale cache from a prior request (9 tokens, base=999).
-    New request: first_step with q_len=12, no AR KV.
+    New request: first_step with q_len=14, no AR KV.
 
     Expected: stale cache is overwritten. New cache = prompt tokens from current request.
     The stale values (999.x) must NOT appear in the new cache.
@@ -421,20 +468,33 @@ def test_cross_request_isolation():
 
     # Simulate leftover from previous request
     leftover_k, leftover_v = _make_known_kv(9, base=999.0)
-    mgr.image_kv_cache_map = (leftover_k, leftover_v)
+    mgr.image_kv_cache_map = (
+        leftover_k.reshape(1, 9, NUM_KV_HEADS, HEAD_DIM),
+        leftover_v.reshape(1, 9, NUM_KV_HEADS, HEAD_DIM),
+    )
     assert mgr._injected_ar_kv is None
 
     # New request first_step
-    bs, seq_len = 1, 12
+    bs = 1
+    prompt_len = 3
+    seq_len = prompt_len + IMAGE_TOKEN_LEN + SUFFIX_TOKEN_LEN
     k_flat, v_flat = _make_known_kv(bs * seq_len, base=1.0)
 
-    _call_mgr(mgr, bs, q_len=seq_len, seq_len=seq_len, key_flat=k_flat, value_flat=v_flat, first_step=True)
+    _call_mgr(
+        mgr,
+        bs,
+        q_len=seq_len,
+        seq_len=seq_len,
+        key_flat=k_flat,
+        value_flat=v_flat,
+        first_step=True,
+        gen_timestep_scatter_index=_gen_timestep_index(bs, prompt_len),
+    )
 
     cached_key, cached_value = mgr.image_kv_cache_map
-    # cached_prompt_len = 12 - 9 = 3
-    assert cached_key.shape[0] == 3
+    assert cached_key.shape == (bs, prompt_len, NUM_KV_HEADS, HEAD_DIM)
     # Must be from current request, not stale
-    assert torch.allclose(cached_key[:3], k_flat[:3])
-    assert torch.allclose(cached_value[:3], v_flat[:3])
+    assert torch.allclose(cached_key[0, :prompt_len], k_flat[:prompt_len])
+    assert torch.allclose(cached_value[0, :prompt_len], v_flat[:prompt_len])
     # Stale values must not be present
     assert not torch.any(cached_key >= 999.0)
diff --git a/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py b/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py
index 7c3256eee..7cb7ea83a 100644
--- a/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py
+++ b/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py
@@ -124,6 +124,58 @@ def test_resolve_stop_token_ids_image_tasks_stop_on_ratio_range():
     assert resolve_stop_token_ids(task="i2t", bot_task=None, tokenizer=tok) == [answer_id]
     assert resolve_stop_token_ids(task="t2t", bot_task=None, tokenizer=tok) == [answer_id]
 
+    # image_size="auto" is equivalent to None (need_ratio=True).
+    for bot in ("think", "recaption", "think_recaption", "vanilla"):
+        assert resolve_stop_token_ids(task="t2i", bot_task=bot, tokenizer=tok, image_size="auto") == expected
+        assert resolve_stop_token_ids(task="it2i", bot_task=bot, tokenizer=tok, image_size="auto") == expected
+
+
+def test_resolve_stop_token_ids_explicit_image_size_stops_on_terminator():
+    """When image_size is an explicit "WxH" string, need_ratio=False: AR
+    already knows the size and stops at the cot terminator instead of the
+    ratio range.
+
+    Mirrors upstream logic where ``need_ratio`` is false when the user
+    supplies a target size, so ``final_stop_tokens`` is set to the
+    recaption/think terminator (not the ratio tokens).
+    """
+    tok = FakeTokenizer()
+
+    end_of_think_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["</think>"]
+    end_of_recaption_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["</recaption>"]
+
+    # think_recaption / recaption -> [end_of_recaption]
+    for bot in ("think_recaption", "recaption"):
+        for task in ("t2i", "it2i"):
+            result = resolve_stop_token_ids(task=task, bot_task=bot, tokenizer=tok, image_size="1024x1024")
+            assert result == [end_of_recaption_id], (
+                f"task={task}, bot_task={bot}, image_size='1024x1024': expected [{end_of_recaption_id}], got {result}"
+            )
+
+    # think -> [end_of_think, end_of_recaption]
+    for task in ("t2i", "it2i"):
+        result = resolve_stop_token_ids(task=task, bot_task="think", tokenizer=tok, image_size="1024x1024")
+        assert result == [end_of_think_id, end_of_recaption_id], (
+            f"task={task}, bot_task='think', image_size='1024x1024': "
+            f"expected [{end_of_think_id}, {end_of_recaption_id}], got {result}"
+        )
+
+    # vanilla -> no early return, falls through to ratio range
+    start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_0>"]
+    end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_32>"]
+    other_start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_33>"]
+    other_end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_36>"]
+    expected_ratio = list(range(start, end + 1)) + list(range(other_start, other_end + 1))
+    for task in ("t2i", "it2i"):
+        result = resolve_stop_token_ids(task=task, bot_task="vanilla", tokenizer=tok, image_size="1024x1024")
+        assert result == expected_ratio, (
+            f"task={task}, bot_task='vanilla', image_size='1024x1024': expected ratio range, got {result}"
+        )
+
+    # Text tasks are unaffected by image_size.
+    answer_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<answer>"]
+    assert resolve_stop_token_ids(task="i2t", bot_task=None, tokenizer=tok, image_size="1024x1024") == [answer_id]
+
 
 @pytest.mark.parametrize(
     "task,bot_task",
diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py
index 665126df7..3930677d3 100644
--- a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py
+++ b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py
@@ -15,12 +15,29 @@ These tests verify:
 import json
 import os
 import tempfile
+from types import SimpleNamespace
 
 import pytest
+import torch
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def _make_ltx23_pipeline(sequence_parallel_size: int = 1):
+    from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+
+    pipeline = object.__new__(LTX23Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.audio_vae_temporal_compression_ratio = 4
+    pipeline.audio_vae_mel_compression_ratio = 4
+    pipeline.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=sequence_parallel_size))
+    pipeline.audio_vae = SimpleNamespace(
+        latents_mean=torch.tensor(0.0),
+        latents_std=torch.tensor(1.0),
+    )
+    return pipeline
+
+
 class TestPipelineIndependence:
     """Verify LTX23Pipeline is fully independent from LTX2Pipeline."""
 
@@ -48,6 +65,313 @@ class TestPipelineIndependence:
 
         assert issubclass(LTX23Pipeline, ProgressBarMixin)
 
+    def test_ltx23_pipeline_has_cfg_parallel_mixin(self):
+        """LTX23Pipeline must use the shared CFG parallel implementation."""
+        from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+
+        assert issubclass(LTX23Pipeline, CFGParallelMixin)
+
+
+class TestCFGParallelHelpers:
+    """Test LTX-2.3 CFG helper math without loading model weights."""
+
+    def test_combine_cfg_noise_matches_x0_space_formula(self):
+        import torch
+
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+
+        pipe = object.__new__(LTX23Pipeline)
+        video_sample = torch.tensor([[[1.0, -2.0]]])
+        audio_sample = torch.tensor([[[0.5, 3.0]]])
+        video_pos = torch.tensor([[[0.2, -0.3]]])
+        video_neg = torch.tensor([[[-0.4, 0.1]]])
+        audio_pos = torch.tensor([[[0.7, -0.2]]])
+        audio_neg = torch.tensor([[[0.1, 0.4]]])
+        video_sigma = torch.tensor(0.25)
+        audio_sigma = torch.tensor(0.5)
+        scale = 4.0
+
+        video_combined, audio_combined = pipe.combine_cfg_noise(
+            (video_pos, audio_pos),
+            (video_neg, audio_neg),
+            scale,
+            video_latents=video_sample,
+            audio_latents=audio_sample,
+            video_sigma=video_sigma,
+            audio_sigma=audio_sigma,
+        )
+
+        x0_video_cond = video_sample - video_pos * video_sigma
+        x0_video_uncond = video_sample - video_neg * video_sigma
+        x0_video_guided = x0_video_cond + (scale - 1) * (x0_video_cond - x0_video_uncond)
+        expected_video = (video_sample - x0_video_guided) / video_sigma
+
+        x0_audio_cond = audio_sample - audio_pos * audio_sigma
+        x0_audio_uncond = audio_sample - audio_neg * audio_sigma
+        x0_audio_guided = x0_audio_cond + (scale - 1) * (x0_audio_cond - x0_audio_uncond)
+        expected_audio = (audio_sample - x0_audio_guided) / audio_sigma
+        assert torch.allclose(video_combined, expected_video)
+        assert torch.allclose(audio_combined, expected_audio)
+
+    def test_two_rank_cfg_parallel_smoke_uses_rank_local_branch_and_x0_formula(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
+
+        pipe = object.__new__(ltx23.LTX23Pipeline)
+        video_sample = torch.tensor([[[1.0, -2.0]]])
+        audio_sample = torch.tensor([[[0.5, 3.0, -1.0]]])
+        video_pos = torch.tensor([[[0.2, -0.3]]])
+        video_neg = torch.tensor([[[-0.4, 0.1]]])
+        audio_pos = torch.tensor([[[0.7, -0.2, 0.3]]])
+        audio_neg = torch.tensor([[[0.1, 0.4, -0.5]]])
+        video_sigma = torch.tensor(0.25)
+        audio_sigma = torch.tensor(0.5)
+        scale = 4.0
+
+        class FakeCfgGroup:
+            def all_gather(self, tensor, separate_tensors=True):
+                assert separate_tensors
+                if tensor.shape == video_pos.shape:
+                    return [video_pos, video_neg]
+                return [audio_pos, audio_neg]
+
+        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2)
+        monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup())
+
+        expected_video = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+            video_sample,
+            video_pos,
+            video_neg,
+            video_sigma,
+            scale,
+        )
+        expected_audio = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+            audio_sample,
+            audio_pos,
+            audio_neg,
+            audio_sigma,
+            scale,
+        )
+
+        for rank, expected_branch in ((0, "positive"), (1, "negative")):
+            calls = []
+            monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda rank=rank: rank)
+
+            def fake_predict_noise(**kwargs):
+                calls.append(kwargs["branch"])
+                if kwargs["branch"] == "positive":
+                    return video_pos, audio_pos
+                return video_neg, audio_neg
+
+            object.__setattr__(pipe, "predict_noise", fake_predict_noise)
+            video_combined, audio_combined = pipe.predict_noise_with_parallel_cfg(
+                true_cfg_scale=scale,
+                positive_kwargs={"branch": "positive"},
+                negative_kwargs={"branch": "negative"},
+                cfg_normalize=False,
+                video_latents=video_sample,
+                audio_latents=audio_sample,
+                video_sigma=video_sigma,
+                audio_sigma=audio_sigma,
+            )
+
+            assert calls == [expected_branch]
+            torch.testing.assert_close(video_combined, expected_video)
+            torch.testing.assert_close(audio_combined, expected_audio)
+
+        assert "_cfg_video_latents" not in pipe.__dict__
+        assert "_cfg_audio_latents" not in pipe.__dict__
+
+
+class TestCFGParallelForwardPath:
+    """Test the LTX-2.3 CFG-parallel denoising path without loading model weights."""
+
+    @pytest.mark.parametrize(("cfg_rank", "expected_prompt_value"), [(0, 1.0), (1, 0.0)])
+    def test_forward_cfg_parallel_steps_video_and_audio_scheduler(
+        self,
+        monkeypatch,
+        cfg_rank,
+        expected_prompt_value,
+    ):
+        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        pipe = object.__new__(ltx23.LTX23Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.device = torch.device("cpu")
+        pipe.tokenizer_max_length = 1
+        pipe.vae_spatial_compression_ratio = 32
+        pipe.vae_temporal_compression_ratio = 1
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+        pipe.audio_sampling_rate = 1
+        pipe.audio_hop_length = 1
+        pipe.audio_vae_temporal_compression_ratio = 1
+        pipe.audio_vae_mel_compression_ratio = 1
+        pipe.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1))
+        pipe.tokenizer = SimpleNamespace(padding_side="left")
+        pipe.vae = SimpleNamespace(
+            latents_mean=torch.zeros(2),
+            latents_std=torch.ones(2),
+            config=SimpleNamespace(scaling_factor=1.0),
+        )
+        pipe.audio_vae = SimpleNamespace(
+            latents_mean=torch.zeros(2),
+            latents_std=torch.ones(2),
+            config=SimpleNamespace(mel_bins=2, latent_channels=1),
+        )
+
+        video_pos = torch.tensor([[[0.2, -0.3]]])
+        video_neg = torch.tensor([[[-0.4, 0.1]]])
+        audio_pos = torch.tensor([[[0.7, -0.2]]])
+        audio_neg = torch.tensor([[[0.1, 0.4]]])
+
+        class FakeCfgGroup:
+            def all_gather(self, tensor, separate_tensors=True):
+                assert separate_tensors
+                if torch.equal(tensor, video_pos) or torch.equal(tensor, video_neg):
+                    return [video_pos, video_neg]
+                if torch.equal(tensor, audio_pos) or torch.equal(tensor, audio_neg):
+                    return [audio_pos, audio_neg]
+                raise AssertionError(f"Unexpected gathered tensor: {tensor}")
+
+        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2)
+        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda: cfg_rank)
+        monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup())
+
+        def fake_retrieve_timesteps(scheduler, num_inference_steps, device, timesteps, sigmas=None, mu=None):
+            scheduler.sigmas = torch.tensor([0.25, 0.25], device=device)
+            return torch.tensor([1.0, 0.5], device=device), 2
+
+        monkeypatch.setattr(ltx23, "retrieve_timesteps", fake_retrieve_timesteps)
+
+        class FakeScheduler:
+            def __init__(self, name="video", calls=None):
+                self.name = name
+                self.calls = [] if calls is None else calls
+                self.config = {
+                    "max_image_seq_len": 4096,
+                    "base_image_seq_len": 1024,
+                    "base_shift": 0.95,
+                    "max_shift": 2.05,
+                }
+                self.sigmas = torch.tensor([0.25, 0.25])
+
+            def __deepcopy__(self, memo):
+                return FakeScheduler("audio", self.calls)
+
+            def step(self, noise_pred, t, latents, return_dict=False, generator=None):
+                self.calls.append((self.name, noise_pred.clone(), t.clone(), latents.clone()))
+                return (latents - noise_pred,)
+
+        class FakeConnectors:
+            def to(self, device):
+                return self
+
+            def __call__(self, prompt_embeds, prompt_attention_mask, padding_side):
+                assert padding_side == "left"
+                assert prompt_embeds.shape[0] == 2
+                return prompt_embeds, prompt_embeds, prompt_attention_mask
+
+        class FakeRope:
+            def prepare_video_coords(self, batch_size, num_frames, height, width, device, fps):
+                return torch.zeros(batch_size, num_frames * height * width, 3, device=device)
+
+            def prepare_audio_coords(self, batch_size, num_frames, device):
+                return torch.zeros(batch_size, num_frames, 1, device=device)
+
+        class FakeTransformer:
+            def __init__(self):
+                self.config = SimpleNamespace(in_channels=2)
+                self.rope = FakeRope()
+                self.audio_rope = FakeRope()
+                self.calls = []
+
+            def __call__(self, **kwargs):
+                self.calls.append(kwargs)
+                expected_prompt = torch.full((1, 1, 1), expected_prompt_value)
+                torch.testing.assert_close(kwargs["encoder_hidden_states"], expected_prompt)
+                torch.testing.assert_close(kwargs["audio_encoder_hidden_states"], expected_prompt)
+                assert kwargs["hidden_states"].shape == (1, 1, 2)
+                assert kwargs["audio_hidden_states"].shape == (1, 1, 2)
+                if cfg_rank == 0:
+                    return video_pos, audio_pos
+                return video_neg, audio_neg
+
+        class DummyProgress:
+            def __enter__(self):
+                return self
+
+            def __exit__(self, exc_type, exc, tb):
+                return False
+
+            def update(self):
+                pass
+
+        pipe.scheduler = FakeScheduler()
+        pipe.connectors = FakeConnectors()
+        pipe.transformer = FakeTransformer()
+        object.__setattr__(pipe, "progress_bar", lambda total: DummyProgress())
+
+        def fake_encode_prompt(**kwargs):
+            return (
+                torch.ones(1, 1, 1),
+                torch.ones(1, 1, dtype=torch.bool),
+                torch.zeros(1, 1, 1),
+                torch.ones(1, 1, dtype=torch.bool),
+            )
+
+        object.__setattr__(pipe, "encode_prompt", fake_encode_prompt)
+
+        video_latents = torch.tensor([[[1.0, -2.0]]])
+        audio_latents = torch.tensor([[[0.5, 3.0]]])
+        req = OmniDiffusionRequest(
+            prompts=[{"prompt": "prompt", "negative_prompt": "negative"}],
+            sampling_params=OmniDiffusionSamplingParams(
+                height=32,
+                width=32,
+                num_frames=1,
+                frame_rate=1.0,
+                num_inference_steps=2,
+                guidance_scale=4.0,
+                latents=video_latents,
+                audio_latents=audio_latents,
+                output_type="latent",
+            ),
+            request_id="ltx23-cfg-parallel-forward-test",
+        )
+
+        output = pipe.forward(req)
+
+        expected_video_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+            video_latents,
+            video_pos,
+            video_neg,
+            pipe.scheduler.sigmas[0],
+            4.0,
+        )
+        expected_audio_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+            audio_latents,
+            audio_pos,
+            audio_neg,
+            pipe.scheduler.sigmas[0],
+            4.0,
+        )
+        scheduler_call_names = [call[0] for call in pipe.scheduler.calls]
+        assert scheduler_call_names == ["video", "audio", "video", "audio"]
+        assert len(pipe.transformer.calls) == 2
+        torch.testing.assert_close(pipe.scheduler.calls[0][1], expected_video_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[1][1], expected_audio_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[2][1], expected_video_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[3][1], expected_audio_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[2][3], video_latents - expected_video_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[3][3], audio_latents - expected_audio_noise)
+
+        video_out, audio_out = output.output
+        torch.testing.assert_close(video_out, (video_latents - 2 * expected_video_noise).reshape(1, 2, 1, 1, 1))
+        torch.testing.assert_close(audio_out, (audio_latents - 2 * expected_audio_noise).reshape(1, 1, 1, 2))
+
 
 class TestRegistryIntegration:
     """Verify all LTX-2.3 pipeline variants are registered."""
@@ -228,3 +552,85 @@ class TestInitExports:
         for name in expected_classes:
             assert hasattr(ltx2, name), f"{name} not exported from ltx2 package"
             assert name in ltx2.__all__, f"{name} not in ltx2.__all__"
+
+
+class TestAudioLatentSPPadding:
+    def test_prepare_audio_latents_pads_generated_dummy_length_for_sp(self):
+        pipeline = _make_ltx23_pipeline(sequence_parallel_size=2)
+
+        latents, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
+            batch_size=1,
+            num_channels_latents=8,
+            num_mel_bins=64,
+            audio_latent_length=1,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+        )
+
+        assert original_num_frames == 1
+        assert padded_num_frames == 2
+        assert latents.shape == (1, 2, 128)
+
+    def test_prepare_audio_latents_pads_provided_packed_sequence_dim_for_sp(self):
+        pipeline = _make_ltx23_pipeline(sequence_parallel_size=4)
+        latents = torch.arange(40, dtype=torch.float32).view(1, 10, 4)
+
+        padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
+            batch_size=1,
+            num_channels_latents=2,
+            num_mel_bins=8,
+            audio_latent_length=10,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+            latents=latents,
+        )
+
+        assert original_num_frames == 10
+        assert padded_num_frames == 12
+        assert padded.shape == (1, 12, 4)
+        torch.testing.assert_close(padded[:, :10], latents)
+        torch.testing.assert_close(padded[:, 10:], torch.zeros(1, 2, 4))
+
+    def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp(self):
+        pipeline = _make_ltx23_pipeline(sequence_parallel_size=4)
+        latents = torch.arange(96, dtype=torch.float32).view(1, 2, 12, 4)
+
+        audio_latent_length = pipeline._resolve_audio_latent_length(10, latents)
+        padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
+            batch_size=1,
+            num_channels_latents=2,
+            num_mel_bins=16,
+            audio_latent_length=audio_latent_length,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+            latents=latents,
+        )
+
+        assert audio_latent_length == 10
+        assert original_num_frames == 10
+        assert padded_num_frames == 12
+        assert padded.shape == (1, 12, 8)
+        torch.testing.assert_close(padded, pipeline._pack_audio_latents(latents))
+
+    def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference(self):
+        pipeline = _make_ltx23_pipeline(sequence_parallel_size=4)
+        latents = torch.zeros(1, 2, 13, 4)
+
+        audio_latent_length = pipeline._resolve_audio_latent_length(10, latents)
+
+        assert audio_latent_length == 13
+
+    def test_prepare_audio_latents_rejects_incompatible_provided_length(self):
+        pipeline = _make_ltx23_pipeline(sequence_parallel_size=4)
+        latents = torch.zeros(1, 11, 4)
+
+        with pytest.raises(ValueError, match="incompatible audio frame count"):
+            pipeline.prepare_audio_latents(
+                batch_size=1,
+                num_channels_latents=2,
+                num_mel_bins=8,
+                audio_latent_length=10,
+                dtype=torch.float32,
+                device=torch.device("cpu"),
+                latents=latents,
+            )
diff --git a/tests/diffusion/models/ovis_image/test_ovis_image.py b/tests/diffusion/models/ovis_image/test_ovis_image.py
index 6e7ed2f58..8e357e6c6 100644
--- a/tests/diffusion/models/ovis_image/test_ovis_image.py
+++ b/tests/diffusion/models/ovis_image/test_ovis_image.py
@@ -167,6 +167,7 @@ def test_basic_generation(ovis_pipeline):
     # Setup request
     req = OmniDiffusionRequest(
         prompts=["A photo of a cat"],
+        request_id="ovis-basic",
         sampling_params=OmniDiffusionSamplingParams(
             height=256,
             width=256,
@@ -206,6 +207,7 @@ def test_guidance_scale(ovis_pipeline, monkeypatch: pytest.MonkeyPatch):
                 "negative_prompt": "bad quality",
             }
         ],
+        request_id="ovis-guidance",
         sampling_params=OmniDiffusionSamplingParams(
             height=256,
             width=256,
@@ -225,6 +227,7 @@ def test_resolution_check(ovis_pipeline):
     # Pass odd resolution
     req = OmniDiffusionRequest(
         prompts=["test"],
+        request_id="ovis-resolution",
         sampling_params=OmniDiffusionSamplingParams(
             height=250,  # Not divisible by 16 (8*2)
             width=250,
diff --git a/tests/diffusion/quantization/test_mxfp8_config.py b/tests/diffusion/quantization/test_mxfp8_config.py
index d831b5d47..cc34fb232 100644
--- a/tests/diffusion/quantization/test_mxfp8_config.py
+++ b/tests/diffusion/quantization/test_mxfp8_config.py
@@ -115,37 +115,64 @@ def test_mxfp8_quant_config_structure():
 # ---------------------------------------------------------------------------
 
 
-def test_get_quant_method_npu_offline(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
-    """Offline (serialized) path on NPU must return NPUMxfp8LinearMethod."""
-    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config, NPUMxfp8LinearMethod
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="Native MXFP8 offline only supported on NPU")
+def test_get_quant_method_offline_npu(mocker: MockerFixture):
+    """Offline (serialized) path must return NPUMxfp8LinearMethod on NPU."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
 
     config = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=True)
     layer = mocker.Mock(spec=LinearBase)
-    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
 
     method = config.get_quant_method(layer, "blocks.0.attn1.to_q")
-    assert isinstance(method, NPUMxfp8LinearMethod)
+    assert type(method).__name__ == "NPUMxfp8LinearMethod"
 
 
-def test_get_quant_method_npu_online(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
-    """Online (BF16 checkpoint) path on NPU must return NPUMxfp8OnlineLinearMethod."""
-    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config, NPUMxfp8OnlineLinearMethod
+@pytest.mark.skipif(not current_omni_platform.is_xpu(), reason="XPU platform not available")
+def test_get_quant_method_offline_xpu_raises(mocker: MockerFixture):
+    """XPU offline mode must raise NotImplementedError (use AutoRound MXFP8 instead)."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    config = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=True)
+    layer = mocker.Mock(spec=LinearBase)
+
+    with pytest.raises(NotImplementedError, match="Native MXFP8 offline mode is not supported on XPU"):
+        config.get_quant_method(layer, "blocks.0.attn1.to_q")
+
+
+@pytest.mark.skipif(
+    not (current_omni_platform.is_npu() or current_omni_platform.is_xpu()), reason="MXFP8 only supported on NPU and XPU"
+)
+def test_get_quant_method_online(mocker: MockerFixture):
+    """Online (BF16 checkpoint) path must return platform-specific method on current platform."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    # Mock the vLLM online method to avoid config dependency for XPU
+    if current_omni_platform.is_xpu():
+        mock_inner = mocker.Mock()
+        mocker.patch(
+            "vllm_omni.quantization.mxfp8_config.VllmMxfp8OnlineLinearMethod.__init__",
+            lambda self: setattr(self, "_inner", mock_inner),
+        )
 
     config = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=False)
     layer = mocker.Mock(spec=LinearBase)
-    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
 
     method = config.get_quant_method(layer, "blocks.0.attn1.to_q")
-    assert isinstance(method, NPUMxfp8OnlineLinearMethod)
+
+    if current_omni_platform.is_npu():
+        assert type(method).__name__ == "NPUMxfp8OnlineLinearMethod"
+    elif current_omni_platform.is_xpu():
+        assert type(method).__name__ == "VllmMxfp8OnlineLinearMethod"
 
 
-def test_get_quant_method_non_npu_raises(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
-    """Non-NPU platform must raise NotImplementedError (CUDA not yet supported)."""
+def test_get_quant_method_unsupported_platform_raises(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
+    """Unsupported platform (CUDA, ROCm) must raise NotImplementedError."""
     from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
 
     config = DiffusionMXFP8Config()
     layer = mocker.Mock(spec=LinearBase)
     monkeypatch.setattr(current_omni_platform, "is_npu", lambda: False)
+    monkeypatch.setattr(current_omni_platform, "is_xpu", lambda: False)
     monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False)
 
     with pytest.raises(NotImplementedError):
@@ -159,6 +186,8 @@ def test_get_quant_method_ignored_layer(mocker: MockerFixture, monkeypatch: pyte
     config = DiffusionMXFP8Config(ignored_layers=["proj_out"])
     layer = mocker.Mock(spec=LinearBase)
     monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+    monkeypatch.setattr(current_omni_platform, "is_xpu", lambda: False)
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False)
 
     method = config.get_quant_method(layer, "proj_out")
     assert isinstance(method, UnquantizedLinearMethod)
diff --git a/tests/diffusion/quantization/test_quantization_quality.py b/tests/diffusion/quantization/test_quantization_quality.py
index 692d83c12..5f0dc4d62 100644
--- a/tests/diffusion/quantization/test_quantization_quality.py
+++ b/tests/diffusion/quantization/test_quantization_quality.py
@@ -71,7 +71,7 @@ class QualityTestConfig:
     prompt: str  # generation prompt
     max_lpips: float  # fail threshold — higher = more lenient
     model: str | None = None  # HF model name
-    quantization: str | None = None  # quantization method, e.g. "fp8"
+    quantization: str | dict[str, object] | None = None  # quantization method/config, e.g. "fp8"
     baseline_model: str | None = None  # explicit BF16/local baseline path
     quantized_model: str | None = None  # explicit quantized/local model path
     height: int = 1024
@@ -118,9 +118,57 @@ QUALITY_CONFIGS = [
     QualityTestConfig(
         id="fp8_z_image",
         model="Tongyi-MAI/Z-Image-Turbo",
-        quantization="fp8",
+        quantization={
+            "method": "fp8",
+            "ignored_layers": [
+                "img_mlp",
+                *[
+                    f"layers.{layer_id}.{suffix}"
+                    for layer_id in range(15, 30)
+                    for suffix in (
+                        "attention.to_qkv",
+                        "attention.to_out.0",
+                        "feed_forward.w13",
+                        "feed_forward.w2",
+                    )
+                ],
+                *[
+                    f"model.layers.{layer_id}.{suffix}"
+                    for layer_id in range(28, 36)
+                    for suffix in (
+                        "self_attn.q_proj",
+                        "self_attn.k_proj",
+                        "self_attn.v_proj",
+                        "self_attn.o_proj",
+                        "mlp.gate_proj",
+                        "mlp.up_proj",
+                        "mlp.down_proj",
+                    )
+                ],
+            ],
+        },
         task="t2i",
-        prompt="a cup of coffee on a wooden table, morning light",
+        prompt=(
+            "A breathtaking twilight scene atop a floating archipelago of crystalline islands suspended in an "
+            "endless nebula-drenched sky, where waterfalls of liquid starlight cascade from the edges of each "
+            "island into the cosmic abyss below. The largest island hosts an ancient, overgrown observatory "
+            "crafted from pearlescent white stone and living silverwood trees whose bioluminescent leaves pulse "
+            "with soft cerulean and violet light. At the observatory's center stands a solitary figure—a young "
+            "astronomer in flowing robes woven from woven moonlight and deep-space fabric, their face illuminated "
+            "by the glow of a holographic star chart hovering above an altar of polished obsidian. The sky swirls "
+            "with the birth of a new galaxy: ribbons of magenta and gold gas clouds twist around nascent stars, "
+            "while distant supernovae bloom like cosmic flowers in the far reaches of the void. Below, the abyss "
+            "reveals faint echoes of forgotten civilizations—ghostly silhouettes of submerged cities shimmering "
+            "in layers of atmospheric haze. Cinematic volumetric lighting cuts through the scene as twin moons—one "
+            "copper-hued, one opalescent—rise on opposite horizons, casting long, intersecting shadows across "
+            "moss-covered ruins and crystalline flora that refract light into prismatic halos. Hyper-detailed, "
+            "photorealistic rendering with the atmospheric depth of Roger Deakins' cinematography, the "
+            "architectural grandeur of Zaha Hadid, and the cosmic wonder of James Jean's illustrations. Shot on "
+            "a mythical 150mm lens with shallow depth of field, 8K resolution, Unreal Engine 5 realism, subsurface "
+            "scattering on organic elements, and ray-traced reflections dancing across every water droplet in the "
+            "starlight waterfalls. Ethereal, melancholic, and transcendent mood—like a moment of quiet revelation "
+            "at the edge of existence."
+        ),
         max_lpips=0.15,
         num_inference_steps=20,
     ),
@@ -391,12 +439,6 @@ _OUTPUT_DIR = Path(os.environ["VLLM_OMNI_QUALITY_OUTPUT_DIR"]) if "VLLM_OMNI_QUA
 
 def _quality_param(c: QualityTestConfig):
     marks = list(_marks)
-    if c.id == "fp8_z_image":
-        marks.append(
-            pytest.mark.skip(
-                reason="Z-Image FP8 quality gate temporarily disabled: https://github.com/vllm-project/vllm-omni/issues/3531"
-            )
-        )
     if c.id == "fp8_qwen_image":
         marks.append(
             pytest.mark.skip(reason="Qwen-Image FP8 quality gate temporarily disabled (see CI / issue tracker).")
diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py
index dca9a001a..2d44f6a68 100644
--- a/tests/diffusion/test_diffusion_engine.py
+++ b/tests/diffusion/test_diffusion_engine.py
@@ -22,15 +22,15 @@ class DiffusionSchedulerOutput:
     num_waiting_reqs: int = 0
 
     @property
-    def scheduled_req_ids(self):
+    def scheduled_request_ids(self):
         ids = [req.request_id for req in self.scheduled_new_reqs]
         if self.scheduled_cached_reqs:
-            ids.extend(self.scheduled_cached_reqs.sched_req_ids)
+            ids.extend(self.scheduled_cached_reqs.request_ids)
         return ids
 
     @property
     def is_empty(self):
-        return len(self.scheduled_req_ids) == 0
+        return len(self.scheduled_request_ids) == 0
 
 
 class MockScheduler:
@@ -71,31 +71,36 @@ async def test_async_add_req_and_wait_for_response():
     engine.scheduler = MockScheduler()
     engine._out_queue = {}
     engine.abort_queue: queue.Queue[str] = queue.Queue()
-    engine._rpc_lock = threading.Lock()
+    engine._rpc_queue = queue.Queue()
+    engine._rpc_lock = threading.RLock()
+    engine._cv = threading.Condition(engine._rpc_lock)
+    engine._init_lock = asyncio.Lock()
+    engine._closed = False
+    engine._loop_started = False
+    engine.main_loop = None
 
     engine._finalize_finished_request = lambda rid, out, err: out.result
 
     def mock_execute_batch(sched_output):
-        req_ids = sched_output.scheduled_req_ids
+        request_ids = sched_output.scheduled_request_ids
 
         time.sleep(1)
 
         class MockRunnerOutput:
             def __init__(self, ids):
-                self.req_id = ids
+                self.request_id = ids
                 self.step_index = 0
                 self.finished = True
                 self._results = {rid: SimpleNamespace(result_data=f"data_{rid}") for rid in ids}
 
-            def get_req_output(self, rid):
+            def get_request_output(self, rid):
                 return SimpleNamespace(result=self._results[rid], step_index=0, finished=True)
 
-        return MockRunnerOutput(req_ids)
+        return MockRunnerOutput(request_ids)
 
     engine.execute_fn = mock_execute_batch
 
-    engine.stop_event = threading.Event()
-    engine.start_background_loop()
+    await engine._check_and_start_background_loop()
 
     async def run_task(rid):
         req = SimpleNamespace(request_id=rid)
@@ -105,10 +110,13 @@ async def test_async_add_req_and_wait_for_response():
 
     task_ids = [f"req_{i}" for i in range(5)]
     tasks = [run_task(rid) for rid in task_ids]
-    results = await asyncio.gather(*tasks)
-
-    engine.stop_event.set()
-    engine.worker_thread.join(timeout=5)
+    try:
+        results = await asyncio.gather(*tasks)
+    finally:
+        with engine._cv:
+            engine.stop_event.set()
+            engine._cv.notify_all()
+        engine.worker_thread.join(timeout=5)
     assert len(results) == 5
     for rid, res, elapsed in results:
         assert rid in res.result_data
diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py
index a7b945eea..a8985f372 100644
--- a/tests/diffusion/test_diffusion_engine_cleanup.py
+++ b/tests/diffusion/test_diffusion_engine_cleanup.py
@@ -17,15 +17,10 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
-def _make_request(
-    request_id: str,
-    *,
-    request_ids: list[str] | None = None,
-) -> OmniDiffusionRequest:
+def _make_request(request_id: str) -> OmniDiffusionRequest:
     return OmniDiffusionRequest(
         prompts=[f"prompt_{request_id}"],
         sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
-        request_ids=request_ids or [request_id],
         request_id=request_id,
     )
 
@@ -79,15 +74,15 @@ def test_handle_finished_requests_ignores_already_drained_waiter() -> None:
     engine._finalize_finished_request.assert_not_called()
 
 
-def test_abort_parent_request_id_aborts_batched_scheduler_request() -> None:
+def test_abort_request_id_aborts_scheduler_request() -> None:
     engine = _make_engine()
-    request = _make_request("batch-parent", request_ids=["batch-parent-0", "batch-parent-1"])
-    sched_req_id = engine.scheduler.add_request(request)
+    request = _make_request("batch-parent")
+    request_id = engine.scheduler.add_request(request)
 
     engine.abort("batch-parent")
     engine._process_aborts_queue()
 
-    state = engine.scheduler.get_request_state(sched_req_id)
+    state = engine.scheduler.get_request_state(request_id)
     assert state is not None
     assert state.status == DiffusionRequestStatus.FINISHED_ABORTED
 
diff --git a/tests/diffusion/test_diffusion_engine_dummy_run.py b/tests/diffusion/test_diffusion_engine_dummy_run.py
new file mode 100644
index 000000000..71cb47512
--- /dev/null
+++ b/tests/diffusion/test_diffusion_engine_dummy_run.py
@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm_omni.diffusion import diffusion_engine
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_dummy_run_num_frames_uses_explicit_model_setting(monkeypatch: pytest.MonkeyPatch) -> None:
+    class JointAudioVideoModel:
+        dummy_run_num_frames = 2
+
+    monkeypatch.setattr(
+        diffusion_engine.DiffusionModelRegistry,
+        "_try_load_model_cls",
+        lambda model_class_name: JointAudioVideoModel,
+    )
+
+    assert diffusion_engine.get_dummy_run_num_frames("joint_audio_video", supports_audio_input=False) == 2
+
+
+def test_dummy_run_num_frames_keeps_audio_output_default(monkeypatch: pytest.MonkeyPatch) -> None:
+    class AudioOutputModel:
+        support_audio_output = True
+
+    monkeypatch.setattr(
+        diffusion_engine.DiffusionModelRegistry,
+        "_try_load_model_cls",
+        lambda model_class_name: AudioOutputModel,
+    )
+
+    assert diffusion_engine.get_dummy_run_num_frames("audio_output", supports_audio_input=False) == 2
+
+
+def test_dummy_run_num_frames_defaults_to_single_frame(monkeypatch: pytest.MonkeyPatch) -> None:
+    class VideoOnlyModel:
+        pass
+
+    monkeypatch.setattr(
+        diffusion_engine.DiffusionModelRegistry,
+        "_try_load_model_cls",
+        lambda model_class_name: VideoOnlyModel,
+    )
+
+    assert diffusion_engine.get_dummy_run_num_frames("video_only", supports_audio_input=False) == 1
+
+
+def test_dummy_run_num_frames_uses_audio_input_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(
+        diffusion_engine.DiffusionModelRegistry,
+        "_try_load_model_cls",
+        lambda model_class_name: None,
+    )
+
+    assert diffusion_engine.get_dummy_run_num_frames("unknown", supports_audio_input=True) == 2
diff --git a/tests/diffusion/test_diffusion_engine_rpc_routing.py b/tests/diffusion/test_diffusion_engine_rpc_routing.py
index b1b63a1a2..852f6dc2e 100644
--- a/tests/diffusion/test_diffusion_engine_rpc_routing.py
+++ b/tests/diffusion/test_diffusion_engine_rpc_routing.py
@@ -89,7 +89,7 @@ class _ConcurrencyTrackingExecutor:
             if method in {"execute_model", "execute_stepwise", "generate"}:
                 # args[0] is the request-like object for execute_model.
                 req = args[0] if args else None
-                tag = req.request_ids[0] if req is not None and hasattr(req, "request_ids") else "unknown"
+                tag = req.request_id if req is not None and hasattr(req, "request_id") else "unknown"
                 return DiffusionOutput(error=f"result_for_{tag}")
             tag = args[0] if args else method
             return DiffusionOutput(error=f"rpc_result_for_{tag}")
@@ -108,7 +108,7 @@ class _ConcurrencyTrackingExecutor:
             exec_all_ranks=True,
         )
         return RunnerOutput(
-            req_id=new_req.sched_req_id,
+            request_id=new_req.request_id,
             step_index=None,
             finished=True,
             result=result,
@@ -120,7 +120,7 @@ class _ConcurrencyTrackingExecutor:
 
 def _make_request(tag: str):
     return SimpleNamespace(
-        request_ids=[tag],
+        request_id=tag,
         prompts=[f"prompt_{tag}"],
         sampling_params=SimpleNamespace(num_inference_steps=1, **_SAMPLING_KEY_DEFAULTS),
     )
@@ -148,6 +148,7 @@ def _make_engine_with_loop(
     engine._rpc_lock = threading.RLock()
     engine._cv = threading.Condition(engine._rpc_lock)
     engine._out_queue = {}
+    engine._closed = False
     engine.abort_queue = queue.Queue()
     engine._rpc_queue = queue.Queue()
 
diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py
index cb12019bb..4205d8280 100644
--- a/tests/diffusion/test_diffusion_model_runner.py
+++ b/tests/diffusion/test_diffusion_model_runner.py
@@ -52,6 +52,7 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
     runner.pipeline = _DummyPipeline(output=SimpleNamespace(output="ok"))
     runner.cache_backend = cache_backend
     runner.offload_backend = None
+    runner.prompt_embed_cache = None
     runner.od_config = SimpleNamespace(
         cache_backend=cache_backend_name,
         enable_cache_dit_summary=enable_cache_dit_summary,
diff --git a/tests/diffusion/test_diffusion_request.py b/tests/diffusion/test_diffusion_request.py
index 56ec95d0c..9087a6e2f 100644
--- a/tests/diffusion/test_diffusion_request.py
+++ b/tests/diffusion/test_diffusion_request.py
@@ -5,7 +5,7 @@ import random
 
 import pytest
 
-from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -15,9 +15,53 @@ def _make_request() -> OmniDiffusionRequest:
     return OmniDiffusionRequest(
         prompts=[{"prompt": "a cup of coffee on a table"}],
         sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        request_id="request-test",
     )
 
 
+def test_request_id_is_required():
+    with pytest.raises(TypeError, match="request_id"):
+        OmniDiffusionRequest(
+            prompts=[{"prompt": "a cup of coffee on a table"}],
+            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        )
+
+
+def test_request_ids_identity_list_is_removed():
+    req = _make_request()
+
+    assert req.request_id == "request-test"
+    assert not hasattr(req, "request_ids")
+
+
+def test_request_id_must_be_non_empty():
+    with pytest.raises(ValueError, match="request_id must be a non-empty string"):
+        OmniDiffusionRequest(
+            prompts=[{"prompt": "a cup of coffee on a table"}],
+            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+            request_id="",
+        )
+
+
+def test_dummy_run_request_is_identified_by_reserved_request_id():
+    req = OmniDiffusionRequest(
+        prompts=[{"prompt": "dummy run"}],
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        request_id=DUMMY_DIFFUSION_REQUEST_ID,
+    )
+
+    assert req.is_dummy_run()
+    assert OmniDiffusionRequest.is_dummy_run_request_id(DUMMY_DIFFUSION_REQUEST_ID)
+
+
+def test_non_dummy_request_is_not_identified_as_dummy_run():
+    req = _make_request()
+
+    assert not req.is_dummy_run()
+    assert not OmniDiffusionRequest.is_dummy_run_request_id(req.request_id)
+    assert not OmniDiffusionRequest.is_dummy_run_request_id(None)
+
+
 def test_tp_seed_same_across_ranks_and_varies_across_requests():
     random.seed(0)
     n_requests = 5
diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py
index ca28b2629..85d2e45c9 100644
--- a/tests/diffusion/test_diffusion_scheduler.py
+++ b/tests/diffusion/test_diffusion_scheduler.py
@@ -31,13 +31,13 @@ def _make_request(req_id: str) -> OmniDiffusionRequest:
     return OmniDiffusionRequest(
         prompts=[f"prompt_{req_id}"],
         sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
-        request_ids=[req_id],
+        request_id=req_id,
     )
 
 
 def _make_request_output(req_id: str, *, error: str | None = None, finished: bool = True):
     return RunnerOutput(
-        req_id=req_id,
+        request_id=req_id,
         step_index=None,
         finished=finished,
         result=DiffusionOutput(output=None, error=error),
@@ -52,7 +52,7 @@ def _make_step_output(
     error: str | None = None,
 ):
     return RunnerOutput(
-        req_id=req_id,
+        request_id=req_id,
         step_index=step_index,
         finished=finished,
         result=DiffusionOutput(output=None, error=error) if error is not None else None,
@@ -73,16 +73,16 @@ def _make_step_request(
             num_inference_steps=num_inference_steps,
             step_index=step_index,
         ),
-        request_ids=[req_id],
+        request_id=req_id,
     )
 
 
 def _new_ids(sched_output) -> list[str]:
-    return [req.sched_req_id for req in sched_output.scheduled_new_reqs]
+    return [req.request_id for req in sched_output.scheduled_new_reqs]
 
 
 def _cached_ids(sched_output) -> list[str]:
-    return list(sched_output.scheduled_cached_reqs.sched_req_ids)
+    return list(sched_output.scheduled_cached_reqs.request_ids)
 
 
 class _StubScheduler(SchedulerInterface):
@@ -90,7 +90,7 @@ class _StubScheduler(SchedulerInterface):
         self._request = request
         self._output = output
         self.initialized_with = None
-        self._sched_req_id = request.request_ids[0]
+        self._request_id = request.request_id
         self._state = None
         self._scheduled = False
         self.max_num_running_reqs = 1
@@ -100,22 +100,22 @@ class _StubScheduler(SchedulerInterface):
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
         assert request is self._request
-        self._state = SimpleNamespace(sched_req_id=self._sched_req_id, req=request)
-        return self._sched_req_id
+        self._state = SimpleNamespace(request_id=self._request_id, req=request)
+        return self._request_id
 
     def schedule(self):
         if self._scheduled or self._state is None:
             return SimpleNamespace(
                 scheduled_new_reqs=[],
                 scheduled_cached_reqs=CachedRequestData.make_empty(),
-                scheduled_req_ids=[],
+                scheduled_request_ids=[],
                 is_empty=True,
             )
         self._scheduled = True
         return SimpleNamespace(
             scheduled_new_reqs=[NewRequestData.from_state(self._state)],
             scheduled_cached_reqs=CachedRequestData.make_empty(),
-            scheduled_req_ids=[self._state.sched_req_id],
+            scheduled_request_ids=[self._state.request_id],
             is_empty=False,
         )
 
@@ -123,30 +123,25 @@ class _StubScheduler(SchedulerInterface):
         del sched_output
         assert output is self._output
         self._state.status = DiffusionRequestStatus.FINISHED_COMPLETED
-        return {self._sched_req_id}
+        return {self._request_id}
 
     def has_requests(self) -> bool:
         return not self._scheduled
 
-    def get_request_state(self, sched_req_id: str):
-        del sched_req_id
+    def get_request_state(self, request_id: str):
+        del request_id
         return self._state
 
-    def get_sched_req_id(self, request_id: str) -> str | None:
-        if request_id in self._request.request_ids:
-            return self._sched_req_id
-        return None
-
-    def pop_request_state(self, sched_req_id: str):
-        del sched_req_id
+    def pop_request_state(self, request_id: str):
+        del request_id
         return self._state
 
-    def preempt_request(self, sched_req_id: str) -> bool:
-        del sched_req_id
+    def preempt_request(self, request_id: str) -> bool:
+        del request_id
         return False
 
-    def finish_requests(self, sched_req_ids, status) -> None:
-        del sched_req_ids, status
+    def finish_requests(self, request_ids, status) -> None:
+        del request_ids, status
         return None
 
     def close(self) -> None:
@@ -171,7 +166,7 @@ class TestGetSamplingParamsKey:
         return OmniDiffusionRequest(
             prompts=["prompt"],
             sampling_params=sp,
-            request_ids=[f"req-{lora_int_id}-{lora_scale}"],
+            request_id=f"req-{lora_int_id}-{lora_scale}",
         )
 
     def test_distinguishes_lora_id(self) -> None:
@@ -291,7 +286,7 @@ class TestRequestScheduler:
             OmniDiffusionRequest(
                 prompts=["prompt_b"],
                 sampling_params=OmniDiffusionSamplingParams(width=768),
-                request_ids=["b"],
+                request_id="b",
             )
         )
         scheduler.add_request(_make_request("c"))
@@ -333,7 +328,7 @@ class TestRequestScheduler:
         assert state_a.status == DiffusionRequestStatus.FINISHED_ABORTED
 
         assert self.scheduler.has_requests() is False
-        assert self.scheduler.schedule().scheduled_req_ids == []
+        assert self.scheduler.schedule().scheduled_request_ids == []
 
     def test_has_requests_state_transition(self) -> None:
         assert self.scheduler.has_requests() is False
@@ -348,60 +343,28 @@ class TestRequestScheduler:
         assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
         assert self.scheduler.has_requests() is False
 
-    def test_request_id_mapping_lifecycle(self) -> None:
+    def test_request_id_is_scheduler_key(self) -> None:
         request = OmniDiffusionRequest(
             prompts=["prompt_map_a", "prompt_map_b"],
             sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
-            request_ids=["map-a", "map-b"],
             request_id="map-parent",
         )
 
-        sched_req_id = self.scheduler.add_request(request)
+        request_id = self.scheduler.add_request(request)
 
-        assert self.scheduler.get_sched_req_id("map-a") == sched_req_id
-        assert self.scheduler.get_sched_req_id("map-b") == sched_req_id
-        assert self.scheduler.get_sched_req_id("map-parent") == sched_req_id
+        assert request_id == "map-parent"
+        state = self.scheduler.get_request_state("map-parent")
+        assert state.request_id == "map-parent"
 
-        self.scheduler.pop_request_state(sched_req_id)
+        self.scheduler.pop_request_state("map-parent")
 
-        assert self.scheduler.get_sched_req_id("map-a") is None
-        assert self.scheduler.get_sched_req_id("map-b") is None
-        assert self.scheduler.get_sched_req_id("map-parent") is None
+        assert self.scheduler.get_request_state("map-parent") is None
 
-    def test_parent_request_id_registration_failure_rolls_back_child_ids(self) -> None:
-        self.scheduler.add_request(
-            OmniDiffusionRequest(
-                prompts=["prompt_existing"],
-                sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
-                request_ids=["existing-child"],
-                request_id="duplicate-parent",
-            )
-        )
+    def test_duplicate_request_id_is_rejected(self) -> None:
+        self.scheduler.add_request(_make_request("dup"))
 
-        colliding_request = OmniDiffusionRequest(
-            prompts=["prompt_unique"],
-            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
-            request_ids=["unique-child"],
-            request_id="duplicate-parent",
-        )
-
-        with pytest.raises(ValueError, match="duplicate-parent"):
-            self.scheduler.add_request(colliding_request)
-
-        assert self.scheduler.get_request_state("unique-child") is None
-        assert self.scheduler.get_sched_req_id("unique-child") is None
-
-        valid_request = OmniDiffusionRequest(
-            prompts=["prompt_unique"],
-            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
-            request_ids=["unique-child"],
-            request_id="unique-parent",
-        )
-        sched_req_id = self.scheduler.add_request(valid_request)
-
-        assert sched_req_id == "unique-child"
-        assert self.scheduler.get_sched_req_id("unique-child") == sched_req_id
-        assert self.scheduler.get_sched_req_id("unique-parent") == sched_req_id
+        with pytest.raises(ValueError, match="request_id 'dup' is already active"):
+            self.scheduler.add_request(_make_request("dup"))
 
 
 class TestDiffusionEngine:
@@ -662,7 +625,7 @@ class TestStepScheduler:
         finished = self.scheduler.update_from_output(
             sched_output,
             RunnerOutput(
-                req_id=req_id,
+                request_id=req_id,
                 step_index=None,
                 finished=True,
                 result=None,
@@ -934,7 +897,7 @@ class TestStepScheduler:
 
         for _ in range(expected_steps - 1):
             sched_output = self.scheduler.schedule()
-            assert sched_output.scheduled_req_ids == [req_id]
+            assert sched_output.scheduled_request_ids == [req_id]
             next_step = request.sampling_params.step_index + 1
             assert (
                 self.scheduler.update_from_output(
@@ -945,7 +908,7 @@ class TestStepScheduler:
             )
 
         final_output = self.scheduler.schedule()
-        assert final_output.scheduled_req_ids == [req_id]
+        assert final_output.scheduled_request_ids == [req_id]
         assert self.scheduler.update_from_output(
             final_output,
             _make_step_output(req_id, step_index=expected_steps, finished=True),
diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py
index 2722de7e3..029fea120 100644
--- a/tests/diffusion/test_diffusion_step_pipeline.py
+++ b/tests/diffusion/test_diffusion_step_pipeline.py
@@ -206,7 +206,7 @@ class _DistributedStepPipeline(CFGParallelMixin):
 def _make_step_request(num_inference_steps: int = 2):
     return SimpleNamespace(
         prompts=["a prompt"],
-        request_ids=["req-1"],
+        request_id="req-1",
         sampling_params=SimpleNamespace(
             generator=None,
             seed=None,
@@ -227,7 +227,7 @@ def _make_engine_request(req_id: str = "req-1", num_inference_steps: int = 2) ->
     return OmniDiffusionRequest(
         prompts=[f"prompt-{req_id}"],
         sampling_params=OmniDiffusionSamplingParams(num_inference_steps=num_inference_steps),
-        request_ids=[req_id],
+        request_id=req_id,
     )
 
 
@@ -274,10 +274,10 @@ def _make_distributed_runner(mode: str, device: torch.device):
     return runner
 
 
-def _make_scheduler_output(req, sched_req_id="req-1", step_id=0, finished_req_ids=None):
+def _make_scheduler_output(req, request_id="req-1", step_id=0, finished_req_ids=None):
     return DiffusionSchedulerOutput(
         step_id=step_id,
-        scheduled_new_reqs=[NewRequestData(sched_req_id=sched_req_id, req=req)],
+        scheduled_new_reqs=[NewRequestData(request_id=request_id, req=req)],
         scheduled_cached_reqs=CachedRequestData.make_empty(),
         finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
         num_running_reqs=1,
@@ -287,7 +287,7 @@ def _make_scheduler_output(req, sched_req_id="req-1", step_id=0, finished_req_id
 
 def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None):
     """Scheduler output for a homogeneous batch (one NewRequestData per req)."""
-    new_reqs = [NewRequestData(sched_req_id=r.request_ids[0], req=r) for r in reqs]
+    new_reqs = [NewRequestData(request_id=r.request_id, req=r) for r in reqs]
     return DiffusionSchedulerOutput(
         step_id=step_id,
         scheduled_new_reqs=new_reqs,
@@ -298,11 +298,11 @@ def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None):
     )
 
 
-def _make_cached_scheduler_output(sched_req_id="req-1", step_id=1, finished_req_ids=None):
+def _make_cached_scheduler_output(request_id="req-1", step_id=1, finished_req_ids=None):
     return DiffusionSchedulerOutput(
         step_id=step_id,
         scheduled_new_reqs=[],
-        scheduled_cached_reqs=CachedRequestData(sched_req_ids=[sched_req_id]),
+        scheduled_cached_reqs=CachedRequestData(request_ids=[request_id]),
         finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
         num_running_reqs=1,
         num_waiting_reqs=0,
@@ -359,7 +359,7 @@ def _distributed_step_worker(local_rank: int, world_size: int, mode: str, master
             runner,
             _make_scheduler_output(_make_step_request(num_inference_steps=1), step_id=0),
         )
-        output = result.get_req_output("req-1")
+        output = result.get_request_output("req-1")
 
         assert output.finished is True
         assert output.result is not None
@@ -384,16 +384,16 @@ class TestRunner:
         monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
 
         result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
-        first = result.get_req_output("req-1")
-        assert first.req_id == "req-1"
+        first = result.get_request_output("req-1")
+        assert first.request_id == "req-1"
         assert first.step_index == 1
         assert first.finished is False
         assert first.result is None
         assert "req-1" in runner.state_cache
 
         result = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1))
-        second = result.get_req_output("req-1")
-        assert second.req_id == "req-1"
+        second = result.get_request_output("req-1")
+        assert second.request_id == "req-1"
         assert second.step_index == 2
         assert second.finished is True
         assert second.result is not None
@@ -410,13 +410,13 @@ class TestRunner:
         runner = _make_runner()
         req_1 = _make_step_request()
         req_2 = _make_step_request()
-        req_2.request_ids = ["req-2"]
+        req_2.request_id = "req-2"
 
         scheduler_output = DiffusionSchedulerOutput(
             step_id=0,
             scheduled_new_reqs=[
-                NewRequestData(sched_req_id="req-1", req=req_1),
-                NewRequestData(sched_req_id="req-2", req=req_2),
+                NewRequestData(request_id="req-1", req=req_1),
+                NewRequestData(request_id="req-2", req=req_2),
             ],
             scheduled_cached_reqs=CachedRequestData.make_empty(),
             finished_req_ids=set(),
@@ -431,7 +431,7 @@ class TestRunner:
         runner = _make_runner()
 
         with pytest.raises(ValueError, match="Missing cached state"):
-            DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(sched_req_id="req-missing"))
+            DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(request_id="req-missing"))
 
     def test_interrupt_marks_request_finished_and_clears_state(self, monkeypatch):
         runner = _make_runner()
@@ -440,8 +440,8 @@ class TestRunner:
         monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
 
         result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
-        output = result.get_req_output("req-1")
-        assert output.req_id == "req-1"
+        output = result.get_request_output("req-1")
+        assert output.request_id == "req-1"
         assert output.step_index == 0
         assert output.finished is True
         assert output.result is not None
@@ -515,7 +515,7 @@ def _make_step_worker(lora_manager=None, *, expected_output=None):
     worker = object.__new__(DiffusionWorker)
     worker.lora_manager = lora_manager
     worker._step_lora_state = {}
-    output = expected_output if expected_output is not None else RunnerOutput(req_id="req-1")
+    output = expected_output if expected_output is not None else RunnerOutput(request_id="req-1")
     worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: output)
     return worker
 
@@ -525,9 +525,9 @@ class TestWorker:
     """DiffusionWorker.execute_stepwise"""
 
     def test_delegates_to_model_runner(self):
-        expected = RunnerOutput(req_id="req-1", step_index=1, finished=False, result=None)
+        expected = RunnerOutput(request_id="req-1", step_index=1, finished=False, result=None)
         worker = _make_step_worker(expected_output=expected)
-        scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), sched_req_id="req-1")
+        scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), request_id="req-1")
 
         output = DiffusionWorker.execute_stepwise(worker, scheduler_output)
 
@@ -536,7 +536,7 @@ class TestWorker:
     def test_deactivates_lora_when_request_has_no_adapter(self):
         manager = _RecordingLoRAManager()
         worker = _make_step_worker(lora_manager=manager)
-        scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), sched_req_id="req-1")
+        scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), request_id="req-1")
 
         DiffusionWorker.execute_stepwise(worker, scheduler_output)
 
@@ -552,7 +552,7 @@ class TestWorker:
 
         manager = _RecordingLoRAManager()
         worker = _make_step_worker(lora_manager=manager)
-        scheduler_output = _make_scheduler_output(request, sched_req_id="req-1")
+        scheduler_output = _make_scheduler_output(request, request_id="req-1")
 
         DiffusionWorker.execute_stepwise(worker, scheduler_output)
 
@@ -568,8 +568,8 @@ class TestWorker:
 
         manager = _RecordingLoRAManager()
         worker = _make_step_worker(lora_manager=manager)
-        first = _make_scheduler_output(request, sched_req_id="req-1")
-        second = _make_cached_scheduler_output(sched_req_id="req-1", step_id=1)
+        first = _make_scheduler_output(request, request_id="req-1")
+        second = _make_cached_scheduler_output(request_id="req-1", step_id=1)
 
         DiffusionWorker.execute_stepwise(worker, first)
         DiffusionWorker.execute_stepwise(worker, second)
@@ -610,10 +610,10 @@ class TestWorker:
         next_request.sampling_params.lora_request = lora_request
 
         worker = _make_step_worker(lora_manager=_RecordingLoRAManager())
-        first = _make_scheduler_output(finishing, sched_req_id="req-1")
+        first = _make_scheduler_output(finishing, request_id="req-1")
         next_batch = _make_scheduler_output(
             next_request,
-            sched_req_id="req-2",
+            request_id="req-2",
             step_id=1,
             finished_req_ids={"req-1"},
         )
@@ -633,11 +633,11 @@ class TestExecutor:
     def test_execute_step_passes_through_runner_output(self, mocker: MockerFixture):
         executor = object.__new__(MultiprocDiffusionExecutor)
         executor._ensure_open = lambda: None
-        expected = RunnerOutput(req_id="req-step", step_index=1, finished=False, result=None)
+        expected = RunnerOutput(request_id="req-step", step_index=1, finished=False, result=None)
         executor.collective_rpc = mocker.Mock(return_value=expected)
 
         request = _make_engine_request("req-step", num_inference_steps=2)
-        scheduler_output = _make_scheduler_output(request, sched_req_id="req-step")
+        scheduler_output = _make_scheduler_output(request, request_id="req-step")
 
         output = MultiprocDiffusionExecutor.execute_step(executor, scheduler_output)
 
@@ -653,7 +653,7 @@ class TestEngine:
         [
             (
                 lambda _: RunnerOutput(
-                    req_id="req-error",
+                    request_id="req-error",
                     step_index=1,
                     finished=True,
                     result=DiffusionOutput(error="boom"),
@@ -688,7 +688,7 @@ class TestEngine:
             call_count["n"] += 1
             finished = call_count["n"] == 2
             return RunnerOutput(
-                req_id="req-step",
+                request_id="req-step",
                 step_index=call_count["n"],
                 finished=finished,
                 result=(DiffusionOutput(output=torch.tensor([2.0])) if finished else None),
@@ -714,7 +714,7 @@ class TestEngine:
             step["n"] += 1
             engine.abort("req-stop")
             return RunnerOutput(
-                req_id="req-stop",
+                request_id="req-stop",
                 step_index=1,
                 finished=False,
                 result=None,
@@ -741,7 +741,7 @@ class TestEngine:
                 assert sched_output == _make_cached_scheduler_output("req-mid", step_id=1)
                 engine.abort("req-mid")
             return RunnerOutput(
-                req_id="req-mid",
+                request_id="req-mid",
                 step_index=step["n"],
                 finished=False,
                 result=None,
@@ -760,7 +760,7 @@ class TestEngine:
         engine = _make_engine(
             scheduler,
             execute_fn=lambda _: RunnerOutput(
-                req_id="req-missing",
+                request_id="req-missing",
                 step_index=1,
                 finished=True,
                 result=None,
@@ -777,7 +777,7 @@ class TestEngine:
 class TestIPC:
     def test_pack_unpack_runner_output_shm(self):
         tensor = torch.zeros(300_000, dtype=torch.float32)
-        output = RunnerOutput(req_id="req-1", finished=True, result=DiffusionOutput(output=tensor))
+        output = RunnerOutput(request_id="req-1", finished=True, result=DiffusionOutput(output=tensor))
 
         packed = pack_diffusion_output_shm(output)
         assert isinstance(packed.result.output, dict)
diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py
index aa2a74ac1..99d7a9c9c 100644
--- a/tests/diffusion/test_multiproc_engine_concurrency.py
+++ b/tests/diffusion/test_multiproc_engine_concurrency.py
@@ -36,7 +36,7 @@ def _tagged_output(tag: str) -> DiffusionOutput:
 def _mock_request(tag: str):
     """Return a lightweight request object identifiable by *tag*."""
     return SimpleNamespace(
-        request_ids=[tag],
+        request_id=tag,
         prompts=[f"prompt_{tag}"],
         sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
     )
@@ -97,8 +97,8 @@ def _start_worker(req_q, res_q, count=2):
             req = req_q.get(timeout=10)
             method = req.get("method", "")
             args = req.get("args", ())
-            if method in {"generate", "execute_model"} and args and hasattr(args[0], "request_ids"):
-                tag = f"result_for_{args[0].request_ids[0]}"
+            if method in {"generate", "execute_model"} and args and hasattr(args[0], "request_id"):
+                tag = f"result_for_{args[0].request_id}"
             elif args:
                 tag = f"result_for_{args[0]}"
             else:
diff --git a/tests/diffusion/test_omni_diffusion_config_master_port.py b/tests/diffusion/test_omni_diffusion_config_master_port.py
new file mode 100644
index 000000000..d65586021
--- /dev/null
+++ b/tests/diffusion/test_omni_diffusion_config_master_port.py
@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for OmniDiffusionConfig master_port resolution (issue #3794)."""
+
+import socket
+
+import pytest
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestOmniDiffusionConfigMasterPort:
+    def test_honors_master_port_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setenv("MASTER_PORT", "43210")
+        config = OmniDiffusionConfig(model="test")
+        assert config.master_port == 43210
+
+    def test_honors_explicit_master_port(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.delenv("MASTER_PORT", raising=False)
+        config = OmniDiffusionConfig(model="test", master_port=40000)
+        assert config.master_port == 40000
+
+    def test_master_port_env_takes_precedence_over_kwarg(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setenv("MASTER_PORT", "43210")
+        config = OmniDiffusionConfig(model="test", master_port=40000)
+        assert config.master_port == 43210
+
+    def test_explicit_master_port_is_stable_without_jitter(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.delenv("MASTER_PORT", raising=False)
+        ports = {OmniDiffusionConfig(model="test", master_port=40123).master_port for _ in range(5)}
+        assert ports == {40123}
+
+    def test_default_master_port_is_available(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.delenv("MASTER_PORT", raising=False)
+        config = OmniDiffusionConfig(model="test")
+        assert config.master_port is not None
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            sock.bind(("", config.master_port))
+
+    def test_settle_port_steps_when_port_is_busy(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.delenv("MASTER_PORT", raising=False)
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            sock.bind(("", 0))
+            busy_port = sock.getsockname()[1]
+            config = OmniDiffusionConfig(model="test", master_port=busy_port)
+            assert config.master_port == busy_port + 37
diff --git a/tests/diffusion/test_stage_diffusion_proc.py b/tests/diffusion/test_stage_diffusion_proc.py
index 36ceb6cfc..0419bb813 100644
--- a/tests/diffusion/test_stage_diffusion_proc.py
+++ b/tests/diffusion/test_stage_diffusion_proc.py
@@ -68,7 +68,7 @@ def test_process_batch_request_preserves_parent_request_id_and_kv_sender_info():
 
         request = captured["request"]
         assert request.request_id == "req-parent"
-        assert request.request_ids == ["req-parent-0", "req-parent-1"]
+        assert not hasattr(request, "request_ids")
         assert request.kv_sender_info == {0: {"host": "10.0.0.2", "zmq_port": 50151}}
         assert result.request_id == "req-parent"
         assert result.images == ["img-1", "img-2"]
diff --git a/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py b/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py
index a308fe5ea..fe705734b 100644
--- a/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py
+++ b/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py
@@ -269,6 +269,7 @@ def _load_mooncake_config(host: str, rpc_port: int, http_port: int) -> str:
 @pytest.mark.core_model
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
+@pytest.mark.skip(reason="Skip failed CI issue 3977: https://github.com/vllm-project/vllm-omni/issues/3977")
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 def test_bagel_text2img_mooncake_connector(run_level):
     """Test Bagel text2img with Mooncake connector for inter-stage communication."""
diff --git a/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py b/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py
index 57c6a79fb..3906611b1 100644
--- a/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py
+++ b/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py
@@ -258,6 +258,7 @@ def _resolve_deploy_config(config_path: str, run_level: str) -> str:
 @pytest.mark.core_model
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
+@pytest.mark.skip(reason="Skip failed CI issue 3977: https://github.com/vllm-project/vllm-omni/issues/3977")
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"})
 def test_bagel_img2img_shared_memory_connector(run_level):
     """Test Bagel img2img with shared memory connector."""
@@ -275,6 +276,7 @@ def test_bagel_img2img_shared_memory_connector(run_level):
 @pytest.mark.core_model
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
+@pytest.mark.skip(reason="Skip failed CI issue 3977: https://github.com/vllm-project/vllm-omni/issues/3977")
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"})
 def test_bagel_text2img_shared_memory_connector(run_level):
     """Test Bagel text2img with shared memory connector."""
diff --git a/tests/distributed/omni_connectors/test_basic_connectors.py b/tests/distributed/omni_connectors/test_basic_connectors.py
index 662d41fe0..4d0661be8 100644
--- a/tests/distributed/omni_connectors/test_basic_connectors.py
+++ b/tests/distributed/omni_connectors/test_basic_connectors.py
@@ -33,6 +33,20 @@ def test_tensor_serialization():
     assert torch.equal(tensor, deserialized)
 
 
+def test_noncontiguous_tensor_serialization():
+    """Test torch.Tensor serialization for strided views."""
+    import torch
+
+    tensor = torch.arange(40, dtype=torch.long).reshape(4, 10).transpose(0, 1)
+    assert not tensor.is_contiguous()
+    assert tensor.stride(-1) != 1
+
+    serialized = OmniSerializer.serialize({"codes": tensor})
+    deserialized = OmniSerializer.deserialize(serialized)
+
+    assert torch.equal(tensor, deserialized["codes"])
+
+
 def test_ndarray_serialization():
     """Test numpy.ndarray serialization."""
     import numpy as np
diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py
index 9f10aea0c..8bfe837e2 100644
--- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py
+++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py
@@ -37,6 +37,7 @@ def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None)
         status=status,
         prompt_token_ids=[],
         num_computed_tokens=0,
+        num_output_placeholders=0,
         additional_information=None,
         is_finished=lambda: status == RequestStatus.FINISHED_STOPPED,
     )
@@ -44,9 +45,16 @@ def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None)
 
 @pytest.fixture
 def build_adapter(monkeypatch, mocker: MockerFixture):
-    def _build(*, stage_id: int = 1, model_mode: str = "ar", max_num_seqs: int = 2):
+    def _build(
+        *,
+        stage_id: int = 1,
+        model_mode: str = "ar",
+        max_num_seqs: int = 2,
+        connector_extra: dict | None = None,
+    ):
         connector = mocker.MagicMock()
         connector.stage_id = stage_id
+        connector.config = {"extra": connector_extra or {}}
         connector.get.return_value = None
         connector.put.return_value = (True, 1, {})
 
@@ -127,6 +135,61 @@ def test_load_poll(build_adapter):
     assert "req-1" not in adapter._pending_load_reqs
 
 
+def test_load_poll_generation_tensor_codes_use_placeholder_prompt(build_adapter):
+    adapter, connector = build_adapter(stage_id=1, model_mode="generation")
+    request = _req("req-tensor", RequestStatus.WAITING, external_req_id="external-tensor")
+
+    codes = torch.tensor([[1, 2], [3, 4]], dtype=torch.long)
+    payload: OmniPayload = {
+        "codes": {"audio": codes},
+        "meta": {
+            "left_context_size": 1,
+            "finished": torch.tensor(False, dtype=torch.bool),
+        },
+    }
+    connector.get.return_value = (payload, 16)
+
+    adapter._poll_single_request(request)
+
+    assert request.prompt_token_ids == [0]
+    assert request.num_computed_tokens == 0
+    assert torch.equal(request.additional_information["codes"]["audio"], codes)
+    assert request.additional_information["meta"]["left_context_size"] == 1
+    assert "finished" not in request.additional_information["meta"]
+    assert "req-tensor" in adapter._finished_load_reqs
+
+
+def test_load_poll_generation_empty_nonterminal_chunk_keeps_polling(build_adapter):
+    adapter, connector = build_adapter(stage_id=1, model_mode="generation")
+    request = _req("req-empty-tensor", RequestStatus.WAITING, external_req_id="external-empty")
+
+    empty_payload: OmniPayload = {
+        "codes": {"audio": torch.empty((4, 0), dtype=torch.long)},
+        "meta": {
+            "left_context_size": 0,
+            "finished": torch.tensor(False, dtype=torch.bool),
+        },
+    }
+    ready_payload: OmniPayload = {
+        "codes": {"audio": torch.tensor([[1, 2]], dtype=torch.long)},
+        "meta": {
+            "left_context_size": 0,
+            "finished": torch.tensor(False, dtype=torch.bool),
+        },
+    }
+    connector.get.side_effect = [(empty_payload, 16), (ready_payload, 16)]
+
+    assert adapter._poll_single_request(request) is False
+    assert request.request_id not in adapter._finished_load_reqs
+    assert request.request_id not in adapter.requests_with_ready_chunks
+    assert adapter.get_req_chunk[request.request_id] == 1
+
+    assert adapter._poll_single_request(request) is True
+    assert request.request_id in adapter._finished_load_reqs
+    assert torch.equal(request.additional_information["codes"]["audio"], ready_payload["codes"]["audio"])
+    assert adapter.get_req_chunk[request.request_id] == 2
+
+
 def test_save_async(build_adapter):
     adapter, _ = build_adapter(stage_id=1)
     request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1")
@@ -140,6 +203,18 @@ def test_save_async(build_adapter):
     assert task["is_finished"] is False
 
 
+def test_save_async_uses_confirmed_tokens_for_async_scheduler_watermark(build_adapter):
+    adapter, _ = build_adapter(stage_id=1)
+    request = _req("req-async", RequestStatus.WAITING, external_req_id="external-async")
+    request.num_computed_tokens = 10
+    request.num_output_placeholders = 2
+
+    adapter.save_async(pooling_output=None, request=request)
+
+    assert adapter.requests_num_chunks_sent["external-async"] == 8
+    assert len(adapter._pending_save_reqs) == 1
+
+
 def test_send_single_request_struct_without_meta_does_not_crash(build_adapter, monkeypatch):
     """Producer may return a struct with ``meta=None`` (e.g. payload that
     carries only ``embed`` or ``codes``). The sender's ``meta is not None``
@@ -155,9 +230,11 @@ def test_send_single_request_struct_without_meta_does_not_crash(build_adapter, m
     cleanup_calls = []
     monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))
 
-    adapter._send_single_request({"pooling_output": None, "request": request, "is_finished": False})
+    adapter._send_single_request(
+        {"pooling_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
+    )
 
-    assert cleanup_calls == []  # no terminal cleanup; meta.finished is unobservable
+    assert cleanup_calls == []  # no terminal cleanup; meta.finished is false
 
 
 def test_send_single_request_empty_struct_goes_on_wire(build_adapter, monkeypatch):
@@ -173,12 +250,50 @@ def test_send_single_request_empty_struct_goes_on_wire(build_adapter, monkeypatc
     adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct()
     monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)
 
-    adapter._send_single_request({"pooling_output": None, "request": request, "is_finished": False})
+    adapter._send_single_request(
+        {"pooling_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
+    )
 
     assert connector.put.called
     sent_payload = connector.put.call_args.kwargs["data"]
     assert isinstance(sent_payload, OmniPayloadStruct)
-    assert sent_payload.meta is None  # confirms it's the empty struct on the wire
+    assert sent_payload.meta.finished.item() is False
+    assert sent_payload.meta.is_segment_finished.item() is False
+
+
+def test_send_single_request_struct_preserves_segment_finished(build_adapter, monkeypatch):
+    adapter, connector = build_adapter(stage_id=1)
+    request = _req("req-segment", RequestStatus.WAITING, external_req_id="ext-segment")
+
+    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct()
+    monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)
+
+    adapter._send_single_request(
+        {"pooling_output": None, "request": request, "is_finished": False, "is_segment_finished": True}
+    )
+
+    sent_payload = connector.put.call_args.kwargs["data"]
+    assert sent_payload.meta.finished.item() is False
+    assert sent_payload.meta.is_segment_finished.item() is True
+
+
+def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapter):
+    adapter, _ = build_adapter(stage_id=1)
+    request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
+    request.resumable = True
+    request.num_computed_tokens = 0
+    adapter.requests_num_chunks_sent["ext-stream"] = 111
+
+    adapter.save_async(pooling_output=None, request=request, is_segment_finished=False)
+
+    assert len(adapter._pending_save_reqs) == 0
+    assert adapter.requests_num_chunks_sent["ext-stream"] == 111
+
+    adapter.requests_num_chunks_sent.pop("ext-stream")
+    adapter.save_async(pooling_output=None, request=request, is_segment_finished=False)
+
+    assert len(adapter._pending_save_reqs) == 1
+    assert adapter.requests_num_chunks_sent["ext-stream"] == 0
 
 
 def test_send_single_request_cleans_up_after_finished_payload(build_adapter, monkeypatch):
@@ -191,7 +306,9 @@ def test_send_single_request_cleans_up_after_finished_payload(build_adapter, mon
     cleanup_calls = []
     monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))
 
-    adapter._send_single_request({"pooling_output": None, "request": request, "is_finished": True})
+    adapter._send_single_request(
+        {"pooling_output": None, "request": request, "is_finished": True, "is_segment_finished": True}
+    )
 
     assert len(cleanup_calls) == 1
     args, _ = cleanup_calls[0]
@@ -458,6 +575,52 @@ def test_finish_requests_restores_status(build_adapter):
     assert req_id not in adapter.requests_origin_status
 
 
+def test_finish_requests_removes_zombies_from_chunk_waiting_deques(build_adapter):
+    adapter, _ = build_adapter(stage_id=1)
+    zombie = _req("req-zombie", RequestStatus.WAITING_FOR_CHUNK)
+    other = _req("req-live", RequestStatus.WAITING_FOR_CHUNK)
+    adapter.waiting_for_chunk_waiting_requests = deque([zombie, other])
+    adapter.waiting_for_chunk_running_requests = deque([other, zombie])
+    adapter.requests_with_ready_chunks.add("req-zombie")
+    adapter.finished_requests.add("req-zombie")
+    requests_map = {
+        "req-zombie": zombie,
+        "req-live": other,
+    }
+
+    adapter.finish_requests(
+        ["req-zombie"],
+        RequestStatus.FINISHED_ABORTED,
+        requests_map,
+    )
+
+    assert [req.request_id for req in adapter.waiting_for_chunk_waiting_requests] == ["req-live"]
+    assert [req.request_id for req in adapter.waiting_for_chunk_running_requests] == ["req-live"]
+    assert "req-zombie" not in adapter.requests_with_ready_chunks
+    assert "req-zombie" not in adapter.finished_requests
+
+
+def test_restore_queues_skips_requests_missing_from_scheduler_requests(build_adapter):
+    adapter, _ = build_adapter(stage_id=1)
+    zombie = _req("req-zombie", RequestStatus.WAITING_FOR_CHUNK)
+    live = _req("req-live", RequestStatus.WAITING_FOR_CHUNK)
+    waiting_queue = DummyWaitingQueue()
+    running_queue = []
+    adapter.waiting_for_chunk_waiting_requests = deque([zombie, live])
+    adapter.waiting_for_chunk_running_requests = deque([zombie, live])
+
+    adapter.restore_queues(
+        waiting_queue,
+        running_queue,
+        scheduler_requests={"req-live": live},
+    )
+
+    assert [req.request_id for req in waiting_queue] == ["req-live"]
+    assert [req.request_id for req in running_queue] == ["req-live"]
+    assert not adapter.waiting_for_chunk_waiting_requests
+    assert not adapter.waiting_for_chunk_running_requests
+
+
 # ---------------------------------------------------------------
 # Scheduler trigger tests
 # ---------------------------------------------------------------
diff --git a/tests/distributed/omni_connectors/test_kv_flow.py b/tests/distributed/omni_connectors/test_kv_flow.py
index cea186019..2cc19f1b3 100644
--- a/tests/distributed/omni_connectors/test_kv_flow.py
+++ b/tests/distributed/omni_connectors/test_kv_flow.py
@@ -384,7 +384,7 @@ def test_manager_reception(kv_config, mock_connector, common_constants):
     req = OmniDiffusionRequest(
         prompts=["test_recv"],
         sampling_params=OmniDiffusionSamplingParams(),
-        request_ids=[req_id],
+        request_id=req_id,
     )
     # req.need_kv_receive = True # Implicitly handled by receive_kv_cache check? No, manager doesn't check it, runner does.
     # But receive_kv_cache in manager checks request_id. Which we need to fix in manager next.
@@ -428,7 +428,6 @@ def test_manager_reception_prefers_parent_request_id_for_batched_request(kv_conf
     req = OmniDiffusionRequest(
         prompts=["prompt-a", "prompt-b"],
         sampling_params=OmniDiffusionSamplingParams(),
-        request_ids=[f"{parent_req_id}-0", f"{parent_req_id}-1"],
         request_id=parent_req_id,
     )
 
@@ -454,7 +453,6 @@ def test_receive_multi_kv_cache_uses_parent_request_id_for_cfg_collection(kv_con
     req = OmniDiffusionRequest(
         prompts=["prompt-a", "prompt-b"],
         sampling_params=OmniDiffusionSamplingParams(),
-        request_ids=["req-parent-0", "req-parent-1"],
         request_id="req-parent",
     )
     req.sampling_params.cfg_kv_request_ids = {"cfg_text": "req-parent__cfg_text"}
@@ -516,7 +514,7 @@ def test_integration_flow(common_constants):
     req = OmniDiffusionRequest(
         prompts=["test_integ"],
         sampling_params=OmniDiffusionSamplingParams(),
-        request_ids=[req_id],
+        request_id=req_id,
     )
 
     # Receive
diff --git a/tests/e2e/accuracy/conftest.py b/tests/e2e/accuracy/conftest.py
index 431a8cc9b..6c0258aa2 100644
--- a/tests/e2e/accuracy/conftest.py
+++ b/tests/e2e/accuracy/conftest.py
@@ -62,6 +62,19 @@ def pytest_addoption(parser):
         default=1200,
         help="Online serving timeout in seconds for Wan2.2 I2V accuracy tests.",
     )
+    group.addoption(
+        "--hunyuanvideo15-i2v-image-source",
+        action="store",
+        default=None,
+        help="Image source for HunyuanVideo-1.5 I2V accuracy tests. Can be local path or remote URL",
+    )
+    group.addoption(
+        "--hunyuanvideo15-online-timeout-seconds",
+        action="store",
+        type=int,
+        default=3600,
+        help="Online serving timeout in seconds for HunyuanVideo-1.5 accuracy tests.",
+    )
 
 
 def _hf_cache_root() -> Path:
@@ -168,6 +181,17 @@ def wan22_i2v_online_timeout_seconds(request: pytest.FixtureRequest) -> int:
     return int(request.config.getoption("wan22_i2v_online_timeout_seconds"))
 
 
+@pytest.fixture(scope="session")
+def hunyuanvideo15_i2v_image_source(request: pytest.FixtureRequest) -> str | None:
+    value = request.config.getoption("hunyuanvideo15_i2v_image_source")
+    return str(value) if value else None
+
+
+@pytest.fixture(scope="session")
+def hunyuanvideo15_online_timeout_seconds(request: pytest.FixtureRequest) -> int:
+    return int(request.config.getoption("hunyuanvideo15_online_timeout_seconds"))
+
+
 @pytest.fixture(scope="session")
 def gebench_samples_per_type(request: pytest.FixtureRequest) -> int:
     return int(request.config.getoption("gebench_samples_per_type"))
diff --git a/tests/e2e/accuracy/helpers.py b/tests/e2e/accuracy/helpers.py
index 38b045dce..593a9960a 100644
--- a/tests/e2e/accuracy/helpers.py
+++ b/tests/e2e/accuracy/helpers.py
@@ -1,7 +1,17 @@
+import json
 import os
+import re
+import shutil
+import subprocess
+import time
+from base64 import b64decode, b64encode
+from fractions import Fraction
+from hashlib import sha1
 from io import BytesIO
+from mimetypes import guess_type
 from pathlib import Path
 from typing import Any
+from urllib.parse import urlparse
 
 import numpy as np
 import pytest
@@ -63,6 +73,274 @@ def model_output_dir(parent_dir: Path, model: str) -> Path:
     return path
 
 
+_SSIM_RE = re.compile(r"All:(?P<score>[0-9.]+)")
+_PSNR_RE = re.compile(r"average:(?P<score>[0-9.]+)")
+
+
+def parse_video_metadata(payload: dict[str, Any]) -> dict[str, int | float]:
+    streams = payload.get("streams")
+    if not isinstance(streams, list) or not streams:
+        raise ValueError(f"ffprobe payload did not include video streams: {payload}")
+
+    stream = streams[0]
+    width = int(stream["width"])
+    height = int(stream["height"])
+    fps = float(Fraction(stream["avg_frame_rate"]))
+    frame_count_value = stream.get("nb_read_frames") or stream.get("nb_frames")
+    if frame_count_value is None:
+        raise ValueError(f"ffprobe payload did not include frame count: {payload}")
+    frame_count = int(frame_count_value)
+    return {
+        "width": width,
+        "height": height,
+        "fps": fps,
+        "frame_count": frame_count,
+    }
+
+
+def parse_ssim_score(output: str) -> float:
+    match = _SSIM_RE.search(output)
+    if match is None:
+        raise ValueError(f"Could not parse SSIM score from ffmpeg output:\n{output}")
+    return float(match.group("score"))
+
+
+def parse_psnr_score(output: str) -> float:
+    match = _PSNR_RE.search(output)
+    if match is None:
+        raise ValueError(f"Could not parse PSNR score from ffmpeg output:\n{output}")
+    return float(match.group("score"))
+
+
+def probe_binary(binary: str, *, test_name: str = "video similarity e2e test") -> str:
+    resolved = shutil.which(binary)
+    if resolved is None:
+        pytest.skip(f"{binary} is required for {test_name}.")
+    return resolved
+
+
+def probe_video(path: Path) -> dict[str, int | float]:
+    probe_binary("ffprobe")
+    result = subprocess.run(
+        [
+            "ffprobe",
+            "-v",
+            "error",
+            "-select_streams",
+            "v:0",
+            "-count_frames",
+            "-show_entries",
+            "stream=width,height,avg_frame_rate,nb_read_frames",
+            "-of",
+            "json",
+            str(path),
+        ],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    return parse_video_metadata(json.loads(result.stdout))
+
+
+def run_ffmpeg_similarity(filter_name: str, first: Path, second: Path) -> str:
+    probe_binary("ffmpeg")
+    result = subprocess.run(
+        [
+            "ffmpeg",
+            "-hide_banner",
+            "-i",
+            str(first),
+            "-i",
+            str(second),
+            "-lavfi",
+            f"[0:v][1:v]{filter_name}",
+            "-f",
+            "null",
+            "-",
+        ],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    return result.stderr
+
+
+def online_timeout_seconds(configured: int | None, *, default: int = 1200) -> int:
+    return int(configured) if configured is not None else default
+
+
+def is_remote_image_source(source: str) -> bool:
+    return source.startswith("http://") or source.startswith("https://")
+
+
+def validate_image_source(source: str) -> None:
+    if is_remote_image_source(source):
+        response = requests.get(source, timeout=60)
+        response.raise_for_status()
+        return
+
+    image_path = Path(source)
+    if not image_path.exists():
+        raise FileNotFoundError(f"Local image source does not exist: {image_path}")
+
+
+def build_online_image_reference(source: str) -> str:
+    if source.startswith("data:image"):
+        return source
+
+    if is_remote_image_source(source):
+        response = requests.get(source, timeout=60)
+        response.raise_for_status()
+        mime_type = response.headers.get("Content-Type", "").split(";", 1)[0].strip()
+        if not mime_type.startswith("image/"):
+            mime_type = guess_type(urlparse(source).path)[0] or "image/png"
+        encoded = b64encode(response.content).decode("ascii")
+        return f"data:{mime_type};base64,{encoded}"
+
+    image_path = Path(source)
+    mime_type = guess_type(image_path.name)[0] or "image/png"
+    encoded = b64encode(image_path.read_bytes()).decode("ascii")
+    return f"data:{mime_type};base64,{encoded}"
+
+
+def materialize_image_source(source: str, output_dir: Path) -> str:
+    if not is_remote_image_source(source):
+        return source
+
+    from diffusers.utils import load_image
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    source_name = Path(urlparse(source).path).name or "input.png"
+    safe_name = re.sub(r"[^A-Za-z0-9_.-]+", "_", source_name)
+    if "." not in safe_name:
+        safe_name = f"{safe_name}.png"
+    image_path = output_dir / safe_name
+    if not image_path.exists():
+        image = load_image(source).convert("RGB")
+        image.save(image_path)
+    return str(image_path)
+
+
+def video_artifact_dir(result_root: Path, source: str) -> Path:
+    if is_remote_image_source(source):
+        source_name = Path(urlparse(source).path).stem or "remote"
+    else:
+        source_name = Path(source).stem or "local"
+
+    safe_name = re.sub(r"[^A-Za-z0-9_.-]+", "_", source_name)
+    digest = sha1(source.encode("utf-8")).hexdigest()[:8]
+    return result_root / f"{safe_name}-{digest}"
+
+
+def send_video_request_with_timeout(
+    openai_client,
+    request_config: dict[str, Any],
+    *,
+    timeout_seconds: int,
+) -> bytes:
+    form_data = request_config.get("form_data")
+    if not isinstance(form_data, dict):
+        raise ValueError("Video request_config must contain 'form_data'")
+
+    if not form_data.get("prompt"):
+        raise ValueError("Video request_config['form_data'] must contain 'prompt'")
+
+    normalized_form_data = {key: str(value) for key, value in form_data.items() if value is not None}
+
+    files: dict[str, tuple[str, BytesIO, str]] = {}
+    image_reference = request_config.get("image_reference")
+    if image_reference:
+        if image_reference.startswith("data:image"):
+            header, encoded = image_reference.split(",", 1)
+            content_type = header.split(";")[0].removeprefix("data:")
+            extension = content_type.split("/")[-1]
+            file_data = b64decode(encoded)
+            files["input_reference"] = (
+                f"reference.{extension}",
+                BytesIO(file_data),
+                content_type,
+            )
+        else:
+            normalized_form_data["image_reference"] = json.dumps({"image_url": image_reference})
+
+    start_time = time.perf_counter()
+    create_url = openai_client._build_url("/v1/videos")
+    response = requests.post(
+        create_url,
+        data=normalized_form_data,
+        files=files,
+        headers={"Accept": "application/json"},
+        timeout=60,
+    )
+    try:
+        response.raise_for_status()
+    except requests.HTTPError as exc:
+        raise requests.HTTPError(f"{exc}; response body: {response.text}", response=response) from exc
+    job_data = response.json()
+    video_id = job_data["id"]
+    openai_client._wait_until_video_completed(video_id, timeout_seconds=timeout_seconds)
+    video_content = openai_client._download_video_content(video_id)
+    print(f"online_video_e2e_latency_s={time.perf_counter() - start_time:.3f}")
+    return video_content
+
+
+def assert_video_metadata(
+    metadata: dict[str, int | float],
+    *,
+    width: int,
+    height: int,
+    fps: int | float,
+    frame_count: int,
+) -> None:
+    assert metadata["width"] == width
+    assert metadata["height"] == height
+    assert metadata["fps"] == float(fps)
+    assert metadata["frame_count"] == frame_count
+
+
+def assert_video_similarity_metrics(
+    *,
+    label: str,
+    online_path: Path,
+    offline_path: Path,
+    ssim_threshold: float,
+    psnr_threshold: float,
+) -> None:
+    ssim_output = run_ffmpeg_similarity("ssim", online_path, offline_path)
+    psnr_output = run_ffmpeg_similarity("psnr", online_path, offline_path)
+    ssim_score = parse_ssim_score(ssim_output)
+    psnr_score = parse_psnr_score(psnr_output)
+
+    print(f"{label} similarity metrics:")
+    print(
+        "  SSIM:"
+        f" value={ssim_score:.6f},"
+        f" threshold>={ssim_threshold:.6f},"
+        " range=[0, 1],"
+        " higher_is_better=True,"
+        " interpretation=structural_similarity"
+    )
+    print(
+        "  PSNR:"
+        f" value={psnr_score:.6f} dB,"
+        f" threshold>={psnr_threshold:.6f} dB,"
+        " range=[0, +inf),"
+        " higher_is_better=True,"
+        " interpretation=pixel_error_in_decibels"
+    )
+    print(f"online_video={online_path}")
+    print(f"offline_video={offline_path}")
+
+    assert ssim_score >= ssim_threshold, (
+        f"SSIM below threshold: got {ssim_score:.6f}, expected >= {ssim_threshold:.6f}. "
+        f"online={online_path} offline={offline_path}"
+    )
+    assert psnr_score >= psnr_threshold, (
+        f"PSNR below threshold: got {psnr_score:.6f}, expected >= {psnr_threshold:.6f}. "
+        f"online={online_path} offline={offline_path}"
+    )
+
+
 def assert_similarity(
     *,
     model_name: str,
diff --git a/tests/e2e/accuracy/hunyuanvideo15_i2v/__init__.py b/tests/e2e/accuracy/hunyuanvideo15_i2v/__init__.py
new file mode 100644
index 000000000..08dfa08ca
--- /dev/null
+++ b/tests/e2e/accuracy/hunyuanvideo15_i2v/__init__.py
@@ -0,0 +1 @@
+"""HunyuanVideo-1.5 I2V accuracy test cases."""
diff --git a/tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py b/tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py
new file mode 100644
index 000000000..84e032770
--- /dev/null
+++ b/tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py
@@ -0,0 +1,262 @@
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+from tests.e2e.accuracy.helpers import (
+    assert_video_metadata,
+    assert_video_similarity_metrics,
+    build_online_image_reference,
+    materialize_image_source,
+    probe_binary,
+    probe_video,
+    send_video_request_with_timeout,
+    validate_image_source,
+    video_artifact_dir,
+)
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+
+MODEL_NAME = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v"
+IMAGE_URL = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG"
+PROMPT = (
+    "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. "
+    "The fluffy-furred feline gazes directly at the camera with a relaxed expression. "
+    "Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, "
+    "and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, "
+    "as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's "
+    "intricate details and the refreshing atmosphere of the seaside."
+)
+WIDTH = 832
+HEIGHT = 480
+FPS = 24
+NUM_FRAMES = 33
+NUM_INFERENCE_STEPS = 50
+GUIDANCE_SCALE = 6.0
+FLOW_SHIFT = 5.0
+SEED = 42
+SSIM_THRESHOLD = 0.94
+PSNR_THRESHOLD = 28.0
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+WORKSPACE_ROOT = REPO_ROOT.parent
+RUNNER_PATH = REPO_ROOT / "examples" / "offline_inference" / "image_to_video" / "image_to_video.py"
+RESULT_ROOT = Path(__file__).parent / "result"
+VIDEO_TIMEOUT_SECONDS = 60 * 60
+SERVER_CASES = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL_NAME,
+            server_args=[
+                "--flow-shift",
+                str(FLOW_SHIFT),
+                "--enforce-eager",
+            ],
+            env_dict={"VLLM_OMNI_STORAGE_PATH": str(RESULT_ROOT / "storage")},
+            use_omni=True,
+        ),
+        id="hunyuanvideo15_i2v_default",
+    )
+]
+
+
+def _runner_env() -> dict[str, str]:
+    env = os.environ.copy()
+    pythonpath_parts = [
+        str(REPO_ROOT),
+        str(WORKSPACE_ROOT / "diffusers" / "src"),
+    ]
+    existing_pythonpath = env.get("PYTHONPATH")
+    if existing_pythonpath:
+        pythonpath_parts.append(existing_pythonpath)
+    env["PYTHONPATH"] = os.pathsep.join(pythonpath_parts)
+    env.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    return env
+
+
+def _resolve_image_source(configured: str | None) -> str:
+    if configured:
+        candidate = Path(configured)
+        if candidate.exists():
+            return str(candidate.resolve())
+        return configured
+    return IMAGE_URL
+
+
+def _artifact_paths(image_source: str) -> tuple[Path, Path]:
+    artifact_dir = video_artifact_dir(RESULT_ROOT, image_source)
+    artifact_dir.mkdir(parents=True, exist_ok=True)
+    return artifact_dir / "online.mp4", artifact_dir / "offline.mp4"
+
+
+def _offline_image_source(image_source: str) -> str:
+    artifact_dir = video_artifact_dir(RESULT_ROOT, image_source)
+    return materialize_image_source(image_source, artifact_dir)
+
+
+def _build_offline_command(*, image_source: str, output_path: Path) -> list[str]:
+    return [
+        sys.executable,
+        str(RUNNER_PATH),
+        "--model",
+        MODEL_NAME,
+        "--image",
+        image_source,
+        "--prompt",
+        PROMPT,
+        "--height",
+        str(HEIGHT),
+        "--width",
+        str(WIDTH),
+        "--num-frames",
+        str(NUM_FRAMES),
+        "--num-inference-steps",
+        str(NUM_INFERENCE_STEPS),
+        "--guidance-scale",
+        str(GUIDANCE_SCALE),
+        "--flow-shift",
+        str(FLOW_SHIFT),
+        "--fps",
+        str(FPS),
+        "--seed",
+        str(SEED),
+        "--enforce-eager",
+        "--output",
+        str(output_path),
+    ]
+
+
+def _generate_offline_video(*, image_source: str) -> Path:
+    _, offline_path = _artifact_paths(image_source)
+    offline_image_source = _offline_image_source(image_source)
+    subprocess.run(
+        _build_offline_command(image_source=offline_image_source, output_path=offline_path),
+        cwd=REPO_ROOT,
+        env=_runner_env(),
+        check=True,
+        timeout=VIDEO_TIMEOUT_SECONDS,
+    )
+    return offline_path
+
+
+def _generate_online_video(
+    *,
+    omni_server,
+    openai_client,
+    image_source: str,
+    timeout_seconds: int,
+) -> Path:
+    online_path, _ = _artifact_paths(image_source)
+    online_image_source = _offline_image_source(image_source)
+    request_config = {
+        "model": omni_server.model,
+        "form_data": {
+            "prompt": PROMPT,
+            "height": HEIGHT,
+            "width": WIDTH,
+            "fps": FPS,
+            "num_frames": NUM_FRAMES,
+            "guidance_scale": GUIDANCE_SCALE,
+            "flow_shift": FLOW_SHIFT,
+            "num_inference_steps": NUM_INFERENCE_STEPS,
+            "seed": SEED,
+        },
+        "image_reference": build_online_image_reference(online_image_source),
+    }
+    online_video_bytes = send_video_request_with_timeout(
+        openai_client,
+        request_config,
+        timeout_seconds=timeout_seconds,
+    )
+    online_path.write_bytes(online_video_bytes)
+    return online_path
+
+
+@pytest.mark.benchmark
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_hunyuanvideo15_i2v_diffusers_offline_generates_video(
+    hunyuanvideo15_i2v_image_source: str | None,
+) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("HunyuanVideo-1.5 I2V offline accuracy test requires CUDA.")
+
+    probe_binary("ffprobe")
+    if not RUNNER_PATH.exists():
+        raise AssertionError(f"Offline runner does not exist: {RUNNER_PATH}")
+
+    image_source = _resolve_image_source(hunyuanvideo15_i2v_image_source)
+    validate_image_source(image_source)
+    offline_path = _generate_offline_video(image_source=image_source)
+    assert offline_path.exists(), f"Expected offline video artifact at {offline_path}"
+    offline_metadata = probe_video(offline_path)
+    assert_video_metadata(offline_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
+
+
+@pytest.mark.benchmark
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True)
+def test_hunyuanvideo15_i2v_online_serving_generates_video(
+    omni_server,
+    openai_client,
+    hunyuanvideo15_i2v_image_source: str | None,
+    hunyuanvideo15_online_timeout_seconds: int,
+) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("HunyuanVideo-1.5 I2V online accuracy test requires CUDA.")
+
+    probe_binary("ffprobe")
+    image_source = _resolve_image_source(hunyuanvideo15_i2v_image_source)
+    validate_image_source(image_source)
+    online_path = _generate_online_video(
+        omni_server=omni_server,
+        openai_client=openai_client,
+        image_source=image_source,
+        timeout_seconds=hunyuanvideo15_online_timeout_seconds,
+    )
+    assert online_path.exists(), f"Expected online video artifact at {online_path}"
+    online_metadata = probe_video(online_path)
+    assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
+
+
+@pytest.mark.benchmark
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_hunyuanvideo15_i2v_serving_matches_offline_video_similarity(
+    hunyuanvideo15_i2v_image_source: str | None,
+) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("HunyuanVideo-1.5 I2V video similarity test requires CUDA.")
+
+    probe_binary("ffmpeg")
+    probe_binary("ffprobe")
+    image_source = _resolve_image_source(hunyuanvideo15_i2v_image_source)
+    validate_image_source(image_source)
+    online_path, offline_path = _artifact_paths(image_source)
+    if not online_path.exists():
+        pytest.skip(f"Missing online artifact from prerequisite test: {online_path}")
+    if not offline_path.exists():
+        pytest.skip(f"Missing offline artifact from prerequisite test: {offline_path}")
+
+    online_metadata = probe_video(online_path)
+    offline_metadata = probe_video(offline_path)
+    assert online_metadata == offline_metadata, (
+        f"Video metadata mismatch:\n"
+        f"online={online_metadata}\n"
+        f"offline={offline_metadata}\n"
+        f"online_path={online_path}\n"
+        f"offline_path={offline_path}"
+    )
+    assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
+    assert_video_similarity_metrics(
+        label="hunyuanvideo15_i2v",
+        online_path=online_path,
+        offline_path=offline_path,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+    )
diff --git a/tests/e2e/accuracy/hunyuanvideo15_t2v/__init__.py b/tests/e2e/accuracy/hunyuanvideo15_t2v/__init__.py
new file mode 100644
index 000000000..bb679db4c
--- /dev/null
+++ b/tests/e2e/accuracy/hunyuanvideo15_t2v/__init__.py
@@ -0,0 +1 @@
+"""HunyuanVideo-1.5 T2V accuracy test cases."""
diff --git a/tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py b/tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py
new file mode 100644
index 000000000..f2c14e2b0
--- /dev/null
+++ b/tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py
@@ -0,0 +1,211 @@
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+from tests.e2e.accuracy.helpers import (
+    assert_video_metadata,
+    assert_video_similarity_metrics,
+    probe_binary,
+    probe_video,
+    send_video_request_with_timeout,
+)
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+
+MODEL_NAME = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
+PROMPT = "A cat walks on the grass, realistic style."
+WIDTH = 832
+HEIGHT = 480
+FPS = 24
+NUM_FRAMES = 33
+NUM_INFERENCE_STEPS = 50
+GUIDANCE_SCALE = 6.0
+FLOW_SHIFT = 5.0
+SEED = 42
+SSIM_THRESHOLD = 0.94
+PSNR_THRESHOLD = 28.0
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+WORKSPACE_ROOT = REPO_ROOT.parent
+RUNNER_PATH = REPO_ROOT / "examples" / "offline_inference" / "text_to_video" / "text_to_video.py"
+RESULT_ROOT = Path(__file__).parent / "result"
+ARTIFACT_DIR = RESULT_ROOT / "cat_grass"
+VIDEO_TIMEOUT_SECONDS = 60 * 60
+SERVER_CASES = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL_NAME,
+            server_args=[
+                "--flow-shift",
+                str(FLOW_SHIFT),
+            ],
+            env_dict={"VLLM_OMNI_STORAGE_PATH": str(RESULT_ROOT / "storage")},
+            use_omni=True,
+        ),
+        id="hunyuanvideo15_t2v_default",
+    )
+]
+
+
+def _runner_env() -> dict[str, str]:
+    env = os.environ.copy()
+    pythonpath_parts = [
+        str(REPO_ROOT),
+        str(WORKSPACE_ROOT / "diffusers" / "src"),
+    ]
+    existing_pythonpath = env.get("PYTHONPATH")
+    if existing_pythonpath:
+        pythonpath_parts.append(existing_pythonpath)
+    env["PYTHONPATH"] = os.pathsep.join(pythonpath_parts)
+    env.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    return env
+
+
+def _artifact_paths() -> tuple[Path, Path]:
+    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
+    return ARTIFACT_DIR / "online.mp4", ARTIFACT_DIR / "offline.mp4"
+
+
+def _build_offline_command(*, output_path: Path) -> list[str]:
+    return [
+        sys.executable,
+        str(RUNNER_PATH),
+        "--model",
+        MODEL_NAME,
+        "--prompt",
+        PROMPT,
+        "--height",
+        str(HEIGHT),
+        "--width",
+        str(WIDTH),
+        "--num-frames",
+        str(NUM_FRAMES),
+        "--num-inference-steps",
+        str(NUM_INFERENCE_STEPS),
+        "--guidance-scale",
+        str(GUIDANCE_SCALE),
+        "--flow-shift",
+        str(FLOW_SHIFT),
+        "--fps",
+        str(FPS),
+        "--seed",
+        str(SEED),
+        "--output",
+        str(output_path),
+    ]
+
+
+def _generate_offline_video() -> Path:
+    _, offline_path = _artifact_paths()
+    subprocess.run(
+        _build_offline_command(output_path=offline_path),
+        cwd=REPO_ROOT,
+        env=_runner_env(),
+        check=True,
+        timeout=VIDEO_TIMEOUT_SECONDS,
+    )
+    return offline_path
+
+
+def _generate_online_video(*, omni_server, openai_client, timeout_seconds: int) -> Path:
+    online_path, _ = _artifact_paths()
+    request_config = {
+        "model": omni_server.model,
+        "form_data": {
+            "prompt": PROMPT,
+            "height": HEIGHT,
+            "width": WIDTH,
+            "fps": FPS,
+            "num_frames": NUM_FRAMES,
+            "guidance_scale": GUIDANCE_SCALE,
+            "flow_shift": FLOW_SHIFT,
+            "num_inference_steps": NUM_INFERENCE_STEPS,
+            "seed": SEED,
+        },
+    }
+    online_video_bytes = send_video_request_with_timeout(
+        openai_client,
+        request_config,
+        timeout_seconds=timeout_seconds,
+    )
+    online_path.write_bytes(online_video_bytes)
+    return online_path
+
+
+@pytest.mark.benchmark
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_hunyuanvideo15_t2v_diffusers_offline_generates_video() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("HunyuanVideo-1.5 T2V offline accuracy test requires CUDA.")
+
+    probe_binary("ffprobe")
+    if not RUNNER_PATH.exists():
+        raise AssertionError(f"Offline runner does not exist: {RUNNER_PATH}")
+
+    offline_path = _generate_offline_video()
+    assert offline_path.exists(), f"Expected offline video artifact at {offline_path}"
+    offline_metadata = probe_video(offline_path)
+    assert_video_metadata(offline_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
+
+
+@pytest.mark.benchmark
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True)
+def test_hunyuanvideo15_t2v_online_serving_generates_video(
+    omni_server,
+    openai_client,
+    hunyuanvideo15_online_timeout_seconds: int,
+) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("HunyuanVideo-1.5 T2V online accuracy test requires CUDA.")
+
+    probe_binary("ffprobe")
+    online_path = _generate_online_video(
+        omni_server=omni_server,
+        openai_client=openai_client,
+        timeout_seconds=hunyuanvideo15_online_timeout_seconds,
+    )
+    assert online_path.exists(), f"Expected online video artifact at {online_path}"
+    online_metadata = probe_video(online_path)
+    assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
+
+
+@pytest.mark.benchmark
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_hunyuanvideo15_t2v_serving_matches_offline_video_similarity() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("HunyuanVideo-1.5 T2V video similarity test requires CUDA.")
+
+    probe_binary("ffmpeg")
+    probe_binary("ffprobe")
+    online_path, offline_path = _artifact_paths()
+    if not online_path.exists():
+        pytest.skip(f"Missing online artifact from prerequisite test: {online_path}")
+    if not offline_path.exists():
+        pytest.skip(f"Missing offline artifact from prerequisite test: {offline_path}")
+
+    online_metadata = probe_video(online_path)
+    offline_metadata = probe_video(offline_path)
+    assert online_metadata == offline_metadata, (
+        f"Video metadata mismatch:\n"
+        f"online={online_metadata}\n"
+        f"offline={offline_metadata}\n"
+        f"online_path={online_path}\n"
+        f"offline_path={offline_path}"
+    )
+    assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
+    assert_video_similarity_metrics(
+        label="hunyuanvideo15_t2v",
+        online_path=online_path,
+        offline_path=offline_path,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+    )
diff --git a/tests/e2e/accuracy/test_diffusers_backend_similarity.py b/tests/e2e/accuracy/test_diffusers_backend_similarity.py
index da05d77ec..26b39ea78 100644
--- a/tests/e2e/accuracy/test_diffusers_backend_similarity.py
+++ b/tests/e2e/accuracy/test_diffusers_backend_similarity.py
@@ -30,10 +30,14 @@ from tests.e2e.accuracy.helpers import (
     env_to_apply_ftfy_mock_in_subproc,
     model_output_dir,
 )
-from tests.e2e.accuracy.wan22_i2v.test_wan22_i2v_video_similarity import (
-    _parse_psnr_score,
-    _parse_ssim_score,
-    _run_ffmpeg_similarity,
+from tests.e2e.accuracy.helpers import (
+    parse_psnr_score as _parse_psnr_score,
+)
+from tests.e2e.accuracy.helpers import (
+    parse_ssim_score as _parse_ssim_score,
+)
+from tests.e2e.accuracy.helpers import (
+    run_ffmpeg_similarity as _run_ffmpeg_similarity,
 )
 from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
 from tests.helpers.mark import hardware_test
diff --git a/tests/e2e/accuracy/test_hunyuan_image3.py b/tests/e2e/accuracy/test_hunyuan_image3.py
index 6541866d4..7ed3c85a0 100644
--- a/tests/e2e/accuracy/test_hunyuan_image3.py
+++ b/tests/e2e/accuracy/test_hunyuan_image3.py
@@ -6,12 +6,15 @@ from __future__ import annotations
 import copy
 import gc
 import importlib
+import json
 import os
 import tempfile
 import time
+from dataclasses import dataclass
 from pathlib import Path
 
 import pytest
+import requests
 import torch
 import yaml
 from PIL import Image
@@ -19,16 +22,18 @@ from PIL import Image
 from tests.e2e.accuracy.helpers import (
     CLIPScorer,
     SemanticSimilarityScorer,
+    assert_similarity,
     compute_image_ssim_psnr,
     download_images,
     model_output_dir,
 )
-from tests.helpers.mark import hardware_marks
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner, OmniServer
 from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens, resolve_stop_token_ids
 
 os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA"
 
-pytestmark = [pytest.mark.local_model, pytest.mark.diffusion, *hardware_marks(res={"cuda": "H100"}, num_cards=8)]
+pytestmark = [pytest.mark.local_model, pytest.mark.diffusion]
 
 # ============================================================================
 # Configurable Parameters
@@ -72,58 +77,133 @@ THRESHOLDS = {
     "text_prefix_match": 10,  # First 10 characters must match exactly
     "cot_semantic_sim": 0.9,  # Full CoT semantic similarity
     # Image comparison
-    "clip_score": 85,  # CLIP image semantic similarity
-    "ssim": 0.20,  # Structural similarity
-    "psnr": 11.0,  # Peak signal-to-noise ratio (dB)
+    "clip_score": 90,  # CLIP image semantic similarity
+    "ssim": 0.26,  # Structural similarity
+    "psnr": 12.5,  # Peak signal-to-noise ratio (dB)
 }
+
+QUANT_PROMPT = "A brown and white dog is running on the grass."
+QUANT_HEIGHT, QUANT_WIDTH = 1024, 1024
+QUANT_PSNR_THRESHOLD = 10.0
+QUANT_SSIM_THRESHOLD = 0.20
+QUANT_CLIP_SCORE_THRESHOLD = 20.0
+QUANT_CLIP_SCORE_DROP_THRESHOLD = float(os.environ.get("HUNYUAN_IMAGE3_QUANT_CLIP_SCORE_DROP_THRESHOLD", "5.0"))
+QUANT_RUN_ENV = "HUNYUAN_IMAGE3_RUN_QUANT_ACCURACY"
+QUANT_BF16_ENV = "HUNYUAN_IMAGE3_BF16_MODEL"
+QUANT_FP8_ENV = "HUNYUAN_IMAGE3_FP8_MODEL"
+QUANT_NVFP4_ENV = "HUNYUAN_IMAGE3_NVFP4_MODEL"
+_TRUE_ENV_VALUES = {"1", "true", "yes", "on"}
 # fmt: off
-_BASE_CONFIG = {
-    "stage_args": [
+_DEPLOY_CONFIG = {
+    "pipeline": "hunyuan_image_3_moe",
+    "async_chunk": False,
+    "trust_remote_code": True,
+    "connectors": {
+        "shared_memory_connector": {
+            "name": "SharedMemoryConnector",
+            "extra": {"shm_threshold_bytes": 65536},
+        },
+    },
+    "stages": [
         {
-            "stage_id": 0, "stage_type": "llm",
-            "runtime": {"process": True, "devices": AR_DEVICES, "max_batch_size": 1, "requires_multimodal_data": True},
-            "engine_args": {
-                "model_stage": "AR", "model_arch": "HunyuanImage3ForCausalMM",
-                "worker_cls": "vllm_omni.worker.gpu_ar_worker.GPUARWorker",
-                "scheduler_cls": "vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler",
-                "gpu_memory_utilization": 0.95, "enforce_eager": True, "trust_remote_code": True,
-                "engine_output_type": "latent", "enable_prefix_caching": False,
-                "max_num_batched_tokens": 32768, "tensor_parallel_size": AR_TP_SIZE, "pipeline_parallel_size": 1,
-                "hf_overrides": {"rope_parameters": {"mrope_section": [0, 32, 32], "rope_type": "default"}},
-            },
-            "is_comprehension": False, "final_output": True, "final_output_type": "text",
-            "default_sampling_params": {
-                "temperature": 0.0, "top_p": 1, "top_k": -1, "max_tokens": 8192,
-                "stop_token_ids": [128025], "detokenize": True, "skip_special_tokens": False,
+            "stage_id": 0,
+            "is_comprehension": False,
+            "final_output": True,
+            "final_output_type": "text",
+            "max_num_seqs": 1,
+            "gpu_memory_utilization": 0.95,
+            "enforce_eager": True,
+            "trust_remote_code": True,
+            "max_num_batched_tokens": 32768,
+            "devices": AR_DEVICES,
+            "tensor_parallel_size": AR_TP_SIZE,
+            "hf_overrides": {
+                "rope_parameters": {"mrope_section": [0, 32, 32], "rope_type": "default"},
             },
             "output_connectors": {"to_stage_1": "shared_memory_connector"},
+            "default_sampling_params": {
+                "temperature": 0.0,
+                "top_p": 1,
+                "top_k": -1,
+                "max_tokens": 8192,
+                "stop_token_ids": [128025],
+                "detokenize": True,
+                "skip_special_tokens": False,
+            },
         },
         {
-            "stage_id": 1, "stage_type": "diffusion",
-            "runtime": {"process": True, "devices": DIT_DEVICES, "max_batch_size": 1, "requires_multimodal_data": True},
-            "engine_args": {
-                "model_stage": "dit", "model_arch": "HunyuanImage3ForCausalMM",
-                "enforce_eager": True, "trust_remote_code": True, "distributed_executor_backend": "mp",
-                "parallel_config": {"tensor_parallel_size": DIT_TP_SIZE, "enable_expert_parallel": True},
-            },
-            "engine_input_source": [0],
-            "custom_process_input_func": "vllm_omni.model_executor.stage_input_processors.hunyuan_image3.ar2diffusion",
-            "final_output": True, "final_output_type": "image",
-            "default_sampling_params": {"num_inference_steps": NUM_INFERENCE_STEPS, "guidance_scale": GUIDANCE_SCALE},
+            "stage_id": 1,
+            "max_num_seqs": 1,
+            "enforce_eager": True,
+            "trust_remote_code": True,
+            "devices": DIT_DEVICES,
+            "distributed_executor_backend": "mp",
+            "parallel_config": {"tensor_parallel_size": DIT_TP_SIZE, "enable_expert_parallel": True},
             "input_connectors": {"from_stage_0": "shared_memory_connector"},
+            "default_sampling_params": {
+                "num_inference_steps": NUM_INFERENCE_STEPS,
+                "guidance_scale": GUIDANCE_SCALE,
+            },
         },
     ],
-    "runtime": {
-        "enabled": True,
-        "connectors": {"shared_memory_connector": {
-            "name": "SharedMemoryConnector",
-            "extra": {"shm_threshold_bytes": 65536}
-        }},
-        "edges": [{"from": 0, "to": 1}],
-    },
+    "edges": [{"from": 0, "to": 1}],
 }
 # fmt: on
 
+_QUANT_DIT_CONFIG = {
+    "pipeline": "hunyuan_image_3_moe",
+    "async_chunk": False,
+    "trust_remote_code": True,
+    "stages": [
+        {
+            "stage_id": 0,
+            "model_stage": "dit",
+            "enforce_eager": True,
+            "trust_remote_code": True,
+            "devices": "0,1",
+            "distributed_executor_backend": "mp",
+            "force_cutlass_fp8": True,
+            "moe_backend": "cutlass",
+            "parallel_config": {
+                "tensor_parallel_size": 2,
+                "enable_expert_parallel": True,
+            },
+            "omni_kv_config": {"need_recv_cache": True},
+            "final_output": True,
+            "final_output_type": "image",
+            "is_comprehension": False,
+            "default_sampling_params": {"seed": SEED},
+        }
+    ],
+}
+
+
+@dataclass(frozen=True)
+class _QuantAccuracyCase:
+    name: str
+    model_env: str
+    nvfp4_backend: str | None = None
+
+
+def _quant_accuracy_cases() -> list[pytest.ParameterSet]:
+    cases = [
+        _QuantAccuracyCase(name="fp8", model_env=QUANT_FP8_ENV),
+        _QuantAccuracyCase(name="mixed_nvfp4", model_env=QUANT_NVFP4_ENV, nvfp4_backend="cutlass"),
+    ]
+    params: list[pytest.ParameterSet] = []
+    run_quant_accuracy = os.environ.get(QUANT_RUN_ENV, "").lower() in _TRUE_ENV_VALUES
+    for case in cases:
+        marks = []
+        if not run_quant_accuracy:
+            marks.append(pytest.mark.skip(reason=f"Set {QUANT_RUN_ENV}=1 to run HunyuanImage3 quant accuracy."))
+        if not os.environ.get(QUANT_BF16_ENV):
+            marks.append(pytest.mark.skip(reason=f"Set {QUANT_BF16_ENV} to run HunyuanImage3 quant accuracy."))
+        if not os.environ.get(case.model_env):
+            marks.append(pytest.mark.skip(reason=f"Set {case.model_env} to run this quant accuracy case."))
+        params.append(pytest.param(case, id=case.name, marks=marks))
+    return params
+
+
 # fmt: off
 COT_REF = ("首先，我分析所有输入图像：图像1是一个圆形的logo，设计现代且抽象。它由不同色调的蓝色（深蓝、中蓝、浅蓝）和白色构成，这些色块以流畅的曲线相互交织，形成一个动态的、类似旋涡或波浪的图案。整个logo是扁平化的矢量图形，背景为纯黑色。图像2展示了四个并排摆放的卡通动物造型冰箱贴，"
            "它们被放置在灰色的织物背景上。这些冰箱贴的关键特征是其材质：它们具有光滑、高光的珐琅或烤漆质感，边缘有明显的金属包边，整体呈现出一种立体的、有厚度的实体感。用户的指令是“基于图一的logo，参考图二中冰箱贴的材质，制作一个新的冰箱贴”。这个指令要求将一个二维的平面设计（logo）"
@@ -138,16 +218,30 @@ COT_REF = ("首先，我分析所有输入图像：图像1是一个圆形的logo
 
 
 def _make_config(enable_kv_reuse: bool, path: Path) -> None:
-    config = copy.deepcopy(_BASE_CONFIG)
-    config["stage_args"][0]["engine_args"]["omni_kv_config"] = {"need_send_cache": enable_kv_reuse}
-    config["stage_args"][1]["engine_args"]["omni_kv_config"] = {"need_recv_cache": enable_kv_reuse}
+    config = copy.deepcopy(_DEPLOY_CONFIG)
+    config["stages"][0]["omni_kv_config"] = {"need_send_cache": enable_kv_reuse}
+    config["stages"][1]["omni_kv_config"] = {"need_recv_cache": enable_kv_reuse}
     path.write_text(yaml.dump(config, default_flow_style=False, sort_keys=False))
 
 
-def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Image, str, float]:
+def _quant_devices() -> str:
+    return os.environ.get("HUNYUAN_IMAGE3_QUANT_DEVICES", "0,1")
+
+
+def _quant_tensor_parallel_size() -> int:
+    return int(os.environ.get("HUNYUAN_IMAGE3_QUANT_TP", str(len(_quant_devices().split(",")))))
+
+
+def _make_quant_dit_config(path: Path) -> None:
+    config = copy.deepcopy(_QUANT_DIT_CONFIG)
+    config["stages"][0]["devices"] = _quant_devices()
+    config["stages"][0]["parallel_config"]["tensor_parallel_size"] = _quant_tensor_parallel_size()
+    path.write_text(yaml.dump(config, default_flow_style=False, sort_keys=False))
+
+
+def _run_offline(deploy_config_path: str, output_path: Path) -> tuple[Image.Image, str, float]:
     from transformers import AutoTokenizer
 
-    from tests.helpers.runtime import OmniRunner
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
     from vllm_omni.platforms import current_omni_platform
 
@@ -163,7 +257,7 @@ def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Imag
     system_prompt_type = result.system_prompt_type
 
     ar_stop_token_ids = resolve_stop_token_ids(task="it2i", bot_task="think_recaption", tokenizer=tokenizer)
-    with OmniRunner(MODEL_PATH, stage_configs_path=stage_configs_path) as runner:
+    with OmniRunner(MODEL_PATH, deploy_config=deploy_config_path) as runner:
         params_list = list(runner.omni.default_sampling_params_list)
         for sp in params_list:
             if isinstance(sp, OmniDiffusionSamplingParams):
@@ -220,6 +314,166 @@ def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Imag
     return image, cot_text, elapsed
 
 
+def _run_online(stage_configs_path: str, output_path: Path) -> tuple[Image.Image, str, float]:
+    from benchmarks.accuracy.common import decode_base64_image, pil_to_png_bytes
+
+    server_args = [
+        "--stage-configs-path",
+        stage_configs_path,
+        "--stage-init-timeout",
+        "300",
+        "--init-timeout",
+        "900",
+    ]
+    try:
+        with OmniServer(MODEL_PATH, server_args, use_omni=True) as omni_server:
+            images = download_images(TEST_IMAGE_URLS)
+            t0 = time.perf_counter()
+            response = requests.post(
+                f"http://{omni_server.host}:{omni_server.port}/v1/images/edits",
+                data={
+                    "model": omni_server.model,
+                    "prompt": PROMPT,
+                    "n": 1,
+                    "response_format": "b64_json",
+                    "num_inference_steps": NUM_INFERENCE_STEPS,
+                    "guidance_scale": GUIDANCE_SCALE,
+                    "seed": SEED,
+                    "sys_type": "en_unified",
+                    "bot_task": "think_recaption",
+                    "size": "1280x720",
+                },
+                files=[
+                    ("image", (f"image_{i}.png", pil_to_png_bytes(img), "image/png")) for i, img in enumerate(images)
+                ],
+                timeout=600,
+            )
+            elapsed = time.perf_counter() - t0
+            if not response.ok:
+                print(f"[ONLINE] HTTP {response.status_code} response body: {response.text}")
+            response.raise_for_status()
+            payload = response.json()
+            assert len(payload["data"]) == 1
+            image = decode_base64_image(payload["data"][0]["b64_json"])
+            image.load()
+            image.save(output_path / "image_online.png")
+            cot_text = payload.get("cot_output") or ""
+            (output_path / "cot_online.txt").write_text(cot_text, encoding="utf-8")
+            return image, cot_text, elapsed
+    finally:
+        gc.collect()
+        if torch.accelerator.is_available():
+            torch.accelerator.empty_cache()
+
+
+@pytest.mark.skipif(
+    torch.accelerator.device_count() < AR_TP_SIZE + DIT_TP_SIZE,
+    reason=f"Needs {AR_TP_SIZE + DIT_TP_SIZE}+ GPUs ({AR_TP_SIZE} AR + {DIT_TP_SIZE} DiT)",
+)
+def test_image_to_image_alignment_online(accuracy_artifact_root: Path, accuracy_assets_root: Path) -> None:
+    """Online API test: same pipeline, same seed as offline → PSNR >= 10 dB."""
+    if importlib.util.find_spec("FlagEmbedding") is None:
+        raise ImportError("Missing dependency: FlagEmbedding\nInstall with: pip install FlagEmbedding")
+    from tabulate import tabulate
+
+    output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME + "-online")
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmp = Path(tmpdir)
+        _make_config(True, tmp / "online.yaml")
+        online_image, online_cot, _ = _run_online(str(tmp / "online.yaml"), output_dir)
+
+    online_cot = online_cot.lstrip("\n")
+    scorer = SemanticSimilarityScorer()
+    clip_scorer = CLIPScorer()
+    cot_results = scorer.text_similarity(online_cot, COT_REF)
+    image_ref = Image.open(str(accuracy_assets_root / "hunyuan_image_ref.png")).convert("RGB")
+    image_clip_score = clip_scorer.image_image_score(online_image, image_ref)
+    ssim_value, psnr_value = compute_image_ssim_psnr(prediction=online_image, reference=image_ref, compare_mode="RGB")
+
+    table = [
+        ["COT similarity to reference", f"{cot_results['cot_semantic_sim']:.4f}", 0.9644],
+        ["COT prefix match", f"{cot_results['text_prefix_match_count']:.4f}", 29],
+        ["Image-Image similarity", f"{image_clip_score:.4f}", 94.5538],
+        ["SSIM", f"{ssim_value:.4f}", 0.242],
+        ["PSNR (dB)", f"{psnr_value:.2f}", 14.1],
+    ]
+    print("[ONLINE] " + tabulate(table, headers=["Metric", "Value", "L20x Reference"], tablefmt="grid"))
+
+    assert cot_results["cot_semantic_sim"] >= THRESHOLDS["cot_semantic_sim"], (
+        f"[ONLINE] COT semantic similarity {cot_results['cot_semantic_sim']:.4f} below threshold {THRESHOLDS['cot_semantic_sim']}"
+    )
+    assert cot_results["text_prefix_match_count"] >= THRESHOLDS["text_prefix_match"], (
+        f"[ONLINE] COT prefix match {cot_results['text_prefix_match_count']} below threshold {THRESHOLDS['text_prefix_match']}"
+    )
+    assert image_clip_score >= THRESHOLDS["clip_score"], (
+        f"[ONLINE] Image-Image similarity {image_clip_score:.4f} below threshold {THRESHOLDS['clip_score']}"
+    )
+    assert ssim_value >= THRESHOLDS["ssim"], f"[ONLINE] SSIM {ssim_value:.4f} below threshold {THRESHOLDS['ssim']}"
+    assert psnr_value >= THRESHOLDS["psnr"], (
+        f"[ONLINE] PSNR {psnr_value:.2f} dB below threshold {THRESHOLDS['psnr']} dB"
+    )
+
+
+def _extract_image(outputs) -> Image.Image:
+    assert outputs, "Pipeline produced no outputs"
+    for output in outputs:
+        images = getattr(output, "images", None)
+        request_output = getattr(output, "request_output", None)
+        if not images and request_output is not None:
+            images = getattr(request_output, "images", None)
+        if images:
+            image = images[0].convert("RGB")
+            image.load()
+            return image
+    raise AssertionError("Pipeline output had no images")
+
+
+def _run_dit_model(
+    model: str,
+    deploy_config_path: str,
+    output_path: Path,
+    *,
+    nvfp4_backend: str | None = None,
+) -> tuple[Image.Image, float]:
+    from tests.helpers.runtime import OmniRunner
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.platforms import current_omni_platform
+
+    old_backend = os.environ.get("VLLM_NVFP4_GEMM_BACKEND")
+    if nvfp4_backend is not None:
+        os.environ["VLLM_NVFP4_GEMM_BACKEND"] = nvfp4_backend
+
+    try:
+        with OmniRunner(model, deploy_config=deploy_config_path, mode="text-to-image") as runner:
+            generator = torch.Generator(device=current_omni_platform.device_type or "cuda").manual_seed(SEED)
+            params = OmniDiffusionSamplingParams(
+                height=QUANT_HEIGHT,
+                width=QUANT_WIDTH,
+                seed=SEED,
+                generator=generator,
+                num_inference_steps=20,
+                guidance_scale=4.0,
+                guidance_scale_provided=True,
+            )
+            t0 = time.perf_counter()
+            outputs = runner.omni.generate({"prompt": QUANT_PROMPT}, params)
+            elapsed = time.perf_counter() - t0
+            image = _extract_image(outputs)
+            image.save(output_path)
+            return image, elapsed
+    finally:
+        if nvfp4_backend is not None:
+            if old_backend is None:
+                os.environ.pop("VLLM_NVFP4_GEMM_BACKEND", None)
+            else:
+                os.environ["VLLM_NVFP4_GEMM_BACKEND"] = old_backend
+        gc.collect()
+        if torch.accelerator.is_available():
+            torch.accelerator.empty_cache()
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=8)
 @pytest.mark.skipif(
     torch.accelerator.device_count() < AR_TP_SIZE + DIT_TP_SIZE,
     reason=f"Needs {AR_TP_SIZE + DIT_TP_SIZE}+ GPUs ({AR_TP_SIZE} AR + {DIT_TP_SIZE} DiT)",
@@ -265,3 +519,88 @@ def test_image_to_image_alignment(accuracy_artifact_root: Path, accuracy_assets_
     )
     assert ssim_value >= THRESHOLDS["ssim"], f"SSIM {ssim_value:.4f} is below threshold {THRESHOLDS['ssim']}"
     assert psnr_value >= THRESHOLDS["psnr"], f"PSNR {psnr_value:.2f} dB is below threshold {THRESHOLDS['psnr']} dB"
+
+
+@pytest.mark.parametrize("case", _quant_accuracy_cases())
+@pytest.mark.skipif(torch.accelerator.device_count() < 2, reason="Needs 2+ GPUs for HunyuanImage3 DiT")
+def test_quantized_dit_matches_bf16_accuracy(
+    case: _QuantAccuracyCase,
+    accuracy_artifact_root: Path,
+) -> None:
+    """Quantized DiT checkpoints should preserve prompt-aligned image quality."""
+    output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME + "-quant")
+    bf16_model = os.environ[QUANT_BF16_ENV]
+    quant_model = os.environ[case.model_env]
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        deploy_config_path = Path(tmpdir) / "hunyuan_image3_quant_dit.yaml"
+        _make_quant_dit_config(deploy_config_path)
+
+        bf16_image, bf16_time = _run_dit_model(
+            bf16_model,
+            str(deploy_config_path),
+            output_dir / "bf16.png",
+        )
+        quant_image, quant_time = _run_dit_model(
+            quant_model,
+            str(deploy_config_path),
+            output_dir / f"{case.name}.png",
+            nvfp4_backend=case.nvfp4_backend,
+        )
+
+    ssim_score, psnr_score = compute_image_ssim_psnr(
+        prediction=quant_image,
+        reference=bf16_image,
+    )
+    assert_similarity(
+        model_name=f"{MODEL_NAME} {case.name} vs bf16",
+        vllm_image=quant_image,
+        diffusers_image=bf16_image,
+        ssim_threshold=QUANT_SSIM_THRESHOLD,
+        psnr_threshold=QUANT_PSNR_THRESHOLD,
+        width=QUANT_WIDTH,
+        height=QUANT_HEIGHT,
+    )
+
+    clip_scorer = CLIPScorer()
+    bf16_clip_score = clip_scorer.score(bf16_image, QUANT_PROMPT)
+    quant_clip_score = clip_scorer.score(quant_image, QUANT_PROMPT)
+    clip_score_drop = bf16_clip_score - quant_clip_score
+
+    metrics = {
+        "case": case.name,
+        "bf16_model": bf16_model,
+        "quant_model": quant_model,
+        "prompt": QUANT_PROMPT,
+        "seed": SEED,
+        "height": QUANT_HEIGHT,
+        "width": QUANT_WIDTH,
+        "num_inference_steps": 20,
+        "guidance_scale": 4.0,
+        "bf16_elapsed_s": bf16_time,
+        "quant_elapsed_s": quant_time,
+        "ssim": ssim_score,
+        "psnr": psnr_score,
+        "bf16_clip_score": bf16_clip_score,
+        "quant_clip_score": quant_clip_score,
+        "clip_score_drop": clip_score_drop,
+    }
+    metrics_path = output_dir / f"{case.name}_metrics.json"
+    metrics_path.write_text(json.dumps(metrics, indent=2, sort_keys=True), encoding="utf-8")
+
+    print(f"\nHunyuanImage3 quant accuracy ({case.name})")
+    print(f"  bf16 model:       {bf16_model}")
+    print(f"  quant model:      {quant_model}")
+    print(f"  BF16 CLIP score:  {bf16_clip_score:.4f}")
+    print(f"  quant CLIP score: {quant_clip_score:.4f} threshold>={QUANT_CLIP_SCORE_THRESHOLD:.4f}")
+    print(f"  CLIP score drop:  {clip_score_drop:.4f} threshold<={QUANT_CLIP_SCORE_DROP_THRESHOLD:.4f}")
+    print(f"  metrics:          {metrics_path}")
+
+    assert quant_clip_score >= QUANT_CLIP_SCORE_THRESHOLD, (
+        f"{case.name} CLIP score below threshold: got {quant_clip_score:.4f}, "
+        f"expected >= {QUANT_CLIP_SCORE_THRESHOLD:.4f}"
+    )
+    assert clip_score_drop <= QUANT_CLIP_SCORE_DROP_THRESHOLD, (
+        f"{case.name} CLIP score drop too large: got {clip_score_drop:.4f}, "
+        f"expected <= {QUANT_CLIP_SCORE_DROP_THRESHOLD:.4f}"
+    )
diff --git a/tests/e2e/accuracy/test_qwen_image_edit.py b/tests/e2e/accuracy/test_qwen_image_edit.py
index ea4a9ff1b..338758cc0 100644
--- a/tests/e2e/accuracy/test_qwen_image_edit.py
+++ b/tests/e2e/accuracy/test_qwen_image_edit.py
@@ -20,7 +20,7 @@ pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
 
 
 SINGLE_MODEL = "Qwen/Qwen-Image-Edit"
-MULTIPLE_MODEL = "Qwen/Qwen-Image-Edit-2509"
+MULTIPLE_MODEL = "Qwen/Qwen-Image-Edit-2511"
 WIDTH = 512
 HEIGHT = 512
 NUM_INFERENCE_STEPS = 20
@@ -253,37 +253,6 @@ def _diffusers_output_single_image(accuracy_artifact_root: Path, qwen_bear_image
     )
 
 
-def _vllm_omni_output_multiple_image(
-    accuracy_artifact_root: Path,
-    qwen_bear_image: Image.Image,
-    rabbit_image: Image.Image,
-) -> Image.Image:
-    output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL)
-    output_path = output_dir / "vllm_omni_multiple.png"
-    with OmniServer(model=MULTIPLE_MODEL, serve_args=SERVER_ARGS) as server:
-        output = _run_vllm_omni_image_edit(
-            omni_server=server,
-            prompt=PROMPT_MULTIPLE_IMAGE,
-            input_images=[qwen_bear_image, rabbit_image],
-            output_path=output_path,
-        )
-    return output
-
-
-def _diffusers_output_multiple_image(
-    accuracy_artifact_root: Path, qwen_bear_image: Image.Image, rabbit_image: Image.Image
-) -> Image.Image:
-    output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL)
-    output_path = output_dir / "diffusers_multiple.png"
-    return _run_diffusers_image_edit(
-        model=MULTIPLE_MODEL,
-        pipeline_class=QwenImageEditPlusPipeline,
-        prompt=PROMPT_MULTIPLE_IMAGE,
-        input_images=[qwen_bear_image, rabbit_image],
-        output_path=output_path,
-    )
-
-
 @pytest.mark.benchmark
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 def test_qwen_image_edit_single_matches_diffusers(
@@ -309,37 +278,6 @@ def test_qwen_image_edit_single_matches_diffusers(
     )
 
 
-@pytest.mark.benchmark
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.skip(
-    reason="Skipping as the second image seems to be ignored by the API. Will come back to this later after #2772 is merged."
-)
-def test_qwen_image_edit_multiple_matches_diffusers(
-    accuracy_artifact_root: Path,
-    qwen_bear_image: Image.Image,
-    rabbit_image: Image.Image,
-) -> None:
-    vllm_image = _vllm_omni_output_multiple_image(
-        accuracy_artifact_root=accuracy_artifact_root,
-        qwen_bear_image=qwen_bear_image,
-        rabbit_image=rabbit_image,
-    )
-    diffusers_image = _diffusers_output_multiple_image(
-        accuracy_artifact_root=accuracy_artifact_root,
-        qwen_bear_image=qwen_bear_image,
-        rabbit_image=rabbit_image,
-    )
-    assert_similarity(
-        model_name=MULTIPLE_MODEL,
-        vllm_image=vllm_image,
-        diffusers_image=diffusers_image,
-        width=WIDTH,
-        height=HEIGHT,
-        ssim_threshold=SSIM_THRESHOLD,
-        psnr_threshold=PSNR_THRESHOLD,
-    )
-
-
 @pytest.mark.benchmark
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 def test_qwen_image_edit_2511_matches_diffusers_pixelwise(accuracy_artifact_root: Path) -> None:
diff --git a/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py b/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py
index 17e4fc0f2..a09d2330a 100644
--- a/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py
+++ b/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py
@@ -1,20 +1,10 @@
 from __future__ import annotations
 
-import json
 import os
-import re
-import shutil
 import subprocess
 import sys
-import time
-from base64 import b64decode, b64encode
-from fractions import Fraction
-from hashlib import sha1
-from io import BytesIO
-from mimetypes import guess_type
+from base64 import b64decode
 from pathlib import Path
-from typing import Any
-from urllib.parse import urlparse
 
 import pytest
 import requests
@@ -22,6 +12,38 @@ import torch
 from diffusers import UniPCMultistepScheduler
 from PIL import Image
 
+from tests.e2e.accuracy.helpers import (
+    assert_video_metadata,
+    assert_video_similarity_metrics,
+    video_artifact_dir,
+)
+from tests.e2e.accuracy.helpers import (
+    build_online_image_reference as _build_online_image_reference,
+)
+from tests.e2e.accuracy.helpers import (
+    online_timeout_seconds as _online_timeout_seconds,
+)
+from tests.e2e.accuracy.helpers import (
+    parse_psnr_score as _parse_psnr_score,
+)
+from tests.e2e.accuracy.helpers import (
+    parse_ssim_score as _parse_ssim_score,
+)
+from tests.e2e.accuracy.helpers import (
+    parse_video_metadata as _parse_video_metadata,
+)
+from tests.e2e.accuracy.helpers import (
+    probe_binary as _probe_binary,
+)
+from tests.e2e.accuracy.helpers import (
+    probe_video as _probe_video,
+)
+from tests.e2e.accuracy.helpers import (
+    send_video_request_with_timeout as _send_video_request_with_timeout,
+)
+from tests.e2e.accuracy.helpers import (
+    validate_image_source as _validate_image_source,
+)
 from tests.e2e.accuracy.wan22_i2v.run_wan22_i2v_diffusers_cp import (
     _configure_scheduler,
     _ensure_wan_ftfy_fallback,
@@ -218,8 +240,6 @@ WORKSPACE_ROOT = REPO_ROOT.parent
 RUNNER_PATH = Path(__file__).with_name("run_wan22_i2v_diffusers_cp.py")
 RESULT_ROOT = Path(__file__).parent / "result"
 VIDEO_TIMEOUT_SECONDS = 60 * 60
-_SSIM_RE = re.compile(r"All:(?P<score>[0-9.]+)")
-_PSNR_RE = re.compile(r"average:(?P<score>[0-9.]+)")
 SERVER_CASES = [
     pytest.param(
         OmniServerParams(
@@ -238,48 +258,6 @@ SERVER_CASES = [
 ]
 
 
-def _parse_video_metadata(payload: dict[str, Any]) -> dict[str, int | float]:
-    streams = payload.get("streams")
-    if not isinstance(streams, list) or not streams:
-        raise ValueError(f"ffprobe payload did not include video streams: {payload}")
-
-    stream = streams[0]
-    width = int(stream["width"])
-    height = int(stream["height"])
-    fps = float(Fraction(stream["avg_frame_rate"]))
-    frame_count_value = stream.get("nb_read_frames") or stream.get("nb_frames")
-    if frame_count_value is None:
-        raise ValueError(f"ffprobe payload did not include frame count: {payload}")
-    frame_count = int(frame_count_value)
-    return {
-        "width": width,
-        "height": height,
-        "fps": fps,
-        "frame_count": frame_count,
-    }
-
-
-def _parse_ssim_score(output: str) -> float:
-    match = _SSIM_RE.search(output)
-    if match is None:
-        raise ValueError(f"Could not parse SSIM score from ffmpeg output:\n{output}")
-    return float(match.group("score"))
-
-
-def _parse_psnr_score(output: str) -> float:
-    match = _PSNR_RE.search(output)
-    if match is None:
-        raise ValueError(f"Could not parse PSNR score from ffmpeg output:\n{output}")
-    return float(match.group("score"))
-
-
-def _probe_binary(binary: str) -> str:
-    resolved = shutil.which(binary)
-    if resolved is None:
-        pytest.skip(f"{binary} is required for Wan2.2 video similarity e2e test.")
-    return resolved
-
-
 def _build_diffusers_command(
     *,
     runner_path: Path,
@@ -343,90 +321,8 @@ def _resolve_image_source(configured: str | None) -> str:
     return configured
 
 
-def _is_remote_image_source(source: str) -> bool:
-    return source.startswith("http://") or source.startswith("https://")
-
-
-def _validate_image_source(source: str) -> None:
-    if _is_remote_image_source(source):
-        response = requests.get(source, timeout=60)
-        response.raise_for_status()
-        return
-
-    image_path = Path(source)
-    if not image_path.exists():
-        raise FileNotFoundError(f"Local image source does not exist: {image_path}")
-
-
-def _build_online_image_reference(source: str) -> str:
-    if _is_remote_image_source(source) or source.startswith("data:image"):
-        return source
-
-    image_path = Path(source)
-    mime_type = guess_type(image_path.name)[0] or "image/png"
-    encoded = b64encode(image_path.read_bytes()).decode("ascii")
-    return f"data:{mime_type};base64,{encoded}"
-
-
-def _probe_video(path: Path) -> dict[str, int | float]:
-    _probe_binary("ffprobe")
-    result = subprocess.run(
-        [
-            "ffprobe",
-            "-v",
-            "error",
-            "-select_streams",
-            "v:0",
-            "-count_frames",
-            "-show_entries",
-            "stream=width,height,avg_frame_rate,nb_read_frames",
-            "-of",
-            "json",
-            str(path),
-        ],
-        capture_output=True,
-        text=True,
-        check=True,
-    )
-    return _parse_video_metadata(json.loads(result.stdout))
-
-
-def _run_ffmpeg_similarity(filter_name: str, first: Path, second: Path) -> str:
-    _probe_binary("ffmpeg")
-    result = subprocess.run(
-        [
-            "ffmpeg",
-            "-hide_banner",
-            "-i",
-            str(first),
-            "-i",
-            str(second),
-            "-lavfi",
-            f"[0:v][1:v]{filter_name}",
-            "-f",
-            "null",
-            "-",
-        ],
-        capture_output=True,
-        text=True,
-        check=True,
-    )
-    return result.stderr
-
-
-def _online_timeout_seconds(configured: int | None) -> int:
-    return int(configured) if configured is not None else 1200
-
-
 def _artifact_dir(image_source: str) -> Path:
-    if _is_remote_image_source(image_source):
-        source_name = Path(urlparse(image_source).path).stem or "remote"
-    else:
-        source_name = Path(image_source).stem or "local"
-
-    safe_name = re.sub(r"[^A-Za-z0-9_.-]+", "_", source_name)
-    digest = sha1(image_source.encode("utf-8")).hexdigest()[:8]
-    return RESULT_ROOT / f"{safe_name}-{digest}"
+    return video_artifact_dir(RESULT_ROOT, image_source)
 
 
 def _artifact_paths(image_source: str) -> tuple[Path, Path, Path]:
@@ -439,55 +335,6 @@ def _artifact_paths(image_source: str) -> tuple[Path, Path, Path]:
     )
 
 
-def _send_video_request_with_timeout(
-    openai_client,
-    request_config: dict[str, Any],
-    *,
-    timeout_seconds: int,
-) -> bytes:
-    form_data = request_config.get("form_data")
-    if not isinstance(form_data, dict):
-        raise ValueError("Video request_config must contain 'form_data'")
-
-    if not form_data.get("prompt"):
-        raise ValueError("Video request_config['form_data'] must contain 'prompt'")
-
-    normalized_form_data = {key: str(value) for key, value in form_data.items() if value is not None}
-
-    files: dict[str, tuple[str, BytesIO, str]] = {}
-    image_reference = request_config.get("image_reference")
-    if image_reference:
-        if image_reference.startswith("data:image"):
-            header, encoded = image_reference.split(",", 1)
-            content_type = header.split(";")[0].removeprefix("data:")
-            extension = content_type.split("/")[-1]
-            file_data = b64decode(encoded)
-            files["input_reference"] = (
-                f"reference.{extension}",
-                BytesIO(file_data),
-                content_type,
-            )
-        else:
-            normalized_form_data["image_reference"] = json.dumps({"image_url": image_reference})
-
-    start_time = time.perf_counter()
-    create_url = openai_client._build_url("/v1/videos")
-    response = requests.post(
-        create_url,
-        data=normalized_form_data,
-        files=files,
-        headers={"Accept": "application/json"},
-        timeout=60,
-    )
-    response.raise_for_status()
-    job_data = response.json()
-    video_id = job_data["id"]
-    openai_client._wait_until_video_completed(video_id, timeout_seconds=timeout_seconds)
-    video_content = openai_client._download_video_content(video_id)
-    print(f"online_video_e2e_latency_s={time.perf_counter() - start_time:.3f}")
-    return video_content
-
-
 def _generate_online_video(
     *,
     omni_server,
@@ -557,10 +404,7 @@ def test_wan22_i2v_diffusers_offline_generates_video(
     assert offline_path.exists(), f"Expected offline video artifact at {offline_path}"
     assert offline_metadata_path.exists(), f"Expected offline metadata artifact at {offline_metadata_path}"
     offline_metadata = _probe_video(offline_path)
-    assert offline_metadata["width"] == WIDTH
-    assert offline_metadata["height"] == HEIGHT
-    assert offline_metadata["fps"] == float(FPS)
-    assert offline_metadata["frame_count"] == NUM_FRAMES
+    assert_video_metadata(offline_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
 
 
 @pytest.mark.benchmark
@@ -586,10 +430,7 @@ def test_wan22_i2v_online_serving_generates_video(
     )
     assert online_path.exists(), f"Expected online video artifact at {online_path}"
     online_metadata = _probe_video(online_path)
-    assert online_metadata["width"] == WIDTH
-    assert online_metadata["height"] == HEIGHT
-    assert online_metadata["fps"] == float(FPS)
-    assert online_metadata["frame_count"] == NUM_FRAMES
+    assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
 
 
 @pytest.mark.benchmark
@@ -627,42 +468,12 @@ def test_wan22_i2v_serving_matches_diffusers_video_similarity(
         f"online_path={online_path}\n"
         f"offline_path={offline_path}"
     )
-    assert online_metadata["width"] == WIDTH
-    assert online_metadata["height"] == HEIGHT
-    assert online_metadata["fps"] == float(FPS)
-    assert online_metadata["frame_count"] == NUM_FRAMES
-
-    ssim_output = _run_ffmpeg_similarity("ssim", online_path, offline_path)
-    psnr_output = _run_ffmpeg_similarity("psnr", online_path, offline_path)
-    ssim_score = _parse_ssim_score(ssim_output)
-    psnr_score = _parse_psnr_score(psnr_output)
-
-    print("wan22_i2v_similarity metrics:")
-    print(
-        "  SSIM:"
-        f" value={ssim_score:.6f},"
-        f" threshold>={SSIM_THRESHOLD:.6f},"
-        " range=[0, 1],"
-        " higher_is_better=True,"
-        " interpretation=structural_similarity"
+    assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
+    assert_video_similarity_metrics(
+        label="wan22_i2v",
+        online_path=online_path,
+        offline_path=offline_path,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
     )
-    print(
-        "  PSNR:"
-        f" value={psnr_score:.6f} dB,"
-        f" threshold>={PSNR_THRESHOLD:.6f} dB,"
-        " range=[0, +inf),"
-        " higher_is_better=True,"
-        " interpretation=pixel_error_in_decibels"
-    )
-    print(f"online_video={online_path}")
-    print(f"offline_video={offline_path}")
     print(f"offline_metadata={offline_metadata_path}")
-
-    assert ssim_score >= SSIM_THRESHOLD, (
-        f"SSIM below threshold: got {ssim_score:.6f}, expected >= {SSIM_THRESHOLD:.6f}. "
-        f"online={online_path} offline={offline_path}"
-    )
-    assert psnr_score >= PSNR_THRESHOLD, (
-        f"PSNR below threshold: got {psnr_score:.6f}, expected >= {PSNR_THRESHOLD:.6f}. "
-        f"online={online_path} offline={offline_path}"
-    )
diff --git a/tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py b/tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py
index bd3f2e099..dc701af7a 100644
--- a/tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py
+++ b/tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py
@@ -21,10 +21,12 @@ Usage:
 
 from __future__ import annotations
 
+import asyncio
 import uuid
 from contextlib import ExitStack
 
 import pytest
+import torch
 
 from tests.helpers.mark import hardware_test
 from vllm_omni.entrypoints.async_omni import AsyncOmni
@@ -150,3 +152,101 @@ async def test_generate_after_list_loras_inline_mode():
         assert last_output is not None
         assert isinstance(last_output, OmniRequestOutput)
         assert last_output.images, "Expected at least one generated image"
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+async def test_sleep_memory_reclaimed_custom_pipeline():
+    """sleep(level=1) must physically reclaim CuMemAllocator-tracked memory for
+    custom_pipeline.
+
+    Regression test for: custom pipelines constructed under ``with target_device:``
+    (CUDA default-device context) caused safetensors >=0.20.0 to use a
+    direct-to-GPU fast path (cudaMalloc via the driver API) that bypasses
+    CuMemAllocator, leaving weights invisible to sleep() and pinned in GPU
+    memory after the call.
+
+    The fix moves custom_pipeline init outside the CUDA context so all weights
+    go through the caching allocator and are therefore fully reclaimed by
+    sleep(level=1).  A non-zero ``CuMemAllocator.get_current_usage()`` after
+    sleep is the direct signal that the bypass is still occurring.
+    """
+    with ExitStack() as after:
+        engine = AsyncOmni(
+            model=MODEL,
+            custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
+            worker_extension_cls=WORKER_EXTENSION_CLASS,
+            enforce_eager=True,
+            enable_sleep_mode=True,
+        )
+        after.callback(engine.shutdown)
+
+        assert not await engine.is_sleeping(), "Engine should be awake after creation"
+
+        # Measure global VRAM before sleep (driver view; includes inline worker
+        # thread since inline mode runs in the same process).
+        torch.accelerator.synchronize()
+        free_before, total = torch.cuda.mem_get_info()
+        used_before_gib = (total - free_before) / 1024**3
+
+        # Measure CuMemAllocator-tracked usage before sleep.  In inline mode
+        # the worker runs in a thread pool inside this process, so the allocator
+        # singleton is shared and can be read directly.
+        allocator = None
+        tracked_before = 0
+        try:
+            from vllm.device_allocator.cumem import CuMemAllocator
+
+            allocator = CuMemAllocator.get_instance()
+            tracked_before = allocator.get_current_usage()
+        except Exception:
+            pass
+
+        # Put the engine to sleep; all weights should be offloaded via the pool.
+        acks = await engine.sleep(level=1)
+        await asyncio.sleep(0.5)  # allow the CUDA driver to settle
+        torch.accelerator.synchronize()
+
+        # Measure after sleep.
+        free_after, _ = torch.cuda.mem_get_info()
+        used_after_gib = (total - free_after) / 1024**3
+        drop_gib = used_before_gib - used_after_gib
+
+        # --- Primary assertion: allocator reports zero tracked memory. ---
+        # If this fails it means weights were allocated outside the CuMem pool
+        # (safetensors direct-to-GPU bypass) — the exact regression this test
+        # is designed to catch.
+        if allocator is not None:
+            tracked_after = allocator.get_current_usage()
+            assert tracked_after == 0, (
+                f"CuMemAllocator still tracks {tracked_after / 1024**3:.3f} GiB "
+                f"after sleep(level=1) on custom_pipeline path "
+                f"(was {tracked_before / 1024**3:.3f} GiB before sleep). "
+                "Weights were allocated outside the CuMem pool via the "
+                "safetensors direct-to-GPU fast path — loader-context fix "
+                "may not be applied."
+            )
+
+        # --- Secondary assertion: physical VRAM or ACK freed_bytes confirms
+        # reclamation at the driver level.
+        total_freed_bytes = sum(
+            (ack.freed_bytes if hasattr(ack, "freed_bytes") else ack.get("freed_bytes", 0))
+            for ack in acks
+            if ack is not None
+        )
+        freed_gib = total_freed_bytes / 1024**3
+        assert freed_gib > 0 or drop_gib > 0, (
+            f"Expected GPU memory to be reclaimed after sleep(level=1) on "
+            f"custom_pipeline + enable_sleep_mode=True. "
+            f"CuMemAllocator tracked before={tracked_before / 1024**3:.3f} GiB, "
+            f"ACK freed={freed_gib:.3f} GiB, global VRAM drop={drop_gib:.3f} GiB."
+        )
+
+        # Engine must report it is sleeping.
+        assert await engine.is_sleeping()
+
+        # Wake up and confirm the engine is functional again.
+        await engine.wake_up()
+        assert not await engine.is_sleeping()
diff --git a/tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py b/tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py
index 653b35d7e..16dddde40 100644
--- a/tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py
+++ b/tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py
@@ -7,6 +7,9 @@ from contextlib import ExitStack
 
 import pytest
 
+from tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob import (
+    QwenImagePipelineWithLogProbForTest,
+)
 from tests.e2e.offline_inference.custom_pipeline.worker_extension import (
     vLLMOmniColocateWorkerExtensionForTest,
 )
@@ -14,9 +17,7 @@ from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 
-CUSTOM_PIPELINE_CLASS = (
-    "tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob.QwenImagePipelineWithLogProbForTest"
-)
+CUSTOM_PIPELINE_CLASS = QwenImagePipelineWithLogProbForTest
 WORKER_EXTENSION_CLASS = (
     "tests.e2e.offline_inference.custom_pipeline.worker_extension.vLLMOmniColocateWorkerExtensionForTest"
 )
diff --git a/tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py b/tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py
new file mode 100644
index 000000000..a3f57064f
--- /dev/null
+++ b/tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py
@@ -0,0 +1,290 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""E2E test for DreamID-Omni offline inference with Cache-Dit.
+
+Uses the official DreamID-Omni ``test_case/twoip`` example assets instead of
+synthetic inputs so the coverage matches the upstream demo structure.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+import os
+import re
+import time
+from functools import lru_cache
+from pathlib import Path
+from typing import NamedTuple
+
+import numpy as np
+import pytest
+import soundfile as sf
+from PIL import Image
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion, pytest.mark.cache, pytest.mark.slow]
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+DEFAULT_MODEL_PATH = REPO_ROOT / "dreamid_omni"
+MODEL_PATH_ENV = "DREAMID_OMNI_MODEL_PATH"
+EXAMPLE_ROOT_ENV = "DREAMID_OMNI_EXAMPLE_ROOT"
+NUM_INFERENCE_STEPS = 45
+
+CACHE_CONFIG = {
+    "Fn_compute_blocks": 1,
+    "Bn_compute_blocks": 0,
+    "max_warmup_steps": 4,
+    "max_cached_steps": 20,
+    "residual_diff_threshold": 0.24,
+    "max_continuous_cached_steps": 3,
+    "enable_taylorseer": False,
+    "taylorseer_order": 1,
+    "scm_steps_mask_policy": None,
+    "scm_steps_policy": "dynamic",
+}
+
+
+class DreamIDGenerationResult(NamedTuple):
+    video: np.ndarray
+    audio: np.ndarray
+    fps: int
+    duration_s: float
+    elapsed_s: float
+
+
+class DreamIDExampleCase(NamedTuple):
+    prompt: str
+    image_paths: tuple[Path, ...]
+    audio_paths: tuple[Path, ...]
+
+
+def _require_model_path() -> str:
+    model_path = os.environ.get(MODEL_PATH_ENV, str(DEFAULT_MODEL_PATH))
+    if not Path(model_path).exists():
+        pytest.skip(f"DreamID-Omni model bundle not found at {model_path}. Set {MODEL_PATH_ENV}.")
+
+    if importlib.util.find_spec("dreamid_omni") is None:
+        pytest.skip(
+            "dreamid_omni dependency package is not importable. "
+            "Run examples/offline_inference/x_to_video_audio/download_dreamid_omni.py first."
+        )
+
+    return model_path
+
+
+def _dreamid_repo_root_candidates() -> list[Path]:
+    candidates: list[Path] = []
+    example_root = os.environ.get(EXAMPLE_ROOT_ENV)
+    if example_root:
+        candidates.append(Path(example_root).expanduser().resolve())
+
+    spec = importlib.util.find_spec("dreamid_omni")
+    if spec is not None and spec.origin:
+        candidates.append(Path(spec.origin).resolve().parents[1] / "test_case" / "twoip")
+
+    return candidates
+
+
+@lru_cache(maxsize=1)
+def _load_example_case() -> DreamIDExampleCase:
+    for candidate_root in _dreamid_repo_root_candidates():
+        caption_path = candidate_root / "captions" / "20.json"
+        image_paths = (
+            candidate_root / "imgs" / "20" / "0.png",
+            candidate_root / "imgs" / "20" / "1.png",
+        )
+        audio_paths = (
+            candidate_root / "audios" / "20" / "0.wav",
+            candidate_root / "audios" / "20" / "1.wav",
+        )
+        if caption_path.exists() and all(path.exists() for path in image_paths + audio_paths):
+            prompt = _load_prompt_from_json(caption_path)
+            return DreamIDExampleCase(
+                prompt=prompt,
+                image_paths=image_paths,
+                audio_paths=audio_paths,
+            )
+
+    searched = ", ".join(str(path) for path in _dreamid_repo_root_candidates()) or "<none>"
+    pytest.skip(
+        "DreamID-Omni official example assets not found. "
+        f"Set {EXAMPLE_ROOT_ENV} to the local 'test_case/twoip' directory. "
+        f"Searched: {searched}"
+    )
+
+
+def _load_prompt_from_json(path: Path) -> str:
+    raw_text = path.read_text(encoding="utf-8").strip()
+    parsed = json.loads(raw_text)
+    if isinstance(parsed, str):
+        prompt = parsed
+    elif isinstance(parsed, dict) and isinstance(parsed.get("content"), str):
+        prompt = parsed["content"]
+    else:
+        raise AssertionError(f"Unexpected DreamID caption format in {path}")
+
+    prompt = re.sub(
+        r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]",
+        "",
+        prompt,
+        flags=re.DOTALL,
+    ).strip()
+    prompt = re.sub(
+        r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]",
+        "",
+        prompt,
+        flags=re.DOTALL,
+    ).strip()
+    prompt = re.sub(r"\[[A-Z_]+\]", "", prompt)
+    prompt = re.sub(r"\n\s*\n", "\n", prompt).strip()
+    return prompt
+
+
+def _load_example_images(example_case: DreamIDExampleCase) -> list[Image.Image]:
+    images: list[Image.Image] = []
+    for path in example_case.image_paths:
+        with Image.open(path) as image:
+            images.append(image.convert("RGB"))
+    return images
+
+
+def _load_example_audios(example_case: DreamIDExampleCase) -> list[np.ndarray]:
+    audios: list[np.ndarray] = []
+    for path in example_case.audio_paths:
+        audio_array, sample_rate = sf.read(str(path), dtype="float32")
+        audio_np = np.asarray(audio_array, dtype=np.float32)
+        if audio_np.ndim == 2:
+            audio_np = np.mean(audio_np, axis=1)
+        audio_np = audio_np[int(sample_rate * 1) : int(sample_rate * 3)]
+        audios.append(audio_np.reshape(-1))
+    return audios
+
+
+def _sampling_params(num_inference_steps: int) -> OmniDiffusionSamplingParams:
+    return OmniDiffusionSamplingParams(
+        height=704,
+        width=1280,
+        num_inference_steps=num_inference_steps,
+        seed=42,
+        extra_args={
+            "solver_name": "unipc",
+            "shift": 5.0,
+        },
+    )
+
+
+def _request_payload(example_case: DreamIDExampleCase) -> dict[str, object]:
+    return {
+        "prompt": example_case.prompt,
+        "video_negative_prompt": "jitter, bad hands, blur, distortion",
+        "audio_negative_prompt": "robotic, muffled, echo, distorted",
+        "multi_modal_data": {
+            "image": _load_example_images(example_case),
+            "audio": _load_example_audios(example_case),
+        },
+    }
+
+
+def _extract_generation(outputs: list[object], elapsed_s: float) -> DreamIDGenerationResult:
+    """Extract video and audio from raw Omni outputs"""
+    assert outputs, "DreamID-Omni returned no outputs"
+
+    result = outputs[0]
+    assert hasattr(result, "images") and result.images, "not return video frames"
+
+    video = result.images[0]
+    assert isinstance(video, np.ndarray), f"Expected video ndarray, got {type(video)}"
+    assert video.ndim == 4, f"Expected 4D DreamID video tensor, got {video.shape}"
+
+    multimodal_output = getattr(result, "multimodal_output", None) or {}
+    assert "audio" in multimodal_output, "not return audio"
+    audio = np.asarray(multimodal_output["audio"], dtype=np.float32).reshape(-1)
+    assert audio.size > 0, "returned empty audio"
+
+    audio_sample_rate = int(multimodal_output.get("audio_sample_rate", 16000))
+    fps = multimodal_output.get("fps", 24)
+    duration_s = float(audio.size) / float(audio_sample_rate)
+
+    generation = DreamIDGenerationResult(
+        video=video.astype(np.float32, copy=False),
+        audio=audio,
+        fps=fps,
+        duration_s=duration_s,
+        elapsed_s=elapsed_s,
+    )
+    return generation
+
+
+def _run_generation(
+    *,
+    model_path: str,
+    example_case: DreamIDExampleCase,
+    use_cache_dit: bool,
+) -> DreamIDGenerationResult:
+    runner_kwargs = {
+        "model_type": "dreamid-omni",
+        "parallel_config": DiffusionParallelConfig(cfg_parallel_size=1),
+    }
+    if use_cache_dit:
+        runner_kwargs["cache_backend"] = "cache_dit"
+        runner_kwargs["cache_config"] = CACHE_CONFIG
+
+    with OmniRunner(model_path, **runner_kwargs) as runner:
+        warmup_payload = _request_payload(example_case)
+        _ = runner.omni.generate(warmup_payload, _sampling_params(2))
+
+        elapsed_samples: list[float] = []
+        extracted_result: DreamIDGenerationResult | None = None
+
+        request_payload = _request_payload(example_case)
+        start = time.perf_counter()
+        outputs = runner.omni.generate(request_payload, _sampling_params(NUM_INFERENCE_STEPS))
+        elapsed_s = time.perf_counter() - start
+        elapsed_samples.append(elapsed_s)
+        if extracted_result is None:
+            extracted_result = _extract_generation(outputs, elapsed_s)
+
+        assert extracted_result is not None
+        median_elapsed_s = float(np.median(np.asarray(elapsed_samples, dtype=np.float64)))
+        return extracted_result._replace(elapsed_s=median_elapsed_s)
+
+
+def _video_mae(a: np.ndarray, b: np.ndarray) -> float:
+    assert a.shape == b.shape, f"Video shapes differ: {a.shape} vs {b.shape}"
+    return float(np.mean(np.abs(a - b)))
+
+
+def _audio_rms(audio: np.ndarray) -> float:
+    audio = np.asarray(audio, dtype=np.float32).reshape(-1)
+    return float(np.sqrt(np.mean(audio**2)))
+
+
+def _audio_rms_ratio(a: DreamIDGenerationResult, b: DreamIDGenerationResult) -> float:
+    rms_a = _audio_rms(a.audio)
+    rms_b = _audio_rms(b.audio)
+    return rms_b / max(rms_a, 1e-6)
+
+
+@hardware_test(res={"cuda": "H100"})
+def test_dreamid_omni_cache_dit_e2e() -> None:
+    model_path = _require_model_path()
+    example_case = _load_example_case()
+
+    baseline = _run_generation(model_path=model_path, example_case=example_case, use_cache_dit=False)
+    cached = _run_generation(model_path=model_path, example_case=example_case, use_cache_dit=True)
+
+    video_mae = _video_mae(baseline.video, cached.video)
+    audio_rms_ratio = _audio_rms_ratio(baseline, cached)
+    speedup = baseline.elapsed_s / cached.elapsed_s if cached.elapsed_s > 0 else float("inf")
+
+    assert video_mae <= 0.15, f"video drift too much: {video_mae:.6f}"
+    assert 0.4 <= audio_rms_ratio <= 2.5, f"audio quality regressed too much: rms_ratio={audio_rms_ratio:.6f}"
+    assert speedup >= 1.05, f"cache-dit improve not enough: {speedup:.3f}x"
diff --git a/tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py
new file mode 100644
index 000000000..2ea4dc641
--- /dev/null
+++ b/tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py
@@ -0,0 +1,238 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E tests for GLM-Image AutoRound W4A16 quantized inference.
+
+These tests cover text-to-image and image-to-image generation with
+the W4A16 quantized GLM-Image model.
+
+Requirements:
+  - 2 CUDA GPUs (H100 or equivalent)
+  - The quantized model checkpoint (Intel/GLM-Image-int4-AutoRound)
+"""
+
+import gc
+import math
+import os
+
+import numpy as np
+import pytest
+from PIL import Image
+from vllm import SamplingParams
+
+from tests.helpers.env import DeviceMemoryMonitor
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import generate_synthetic_image
+from tests.helpers.runtime import OmniRunnerHandler
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+QUANTIZED_MODEL = os.environ.get("GLM_IMAGE_AUTOROUND_MODEL", "Intel/GLM-Image-int4-AutoRound")
+
+# Small resolution to keep GPU memory & time manageable
+HEIGHT = 256
+WIDTH = 256
+NUM_STEPS = 2  # minimal for smoke-test
+
+# GLM-Image AR generation config (from generation_config.json)
+GLM_IMAGE_EOS_TOKEN_ID = 16385
+GLM_IMAGE_VISION_VOCAB_SIZE = 16512
+
+_CI_DEPLOY = get_deploy_config_path("glm_image.yaml")
+
+
+def _get_stage_config():
+    """Build a CI-friendly stage config with eager mode for testing."""
+    return modify_stage_config(
+        _CI_DEPLOY,
+        updates={
+            "stages": {
+                0: {"enforce_eager": True},
+                1: {"enforce_eager": True},
+            },
+        },
+    )
+
+
+stage_config = _get_stage_config()
+
+# (model, stage_config_path) for ``omni_runner`` indirect parametrize
+_OMNI_RUNNER_PARAM = (QUANTIZED_MODEL, stage_config)
+
+
+def compute_max_tokens(height: int, width: int, factor: int = 32) -> int:
+    """Compute max_new_tokens for GLM-Image AR text-to-image generation."""
+    token_h = height // factor
+    token_w = width // factor
+    large_tokens = token_h * token_w
+
+    ratio = token_h / token_w if token_w > 0 else 1.0
+    small_token_h = max(1, int(math.sqrt(ratio) * (factor // 2)))
+    small_token_w = max(1, int(math.sqrt(1 / ratio) * (factor // 2)))
+    small_tokens = small_token_h * small_token_w
+
+    return small_tokens + large_tokens + 1
+
+
+def _ar_sampling_params(max_tokens: int, height: int, width: int, seed: int = 42) -> SamplingParams:
+    """Build AR stage SamplingParams for GLM-Image."""
+    return SamplingParams(
+        temperature=0.9,
+        top_p=0.75,
+        top_k=GLM_IMAGE_VISION_VOCAB_SIZE,
+        max_tokens=max_tokens,
+        stop_token_ids=[GLM_IMAGE_EOS_TOKEN_ID],
+        seed=seed,
+        detokenize=False,
+        extra_args={
+            "target_h": height,
+            "target_w": width,
+        },
+    )
+
+
+def _diffusion_sampling_params(
+    height: int = HEIGHT,
+    width: int = WIDTH,
+    num_steps: int = NUM_STEPS,
+    seed: int = 42,
+) -> OmniDiffusionSamplingParams:
+    """Build Diffusion stage OmniDiffusionSamplingParams."""
+    return OmniDiffusionSamplingParams(
+        height=height,
+        width=width,
+        num_inference_steps=num_steps,
+        guidance_scale=0.0,
+        seed=seed,
+    )
+
+
+pytestmark = [
+    pytest.mark.full_model,
+    pytest.mark.diffusion,
+]
+
+
+# ------------------------------------------------------------------
+# Test: text-to-image generation produces a valid image (quantized)
+# ------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True)
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+def test_glm_image_autoround_w4a16_generates_image(omni_runner_handler: OmniRunnerHandler):
+    """Load the W4A16 quantized GLM-Image model and verify it produces a valid image."""
+    gc.collect()
+    current_omni_platform.empty_cache()
+    device_index = current_omni_platform.current_device()
+    current_omni_platform.reset_peak_memory_stats()
+    monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
+    monitor.start()
+
+    prompt_dict = {
+        "prompt": "A photo of a cat sitting on a laptop keyboard",
+        "height": HEIGHT,
+        "width": WIDTH,
+        "mm_processor_kwargs": {
+            "target_h": HEIGHT,
+            "target_w": WIDTH,
+        },
+    }
+    ar_params = _ar_sampling_params(
+        max_tokens=compute_max_tokens(HEIGHT, WIDTH),
+        height=HEIGHT,
+        width=WIDTH,
+        seed=42,
+    )
+    diffusion_params = _diffusion_sampling_params(
+        height=HEIGHT,
+        width=WIDTH,
+        num_steps=NUM_STEPS,
+        seed=42,
+    )
+
+    outputs = omni_runner_handler.runner.generate(
+        [prompt_dict],
+        [ar_params, diffusion_params],
+    )
+
+    monitor.stop()
+
+    first_output = outputs[0]
+    assert first_output.final_output_type == "image"
+    req_out = first_output.request_output
+    assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
+    images = req_out.images
+
+    assert len(images) >= 1, "Expected at least one generated image"
+    img = images[0]
+    assert isinstance(img, Image.Image)
+    assert img.width == WIDTH, f"Expected width {WIDTH}, got {img.width}"
+    assert img.height == HEIGHT, f"Expected height {HEIGHT}, got {img.height}"
+
+    # Sanity: image should not be blank (all-zero)
+    arr = np.array(img)
+    assert arr.std() > 1.0, "Generated image appears blank (std ≈ 0)"
+
+    gc.collect()
+    current_omni_platform.empty_cache()
+
+
+# ------------------------------------------------------------------
+# Test: image-to-image generation (quantized)
+# ------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True)
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+def test_glm_image_autoround_w4a16_image_to_image(omni_runner_handler: OmniRunnerHandler):
+    """Load the W4A16 quantized GLM-Image and verify image-to-image generation works."""
+    ref_image_arr = generate_synthetic_image(WIDTH, HEIGHT)["np_array"]
+
+    gc.collect()
+    current_omni_platform.empty_cache()
+    current_omni_platform.reset_peak_memory_stats()
+
+    prompt_dict = {
+        "prompt": "Make it look like winter",
+        "multi_modal_data": {"image": ref_image_arr},
+        "height": HEIGHT,
+        "width": WIDTH,
+        "mm_processor_kwargs": {
+            "target_h": HEIGHT,
+            "target_w": WIDTH,
+        },
+    }
+    ar_params = _ar_sampling_params(
+        max_tokens=compute_max_tokens(HEIGHT, WIDTH),
+        height=HEIGHT,
+        width=WIDTH,
+        seed=42,
+    )
+    diffusion_params = _diffusion_sampling_params(
+        height=HEIGHT,
+        width=WIDTH,
+        num_steps=NUM_STEPS,
+        seed=42,
+    )
+
+    outputs = omni_runner_handler.runner.generate(
+        [prompt_dict],
+        [ar_params, diffusion_params],
+    )
+
+    first_output = outputs[0]
+    assert first_output.final_output_type == "image"
+    req_out = first_output.request_output
+    assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
+    images = req_out.images
+
+    assert len(images) >= 1, "Expected at least one generated image"
+    img = images[0]
+    assert isinstance(img, Image.Image)
+    assert img.width == WIDTH
+    assert img.height == HEIGHT
+
+    gc.collect()
+    current_omni_platform.empty_cache()
diff --git a/tests/e2e/offline_inference/test_glm_tts.py b/tests/e2e/offline_inference/test_glm_tts.py
new file mode 100644
index 000000000..0b902c6bc
--- /dev/null
+++ b/tests/e2e/offline_inference/test_glm_tts.py
@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E Offline inference tests for GLM-TTS voice cloning.
+
+GLM-TTS is a zero-shot voice cloning model that always requires ref_audio.
+There is no text-only / non-clone inference path in the official implementation.
+"""
+
+import os
+import tempfile
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+import numpy as np
+import pytest
+import soundfile as sf
+
+from tests.helpers.assertions import assert_audio_speech_response
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import convert_audio_file_to_text, get_asset_path
+from tests.helpers.runtime import OmniResponse
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from vllm_omni.model_executor.models.glm_tts.glm_tts import build_glm_tts_prefill_metadata
+
+MODEL = os.environ.get("GLM_TTS_MODEL_PATH", "zai-org/GLM-TTS")
+REF_TEXT = "他当时还跟线下其他的站姐吵架，然后，打架进局子了。"
+
+DEPLOY_CONFIG = get_deploy_config_path("glm_tts.yaml")
+
+REFERENCE_PROMPT_WAV_PATH = get_asset_path("glm_tts/jiayan_zh.wav")
+
+ASYNC_CHUNK_MODES = [
+    pytest.param(False, id="sync"),
+    pytest.param(True, id="async_chunk"),
+]
+
+
+def _load_ref_audio() -> tuple[np.ndarray, int]:
+    audio, sr = sf.read(str(REFERENCE_PROMPT_WAV_PATH), dtype="float32", always_2d=False)
+    if isinstance(audio, np.ndarray) and audio.ndim > 1:
+        audio = np.mean(audio, axis=-1)
+    return np.asarray(audio, dtype=np.float32), int(sr)
+
+
+def _concat_audio(audio_val) -> np.ndarray:
+    import torch
+
+    if isinstance(audio_val, list):
+        tensors = [torch.as_tensor(t).float().reshape(-1) for t in audio_val if t is not None]
+        if not tensors:
+            return np.zeros((0,), dtype=np.float32)
+        return torch.cat(tensors, dim=-1).cpu().numpy().astype(np.float32, copy=False)
+    if isinstance(audio_val, torch.Tensor):
+        return audio_val.float().cpu().numpy().reshape(-1)
+    return np.asarray(audio_val, dtype=np.float32).reshape(-1)
+
+
+def _extract_sample_rate(audio_mm: dict) -> int:
+    sr_raw = audio_mm.get("sr", 24000)
+    if isinstance(sr_raw, list):
+        sr_raw = sr_raw[-1] if sr_raw else 24000
+    if hasattr(sr_raw, "item"):
+        return int(sr_raw.item())
+    return int(sr_raw)
+
+
+def _get_deploy_config(*, async_chunk: bool) -> str:
+    """Build deploy config with explicit sync/async mode and eager execution."""
+    return modify_stage_config(
+        DEPLOY_CONFIG,
+        updates={
+            "async_chunk": async_chunk,
+            "stages": {
+                0: {
+                    "enforce_eager": True,
+                    "async_scheduling": bool(async_chunk),
+                },
+                1: {
+                    "enforce_eager": True,
+                },
+            },
+        },
+    )
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("async_chunk", ASYNC_CHUNK_MODES)
+def test_offline_voice_clone_zh(async_chunk: bool) -> None:
+    """
+    Test voice cloning offline inference.
+    Deploy Setting: glm_tts.yaml with explicit sync/async mode, enforce_eager
+    Input Modal: text + ref_audio + ref_text
+    Output Modal: audio
+
+    Uses the official jiayan_zh.wav reference audio from the upstream
+    GLM-TTS repository (real human speech matching REF_TEXT).
+    """
+    synth_text = "今天天气真不错，适合出去散散步。"
+    prompt_audio = _load_ref_audio()
+    from tests.helpers.runtime import OmniRunner
+
+    with OmniRunner(
+        MODEL,
+        stage_configs_path=_get_deploy_config(async_chunk=async_chunk),
+        stage_init_timeout=600,
+    ) as omni_runner:
+        outputs = omni_runner.omni.generate(
+            [
+                {
+                    "prompt": synth_text,
+                    "multi_modal_data": {"audio": prompt_audio},
+                    "modalities": ["audio"],
+                    "mm_processor_kwargs": {"prompt_text": REF_TEXT},
+                    "additional_information": build_glm_tts_prefill_metadata(
+                        MODEL,
+                        synth_text,
+                        REF_TEXT,
+                        trust_remote_code=True,
+                    ),
+                }
+            ],
+            omni_runner.get_default_sampling_params_list(),
+        )
+
+        assert outputs, "No outputs returned"
+        audio_mm = outputs[0].multimodal_output
+        assert "audio" in audio_mm, "No audio output found"
+        audio = _concat_audio(audio_mm["audio"])
+        assert audio.size > 0, "Generated audio is empty"
+
+        with tempfile.NamedTemporaryFile(suffix=".wav") as out_file:
+            sf.write(out_file.name, audio, samplerate=_extract_sample_rate(audio_mm), format="WAV")
+            transcript = convert_audio_file_to_text(out_file.name)
+            audio_bytes = open(out_file.name, "rb").read()
+
+        assert_audio_speech_response(
+            OmniResponse(
+                success=True,
+                audio_bytes=audio_bytes,
+                audio_content=transcript,
+                audio_format="audio/wav",
+            ),
+            {"input": synth_text, "response_format": "wav"},
+            run_level="advanced_model",
+        )
diff --git a/tests/e2e/offline_inference/test_higgs_audio_v2.py b/tests/e2e/offline_inference/test_higgs_audio_v2.py
new file mode 100644
index 000000000..fdeccc924
--- /dev/null
+++ b/tests/e2e/offline_inference/test_higgs_audio_v2.py
@@ -0,0 +1,231 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+End-to-end offline inference tests for higgs-audio v2.
+
+Mirrors examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py:
+build Stage-0 prompt_token_ids with the model's own processor + plain-text
+chat template, drive the engine via ``Omni.generate``, and inspect the
+Stage-1 PCM payload directly. The OmniRunnerHandler.send_audio_speech_request
+helper applies a Qwen-style ChatML wrap that is not valid for higgs, so this
+test bypasses it (same approach as test_voxtral_tts.py).
+
+v1 scope is plain text -> 24 kHz speech; voice cloning, multi-speaker prompts,
+language overrides, and task_type selection are out of scope.
+"""
+
+from __future__ import annotations
+
+import os
+
+# Match the example script: keep DeepGEMM FP8 off so engine init does not
+# require the optional ``deep_gemm`` backend.
+os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
+os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
+
+import numpy as np
+import pytest
+import torch
+from vllm import SamplingParams
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import get_asset_path
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = "bosonai/higgs-audio-v2-generation-3B-base"
+STAGE_CONFIG = get_deploy_config_path("higgs_audio_v2.yaml")
+SAMPLE_RATE = 24_000
+# A short sentence at 24 kHz must produce at least ~0.5 s of audio.
+MIN_AUDIO_SAMPLES = 12_000
+TEST_TEXT = "Hello world."
+
+# Reuse the qwen3_tts ref clip + transcript for voice clone tests (see
+# tests/e2e/online_serving/test_qwen3_tts_base.py:27-28 for the source).
+_REF_AUDIO_ASSET = "qwen3_tts/clone_2.wav"
+_REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
+
+# (model, stage_config_path, extra_omni_kwargs) — same shape as test_voxtral_tts.py.
+# Stage configs already pin enforce_eager=true per stage, so no Omni-level override.
+_OMNI_RUNNER_PARAM = (MODEL, STAGE_CONFIG, {"trust_remote_code": True})
+
+
+def _compose_request(model_name: str, text: str) -> dict:
+    """Build a higgs prompt dict via the upstream processor + plain-text template.
+
+    The serving layer's ``_build_higgs_audio_v2_params`` does the same thing
+    (see vllm_omni/entrypoints/openai/serving_speech.py:1520) — keeping the
+    offline and online prompt-construction paths identical is what guarantees
+    input-token parity with the HF reference.
+    """
+    from transformers import AutoProcessor
+
+    from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
+        build_plain_text_prompt,
+        input_ids_to_python_list,
+    )
+
+    processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
+    inputs = build_plain_text_prompt(processor, text)
+    return {"prompt_token_ids": input_ids_to_python_list(inputs)}
+
+
+def _extract_pcm(outputs) -> torch.Tensor:
+    """Pull the concatenated PCM tensor out of the Stage-1 multimodal_output."""
+    for stage_output in outputs:
+        mm = getattr(stage_output, "multimodal_output", None)
+        if not mm:
+            continue
+        audio = mm.get("model_outputs") if "model_outputs" in mm else mm.get("audio")
+        if audio is None:
+            continue
+        if isinstance(audio, list):
+            chunks = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
+            if not chunks:
+                continue
+            return torch.cat(chunks, dim=0) if len(chunks) > 1 else chunks[0]
+        return torch.as_tensor(audio).float().cpu().reshape(-1)
+    raise AssertionError(f"no audio payload in any stage output (got {len(outputs)} stages)")
+
+
+@pytest.mark.core_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True)
+def test_higgs_audio_v2_offline_plain_text(omni_runner_function: OmniRunner) -> None:
+    """
+    Offline plain-text -> 24 kHz PCM via the bundled higgs_audio_v2 deploy.
+
+    Deploy Setting: vllm_omni/deploy/higgs_audio_v2.yaml (codec_streaming=true,
+        async_chunk=false, enforce_eager=true per stage).
+    Input Modal: text
+    Output Modal: audio (PCM tensor, 24 kHz mono float)
+    Input Setting: single prompt, default sampling
+    """
+    omni = omni_runner_function.omni
+    inputs = _compose_request(MODEL, TEST_TEXT)
+
+    # max_tokens caps Stage-0 codec frames; Stage 1 max_tokens lifts the
+    # per-codebook prefill cap. Mirror the deploy yaml's defaults.
+    stage_sampling = [
+        SamplingParams(temperature=1.0, top_p=0.95, top_k=50, max_tokens=1024, seed=42),
+        SamplingParams(temperature=0.0, top_p=1.0, top_k=-1, max_tokens=65536, seed=42),
+    ]
+
+    outputs = list(omni.generate([inputs], stage_sampling))
+    assert outputs, "no outputs returned from Omni.generate"
+
+    pcm = _extract_pcm(outputs)
+    arr = pcm.numpy()
+
+    assert arr.size > MIN_AUDIO_SAMPLES, (
+        f"audio too short: {arr.size} samples, expected > {MIN_AUDIO_SAMPLES} "
+        f"({MIN_AUDIO_SAMPLES / SAMPLE_RATE:.2f} s @ {SAMPLE_RATE} Hz)"
+    )
+    # Threshold 5e-4 (not 1e-3) keeps us above true silence while tolerating
+    # the ``load_format=dummy`` Stage-1 codec — random RVQ/DAC weights collapse
+    # the output to a near-constant peak around 7.4e-4. Real-weight runs reach
+    # 0.1+ peak so this still catches genuine silence.
+    assert np.max(np.abs(arr)) > 5e-4, "audio output is silent (peak abs amplitude < 5e-4)"
+    assert np.isfinite(arr).all(), "audio contains non-finite samples (nan/inf)"
+
+
+@pytest.mark.core_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True)
+def test_higgs_audio_v2_offline_batch_two_prompts(omni_runner_function: OmniRunner) -> None:
+    """
+    Submit two prompts in one engine call and verify both yield non-trivial PCM.
+
+    Guards the Stage-0 talker's per-slot audio state under a batched call —
+    the prior example ran prompts one at a time citing 'request-scoped audio
+    state'; this test pins the batched-2 behavior the engine path actually
+    needs to support for /v1/audio/speech under concurrent traffic.
+    """
+    omni = omni_runner_function.omni
+    prompts = [
+        _compose_request(MODEL, "Hello world."),
+        _compose_request(MODEL, "The quick brown fox jumps over the lazy dog."),
+    ]
+    stage_sampling = [
+        SamplingParams(temperature=1.0, top_p=0.95, top_k=50, max_tokens=1024, seed=42),
+        SamplingParams(temperature=0.0, top_p=1.0, top_k=-1, max_tokens=65536, seed=42),
+    ]
+
+    outputs = list(omni.generate(prompts, stage_sampling))
+    assert outputs, "no outputs returned from Omni.generate"
+
+    pcm = _extract_pcm(outputs)
+    # Whether the engine returns one concatenated PCM or per-request tensors is
+    # implementation-defined; for either shape, total samples must clear the
+    # floor for at least one of the two prompts.
+    assert pcm.numel() > MIN_AUDIO_SAMPLES, (
+        f"batched audio too short: {pcm.numel()} samples, expected > {MIN_AUDIO_SAMPLES}"
+    )
+    # See the silence-threshold note in test_higgs_audio_v2_offline_plain_text.
+    assert np.max(np.abs(pcm.numpy())) > 5e-4, "batched audio output is silent"
+
+
+def _compose_voice_clone_request(model_name: str, text: str, ref_text: str, ref_wav_path) -> dict:
+    """Build the voice-clone prompt + additional_information payload.
+
+    Calls the same ``build_voice_clone_prompt`` the serving layer uses online —
+    the HF processor encodes the reference clip on the spot via the bundled
+    HiggsAudioV2TokenizerModel. The talker's
+    ``_maybe_apply_ref_audio_substitution`` consumes the tensors at prefill.
+    """
+    import soundfile as sf
+    from transformers import AutoProcessor
+
+    from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
+        build_voice_clone_prompt,
+    )
+
+    wav, sr = sf.read(str(ref_wav_path), always_2d=False)
+    if wav.ndim == 2:
+        wav = wav.mean(axis=1)  # downmix any stereo asset to mono
+    processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
+    out = build_voice_clone_prompt(processor, text, wav, int(sr), ref_text)
+    return {
+        "prompt_token_ids": out["prompt_token_ids"],
+        # Pass tensors bare (see serving_speech._build_higgs_audio_v2_params
+        # — list-wrapped tensors get dropped by the msgspec serializer).
+        "additional_information": {
+            "audio_input_ids": out["audio_input_ids"],
+            "audio_input_ids_mask": out["audio_input_ids_mask"],
+        },
+    }
+
+
+@pytest.mark.advanced_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True)
+def test_higgs_audio_v2_offline_voice_clone(omni_runner_function: OmniRunner) -> None:
+    """
+    Shallow voice clone via ref_audio + ref_text, offline.
+
+    Deploy Setting: vllm_omni/deploy/higgs_audio_v2.yaml.
+    Input Modal: text + reference audio (qwen3_tts/clone_2.wav vendored asset)
+    Output Modal: audio (PCM tensor, 24 kHz mono float)
+    Input Setting: single prompt
+    """
+    omni = omni_runner_function.omni
+    ref_path = get_asset_path(_REF_AUDIO_ASSET)
+    inputs = _compose_voice_clone_request(MODEL, TEST_TEXT, _REF_TEXT, ref_path)
+
+    stage_sampling = [
+        SamplingParams(temperature=1.0, top_p=0.95, top_k=50, max_tokens=1024, seed=42),
+        SamplingParams(temperature=0.0, top_p=1.0, top_k=-1, max_tokens=65536, seed=42),
+    ]
+
+    outputs = list(omni.generate([inputs], stage_sampling))
+    assert outputs, "no outputs returned from Omni.generate for voice clone"
+
+    pcm = _extract_pcm(outputs)
+    arr = pcm.numpy()
+    assert arr.size > MIN_AUDIO_SAMPLES, f"cloned audio too short: {arr.size} samples, expected > {MIN_AUDIO_SAMPLES}"
+    assert np.max(np.abs(arr)) > 5e-4, "voice-clone output is silent"
+    assert np.isfinite(arr).all(), "voice-clone audio contains non-finite samples"
diff --git a/tests/e2e/offline_inference/test_sensenova_u1_img2img.py b/tests/e2e/offline_inference/test_sensenova_u1_img2img.py
index 10f5f05dc..fedd8e90e 100644
--- a/tests/e2e/offline_inference/test_sensenova_u1_img2img.py
+++ b/tests/e2e/offline_inference/test_sensenova_u1_img2img.py
@@ -167,3 +167,18 @@ def test_sensenova_u1_img2img(run_level):
         generated_image = _generate_sensenova_u1_img2img(runner.omni, input_image)
         if run_level == "advanced_model":
             _validate_pixels(generated_image)
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cache
+@hardware_test(res={"cuda": "H100"})
+def test_sensenova_u1_img2img_cache_dit():
+    """Test SenseNova-U1 img2img with Cache-DiT enabled."""
+    input_image = _load_input_image()
+    with OmniRunner(
+        "SenseNova/SenseNova-U1-8B-MoT",
+        stage_configs_path=None,
+        cache_backend="cache_dit",
+    ) as runner:
+        _generate_sensenova_u1_img2img(runner.omni, input_image)
diff --git a/tests/e2e/offline_inference/test_sensenova_u1_text2img.py b/tests/e2e/offline_inference/test_sensenova_u1_text2img.py
index 334381067..fc9b5411e 100644
--- a/tests/e2e/offline_inference/test_sensenova_u1_text2img.py
+++ b/tests/e2e/offline_inference/test_sensenova_u1_text2img.py
@@ -156,3 +156,17 @@ def test_sensenova_u1_text2img(run_level):
         generated_image = _generate_sensenova_u1_image(runner.omni)
         if run_level == "advanced_model":
             _validate_pixels(generated_image)
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cache
+@hardware_test(res={"cuda": "H100"})
+def test_sensenova_u1_text2img_cache_dit():
+    """Test SenseNova-U1 text2img with Cache-DiT enabled."""
+    with OmniRunner(
+        "SenseNova/SenseNova-U1-8B-MoT",
+        stage_configs_path=None,
+        cache_backend="cache_dit",
+    ) as runner:
+        _generate_sensenova_u1_image(runner.omni)
diff --git a/tests/e2e/online_serving/test_bagel_expansion.py b/tests/e2e/online_serving/test_bagel_expansion.py
index 7152b6986..2feaa3f07 100644
--- a/tests/e2e/online_serving/test_bagel_expansion.py
+++ b/tests/e2e/online_serving/test_bagel_expansion.py
@@ -18,6 +18,7 @@ import pytest
 
 from tests.helpers.mark import hardware_marks
 from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
@@ -25,7 +26,18 @@ PROMPT = "A futuristic city skyline at twilight, cyberpunk style, ultra-detailed
 NEGATIVE_PROMPT = "low quality, blurry, distorted, deformed, watermark"
 
 SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
-PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+PARALLEL_2_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+HSDP_2_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+
+BAGEL_CI_DEPLOY = get_deploy_config_path("ci/bagel.yaml")
+BAGEL_PARALLEL_2_DEPLOY = modify_stage_config(
+    BAGEL_CI_DEPLOY,
+    updates={"stages": {0: {"devices": "0"}, 1: {"devices": "0,1"}}},
+)
+BAGEL_HSDP_2_DEPLOY = modify_stage_config(
+    BAGEL_CI_DEPLOY,
+    updates={"stages": {0: {"devices": "0"}, 1: {"devices": "0,1"}}},
+)
 
 
 def _get_diffusion_feature_cases(model: str):
@@ -63,6 +75,7 @@ def _get_diffusion_feature_cases(model: str):
         pytest.param(
             OmniServerParams(
                 model=model,
+                stage_config_path=BAGEL_PARALLEL_2_DEPLOY,
                 server_args=[
                     "--cache-backend",
                     "tea_cache",
@@ -71,31 +84,33 @@ def _get_diffusion_feature_cases(model: str):
                 ],
             ),
             id="parallel_cfg_2",
-            marks=PARALLEL_FEATURE_MARKS,
+            marks=PARALLEL_2_FEATURE_MARKS,
         ),
         # Ulysses-SP degree=2 (2 GPUs)
         pytest.param(
             OmniServerParams(
                 model=model,
+                stage_config_path=BAGEL_PARALLEL_2_DEPLOY,
                 server_args=[
                     "--usp",
                     "2",
                 ],
             ),
             id="sp_ulysses_2",
-            marks=PARALLEL_FEATURE_MARKS,
+            marks=PARALLEL_2_FEATURE_MARKS,
         ),
         # Ring-Attention degree=2 (2 GPUs)
         pytest.param(
             OmniServerParams(
                 model=model,
+                stage_config_path=BAGEL_PARALLEL_2_DEPLOY,
                 server_args=[
                     "--ring",
                     "2",
                 ],
             ),
             id="sp_ring_2",
-            marks=PARALLEL_FEATURE_MARKS,
+            marks=PARALLEL_2_FEATURE_MARKS,
         ),
         # Layerwise Offloading (single-card)
         pytest.param(
@@ -106,18 +121,19 @@ def _get_diffusion_feature_cases(model: str):
             id="single_card_layerwise_offload",
             marks=SINGLE_CARD_FEATURE_MARKS,
         ),
-        # Hybrid Sharded Data Parallel (4 GPUs)
+        # Hybrid Sharded Data Parallel (2 GPUs)
         pytest.param(
             OmniServerParams(
                 model=model,
+                stage_config_path=BAGEL_HSDP_2_DEPLOY,
                 server_args=[
                     "--use-hsdp",
                     "--hsdp-shard-size",
-                    "4",
+                    "2",
                 ],
             ),
-            id="parallel_hsdp_4",
-            marks=PARALLEL_FEATURE_MARKS,
+            id="parallel_hsdp_2",
+            marks=HSDP_2_FEATURE_MARKS,
         ),
     ]
 
@@ -140,7 +156,7 @@ def test_bagel(
     - Ulysses-SP (degree=2)
     - Ring-Attention (degree=2)
     - Layerwise Offloading
-    - Hybrid Sharded Data Parallel (size=4)
+    - Hybrid Sharded Data Parallel (size=2)
 
     Validation is delegated to assert_diffusion_response in tests/helpers/assertions.py,
     which checks output dimensions and basic correctness.
@@ -158,6 +174,7 @@ def test_bagel(
             # Enable CFG for models that use classifier-free guidance
             "negative_prompt": NEGATIVE_PROMPT,
             "true_cfg_scale": 4.0,
+            "cfg_img_scale": 1.0,
             "seed": 42,
         },
     }
diff --git a/tests/e2e/online_serving/test_flux_2_dev_expansion.py b/tests/e2e/online_serving/test_flux_2_dev_expansion.py
index 953cb448a..aaf118112 100644
--- a/tests/e2e/online_serving/test_flux_2_dev_expansion.py
+++ b/tests/e2e/online_serving/test_flux_2_dev_expansion.py
@@ -3,6 +3,8 @@ End-to-end diffusion coverage for FLUX.2-dev in online serving mode.
 
 Coverage:
 - CPU offload
+- Ulysses sequence parallelism
+- Ring sequence parallelism
 
 This test verifies that FLUX.2-dev can be launched with CPU offload enabled,
 accepts text-to-image requests through the OpenAI-compatible API, and returns
@@ -28,7 +30,7 @@ PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
 
 
 def _get_flux_2_dev_feature_cases(model: str):
-    """Return FLUX.2-dev diffusion feature cases for CPU offload."""
+    """Return FLUX.2-dev diffusion feature cases for CPU offload and SP."""
 
     return [
         pytest.param(
@@ -53,6 +55,30 @@ def _get_flux_2_dev_feature_cases(model: str):
             id="parallel_cfg_2",
             marks=PARALLEL_FEATURE_MARKS,
         ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--enable-cpu-offload",
+                    "--ulysses-degree",
+                    "2",
+                ],
+            ),
+            id="ulysses_2",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--enable-cpu-offload",
+                    "--ring-degree",
+                    "2",
+                ],
+            ),
+            id="ring_2",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
     ]
 
 
diff --git a/tests/e2e/online_serving/test_glm_tts.py b/tests/e2e/online_serving/test_glm_tts.py
new file mode 100644
index 000000000..7c6ad9764
--- /dev/null
+++ b/tests/e2e/online_serving/test_glm_tts.py
@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E Online tests for GLM-TTS model with text input and audio output.
+
+These tests verify the /v1/audio/speech endpoint works correctly with
+the GLM-TTS two-stage pipeline (AR + DiT).
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = os.environ.get("GLM_TTS_MODEL_PATH", "zai-org/GLM-TTS")
+REF_TEXT = "他当时还跟线下其他的站姐吵架，然后，打架进局子了。"
+
+DEPLOY_CONFIG = get_deploy_config_path("glm_tts.yaml")
+
+# Official GLM-TTS reference audio and transcript. Vendored under tests/assets/
+# so Buildkite does not depend on raw.githubusercontent.com during E2E.
+REF_AUDIO_URL = load_test_audio_data_url("glm_tts/jiayan_zh.wav")
+
+
+SYNC_EXTRA_ARGS = [
+    "--trust-remote-code",
+    "--disable-log-stats",
+    "--no-async-chunk",
+]
+
+ASYNC_EXTRA_ARGS = [
+    "--trust-remote-code",
+    "--disable-log-stats",
+]
+
+tts_sync_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=DEPLOY_CONFIG,
+            server_args=SYNC_EXTRA_ARGS,
+            stage_init_timeout=600,
+        ),
+        id="glm_tts_sync",
+    )
+]
+
+tts_async_chunk_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=DEPLOY_CONFIG,
+            server_args=ASYNC_EXTRA_ARGS,
+            stage_init_timeout=600,
+        ),
+        id="glm_tts_async_chunk",
+    )
+]
+
+
+# GLM-TTS is a zero-shot voice cloning model (always requires ref_audio).
+# There is no text-only / non-clone path — the official implementation
+# asserts prompt_feat.shape[1] != 0 in token2wav_with_cache and every
+# example prompt includes a reference audio.  CosyVoice3 follows the same
+# pattern: all tests are voice_clone_*.
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_sync_server_params, indirect=True)
+def test_voice_clone_zh_sync(omni_server, openai_client) -> None:
+    """
+    Test voice cloning TTS with Chinese text.
+    Deploy Setting: glm_tts.yaml with ``--no-async-chunk`` (sync two-stage)
+    Input Modal: text + ref_audio + ref_text
+    Output Modal: audio
+    Input Setting: stream=False
+
+    Uses the official jiayan_zh.wav reference audio from the upstream
+    GLM-TTS repository (real human speech matching REF_TEXT).
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": "今天天气真不错，适合出去散散步。",
+        "stream": False,
+        "response_format": "wav",
+        "ref_audio": REF_AUDIO_URL,
+        "ref_text": REF_TEXT,
+    }
+    openai_client.send_audio_speech_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_async_chunk_server_params, indirect=True)
+def test_voice_clone_zh_async_chunk(omni_server, openai_client) -> None:
+    """
+    Test voice cloning TTS with Chinese text via async_chunk streaming.
+    Deploy Setting: glm_tts.yaml default ``async_chunk: true``
+    Input Modal: text + ref_audio + ref_text
+    Output Modal: audio
+    Input Setting: stream=True
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": "今天天气真不错，适合出去散散步。",
+        "stream": True,
+        "response_format": "wav",
+        "ref_audio": REF_AUDIO_URL,
+        "ref_text": REF_TEXT,
+    }
+    openai_client.send_audio_speech_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_sync_server_params, indirect=True)
+def test_models_endpoint(omni_server, openai_client) -> None:
+    """Test the /v1/models endpoint returns loaded model."""
+    models = openai_client.client.models.list()
+    assert len(models.data) > 0
diff --git a/tests/e2e/online_serving/test_hidream_i1_full_expansion.py b/tests/e2e/online_serving/test_hidream_i1_full_expansion.py
new file mode 100644
index 000000000..f84cc08ad
--- /dev/null
+++ b/tests/e2e/online_serving/test_hidream_i1_full_expansion.py
@@ -0,0 +1,69 @@
+"""
+End-to-end diffusion coverage for HiDream-I1-Full in online serving mode.
+
+This test verifies that HiDream-I1-Full can be launched,
+accepts text-to-image requests through the OpenAI-compatible API, and returns
+valid generated images with the requested resolution.
+
+assert_diffusion_response validates successful generation and the expected
+image resolution.
+"""
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+
+MODEL = "HiDream-ai/HiDream-I1-Full"
+PROMPT = "A cinematic mountain landscape at sunrise, dramatic clouds, ultra-detailed, realistic photography."
+NEGATIVE_PROMPT = "low quality, blurry, distorted, deformed, watermark"
+
+SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+
+
+def _get_hidream_i1_image_feature_cases(model: str):
+    """Return HiDream-I1-Full diffusion feature cases."""
+
+    return [
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--auxiliary-text-encoder",
+                    "meta-llama/Meta-Llama-3.1-8B-Instruct",
+                ],
+            ),
+            id="default",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_hidream_i1_image_feature_cases(MODEL),
+    indirect=True,
+)
+def test_hidream_i1_image(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+):
+    """Validate HiDream-I1-Full online serving with CPU offload."""
+
+    messages = dummy_messages_from_mix_data(content_text=PROMPT)
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "height": 512,
+            "width": 512,
+            "num_inference_steps": 2,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "seed": 42,
+        },
+    }
+
+    openai_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_higgs_audio_v2.py b/tests/e2e/online_serving/test_higgs_audio_v2.py
new file mode 100644
index 000000000..92688ad0b
--- /dev/null
+++ b/tests/e2e/online_serving/test_higgs_audio_v2.py
@@ -0,0 +1,299 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+End-to-end online tests for higgs-audio v2 against /v1/audio/speech.
+
+v1 scope is plain text -> 24 kHz speech plus shallow voice clone via
+ref_audio + ref_text (inline) or voice=<name> (after POST /v1/audio/voices).
+The model-aware request validator
+(vllm_omni/entrypoints/openai/serving_speech.py::_validate_higgs_audio_v2_request)
+rejects multi-speaker tags, language overrides, task_type, and bare
+voice=<name> for names that do not match an uploaded speaker — this suite
+exercises both the happy path (plain text in, audio bytes out) and the
+validator rejections, so a regression that loosens the schema will fail
+loudly.
+"""
+
+from __future__ import annotations
+
+import os
+
+os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+# Match run_server.sh: DeepGEMM FP8 kernels are optional and trip warmup on
+# images without the deep_gemm backend, so disable them by default.
+os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
+os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = "bosonai/higgs-audio-v2-generation-3B-base"
+STAGE_CONFIG = get_deploy_config_path("higgs_audio_v2.yaml")
+SERVER_ARGS = ["--trust-remote-code", "--disable-log-stats"]
+# DeepGEMM warmup is optional; mirror run_server.sh and switch it off in the
+# server subprocess env too (parent-process os.environ above only affects this
+# test driver; the engine subprocesses inherit through env_dict).
+SERVER_ENV = {"VLLM_USE_DEEP_GEMM": "0", "VLLM_MOE_USE_DEEP_GEMM": "0"}
+
+TEST_PARAMS = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=STAGE_CONFIG,
+            server_args=SERVER_ARGS,
+            env_dict=SERVER_ENV,
+        ),
+        id="higgs_audio_v2_plain_text",
+    )
+]
+
+DEFAULT_SPEECH_TIMEOUT_S = 180.0
+# Floor for ~0.5 s of 24 kHz mono PCM_16: 24000 * 0.5 * 2 bytes ≈ 24 KiB.
+# A WAV header adds 44 bytes; pick a conservative floor that catches truncated
+# / silence-only outputs without flagging short legitimate clips.
+_MIN_AUDIO_BYTES = 20_000
+
+# Reuse the qwen3_tts vendored reference clip (clean ~5 s 24 kHz mono human
+# speech) + its transcript. See tests/e2e/online_serving/test_qwen3_tts_base.py
+# for the asset rationale — keeping a single shared reference clip across TTS
+# tests avoids duplicating WAVs in the repo.
+_REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
+_REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
+
+
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+class TestHiggsAudioV2OnlineHappyPath:
+    """Plain-text -> audio happy paths against the live HTTP server."""
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_plain_text_wav(self, omni_server, openai_client) -> None:
+        """Single non-streaming WAV request — covers the canonical TTS happy path."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_plain_text_pcm_streaming(self, omni_server, openai_client) -> None:
+        """Streaming PCM via the shared-memory connector's codec_streaming path.
+
+        higgs_audio_v2.yaml pins ``codec_streaming: true`` + ``async_chunk: false``,
+        so the only streaming surface exposed to clients is the WAV/PCM bytes
+        served chunk-by-chunk from Stage 1 — exercise it directly.
+        """
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Streaming the quick brown fox over the lazy dog.",
+                "stream": True,
+                "response_format": "pcm",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_plain_text_with_max_new_tokens(self, omni_server, openai_client) -> None:
+        """max_new_tokens is one of the few extra fields the higgs validator accepts."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Innovation distinguishes between a leader and a follower.",
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "max_new_tokens": 500,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_concurrent_plain_text(self, omni_server, openai_client) -> None:
+        """Three concurrent non-streaming requests — guards the per-slot audio state."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "It was the night before my birthday.",
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            },
+            request_num=3,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Validator rejections served over HTTP. Each case targets one out-of-scope
+# field; the validator returns 4xx via the OpenAI error path. We keep the
+# error-message substring loose so a phrasing tweak in the validator does not
+# break CI, but tight enough to catch a regression that silently accepts the
+# field.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+class TestHiggsAudioV2OnlineValidatorRejections:
+    """Out-of-scope fields must come back as 4xx with a higgs-named message."""
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_ref_audio_without_ref_text(self, omni_server, openai_client) -> None:
+        """Voice clone needs the transcript too — half-supplied is a 4xx."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "ref_audio": _REF_AUDIO_URL,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "status_code": (400, 422),
+                "err_message": "ref_text",
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_ref_text_without_ref_audio(self, omni_server, openai_client) -> None:
+        """Symmetric guard: ref_text alone is not enough — must come with ref_audio."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "ref_text": "some transcript",
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "status_code": (400, 422),
+                "err_message": "ref_audio",
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_task_type(self, omni_server, openai_client) -> None:
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "task_type": "Base",
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "status_code": (400, 422),
+                "err_message": "task_type",
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_language_override(self, omni_server, openai_client) -> None:
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "language": "Chinese",
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "status_code": (400, 422),
+                "err_message": "language",
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_multi_speaker_tag_in_text(self, omni_server, openai_client) -> None:
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "[SPEAKER0] hi",
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "status_code": (400, 422),
+                "err_message": "multi-speaker",
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_empty_input(self, omni_server, openai_client) -> None:
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "   ",
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "status_code": (400, 422),
+                "err_message": "empty",
+            }
+        )
+
+
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+class TestHiggsAudioV2OnlineVoiceClone:
+    """Shallow voice clone: ref_audio + ref_text -> speech in the cloned voice.
+
+    The HF processor that the serving layer calls (lazy-loaded at the first
+    request) ships with the bundled HiggsAudioV2TokenizerModel and encodes
+    the reference clip in-process. The talker substitutes the encoded codes
+    at the prompt-side audio placeholders via
+    :meth:`HiggsAudioV2TalkerForConditionalGeneration._maybe_apply_ref_audio_substitution`.
+    """
+
+    @pytest.mark.advanced_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_voice_clone_basic(self, omni_server, openai_client) -> None:
+        """Single non-streaming WAV request driven by the qwen3_tts ref clip."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "ref_audio": _REF_AUDIO_URL,
+                "ref_text": _REF_TEXT,
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+    @pytest.mark.advanced_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_voice_clone_pcm_streaming(self, omni_server, openai_client) -> None:
+        """Voice clone over the streaming PCM path."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Innovation distinguishes a leader from a follower.",
+                "ref_audio": _REF_AUDIO_URL,
+                "ref_text": _REF_TEXT,
+                "stream": True,
+                "response_format": "pcm",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
diff --git a/tests/e2e/online_serving/test_lance.py b/tests/e2e/online_serving/test_lance.py
index 197c35746..9c1a40be9 100644
--- a/tests/e2e/online_serving/test_lance.py
+++ b/tests/e2e/online_serving/test_lance.py
@@ -9,7 +9,7 @@ requests via the OpenAI-compatible chat completions API exposed by
 Equivalent to running:
 
     vllm-omni serve "bytedance-research/Lance" --omni \\
-        --deploy-config vllm_omni/deploy/lance.yaml --port 8091
+        --pipeline lance --enforce-eager --trust-remote-code --port 8091
 
     # text2img
     python3 examples/online_serving/lance/openai_chat_client.py \\
@@ -25,7 +25,6 @@ from vllm.assets.image import ImageAsset
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniServerParams
-from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni.diffusion.models.lance.prompts import (
     VIDEO_PAD,
     VISION_END,
@@ -36,7 +35,6 @@ from vllm_omni.diffusion.models.lance.prompts import (
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 MODEL = "bytedance-research/Lance"
-STAGE_CONFIGS_PATH = get_deploy_config_path("ci/lance.yaml")
 
 TEXT2IMG_PROMPT = "A cute corgi astronaut on the moon, cinematic"
 IMG2IMG_PROMPT = "Convert this into a vibrant cartoon-style illustration"
@@ -44,10 +42,26 @@ IMG2IMG_PROMPT = "Convert this into a vibrant cartoon-style illustration"
 # Lance vision block: <|vision_start|><|video_pad|><|vision_end|>
 _VISION_BLOCK = f"{VISION_START}{VIDEO_PAD}{VISION_END}"
 
+# Lance is single-stage diffusion — no deploy YAML.  Pass the engine
+# knobs that used to live in ``vllm_omni/deploy/lance.yaml`` via CLI
+# flags so ``create_default_diffusion`` builds the stage config.
+_LANCE_SERVE_ARGS = [
+    "--pipeline",
+    "lance",
+    "--max-num-batched-tokens",
+    "32768",
+    "--max-num-seqs",
+    "1",
+    "--enforce-eager",
+    "--trust-remote-code",
+    "--no-enable-prefix-caching",
+    "--no-async-chunk",
+]
+
 test_params = [
     OmniServerParams(
         model=MODEL,
-        stage_config_path=STAGE_CONFIGS_PATH,
+        server_args=_LANCE_SERVE_ARGS,
         stage_init_timeout=300,
     ),
 ]
diff --git a/tests/e2e/online_serving/test_omnivoice_expansion.py b/tests/e2e/online_serving/test_omnivoice_expansion.py
index 631c799cd..ff41ee413 100644
--- a/tests/e2e/online_serving/test_omnivoice_expansion.py
+++ b/tests/e2e/online_serving/test_omnivoice_expansion.py
@@ -74,6 +74,41 @@ class TestOmniVoiceTTS:
         openai_client.send_audio_speech_request(request_config)
 
 
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+class TestOmniVoiceSeed:
+    """E2E tests for OmniVoice Seed Params."""
+
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    def test_speech_auto_voice_seed_deterministic(self, omni_server, openai_client) -> None:
+        cfg = {
+            "model": omni_server.model,
+            "input": get_prompt("text"),
+            "response_format": "wav",
+            "seed": 42,
+            "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES,
+        }
+
+        r1 = openai_client.send_audio_speech_request(cfg)[0]
+        r2 = openai_client.send_audio_speech_request(cfg)[0]
+        assert r1.audio_bytes == r2.audio_bytes
+
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    def test_speech_auto_voice_seed_non_deterministic(self, omni_server, openai_client) -> None:
+        cfg = {
+            "model": omni_server.model,
+            "input": get_prompt("text"),
+            "response_format": "wav",
+            "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES,
+        }
+
+        cfg1 = {**cfg, "seed": 42}
+        cfg2 = {**cfg, "seed": 43}
+
+        r1 = openai_client.send_audio_speech_request(cfg1)[0]
+        r2 = openai_client.send_audio_speech_request(cfg2)[0]
+        assert r1.audio_bytes != r2.audio_bytes
+
+
 @pytest.mark.skipif(not _HAS_VOICE_CLONE, reason="Voice cloning requires transformers>=5.3.0")
 @pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
 class TestOmniVoiceVoiceCloning:
diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py
index 90d5a7129..aa176c077 100644
--- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py
+++ b/tests/e2e/online_serving/test_qwen3_omni_expansion.py
@@ -61,7 +61,11 @@ test_params = [
             model=model,
             stage_config_path=default_path,
             use_stage_cli=True,
-            server_args=["--no-async-chunk"],
+            server_args=[
+                "--no-async-chunk",
+                "--stage-overrides",
+                '{"0": {"enable_prefix_caching": true}, "1": {"enable_prefix_caching": true}}',
+            ],
         ),
         id="default",
     ),
@@ -539,3 +543,24 @@ def test_language_001(omni_server, openai_client) -> None:
     }
 
     openai_client.send_omni_request(request_config)
+
+
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_to_audio_long_output_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: text only (long-form generation prompt).
+    Output Modal: text, audio (default ``modalities``);
+    Input Setting: stream=True
+    Datasets: single request
+    """
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        content_text="Tell a 300-word story.",
+    )
+
+    request_config = {"model": omni_server.model, "messages": messages, "stream": True}
+    responses = openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+    text = responses[0].text_content if responses else ""
+    word_count = len(text.split())
+    assert word_count >= 200, f"Expected at least 200 words in long output, got {word_count}"
diff --git a/tests/e2e/online_serving/test_qwen_image_edit.py b/tests/e2e/online_serving/test_qwen_image_edit.py
index bdf8c9122..738a2e50c 100644
--- a/tests/e2e/online_serving/test_qwen_image_edit.py
+++ b/tests/e2e/online_serving/test_qwen_image_edit.py
@@ -5,8 +5,8 @@
 Online serving tests for the Qwen-Image-Edit family (image-to-image via chat completions).
 
 - ``test_single_image_to_image_001``: ``Qwen/Qwen-Image-Edit`` — one reference image, fixed 512×512.
-- ``test_multi_images_to_image_001``: ``Qwen/Qwen-Image-Edit-2509`` — two reference images, fixed 512×512.
-- ``test_different_sizes_001``: ``Qwen/Qwen-Image-Edit-2509`` only, ``advanced_model`` — mixed input
+- ``test_multi_images_to_image_001``: ``Qwen/Qwen-Image-Edit-2511`` — two reference images, fixed 512×512.
+- ``test_different_sizes_001``: ``Qwen/Qwen-Image-Edit-2511`` only, ``advanced_model`` — mixed input
   resolutions; ``extra_body`` uses per-output ``width``/``height`` lists; the test client sends one
   scalar-size request per list index in parallel and merges images (see ``OpenAIClientHandler.send_diffusion_request``).
 
@@ -53,7 +53,6 @@ def _get_diffusion_feature_cases(model: str):
 
 
 @pytest.mark.advanced_model
-@pytest.mark.core_model
 @pytest.mark.diffusion
 @pytest.mark.parametrize(
     "omni_server",
@@ -83,15 +82,16 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA
     openai_client.send_diffusion_request(request_config)
 
 
+@pytest.mark.core_model
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"),
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
     indirect=True,
 )
 def test_multi_images_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    """Two-reference edit smoke for ``Qwen/Qwen-Image-Edit-2509``."""
+    """Two-reference edit smoke for ``Qwen/Qwen-Image-Edit-2511``."""
 
     image_data_url_list = [f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" for _ in range(2)]
 
@@ -118,7 +118,7 @@ def test_multi_images_to_image_001(omni_server: OmniServer, openai_client: OpenA
 @pytest.mark.diffusion
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"),
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
     indirect=True,
 )
 def test_different_sizes_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
diff --git a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py b/tests/e2e/online_serving/test_qwen_image_edit_expansion.py
index 9b30ef4eb..1457ad636 100644
--- a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py
+++ b/tests/e2e/online_serving/test_qwen_image_edit_expansion.py
@@ -2,7 +2,7 @@
 Comprehensive tests of diffusion features that are available in online serving mode
 and are supported by the following models:
 - Qwen-Image-Edit: single image input
-- Qwen-Image-Edit-2509: single image input and two image inputs
+- Qwen-Image-Edit-2511: single image input and two image inputs
 """
 
 import pytest
@@ -14,7 +14,7 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand
 pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
 EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting."
-SINGLE_EDIT_PROMPT_2509 = "Restyle this image into a Vincent van Gogh style impressionist painting."
+SINGLE_EDIT_PROMPT_2511 = "Restyle this image into a Vincent van Gogh style impressionist painting."
 MULTI_EDIT_PROMPT = (
     "Transform the first image into a Dadaism collage art. "
     "Transform the second image into a Vincent van Gogh style painting. "
@@ -34,7 +34,7 @@ def _get_diffusion_feature_cases(model: str):
                 model=model,
                 server_args=[
                     "--cache-backend",
-                    "tea_cache",  # [TODO] may consider changing to cache_dit after #1779 is resolved. Currently cache_dit and layerwise offload cannot work together.
+                    "cache_dit",
                     "--enable-layerwise-offload",
                 ],
             ),
@@ -59,7 +59,7 @@ def _get_diffusion_feature_cases(model: str):
                 model=model,
                 server_args=[
                     "--cache-backend",
-                    "cache_dit",
+                    "tea_cache",
                     "--ring",
                     "2",
                 ],
@@ -141,11 +141,11 @@ def test_qwen_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHan
 
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"),
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
     indirect=True,
 )
-def test_qwen_image_edit_2509_single_image(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    """Test Qwen-Image-Edit-2509 with a single image input.
+def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+    """Test Qwen-Image-Edit-2511 with a single image input.
 
     Regression: with tea_cache enabled and zero_cond_t=True, the TeaCache
     postprocess closure used the doubled temb (shape 2*batch) without halving
@@ -156,7 +156,7 @@ def test_qwen_image_edit_2509_single_image(omni_server: OmniServer, openai_clien
     """
     image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
 
-    messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=SINGLE_EDIT_PROMPT_2509)
+    messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=SINGLE_EDIT_PROMPT_2511)
 
     request_config = {
         "model": omni_server.model,
@@ -176,11 +176,11 @@ def test_qwen_image_edit_2509_single_image(omni_server: OmniServer, openai_clien
 
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"),
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
     indirect=True,
 )
-def test_qwen_image_edit_2509_two_images(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    """Test Qwen-Image-Edit-2509 with two image inputs."""
+def test_qwen_image_edit_2511_two_images(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+    """Test Qwen-Image-Edit-2511 with two image inputs."""
     image_data_url_1 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
     image_data_url_2 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
 
diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py
index b423aa407..9d474e81c 100644
--- a/tests/engine/test_async_omni_engine_stage_init.py
+++ b/tests/engine/test_async_omni_engine_stage_init.py
@@ -1,8 +1,10 @@
+import concurrent.futures
 import importlib
 import os
 import threading
 import time
 import types
+from contextlib import contextmanager
 
 import pytest
 
@@ -189,10 +191,12 @@ def test_initialize_diffusion_replica_restores_device_visibility_after_local_ini
     old_env = os.environ.get(env_var)
     os.environ[env_var] = "0,1"
 
-    def _fake_setup_stage_devices(_stage_id, _runtime_cfg):
+    @contextmanager
+    def _fake_stage_runtime_setup(_stage_id, _runtime_cfg):
         current_omni_platform.set_device_control_env_var("1")
+        yield
 
-    monkeypatch.setattr(engine_mod, "setup_stage_devices", _fake_setup_stage_devices)
+    monkeypatch.setattr(engine_mod, "stage_runtime_setup", _fake_stage_runtime_setup)
     monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None)
     monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", lambda *_, **__: types.SimpleNamespace())
 
@@ -221,7 +225,11 @@ def test_initialize_diffusion_replica_passes_stage_init_timeout_and_inline_flag(
 
     captured: dict[str, object] = {}
 
-    monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None)
+    @contextmanager
+    def _noop_stage_runtime_setup(*_):
+        yield
+
+    monkeypatch.setattr(engine_mod, "stage_runtime_setup", _noop_stage_runtime_setup)
     monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None)
 
     def _capture_initialize_diffusion_stage(
@@ -372,6 +380,44 @@ def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatc
     assert all(replica.num_replicas == 3 for replica in stage_plans[1].replicas)
 
 
+def test_initialize_stage_replicas_reuses_long_lived_executor(monkeypatch):
+    engine = object.__new__(AsyncOmniEngine)
+    engine._stage_init_executor = None
+
+    cfg = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    stage_plans = [_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg, num_replicas=1)]
+
+    created_executors = []
+
+    class FakeExecutor:
+        def __init__(self, *args, **kwargs):
+            self.shutdown_calls = []
+            created_executors.append(self)
+
+        def submit(self, fn, *args, **kwargs):
+            future = concurrent.futures.Future()
+            future.set_result(fn(*args, **kwargs))
+            return future
+
+        def shutdown(self, wait=False, cancel_futures=False):
+            self.shutdown_calls.append((wait, cancel_futures))
+
+    monkeypatch.setattr(concurrent.futures, "ThreadPoolExecutor", FakeExecutor)
+    monkeypatch.setattr(
+        engine,
+        "_initialize_replica",
+        lambda plan, _stage_init_timeout, _stage_launch_lock: f"client-{plan.replica_id}",
+    )
+
+    first = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=1)
+    second = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=1)
+
+    assert first == {0: ["client-0"]}
+    assert second == {0: ["client-0"]}
+    assert len(created_executors) == 1
+    assert engine._stage_init_executor is created_executors[0]
+
+
 def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monkeypatch):
     engine = object.__new__(AsyncOmniEngine)
 
@@ -403,6 +449,42 @@ def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monk
     }
 
 
+def test_shutdown_closes_long_lived_stage_init_executor():
+    engine = object.__new__(AsyncOmniEngine)
+    engine._shutdown_called = False
+    engine._weak_finalizer = None
+    engine._omni_master_server = None
+    engine._coordinator_runtime = None
+    engine.orchestrator_thread = types.SimpleNamespace(join=lambda timeout: None, is_alive=lambda: False)
+
+    class FakeQueue:
+        def __init__(self):
+            self.closed = False
+            self.sync_q = types.SimpleNamespace(put_nowait=lambda _msg: None)
+
+        def close(self):
+            self.closed = True
+
+    class FakeExecutor:
+        def __init__(self):
+            self.shutdown_calls = []
+
+        def shutdown(self, wait=False, cancel_futures=False):
+            self.shutdown_calls.append((wait, cancel_futures))
+
+    engine.request_queue = FakeQueue()
+    engine.output_queue = FakeQueue()
+    engine.rpc_output_queue = FakeQueue()
+    engine._stage_init_executor = FakeExecutor()
+
+    engine.shutdown()
+
+    assert engine._stage_init_executor is None
+    assert engine.request_queue.closed is True
+    assert engine.output_queue.closed is True
+    assert engine.rpc_output_queue.closed is True
+
+
 def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_replica_failure(monkeypatch):
     import vllm_omni.engine.async_omni_engine as engine_mod
 
@@ -527,7 +609,11 @@ def test_initialize_llm_replica_passes_stage_init_timeout_to_complete_stage_hand
     prev_device_env = os.environ.get(device_env_var)
     os.environ[device_env_var] = "0"
 
-    monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None)
+    @contextmanager
+    def _noop_stage_runtime_setup(*_):
+        yield
+
+    monkeypatch.setattr(engine_mod, "stage_runtime_setup", _noop_stage_runtime_setup)
     monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {})
     monkeypatch.setattr(engine_mod, "acquire_device_locks", lambda *_: [])
     monkeypatch.setattr(engine_mod, "spawn_stage_core", lambda **_: (fake_addresses, fake_proc, "ipc://handshake"))
diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py
index 9e050a3f8..55e3edc41 100644
--- a/tests/engine/test_orchestrator.py
+++ b/tests/engine/test_orchestrator.py
@@ -964,6 +964,7 @@ async def test_handle_streaming_update_passes_prompt_text_to_stage_pool() -> Non
 
     pool = RecordingPool()
     orchestrator = object.__new__(Orchestrator)
+    orchestrator.async_chunk = False
     orchestrator.request_states = {
         "req-stream": OrchestratorRequestState(
             request_id="req-stream",
diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py
index 52187c341..04ddbe22d 100644
--- a/tests/engine/test_single_stage_mode.py
+++ b/tests/engine/test_single_stage_mode.py
@@ -4,7 +4,7 @@ from __future__ import annotations
 
 import os
 import threading
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from types import SimpleNamespace
 from typing import Any
 
@@ -835,7 +835,7 @@ class TestSingleStageReplicaInitialization:
         prev_device_env = os.environ.get(device_env_var)
         os.environ[device_env_var] = "0"
 
-        mocker.patch.object(engine_mod, "setup_stage_devices")
+        mocker.patch.object(engine_mod, "stage_runtime_setup", return_value=nullcontext())
         mocker.patch.object(engine_mod, "build_engine_args_dict", return_value={})
         mocker.patch.object(engine_mod, "acquire_device_locks", return_value=[])
         mocker.patch.object(engine_mod, "release_device_locks")
@@ -909,7 +909,7 @@ class TestSingleStageReplicaInitialization:
         prev_device_env = os.environ.get(device_env_var)
         os.environ[device_env_var] = "0"
 
-        mocker.patch.object(engine_mod, "setup_stage_devices")
+        mocker.patch.object(engine_mod, "stage_runtime_setup", return_value=nullcontext())
         mocker.patch.object(engine_mod, "inject_kv_stage_info")
         mocker.patch.object(engine_mod, "build_diffusion_config", return_value="diffusion-config")
         mock_register = mocker.patch.object(
@@ -988,7 +988,7 @@ class TestSingleStageReplicaInitialization:
         prev_device_env = os.environ.get(device_env_var)
         os.environ[device_env_var] = "0"
 
-        mocker.patch.object(engine_mod, "setup_stage_devices")
+        mocker.patch.object(engine_mod, "stage_runtime_setup", return_value=nullcontext())
         mocker.patch.object(engine_mod, "inject_kv_stage_info")
         mocker.patch.object(engine_mod, "build_diffusion_config", return_value="diffusion-config")
         mocker.patch.object(
diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py
index 40adb7a91..2556ac71d 100644
--- a/tests/entrypoints/openai_api/test_image_server.py
+++ b/tests/entrypoints/openai_api/test_image_server.py
@@ -9,6 +9,7 @@ OpenAI-compatible async text-to-image generation API endpoints in api_server.py.
 
 import base64
 import io
+import json
 from argparse import Namespace
 from types import SimpleNamespace
 
@@ -19,6 +20,7 @@ from PIL import Image
 from pytest_mock import MockerFixture
 from vllm import SamplingParams
 from vllm.entrypoints.openai.models.protocol import BaseModelPath
+from vllm.sampling_params import RequestOutputKind
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.api_server import _DiffusionServingModels, router
@@ -120,6 +122,27 @@ class MockGenerationResult:
         self.peak_memory_mb = 0.0
 
 
+class MockStageResult:
+    """Mock multi-stage output for streaming image edit tests."""
+
+    def __init__(self, *, stage_id, final_output_type, text="", texts=None, images=None):
+        self.stage_id = stage_id
+        self.final_output_type = final_output_type
+        self.images = images or []
+        if texts is not None:
+            outputs = [SimpleNamespace(text=item, index=index) for index, item in enumerate(texts)]
+        elif text:
+            outputs = [SimpleNamespace(text=text, index=0)]
+        else:
+            outputs = []
+        self.request_output = SimpleNamespace(
+            outputs=outputs,
+            images=self.images,
+        )
+        self.stage_durations = {}
+        self.peak_memory_mb = 0.0
+
+
 class FakeAsyncOmni:
     """Fake AsyncOmni that yields a single diffusion output."""
 
@@ -391,6 +414,72 @@ def async_omni_stage_configs_only_client():
     return TestClient(app)
 
 
+@pytest.fixture
+def streaming_image_edit_client():
+    """Create a multi-stage client whose engine yields AR text before image output."""
+    from fastapi import FastAPI
+
+    from vllm_omni.entrypoints.async_omni import AsyncOmni
+    from vllm_omni.entrypoints.openai.api_server import router
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    class FakeAsyncOmniClass(AsyncOmni):
+        def __init__(self):
+            stage_configs = [
+                SimpleNamespace(stage_type="llm", is_comprehension=True),
+                SimpleNamespace(stage_type="diffusion", is_comprehension=False),
+            ]
+            default_sampling_params_list = [
+                SamplingParams(temperature=0.1),
+                OmniDiffusionSamplingParams(),
+            ]
+            self.engine = SimpleNamespace(
+                stage_configs=stage_configs,
+                default_sampling_params_list=default_sampling_params_list,
+            )
+            self.default_sampling_params_list = default_sampling_params_list
+            self.captured_sampling_params_list = None
+            self.captured_prompt = None
+            self.od_config = SimpleNamespace(supports_multimodal_inputs=True)
+
+        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None):
+            self.captured_prompt = prompt
+            self.captured_sampling_params_list = sampling_params_list or [sampling_params]
+            assert self.captured_sampling_params_list[0].output_kind == RequestOutputKind.DELTA
+            yield MockStageResult(stage_id=0, final_output_type="text", text="recap")
+            yield MockStageResult(stage_id=0, final_output_type="text", text=" done")
+            yield MockStageResult(
+                stage_id=1,
+                final_output_type="image",
+                images=[Image.new("RGB", (32, 24), color="purple")],
+            )
+
+        def __class_getitem__(cls, item):
+            return cls
+
+        def get_diffusion_od_config(self):
+            return self.od_config
+
+    app = FastAPI()
+    app.include_router(router)
+
+    engine = FakeAsyncOmniClass()
+    chat_handler = object.__new__(OmniOpenAIServingChat)
+    chat_handler.engine_client = engine
+    chat_handler._diffusion_engine = None
+    app.state.openai_serving_chat = chat_handler
+    app.state.engine_client = engine
+    app.state.stage_configs = [
+        SimpleNamespace(stage_type="llm"),
+        SimpleNamespace(stage_type="diffusion"),
+    ]
+    app.state.args = Namespace(
+        default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}',
+        max_generated_image_size=1024 * 1792,
+    )
+    return TestClient(app)
+
+
 def test_health_endpoint(test_client):
     """Test health check endpoint for diffusion mode"""
     response = test_client.get("/health")
@@ -680,6 +769,154 @@ def test_image_edits_async_omni_stage_configs_only(async_omni_stage_configs_only
     assert len(captured) == 2
 
 
+def _parse_sse_payloads(body: str):
+    payloads = []
+    for line in body.splitlines():
+        if not line.startswith("data: "):
+            continue
+        data = line[len("data: ") :]
+        payloads.append(data if data == "[DONE]" else json.loads(data))
+    return payloads
+
+
+def test_image_edits_streaming_returns_ar_delta_then_image(streaming_image_edit_client):
+    img_bytes = make_test_image_bytes((16, 16))
+    response = streaming_image_edit_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes)],
+        data={
+            "prompt": "edit me",
+            "size": "auto",
+            "stream": "true",
+            "output_format": "png",
+        },
+    )
+
+    assert response.status_code == 200
+    assert response.headers["content-type"].startswith("text/event-stream")
+
+    payloads = _parse_sse_payloads(response.text)
+    assert [p["type"] if isinstance(p, dict) and "type" in p else p for p in payloads] == [
+        "ar_delta",
+        "ar_delta",
+        "image",
+        "[DONE]",
+    ]
+    assert payloads[0]["delta"] == "recap"
+    assert payloads[0]["index"] == 0
+    assert payloads[1]["delta"] == " done"
+    assert payloads[2]["output_format"] == "png"
+    assert payloads[2]["size"] == "16x16"
+
+    image_payload = payloads[2]["data"][0]
+    img = Image.open(io.BytesIO(base64.b64decode(image_payload["b64_json"])))
+    assert img.size == (32, 24)
+
+
+def test_image_edits_streaming_ar_delta_chunks_include_index(streaming_image_edit_client):
+    async def generate_multi_output_delta(prompt, request_id, sampling_params=None, sampling_params_list=None):
+        yield MockStageResult(stage_id=0, final_output_type="text", texts=["first", "second"])
+        yield MockStageResult(
+            stage_id=1,
+            final_output_type="image",
+            images=[Image.new("RGB", (32, 24), color="purple")],
+        )
+
+    streaming_image_edit_client.app.state.engine_client.generate = generate_multi_output_delta
+    img_bytes = make_test_image_bytes((16, 16))
+    response = streaming_image_edit_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes)],
+        data={
+            "prompt": "edit me",
+            "stream": "true",
+        },
+    )
+
+    assert response.status_code == 200
+    payloads = _parse_sse_payloads(response.text)
+    assert [(payloads[0]["index"], payloads[0]["delta"]), (payloads[1]["index"], payloads[1]["delta"])] == [
+        (0, "first"),
+        (1, "second"),
+    ]
+
+
+def test_image_edits_streaming_errors_without_final_image(streaming_image_edit_client):
+    async def generate_without_image(prompt, request_id, sampling_params=None, sampling_params_list=None):
+        yield MockStageResult(stage_id=0, final_output_type="text", text="recap")
+
+    streaming_image_edit_client.app.state.engine_client.generate = generate_without_image
+    img_bytes = make_test_image_bytes((16, 16))
+    response = streaming_image_edit_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes)],
+        data={
+            "prompt": "edit me",
+            "stream": "true",
+        },
+    )
+
+    assert response.status_code == 200
+    payloads = _parse_sse_payloads(response.text)
+    assert payloads[0]["type"] == "ar_delta"
+    assert payloads[1]["object"] == "error"
+    assert "without a final image" in payloads[1]["error"]["message"]
+    assert payloads[2] == "[DONE]"
+
+
+def test_image_edits_streaming_errors_on_empty_final_image(streaming_image_edit_client):
+    async def generate_empty_image(prompt, request_id, sampling_params=None, sampling_params_list=None):
+        yield MockStageResult(stage_id=0, final_output_type="text", text="recap")
+        yield MockStageResult(stage_id=1, final_output_type="image", images=[])
+
+    streaming_image_edit_client.app.state.engine_client.generate = generate_empty_image
+    img_bytes = make_test_image_bytes((16, 16))
+    response = streaming_image_edit_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes)],
+        data={
+            "prompt": "edit me",
+            "stream": "true",
+        },
+    )
+
+    assert response.status_code == 200
+    payloads = _parse_sse_payloads(response.text)
+    assert payloads[0]["type"] == "ar_delta"
+    assert payloads[1]["object"] == "error"
+    assert "empty final image" in payloads[1]["error"]["message"]
+    assert payloads[2] == "[DONE]"
+
+
+def test_image_edits_streaming_rejects_single_stage(test_client):
+    img_bytes = make_test_image_bytes((16, 16))
+    response = test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes)],
+        data={
+            "prompt": "edit me",
+            "stream": "true",
+        },
+    )
+
+    assert response.status_code == 400
+    assert "multi-stage" in response.json()["detail"]
+
+
+def test_image_edits_streaming_rejects_single_stage_before_loading_url(test_client):
+    response = test_client.post(
+        "/v1/images/edits",
+        data={
+            "prompt": "edit me",
+            "url": "https://example.invalid/not-fetched.png",
+            "stream": "true",
+        },
+    )
+
+    assert response.status_code == 400
+    assert "multi-stage" in response.json()["detail"]
+
+
 def test_generate_images_max_size_rejected(async_omni_test_client):
     """Test that a size exceeding max_generated_image_size returns 400."""
     response = async_omni_test_client.post(
diff --git a/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py b/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py
new file mode 100644
index 000000000..cf7d6e234
--- /dev/null
+++ b/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py
@@ -0,0 +1,289 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import dataclasses
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@dataclasses.dataclass
+class FakeAck:
+    stage_id: int
+    result: str = "ok"
+
+
+def _make_app(engine_client):
+    from vllm_omni.entrypoints.openai.api_server import router
+
+    app = FastAPI()
+    app.include_router(router)
+    app.state.engine_client = engine_client
+    app.state.sleeping_stages = set()
+    return app
+
+
+@pytest.fixture
+def sleep_capable_engine(mocker):
+    engine = mocker.MagicMock()
+    engine.sleep = mocker.AsyncMock(return_value=[FakeAck(stage_id=0), FakeAck(stage_id=1)])
+    engine.wake_up = mocker.AsyncMock(return_value=[FakeAck(stage_id=0), FakeAck(stage_id=1)])
+    return engine
+
+
+@pytest.fixture
+def sleep_incapable_engine(mocker):
+    return mocker.MagicMock(spec=[])  # no 'sleep' or 'wake_up' attributes
+
+
+# ---------------------------------------------------------------------------
+# /v1/omni/sleep
+# ---------------------------------------------------------------------------
+
+
+def test_sleep_success(sleep_capable_engine):
+    app = _make_app(sleep_capable_engine)
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/sleep", json={"stage_ids": [0, 1], "level": 2})
+
+    assert response.status_code == 200
+    data = response.json()
+    assert data["status"] == "SUCCESS"
+    assert [ack["stage_id"] for ack in data["acks"]] == [0, 1]
+    assert app.state.sleeping_stages == {0, 1}
+    sleep_capable_engine.sleep.assert_awaited_once_with(stage_ids=[0, 1], level=2)
+
+
+def test_sleep_default_level(sleep_capable_engine):
+    app = _make_app(sleep_capable_engine)
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/sleep", json={"stage_ids": [0]})
+
+    assert response.status_code == 200
+    sleep_capable_engine.sleep.assert_awaited_once_with(stage_ids=[0], level=2)
+
+
+def test_sleep_empty_stage_ids(sleep_capable_engine, mocker):
+    """Empty stage_ids: engine is still called and returns SUCCESS with no acks."""
+    sleep_capable_engine.sleep = mocker.AsyncMock(return_value=[])
+    app = _make_app(sleep_capable_engine)
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/sleep", json={"stage_ids": [], "level": 2})
+
+    assert response.status_code == 200
+    data = response.json()
+    assert data["status"] == "SUCCESS"
+    assert data["acks"] == []
+    assert app.state.sleeping_stages == set()
+    sleep_capable_engine.sleep.assert_awaited_once_with(stage_ids=[], level=2)
+
+
+def test_sleep_updates_sleeping_set(sleep_capable_engine):
+    app = _make_app(sleep_capable_engine)
+    app.state.sleeping_stages = {5}  # pre-existing entry
+    client = TestClient(app)
+
+    client.post("/v1/omni/sleep", json={"stage_ids": [0, 1], "level": 2})
+
+    assert app.state.sleeping_stages == {5, 0, 1}
+
+
+def test_sleep_engine_not_support(sleep_incapable_engine):
+    app = _make_app(sleep_incapable_engine)
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/sleep", json={"stage_ids": [0], "level": 1})
+
+    assert response.status_code == 501
+    assert "sleep" in response.json()["detail"]
+
+
+# ---------------------------------------------------------------------------
+# /v1/omni/wakeup
+# ---------------------------------------------------------------------------
+
+
+def test_wakeup_success(sleep_capable_engine):
+    app = _make_app(sleep_capable_engine)
+    app.state.sleeping_stages = {0, 1}
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/wakeup", json={"stage_ids": [0, 1]})
+
+    assert response.status_code == 200
+    data = response.json()
+    assert data["status"] == "SUCCESS"
+    assert [ack["stage_id"] for ack in data["acks"]] == [0, 1]
+    assert app.state.sleeping_stages == set()
+    sleep_capable_engine.wake_up.assert_awaited_once_with(stage_ids=[0, 1])
+
+
+def test_wakeup_skipped_when_not_sleeping(sleep_capable_engine):
+    app = _make_app(sleep_capable_engine)
+    app.state.sleeping_stages = set()  # nothing sleeping
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/wakeup", json={"stage_ids": [0, 1]})
+
+    assert response.status_code == 200
+    data = response.json()
+    assert data["status"] == "SKIPPED"
+    sleep_capable_engine.wake_up.assert_not_awaited()
+
+
+def test_wakeup_partial_sleeping(sleep_capable_engine):
+    """Only stage 0 is sleeping; stage 1 is not. Should still proceed (partial match)."""
+    app = _make_app(sleep_capable_engine)
+    app.state.sleeping_stages = {0}
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/wakeup", json={"stage_ids": [0, 1]})
+
+    assert response.status_code == 200
+    data = response.json()
+    assert data["status"] == "SUCCESS"
+    assert 0 not in app.state.sleeping_stages
+
+
+def test_wakeup_empty_stage_ids(sleep_capable_engine):
+    """Empty stage_ids: any() on empty iterable is False → SKIPPED, engine not called."""
+    app = _make_app(sleep_capable_engine)
+    app.state.sleeping_stages = {0, 1}
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/wakeup", json={"stage_ids": []})
+
+    assert response.status_code == 200
+    assert response.json()["status"] == "SKIPPED"
+    sleep_capable_engine.wake_up.assert_not_awaited()
+
+
+def test_wakeup_engine_not_support(sleep_incapable_engine):
+    app = _make_app(sleep_incapable_engine)
+    app.state.sleeping_stages = {0}
+    client = TestClient(app)
+
+    response = client.post("/v1/omni/wakeup", json={"stage_ids": [0]})
+
+    assert response.status_code == 501
+    assert "wake_up" in response.json()["detail"]
+
+
+def test_wakeup_removes_only_requested_stages(sleep_capable_engine):
+    app = _make_app(sleep_capable_engine)
+    app.state.sleeping_stages = {0, 1, 2}
+    client = TestClient(app)
+
+    client.post("/v1/omni/wakeup", json={"stage_ids": [0]})
+
+    assert app.state.sleeping_stages == {1, 2}
+
+
+def test_sleep_then_wakeup_roundtrip(sleep_capable_engine):
+    app = _make_app(sleep_capable_engine)
+    client = TestClient(app)
+
+    sleep_resp = client.post("/v1/omni/sleep", json={"stage_ids": [0, 1], "level": 2})
+    assert sleep_resp.status_code == 200
+    assert app.state.sleeping_stages == {0, 1}
+
+    wakeup_resp = client.post("/v1/omni/wakeup", json={"stage_ids": [0, 1]})
+    assert wakeup_resp.status_code == 200
+    assert app.state.sleeping_stages == set()
+
+
+# ---------------------------------------------------------------------------
+# pure-diffusion branch: omni_init_app_state initialises sleeping_stages
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def pure_diffusion_engine(mocker):
+    engine = mocker.MagicMock()
+    engine.stage_configs = [{"stage_type": "diffusion"}]
+    engine.sleep = mocker.AsyncMock(return_value=[FakeAck(stage_id=0)])
+    engine.wake_up = mocker.AsyncMock(return_value=[FakeAck(stage_id=0)])
+    # Remove attributes that would make _get_vllm_config return a config
+    del engine.get_vllm_config
+    del engine.vllm_config
+    return engine
+
+
+@pytest.fixture
+def pure_diffusion_app(pure_diffusion_engine, mocker):
+    """App whose state was initialised via omni_init_app_state (pure diffusion path)."""
+    import asyncio
+    from argparse import Namespace
+
+    from vllm_omni.entrypoints.openai.api_server import omni_init_app_state, router
+
+    args = Namespace(
+        served_model_name=None,
+        model="fake-diffusion-model",
+        enable_log_requests=False,
+        disable_log_stats=True,
+        max_log_len=0,
+        enable_server_load_tracking=False,
+    )
+
+    app = FastAPI()
+    app.include_router(router)
+
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.api_server.OmniOpenAIServingChat"
+    ).for_diffusion.return_value = mocker.MagicMock()
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.api_server.OmniOpenAIServingVideo"
+    ).for_diffusion.return_value = mocker.MagicMock()
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.api_server.OmniOpenAIServingSpeech"
+    ).for_diffusion.return_value = mocker.MagicMock()
+    mocker.patch("vllm_omni.entrypoints.openai.api_server._DiffusionServingModels")
+
+    loop = asyncio.new_event_loop()
+    try:
+        loop.run_until_complete(omni_init_app_state(pure_diffusion_engine, app.state, args))
+    finally:
+        loop.close()
+
+    return app
+
+
+def test_pure_diffusion_sleep_success(pure_diffusion_app, pure_diffusion_engine):
+    client = TestClient(pure_diffusion_app)
+
+    response = client.post("/v1/omni/sleep", json={"stage_ids": [0], "level": 2})
+
+    assert response.status_code == 200
+    assert response.json()["status"] == "SUCCESS"
+    assert pure_diffusion_app.state.sleeping_stages == {0}
+    pure_diffusion_engine.sleep.assert_awaited_once_with(stage_ids=[0], level=2)
+
+
+def test_pure_diffusion_wakeup_skipped(pure_diffusion_app, pure_diffusion_engine):
+    # sleeping_stages is empty after init — nothing is sleeping
+    client = TestClient(pure_diffusion_app)
+
+    response = client.post("/v1/omni/wakeup", json={"stage_ids": [0]})
+
+    assert response.status_code == 200
+    assert response.json()["status"] == "SKIPPED"
+    pure_diffusion_engine.wake_up.assert_not_awaited()
+
+
+def test_pure_diffusion_sleep_then_wakeup(pure_diffusion_app, pure_diffusion_engine):
+    client = TestClient(pure_diffusion_app)
+
+    sleep_resp = client.post("/v1/omni/sleep", json={"stage_ids": [0], "level": 2})
+    assert sleep_resp.status_code == 200
+    assert pure_diffusion_app.state.sleeping_stages == {0}
+
+    wakeup_resp = client.post("/v1/omni/wakeup", json={"stage_ids": [0]})
+    assert wakeup_resp.status_code == 200
+    assert pure_diffusion_app.state.sleeping_stages == set()
diff --git a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py
index 00a6565d9..39b27103e 100644
--- a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py
+++ b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py
@@ -44,7 +44,6 @@ realtime_server_params = [
             model=MODEL,
             stage_config_path=default_stage_config,
             use_stage_cli=True,
-            server_args=["--no-async-chunk"],
         ),
         id="default",
     ),
diff --git a/tests/entrypoints/openai_api/test_qwen3_tts_speaker_embedding.py b/tests/entrypoints/openai_api/test_qwen3_tts_speaker_embedding.py
index 1f71cf8fb..d82e1a68b 100644
--- a/tests/entrypoints/openai_api/test_qwen3_tts_speaker_embedding.py
+++ b/tests/entrypoints/openai_api/test_qwen3_tts_speaker_embedding.py
@@ -28,7 +28,7 @@ MODEL_BASE = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
 MODEL_BASE_1_7B = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
 
 # A synthetic 1024-dim speaker embedding (all 0.1 — not a real voice, but
-# exercises the full code path through the talker's _build_prompt_embeds).
+# exercises the full code path through Qwen3TTSPromptEmbedsBuilder.build_prompt_embeds).
 DUMMY_EMBEDDING_1024 = [0.1] * 1024
 DUMMY_EMBEDDING_2048 = [0.1] * 2048
 
diff --git a/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py b/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py
index 4b63588ba..c8bb050e7 100644
--- a/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py
+++ b/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py
@@ -363,3 +363,104 @@ def test_build_multistage_generation_inputs_custom_system_prompt(serving_chat):
         f"custom system_prompt content must reach the rendered prompt; "
         f"marker {marker!r} not found in prompt of length {len(out['prompt'])}"
     )
+
+
+def test_build_multistage_generation_inputs_sets_ar_stop_token_ids_with_explicit_size(serving_chat):
+    """When height+width are provided with bot_task, the AR (llm) stage
+    must receive stop_token_ids from resolve_stop_token_ids so AR stops
+    at the correct terminator instead of generating until max_tokens.
+
+    Without this, AR would generate past the cot boundary and produce
+    garbage that the DiT bridge cannot parse.
+    """
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    class FakeTokenizer:
+        SPECIAL = {
+            "<|startoftext|>": 1,
+            "<img>": 2,
+            "<recaption>": 4,
+        }
+
+        def convert_tokens_to_ids(self, tok: str) -> int:
+            return self.SPECIAL.get(tok, 0)
+
+        def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
+            return list(range(100, 100 + len(text)))
+
+    engine = SimpleNamespace(
+        stage_configs=[
+            SimpleNamespace(stage_type="llm", is_comprehension=True),
+            SimpleNamespace(stage_type="diffusion", is_comprehension=False),
+        ],
+        default_sampling_params_list=[
+            SamplingParams(temperature=0.0),
+            OmniDiffusionSamplingParams(),
+        ],
+    )
+    images = [Image.new("RGB", (32, 32), color="red")]
+
+    # With height+width and bot_task=think, ar_image_size="1024x768" ->
+    # need_ratio=False -> stop_token_ids=[end_of_think, end_of_recaption]
+    _, sampling_params_list = OmniOpenAIServingChat._build_multistage_generation_inputs(
+        serving_chat,
+        engine=engine,
+        prompt="draw a cat",
+        extra_body={"bot_task": "think"},
+        reference_images=images,
+        gen_params=OmniDiffusionSamplingParams(height=768, width=1024),
+        tokenizer=FakeTokenizer(),
+    )
+
+    # AR stage (index 0) must have stop_token_ids set.
+    ar_params = sampling_params_list[0]
+    assert ar_params.stop_token_ids is not None, (
+        "AR stage must have stop_token_ids set when height+width and bot_task are provided"
+    )
+    assert len(ar_params.stop_token_ids) > 0, "stop_token_ids must be non-empty"
+
+    # Diffusion stage (index 1) must NOT have stop_token_ids set.
+    diff_params = sampling_params_list[1]
+    assert getattr(diff_params, "stop_token_ids", None) is None, "Diffusion stage must not have stop_token_ids set"
+
+
+def test_build_multistage_generation_inputs_no_stop_token_ids_without_size(serving_chat):
+    """Without height+width, ar_image_size=None -> need_ratio=True ->
+    stop_token_ids is set to the ratio range (AR predicts ratio).
+    This still assigns stop_token_ids, but the set is the ratio range,
+    not the terminator.
+
+    Without bot_task at all, ar_stop_token_ids stays None and
+    stop_token_ids is not set on any stage.
+    """
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    engine = SimpleNamespace(
+        stage_configs=[
+            SimpleNamespace(stage_type="llm", is_comprehension=True),
+            SimpleNamespace(stage_type="diffusion", is_comprehension=False),
+        ],
+        default_sampling_params_list=[
+            SamplingParams(temperature=0.0),
+            OmniDiffusionSamplingParams(),
+        ],
+    )
+    images = [Image.new("RGB", (32, 32), color="red")]
+
+    # No height/width, no bot_task -> ar_stop_token_ids=None ->
+    # stop_token_ids not set on any stage.
+    _, sampling_params_list = OmniOpenAIServingChat._build_multistage_generation_inputs(
+        serving_chat,
+        engine=engine,
+        prompt="draw a cat",
+        extra_body={},
+        reference_images=images,
+        gen_params=OmniDiffusionSamplingParams(),
+    )
+
+    # SamplingParams defaults stop_token_ids=[], not None.
+    # The key contract: it was NOT set by resolve_stop_token_ids,
+    # so it stays as the SamplingParams default (empty list).
+    assert sampling_params_list[0].stop_token_ids == [], (
+        "Without bot_task, AR stage stop_token_ids must be the default empty list"
+    )
diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py
index 2440fe586..872410aaa 100644
--- a/tests/entrypoints/openai_api/test_serving_speech.py
+++ b/tests/entrypoints/openai_api/test_serving_speech.py
@@ -2335,6 +2335,84 @@ class TestCosyVoice3Serving:
         cosyvoice3_server._build_cosyvoice3_prompt.assert_awaited_once()
 
 
+# ---- GLM-TTS Serving Tests ----
+
+
+@pytest.fixture
+def glm_tts_server(mocker: MockerFixture):
+    mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
+    mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
+
+    mock_engine_client = mocker.MagicMock()
+    mock_engine_client.errored = False
+    mock_engine_client.model_config = mocker.MagicMock(
+        model="zai-org/GLM-TTS",
+        hf_config=SimpleNamespace(min_token_text_ratio=2, max_token_text_ratio=20),
+    )
+    mock_engine_client.default_sampling_params_list = [
+        SimpleNamespace(max_tokens=2048, min_tokens=None, extra_args=None)
+    ]
+    mock_engine_client.tts_batch_max_items = 32
+    mock_engine_client.generate = mocker.MagicMock(return_value="generator")
+    mock_engine_client.stage_configs = [
+        SimpleNamespace(
+            engine_args=SimpleNamespace(model_stage="glm_tts"),
+            tts_args={},
+        )
+    ]
+
+    mock_models = mocker.MagicMock()
+    mock_models.is_base_model.return_value = True
+
+    return OmniOpenAIServingSpeech(
+        engine_client=mock_engine_client,
+        models=mock_models,
+        request_logger=mocker.MagicMock(),
+    )
+
+
+class TestGLMTTSServing:
+    def test_validate_glm_tts_requires_ref_audio(self, glm_tts_server):
+        request = OpenAICreateSpeechRequest(input="Hello", ref_text="Reference transcript")
+        error = glm_tts_server._validate_glm_tts_request(request)
+        assert error is not None
+        assert "ref_audio" in error
+
+    def test_validate_glm_tts_requires_ref_text(self, glm_tts_server):
+        request = OpenAICreateSpeechRequest(input="Hello", ref_audio="data:audio/wav;base64,abc")
+        error = glm_tts_server._validate_glm_tts_request(request)
+        assert error is not None
+        assert "ref_text" in error
+
+    def test_estimate_glm_tts_target_text_len_uses_tokenizer_tokens(
+        self,
+        glm_tts_server,
+        mocker: MockerFixture,
+    ):
+        class FakeTokenizer:
+            def encode(self, text):
+                assert text == "normalized target"
+                return [10, 20, 30]
+
+        mocker.patch(
+            "vllm_omni.model_executor.models.glm_tts.glm_tts.resolve_glm_tts_tokenizer_path",
+            return_value="resolved-tokenizer",
+        )
+        load_tokenizer = mocker.patch(
+            "vllm_omni.model_executor.models.glm_tts.glm_tts.load_glm_tts_tokenizer",
+            return_value=FakeTokenizer(),
+        )
+        mocker.patch(
+            "vllm_omni.model_executor.models.glm_tts.text_frontend.GLMTTSTextFrontend.text_normalize",
+            return_value="normalized target",
+        )
+
+        text_token_len = glm_tts_server._estimate_glm_tts_text_token_len("abcdef")
+
+        assert text_token_len == 3
+        load_tokenizer.assert_called_once()
+
+
 class TestTTSAsyncOffloading:
     """Tests for event-loop-safe offloading of blocking TTS operations."""
 
diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py
index c8ef5eddf..5d7176490 100644
--- a/tests/entrypoints/test_async_omni.py
+++ b/tests/entrypoints/test_async_omni.py
@@ -1,12 +1,19 @@
 import asyncio
+import re
 from types import SimpleNamespace
 
 import pytest
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 
+from tests.helpers.mark import hardware_test
+from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model]
+
+DIFFUSION_MODEL = "riverclouds/qwen_image_random"
+OMNI_MODEL = "Qwen/Qwen2.5-Omni-7B"
+OMNI_STAGE_CONFIG = get_deploy_config_path("ci/qwen2_5_omni_thinker_only.yaml")
 
 
 async def _noop(**kw):
@@ -59,6 +66,89 @@ def get_async_omni_instance(fake_add_request=_noop, fake_abort_request=_noop) ->
     return omni
 
 
+@pytest.mark.cpu
+def test_generate_submits_randomized_id_to_engine():
+    """Ensure the engine receives a UUID-suffixed ID, not the raw request ID"""
+
+    async def run():
+        submitted_ids = []
+        omni = get_async_omni_instance(fake_add_request=get_fake_add_request(submitted_ids))
+
+        req_id = "my-req-1"
+        async for _ in omni.generate(
+            prompt={"prompt": "test"},
+            request_id=req_id,
+            sampling_params_list=[SimpleNamespace()],
+            output_modalities=["text"],
+        ):
+            pass
+
+        assert len(submitted_ids) == 1
+        assert submitted_ids[0] != req_id
+        assert submitted_ids[0].startswith(f"{req_id}-")
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "req_ids,cancel_prefix,expected_cancel_count",
+    [
+        (["cancel-me"], "cancel-me", 1),
+        (["cancel-me", "cancel-me"], "cancel-me", 2),
+        (["cancel-hello", "cancel-hello-world"], "cancel-hello", 1),
+    ],
+)
+def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefix: str, expected_cancel_count: int):
+    """Ensure that abort() with the user-visible ID resolves correctly.
+
+    NOTE: In the case of concurrent / colliding request(s), all requests matching the
+    user provided request ID will be aborted."""
+
+    async def run():
+        aborted_batches = []
+        omni = get_async_omni_instance(
+            fake_abort_request=get_fake_abort(aborted_batches),
+        )
+
+        async def exhaust(agen):
+            async for _ in agen:
+                pass
+
+        tasks = []
+        for user_request_id in req_ids:
+            t = asyncio.create_task(
+                exhaust(
+                    omni.generate(
+                        prompt={"prompt": "test"},
+                        request_id=user_request_id,
+                        sampling_params_list=[SimpleNamespace()],
+                    )
+                )
+            )
+            await asyncio.sleep(0)
+            tasks.append(t)
+
+        assert len(omni.request_states) == len(req_ids)
+        await omni.abort(cancel_prefix)
+
+        assert len(aborted_batches) == 1
+        aborted_ids = aborted_batches[0]
+        # Aborted requests will have fmt {ext_id}-{UUID} to avoid collisions
+        for rid in aborted_ids:
+            assert re.fullmatch(rf"{re.escape(cancel_prefix)}-[0-9a-f]+", rid)
+        assert len(aborted_ids) == expected_cancel_count
+        assert len(set(aborted_ids)) == expected_cancel_count
+        assert len(omni.request_states) == len(req_ids) - expected_cancel_count
+
+        for t in tasks:
+            t.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
 def test_generate_accepts_request_after_repeated_cancellations():
     async def run_test():
         submitted_request_ids = []
@@ -91,22 +181,23 @@ def test_generate_accepts_request_after_repeated_cancellations():
                 await task
 
         assert len(await collect_outputs("after-cancel")) == 1
-        assert submitted_request_ids == [
-            "baseline",
-            "cancel-0",
-            "cancel-1",
-            "cancel-2",
-            "after-cancel",
-        ]
-        assert aborted_request_batches == [
-            ["cancel-0"],
-            ["cancel-1"],
-            ["cancel-2"],
-        ]
+
+        # Check prefixes instead of equality since generate will add
+        # a random numeric suffix to ensure the request ID is unique
+        expected_prefixes = ["baseline-", "cancel-0-", "cancel-1-", "cancel-2-", "after-cancel-"]
+        assert len(submitted_request_ids) == len(expected_prefixes)
+        for submitted, prefix in zip(submitted_request_ids, expected_prefixes):
+            assert submitted.startswith(prefix)
+
+        assert len(aborted_request_batches) == 3
+        for batch, prefix in zip(aborted_request_batches, ["cancel-0-", "cancel-1-", "cancel-2-"]):
+            assert len(batch) == 1
+            assert batch[0].startswith(prefix)
 
     asyncio.run(run_test())
 
 
+@pytest.mark.cpu
 @pytest.mark.parametrize(
     "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY, RequestOutputKind.CUMULATIVE]
 )
@@ -132,3 +223,79 @@ def test_output_kind_is_preserved_with_explicit_sampling_params(output_kind):
 
     asyncio.run(run())
     assert captured_params[0].output_kind == output_kind
+
+
+# End to end tests for ensuring internal manipulation of request ID
+# in diffusion / Omni models don't leak back to the user.
+#
+# NOTE: It seems like we currently need the shutdowns here, otherwise
+# running the tests sequentially seems to leave a zombie process in diffusion
+# that can OOM the Omni tests.
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.omni
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "req_id",
+    ["my-req-1", "img_gen-abc123", "chatcmpl-xyz"],
+    ids=["plain", "prefixed-img", "prefixed-chat"],
+)
+async def test_diffusion_generate_preserves_request_id(req_id):
+    """Ensure diffusion model requests don't leak internal UUID-suffixed req id back to user."""
+    engine = AsyncOmni(model=DIFFUSION_MODEL)
+    try:
+        async for output in engine.generate("a white cat", request_id=req_id):
+            assert output.request_id == req_id
+    finally:
+        engine.shutdown()
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.omni
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "req_id",
+    ["my-req-1", "img_gen-abc123", "chatcmpl-xyz"],
+    ids=["plain", "prefixed-img", "prefixed-chat"],
+)
+async def test_omni_generate_preserves_request_id(req_id):
+    """Ensure omni model requests don't leak internal UUID-suffixed req id back to user."""
+    engine = AsyncOmni(model=OMNI_MODEL, stage_configs_path=OMNI_STAGE_CONFIG)
+    try:
+        async for output in engine.generate(
+            "Say hello in one word.",
+            request_id=req_id,
+            output_modalities=["text"],
+        ):
+            assert output.request_id == req_id
+    finally:
+        engine.shutdown()
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.omni
+@pytest.mark.asyncio
+async def test_diffusion_generate_empty_request_id():
+    """Empty request_id should get a generated internal ID, not stay empty (diffusion)"""
+    engine = AsyncOmni(model=DIFFUSION_MODEL)
+    try:
+        async for output in engine.generate("a white cat", request_id=""):
+            assert output.request_id != ""
+    finally:
+        engine.shutdown()
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.omni
+@pytest.mark.asyncio
+async def test_omni_generate_empty_request_id():
+    """Empty request_id should get a generated internal ID, not stay empty (omni)"""
+    engine = AsyncOmni(model=OMNI_MODEL, stage_configs_path=OMNI_STAGE_CONFIG)
+    try:
+        async for output in engine.generate(
+            "Say hello in one word.",
+            request_id="",
+            output_modalities=["text"],
+        ):
+            assert output.request_id != ""
+    finally:
+        engine.shutdown()
diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py
index 3612020d4..bbef876d4 100644
--- a/tests/entrypoints/test_omni_entrypoints.py
+++ b/tests/entrypoints/test_omni_entrypoints.py
@@ -16,6 +16,7 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 from vllm_omni.engine.async_omni_engine import StageRuntimeInfo
 from vllm_omni.engine.messages import ErrorMessage, OutputMessage
 from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.entrypoints.omni_base import OmniEngineDeadError
 from vllm_omni.outputs import OmniRequestOutput
@@ -165,6 +166,14 @@ class FakeAsyncOmniEngine:
 def _patch_engine(monkeypatch: pytest.MonkeyPatch, engine: FakeAsyncOmniEngine) -> None:
     monkeypatch.setattr("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", lambda *args, **kwargs: engine)
     monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model)
+    # Don't add random UUIDs to requests calling .generate since we usually
+    # just want to check for present requests anyway, and would need to just
+    # strip the UUID. Explicit checks against the mapping are in tests for
+    # AsyncOmni, or explicitly set the client req states.
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.async_omni.AsyncOmni._get_unique_request_id",
+        staticmethod(lambda request_id: request_id),
+    )
 
 
 def test_direct_omni_with_nullified_parser_only_nulls_untyped_override_fields(
@@ -623,16 +632,23 @@ async def test_async_omni_llm_diffusion_yields_text_stream_then_image(monkeypatc
 async def test_async_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch):
     engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META)
     _patch_engine(monkeypatch, engine)
-
     app = AsyncOmni("dummy-model")
+
+    # Requests internally have a random UUID appended to the
+    # external ID to avoid collisions, so this also tests mapping
+    external_req_id = "req-1"
+    req_id = "req-1-12345678"
     try:
-        app.request_states["req-1"] = object()
-        await app.abort("req-1")
+        app.request_states[req_id] = ClientRequestState(
+            request_id=req_id,
+            external_request_id=external_req_id,
+        )
+        await app.abort(external_req_id)
     finally:
         app.shutdown()
 
-    assert engine.aborted == [["req-1"]]
-    assert "req-1" not in app.request_states
+    assert engine.aborted == [[req_id]]
+    assert external_req_id not in app.request_states
 
 
 @pytest.mark.asyncio
diff --git a/tests/examples/offline_inference/test_qwen3_tts_estimator.py b/tests/examples/offline_inference/test_qwen3_tts_estimator.py
index fde638ba5..4f94eef0d 100644
--- a/tests/examples/offline_inference/test_qwen3_tts_estimator.py
+++ b/tests/examples/offline_inference/test_qwen3_tts_estimator.py
@@ -57,11 +57,11 @@ def wav_file(tmp_path: Path) -> Path:
 def estimator_cache(monkeypatch):
     """
     Bypass model loading by pre-populating the estimator cache, and shortcut
-    Qwen3TTSTalker.estimate_prompt_len_from_additional_information so that
-    the test directly observes the real _estimate_ref_code_len closure.
+    Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information
+    so that the test directly observes the real _estimate_ref_code_len closure.
     """
-    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
-        Qwen3TTSTalkerForConditionalGeneration,
+    from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
+        Qwen3TTSPromptEmbedsBuilder,
     )
 
     def _shortcut(**kwargs):
@@ -69,7 +69,7 @@ def estimator_cache(monkeypatch):
         return kwargs["estimate_ref_code_len"](ref_audio)
 
     monkeypatch.setattr(
-        Qwen3TTSTalkerForConditionalGeneration,
+        Qwen3TTSPromptEmbedsBuilder,
         "estimate_prompt_len_from_additional_information",
         staticmethod(_shortcut),
     )
diff --git a/tests/examples/offline_inference/test_vace_video_generation.py b/tests/examples/offline_inference/test_vace_video_generation.py
new file mode 100644
index 000000000..11746e58a
--- /dev/null
+++ b/tests/examples/offline_inference/test_vace_video_generation.py
@@ -0,0 +1,35 @@
+"""
+Offline inference tests: vace video generation.
+See examples/offline_inference/vace/vace_video_generation.md
+"""
+
+from pathlib import Path
+
+import pytest
+
+from tests.examples.helpers import EXAMPLES, ExampleRunner, ReadmeSnippet
+from tests.helpers.assertions import assert_video_valid
+from tests.helpers.mark import hardware_marks
+
+pytestmark = [
+    pytest.mark.usefixtures("clean_gpu_memory_between_tests"),
+    pytest.mark.full_model,
+    pytest.mark.example,
+    *hardware_marks(res={"cuda": "H100"}),
+]
+
+VACE_SCRIPT = EXAMPLES / "offline_inference" / "vace" / "vace_video_generation.py"
+README_PATH = VACE_SCRIPT.with_name("vace_video_generation.md")
+EXAMPLE_OUTPUT_SUBFOLDER = "example_offline_vace"
+
+README_SNIPPETS = ReadmeSnippet.extract_readme_snippets(README_PATH)
+
+
+@pytest.mark.parametrize("snippet", README_SNIPPETS, ids=lambda snippet: snippet.test_id)
+def test_vace_video_generation(snippet: ReadmeSnippet, example_runner: ExampleRunner):
+    should_skip, reason = snippet.skip
+    if should_skip:
+        pytest.skip(reason)
+    result = example_runner.run(snippet, output_subfolder=Path(EXAMPLE_OUTPUT_SUBFOLDER))
+    for asset in result.assets:
+        assert_video_valid(asset)
diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py
index 5898b624b..f0d515d01 100644
--- a/tests/helpers/assertions.py
+++ b/tests/helpers/assertions.py
@@ -477,6 +477,7 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve
         # Verify similarity (Whisper transcript vs streamed/detokenized text)
         if "audio" in modalities:
             audio_ref_text = request_config.get("audio_ref_text")
+            similarity_threshold = request_config.get("similarity_threshold", 0.8)
             if "text" in modalities:
                 transcript = (response.audio_content or "").strip()
                 text_output = (response.text_content or "").strip()
@@ -484,14 +485,14 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve
                     transcript.lower(),
                     text_output.lower(),
                 )
-                assert similarity > 0.9, "The audio content is not same as the text"
                 print(f"similarity is: {similarity}")
+                assert similarity > similarity_threshold, "The audio content is not same as the text"
             if audio_ref_text:
                 audio_similarity = cosine_similarity_text(
                     response.audio_content.lower(),
                     str(audio_ref_text).lower(),
                 )
-                assert audio_similarity > 0.9, (
+                assert audio_similarity > similarity_threshold, (
                     f"The audio content does not match reference text: similarity={audio_similarity:.3f}"
                 )
 
diff --git a/tests/helpers/fixtures/env.py b/tests/helpers/fixtures/env.py
index 452ae2d2d..2d38e4907 100644
--- a/tests/helpers/fixtures/env.py
+++ b/tests/helpers/fixtures/env.py
@@ -2,6 +2,7 @@ import os
 
 import pytest
 import torch
+from vllm.distributed import parallel_state
 
 
 @pytest.fixture(scope="session", autouse=True)
@@ -59,3 +60,15 @@ def default_vllm_config():
 
     with set_current_vllm_config(VllmConfig(device_config=device_config)):
         yield
+
+
+@pytest.fixture()
+def init_fake_tp_group(mocker):
+    """Provide a fake TP group so vllm linear layers can be instantiated."""
+    mock_tp = mocker.MagicMock()
+    mock_tp.world_size = 1
+    mock_tp.rank_in_group = 0
+    old = parallel_state._TP
+    parallel_state._TP = mock_tp
+    yield
+    parallel_state._TP = old
diff --git a/tests/helpers/media.py b/tests/helpers/media.py
index daf3a2ba4..5c76a9e1e 100644
--- a/tests/helpers/media.py
+++ b/tests/helpers/media.py
@@ -623,9 +623,11 @@ def _whisper_transcribe_in_current_process(output_path: str) -> str:
 
     if current_omni_platform.is_available():
         n = current_omni_platform.get_device_count()
-        if n == 1:
-            device_index = 0
-        elif n > 1:
+        # Single-GPU runners (e.g. the L4 nightly): the model server already
+        # occupies device 0. Loading Whisper there, once per concurrent
+        # request, competes for VRAM and OOMs. Only borrow an accelerator
+        # when a spare device exists; otherwise validate on CPU.
+        if n > 1:
             device_index = n - 1
 
     if device_index is not None:
diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py
index 8b007dc16..a3d337e1b 100644
--- a/tests/helpers/runtime.py
+++ b/tests/helpers/runtime.py
@@ -531,10 +531,6 @@ class OmniServerStageCli(OmniServer):
         if self.env_dict is not None:
             env.update(self.env_dict)
 
-        devices = self.stage_runtime_devices.get(stage_id)
-        if devices:
-            self._set_stage_device_env(stage_id, env, devices, replica_id=replica_id)
-
         cmd = self._build_stage_cmd(stage_id, headless=headless, replica_id=replica_id)
         print(f"Launching OmniServerStageCli stage {stage_id} replica {replica_id}: {' '.join(cmd)}")
         # Capture each subprocess's stdout+stderr to a per-stage log file so
@@ -1671,7 +1667,7 @@ class OpenAIClientHandler:
         # Qwen3-TTS custom fields, forwarded via extra_body.
         extra_body: dict[str, Any] = {}
         # Keep this list aligned with vllm_omni.entrypoints.openai.protocol.audio params.
-        for key in ("task_type", "ref_text", "ref_audio", "language", "max_new_tokens"):
+        for key in ("task_type", "ref_text", "ref_audio", "language", "max_new_tokens", "seed"):
             if key in request_config:
                 extra_body[key] = request_config[key]
 
@@ -2147,8 +2143,8 @@ class OmniRunner:
         _cache = self._prompt_len_estimate_cache
         try:
             from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
-            from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
-                Qwen3TTSTalkerForConditionalGeneration,
+            from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
+                Qwen3TTSPromptEmbedsBuilder,
             )
 
             if model_name not in _cache:
@@ -2160,7 +2156,7 @@ class OmniRunner:
 
             tok, tcfg = _cache[model_name]
             task_type = (additional_information.get("task_type") or ["CustomVoice"])[0]
-            return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
+            return Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
                 additional_information=additional_information,
                 task_type=task_type,
                 tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
diff --git a/tests/helpers/stage_config.py b/tests/helpers/stage_config.py
index 265c42821..e34680be0 100644
--- a/tests/helpers/stage_config.py
+++ b/tests/helpers/stage_config.py
@@ -332,15 +332,14 @@ _CI_OVERLAYS: dict[str, dict[str, Any]] = {
             },
             {
                 "stage_id": 1,
-                "gpu_memory_utilization": 0.5,
                 "max_num_seqs": 5,
+                "gpu_memory_utilization": 0.5,
                 "max_model_len": 32768,
                 "default_sampling_params": {"max_tokens": 1000},
             },
             {
                 "stage_id": 2,
                 "max_num_seqs": 5,
-                "max_num_batched_tokens": 100000,
                 "default_sampling_params": {"max_tokens": 2000},
             },
         ],
@@ -495,15 +494,6 @@ _CI_OVERLAYS: dict[str, dict[str, Any]] = {
             },
         ],
     },
-    "lance": {
-        "base_config": "lance.yaml",
-        "stages": [
-            {
-                "stage_id": 0,
-                "max_num_seqs": 1,
-            },
-        ],
-    },
     "bagel_mooncake": {
         "base_config": "bagel.yaml",
         "stages": [
diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py
index 0e071f724..bf2261cb9 100644
--- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py
+++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py
@@ -293,7 +293,13 @@ def test_code2wav_forward_finalizes_hift_tail():
     model = object.__new__(CosyVoice3Code2Wav)
     nn.Module.__init__(model)
     model.hift = DummyHiFT()
-    model._forward_mel = lambda **_: torch.ones((1, 80, 8), dtype=torch.float32)
+    forward_mel_calls = []
+
+    def fake_forward_mel(**kwargs):
+        forward_mel_calls.append(kwargs)
+        return torch.ones((1, 80, 8), dtype=torch.float32)
+
+    model._forward_mel = fake_forward_mel
 
     out = model.forward(
         token=torch.tensor([[1, 2, 3]], dtype=torch.int32),
@@ -304,3 +310,4 @@ def test_code2wav_forward_finalizes_hift_tail():
 
     assert out.shape == (1, 1, 8)
     assert model.hift.finalize_calls == [True]
+    assert forward_mel_calls[0]["token_offset_tokens"] == 0
diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py
index bc6a26466..b0afc9592 100644
--- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py
+++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py
@@ -138,27 +138,6 @@ def _make_sampling_metadata(
     )
 
 
-def test_split_request_ids_uses_seq_token_counts():
-    CosyVoice3Model, _ = _cosyvoice3_model_and_runner()
-    ids = torch.tensor([10, 11, 12, 13, 14], dtype=torch.long)
-    chunks = CosyVoice3Model._split_request_ids(ids, [2, 2, 2])
-    assert [c.tolist() for c in chunks] == [[10, 11], [12, 13], [14]]
-
-
-def test_split_request_ids_honors_single_request_seq_token_counts():
-    CosyVoice3Model, _ = _cosyvoice3_model_and_runner()
-    ids = torch.tensor([10, 11, 12, 13, 14], dtype=torch.long)
-    chunks = CosyVoice3Model._split_request_ids(ids, [3])
-    assert [c.tolist() for c in chunks] == [[10, 11, 12]]
-
-
-def test_sanitize_codec_tokens_filters_out_of_range():
-    model = _make_code2wav_model()
-    raw = torch.tensor([-1, 0, 3, 4, 99], dtype=torch.long)
-    clean = model._sanitize_codec_tokens(raw)
-    assert clean.tolist() == [0, 3]
-
-
 def test_forward_prefers_token_offset_when_present():
     model = _make_code2wav_model()
 
@@ -265,6 +244,31 @@ def test_forward_uses_non_stream_decode_without_chunk_metadata():
     assert len(model.code2wav.forward_streaming_calls) == 0
     call = model.code2wav.forward_calls[0]
     assert call["token"].tolist() == [[0, 1, 2]]
+    assert call["token_offset_tokens"] == 0
+
+
+def test_forward_uses_non_stream_talker_prefill_offset():
+    model = _make_code2wav_model()
+
+    runtime_info = [
+        {
+            "embed": {
+                "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+                "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+                "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            },
+            "meta": {"talker_prefill_offset": 3},
+        }
+    ]
+
+    model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+
+    assert model.code2wav.forward_calls[0]["token_offset_tokens"] == 3
 
 
 def test_forward_reuses_streaming_cache_state_between_chunks():
@@ -406,6 +410,9 @@ def test_gpu_ar_model_runner_prefers_model_sampler_when_opted_in():
             self.updated = False
 
         def update_async_output_token_ids(self):
+            # After PR 3681 fix, update_async_output_token_ids is called
+            # BEFORE model sampler path to ensure async placeholder repair
+            # runs for all sampling paths
             self.updated = True
 
     _, GPUARModelRunner = _cosyvoice3_model_and_runner()
@@ -420,7 +427,7 @@ def test_gpu_ar_model_runner_prefers_model_sampler_when_opted_in():
     out = runner._sample(torch.tensor([[0.1, 0.2]], dtype=torch.float32), spec_decode_metadata=None)
 
     assert out is expected
-    assert runner.input_batch.updated is False
+    assert runner.input_batch.updated is True
     assert len(calls) == 1
 
 
@@ -436,9 +443,13 @@ def test_gpu_ar_model_runner_supplies_req_output_history_to_model_sampler():
             self.sampled_token_ids_cpu = None
             self.async_copy_ready_event = None
             self.prev_req_id_to_index = None
+            self.update_async_called = False
 
         def update_async_output_token_ids(self):
-            raise AssertionError("fallback async repair should not run for model sampler path")
+            # After PR 3681 fix, update_async_output_token_ids is called
+            # BEFORE model sampler path to ensure async placeholder repair
+            # runs for all sampling paths
+            self.update_async_called = True
 
     _, GPUARModelRunner = _cosyvoice3_model_and_runner()
     runner = object.__new__(GPUARModelRunner)
@@ -454,6 +465,7 @@ def test_gpu_ar_model_runner_supplies_req_output_history_to_model_sampler():
 
     runner._sample(torch.tensor([[0.1, 0.2]], dtype=torch.float32), spec_decode_metadata=None)
 
+    assert runner.input_batch.update_async_called is True
     assert seen_histories == [[[1, 2, 3]]]
 
 
@@ -476,9 +488,13 @@ def test_gpu_ar_model_runner_repairs_async_placeholders_for_model_sampler():
             self.sampled_token_ids_cpu = torch.tensor([[29]], dtype=torch.int32)
             self.async_copy_ready_event = _ReadyEvent()
             self.prev_req_id_to_index = {"rid-1": 0}
+            self.update_async_called = False
 
         def update_async_output_token_ids(self):
-            raise AssertionError("fallback async repair should not run for model sampler path")
+            # After PR 3681 fix, update_async_output_token_ids is called
+            # BEFORE model sampler path to ensure async placeholder repair
+            # runs for all sampling paths (model sampler + fallback sampler)
+            self.update_async_called = True
 
     _, GPUARModelRunner = _cosyvoice3_model_and_runner()
     runner = object.__new__(GPUARModelRunner)
@@ -495,4 +511,5 @@ def test_gpu_ar_model_runner_repairs_async_placeholders_for_model_sampler():
     runner._sample(torch.tensor([[0.1, 0.2]], dtype=torch.float32), spec_decode_metadata=None)
 
     assert runner.input_batch.async_copy_ready_event.synced is True
+    assert runner.input_batch.update_async_called is True
     assert seen_histories == [[[11, 29]]]
diff --git a/tests/model_executor/models/fish_speech/test_fish_speech_fast_ar_runtime_buffers.py b/tests/model_executor/models/fish_speech/test_fish_speech_fast_ar_runtime_buffers.py
new file mode 100644
index 000000000..41a29ca63
--- /dev/null
+++ b/tests/model_executor/models/fish_speech/test_fish_speech_fast_ar_runtime_buffers.py
@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+import torch.nn as nn
+
+from vllm_omni.model_executor.models.fish_speech.configuration_fish_speech import (
+    FishSpeechFastARConfig,
+    FishSpeechSlowARConfig,
+)
+from vllm_omni.model_executor.models.fish_speech.fish_speech_fast_ar import FishSpeechFastAR
+from vllm_omni.model_executor.models.fish_speech.fish_speech_slow_ar import (
+    FishSpeechSlowARForConditionalGeneration,
+)
+
+
+def test_fast_ar_reuses_dense_position_id_buffer(monkeypatch):
+    fast_config = FishSpeechFastARConfig(
+        vocab_size=16,
+        num_codebooks=4,
+        dim=8,
+        n_head=2,
+        n_local_heads=1,
+        head_dim=4,
+        n_layer=1,
+        intermediate_size=16,
+        max_seq_len=5,
+    )
+    slow_config = FishSpeechSlowARConfig(
+        vocab_size=32,
+        dim=8,
+        n_head=2,
+        n_local_heads=1,
+        head_dim=4,
+        n_layer=1,
+        intermediate_size=16,
+        codebook_size=16,
+        num_codebooks=4,
+        semantic_begin_id=4,
+        semantic_end_id=19,
+    )
+    fast_ar = object.__new__(FishSpeechFastAR)
+    nn.Module.__init__(fast_ar)
+    fast_ar.config = fast_config
+    fast_ar.slow_ar_config = slow_config
+    fast_ar.fast_project_in = nn.Identity()
+    fast_ar.fast_embeddings = nn.Embedding(fast_config.vocab_size, fast_config.hidden_size)
+    fast_ar.fast_output = nn.Linear(fast_config.hidden_size, fast_config.vocab_size, bias=False)
+    fast_ar.fast_norm = nn.Identity()
+    fast_ar._num_codebooks = fast_config.num_codebooks
+    fast_ar._fast_dim = fast_config.hidden_size
+    fast_ar._embed_buf = None
+    fast_ar._pos_ids = None
+    fast_ar._k_cache = None
+    fast_ar._v_cache = None
+    fast_ar._compiled_model_fwd = None
+    fast_ar._compile_attempted = True
+    fast_ar._compile_failed = False
+    fast_ar._disable_compile_for_graph = False
+
+    seen_position_ids: list[torch.Tensor] = []
+
+    def fake_run_model_one(step_input, step_pos_ids, cache_pos):
+        seen_position_ids.append(step_pos_ids)
+        return torch.zeros(
+            step_input.shape[0],
+            fast_config.hidden_size,
+            dtype=step_input.dtype,
+            device=step_input.device,
+        )
+
+    monkeypatch.setattr(fast_ar, "_run_model_one", fake_run_model_one)
+
+    hidden = torch.zeros(3, slow_config.hidden_size, dtype=torch.float32)
+    semantic = torch.full((3,), slow_config.semantic_begin_id, dtype=torch.long)
+    fast_ar(hidden, semantic, do_sample=False)
+
+    assert fast_ar._pos_ids is not None
+    assert fast_ar._pos_ids.shape == (3, fast_config.num_codebooks + 1)
+    assert len(seen_position_ids) == fast_config.num_codebooks
+    for step, step_pos_ids in enumerate(seen_position_ids):
+        assert step_pos_ids.shape == (3,)
+        assert step_pos_ids.untyped_storage().data_ptr() == fast_ar._pos_ids.untyped_storage().data_ptr()
+        assert step_pos_ids.stride() == (fast_ar._pos_ids.stride(0),)
+        expected = torch.full((3,), step)
+        assert torch.equal(step_pos_ids.cpu(), expected)
+
+
+def test_talker_mtp_does_not_mutate_input():
+    model = object.__new__(FishSpeechSlowARForConditionalGeneration)
+    nn.Module.__init__(model)
+    model._semantic_begin_id = 4
+    model._semantic_end_id = 19
+    model._codebook_size = 16
+    model._num_codebooks = 4
+    model.codebook_embeddings = nn.Embedding(model._num_codebooks * model._codebook_size, 8)
+    model.fast_ar = lambda **_: torch.tensor(
+        [
+            [0, 1, 2, 3],
+            [15, 14, 13, 12],
+        ],
+        dtype=torch.long,
+    )
+    input_ids = torch.tensor([4, 2], dtype=torch.long)
+    input_embeds = torch.randn(2, 8)
+    input_embeds_before = input_embeds.clone()
+    last_hidden = torch.randn(2, 8, dtype=torch.bfloat16)
+    text_step = torch.zeros(2, 8, dtype=torch.bfloat16)
+
+    out1, _ = FishSpeechSlowARForConditionalGeneration.talker_mtp(
+        model,
+        input_ids,
+        input_embeds,
+        last_hidden,
+        text_step,
+    )
+
+    assert torch.equal(input_embeds, input_embeds_before)
+    assert out1.untyped_storage().data_ptr() != input_embeds.untyped_storage().data_ptr()
+
+
+def test_talker_mtp_forwards_sampling_params():
+    model = object.__new__(FishSpeechSlowARForConditionalGeneration)
+    nn.Module.__init__(model)
+    model._semantic_begin_id = 4
+    model._semantic_end_id = 19
+    model._codebook_size = 16
+    model._num_codebooks = 4
+    model.codebook_embeddings = nn.Embedding(model._num_codebooks * model._codebook_size, 8)
+    seen = {}
+
+    def fake_fast_ar(**kwargs):
+        seen.update(kwargs)
+        return torch.tensor([[0, 1, 2, 3]], dtype=torch.long)
+
+    model.fast_ar = fake_fast_ar
+    input_ids = torch.tensor([4], dtype=torch.long)
+    input_embeds = torch.randn(1, 8)
+    last_hidden = torch.randn(1, 8, dtype=torch.bfloat16)
+    text_step = torch.zeros(1, 8, dtype=torch.bfloat16)
+
+    FishSpeechSlowARForConditionalGeneration.talker_mtp(
+        model,
+        input_ids,
+        input_embeds,
+        last_hidden,
+        text_step,
+        do_sample=False,
+        temperature=0.1,
+        top_k=5,
+        top_p=0.7,
+    )
+
+    assert seen["do_sample"] is False
+    assert seen["temperature"] == 0.1
+    assert seen["top_k"] == 5
+    assert seen["top_p"] == 0.7
diff --git a/tests/model_executor/models/fish_speech/test_fish_speech_regressions.py b/tests/model_executor/models/fish_speech/test_fish_speech_regressions.py
index dbb0150e8..b21ff3416 100644
--- a/tests/model_executor/models/fish_speech/test_fish_speech_regressions.py
+++ b/tests/model_executor/models/fish_speech/test_fish_speech_regressions.py
@@ -1,5 +1,6 @@
 import functools
 import math
+from types import SimpleNamespace
 
 import pytest
 import torch
@@ -26,6 +27,39 @@ class _FakeCodec:
         return wav, audio_lengths
 
 
+class _RecordingCodec:
+    def __init__(self):
+        self.calls = []
+
+    def decode(self, codes_bqf: torch.Tensor, feature_lengths: torch.Tensor):
+        self.calls.append((tuple(codes_bqf.shape), tuple(int(x) for x in feature_lengths.tolist())))
+        batch = codes_bqf.shape[0]
+        wav_len = int(codes_bqf.shape[-1]) * 10
+        wav = torch.arange(batch * wav_len, dtype=torch.float32).view(batch, 1, wav_len)
+        audio_lengths = feature_lengths.to(torch.long) * 10
+        return wav, audio_lengths
+
+
+class _FakeOneArgCodec:
+    def decode(self, codes_bqf: torch.Tensor):
+        batch = codes_bqf.shape[0]
+        return torch.arange(batch * 100, dtype=torch.float32).view(batch, 1, 100)
+
+
+class _FakeFromIndicesCodec:
+    def __init__(self):
+        self.from_indices_calls = 0
+
+    def from_indices(self, codes_bqf: torch.Tensor):
+        self.from_indices_calls += 1
+        batch = codes_bqf.shape[0]
+        return torch.arange(batch * 100, dtype=torch.float32).view(batch, 1, 100)
+
+    def decode(self, codes_bqf: torch.Tensor):
+        del codes_bqf
+        raise AssertionError("decode() should not be used for index tensors when from_indices() exists")
+
+
 class _FakeTokenizer:
     def __init__(self, mapping, unk_token_id=-1):
         self._mapping = mapping
@@ -40,6 +74,7 @@ def test_dac_decoder_mixed_batch_empty_request_does_not_misalign_indices():
     decoder = object.__new__(FishSpeechDACDecoder)
     torch.nn.Module.__init__(decoder)
     decoder._codec = _FakeCodec()
+    decoder._codec_decode_takes_lengths = True
     decoder._num_codebooks = 10
     decoder._output_sample_rate = 44100
     decoder._hop_length = 512
@@ -62,6 +97,132 @@ def test_dac_decoder_mixed_batch_empty_request_does_not_misalign_indices():
     assert audios[1].shape[0] == 50
 
 
+def test_dac_decoder_uses_tensor_codes_from_runtime_info():
+    _, FishSpeechDACDecoder, _ = _fish_speech_regression_modules()
+    decoder = object.__new__(FishSpeechDACDecoder)
+    torch.nn.Module.__init__(decoder)
+    decoder._codec = _FakeCodec()
+    decoder._codec_decode_takes_lengths = True
+    decoder._num_codebooks = 2
+    decoder._output_sample_rate = 44100
+    decoder._hop_length = 512
+    decoder._logged_codec_stats = False
+    decoder._ensure_codec_loaded = lambda: None
+
+    out = decoder.forward(
+        input_ids=torch.tensor([0], dtype=torch.long),
+        runtime_additional_information=[
+            {
+                "codes": {"audio": torch.tensor([[1, 2], [3, 4]], dtype=torch.long)},
+                "meta": {"left_context_size": 1},
+            }
+        ],
+    )
+
+    audios = out.multimodal_outputs["model_outputs"]
+    assert len(audios) == 1
+    # 2 total frames with 1 frame of left context => proportional trim removes half the samples.
+    assert audios[0].shape[0] == 50
+
+
+def test_dac_decoder_supports_codec_decode_without_lengths():
+    _, FishSpeechDACDecoder, _ = _fish_speech_regression_modules()
+    decoder = object.__new__(FishSpeechDACDecoder)
+    torch.nn.Module.__init__(decoder)
+    decoder._codec = _FakeOneArgCodec()
+    decoder._codec_decode_takes_lengths = False
+    decoder._num_codebooks = 10
+    decoder._output_sample_rate = 44100
+    decoder._hop_length = 512
+    decoder._logged_codec_stats = False
+    decoder._ensure_codec_loaded = lambda: None
+    decoder._split_request_ids = lambda ids, seq_token_counts=None: [
+        torch.arange(10, dtype=torch.long),
+        torch.arange(20, dtype=torch.long),
+    ]
+
+    out = decoder.forward(
+        input_ids=torch.arange(30, dtype=torch.long),
+        runtime_additional_information=[{}, {}],
+    )
+
+    audios = out.multimodal_outputs["model_outputs"]
+    assert len(audios) == 2
+    assert audios[0].shape[0] == 50
+    assert audios[1].shape[0] == 100
+
+
+def test_dac_decoder_prefers_from_indices_for_one_arg_codec():
+    _, FishSpeechDACDecoder, _ = _fish_speech_regression_modules()
+    codec = _FakeFromIndicesCodec()
+    decoder = object.__new__(FishSpeechDACDecoder)
+    torch.nn.Module.__init__(decoder)
+    decoder._codec = codec
+    decoder._codec_decode_takes_lengths = False
+    decoder._num_codebooks = 10
+    decoder._output_sample_rate = 44100
+    decoder._hop_length = 512
+    decoder._logged_codec_stats = False
+    decoder._ensure_codec_loaded = lambda: None
+    decoder._split_request_ids = lambda ids, seq_token_counts=None: [
+        torch.arange(10, dtype=torch.long),
+        torch.arange(20, dtype=torch.long),
+    ]
+
+    out = decoder.forward(
+        input_ids=torch.arange(30, dtype=torch.long),
+        runtime_additional_information=[{}, {}],
+    )
+
+    audios = out.multimodal_outputs["model_outputs"]
+    assert codec.from_indices_calls == 1
+    assert len(audios) == 2
+    assert audios[0].shape[0] == 50
+    assert audios[1].shape[0] == 100
+
+
+def test_dac_decoder_splits_groups_by_padded_frame_budget():
+    _, FishSpeechDACDecoder, _ = _fish_speech_regression_modules()
+    codec = _RecordingCodec()
+    decoder = object.__new__(FishSpeechDACDecoder)
+    torch.nn.Module.__init__(decoder)
+    decoder._codec = codec
+    decoder._codec_decode_takes_lengths = True
+    decoder._num_codebooks = 2
+    decoder._output_sample_rate = 44100
+    decoder._hop_length = 512
+    decoder._logged_codec_stats = False
+    decoder._decode_batch_max_padded_frames = 100
+    decoder._decode_batch_max_batch = 0
+    decoder._ensure_codec_loaded = lambda: None
+    decoder._split_request_ids = lambda ids, seq_token_counts=None: [
+        torch.arange(2 * 46, dtype=torch.long),
+        torch.arange(2 * 45, dtype=torch.long),
+        torch.arange(2 * 47, dtype=torch.long),
+    ]
+
+    decoder.forward(input_ids=torch.arange(2 * (46 + 45 + 47), dtype=torch.long))
+
+    assert codec.calls == [
+        ((2, 2, 46), (45, 46)),
+        ((1, 2, 47), (47,)),
+    ]
+
+
+def test_dac_decoder_dtype_can_be_configured():
+    _, FishSpeechDACDecoder, _ = _fish_speech_regression_modules()
+    cfg = SimpleNamespace(
+        model_config=SimpleNamespace(
+            model="unused",
+            stage_connector_config={"extra": {"fish_speech_dac_dtype": "fp16"}},
+        )
+    )
+
+    decoder = FishSpeechDACDecoder(vllm_config=cfg)
+
+    assert decoder._dac_dtype == torch.float16
+
+
 def test_structured_voice_clone_prefill_adds_full_codebooks_with_decode_scale(monkeypatch):
     slow_ar_module, _, FishSpeechSlowARForConditionalGeneration = _fish_speech_regression_modules()
     model = object.__new__(FishSpeechSlowARForConditionalGeneration)
diff --git a/tests/model_executor/models/glm_image/test_glm_image_init.py b/tests/model_executor/models/glm_image/test_glm_image_init.py
new file mode 100644
index 000000000..eb42e9649
--- /dev/null
+++ b/tests/model_executor/models/glm_image/test_glm_image_init.py
@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for GLM-Image model __init__.py lazy import pattern.
+
+The ``__init__.py`` uses ``__getattr__`` for lazy loading to avoid importing
+``transformers.models.glm_image`` at module init, which may not be available
+in all environments.
+"""
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestLazyImport:
+    """Test the __getattr__ lazy import pattern in glm_image __init__.py."""
+
+    def test_getattr_exists_and_is_callable(self):
+        """Verify __getattr__ exists and is callable."""
+        import vllm_omni.model_executor.models.glm_image as glm_image_pkg
+
+        assert hasattr(glm_image_pkg, "__getattr__")
+        assert callable(glm_image_pkg.__getattr__)
+
+    def test_getattr_returns_class_for_known_attribute(self):
+        """Verify __getattr__ returns GlmImageForConditionalGeneration for known attribute."""
+        import vllm_omni.model_executor.models.glm_image as glm_image_pkg
+
+        # Call __getattr__ directly to test the lazy import logic
+        result = glm_image_pkg.__getattr__("GlmImageForConditionalGeneration")
+
+        # Verify we get a class (the actual GlmImageForConditionalGeneration)
+        assert result is not None
+        assert isinstance(result, type)
+        assert result.__name__ == "GlmImageForConditionalGeneration"
+
+    def test_getattr_raises_for_unknown_attribute(self):
+        """Verify __getattr__ raises AttributeError for unknown attributes."""
+        import vllm_omni.model_executor.models.glm_image as glm_image_pkg
+
+        # Test unknown attribute via __getattr__ directly
+        with pytest.raises(AttributeError, match="has no attribute"):
+            glm_image_pkg.__getattr__("UnknownClass")
+
+    def test___all___exports_correct_symbols(self):
+        """Verify __all__ contains the expected exported symbols."""
+        import vllm_omni.model_executor.models.glm_image as glm_image_pkg
+
+        assert hasattr(glm_image_pkg, "__all__")
+        assert "GlmImageForConditionalGeneration" in glm_image_pkg.__all__
+
+
+class TestLazyImportDoesNotImportTransformersAtInit:
+    """Verify that importing the package does not eagerly load transformers."""
+
+    def test_glm_image_module_has_getattr(self):
+        """Test that the module has __getattr__ for lazy loading."""
+        import vllm_omni.model_executor.models.glm_image as glm_image_pkg
+
+        # The module should have __getattr__
+        assert hasattr(glm_image_pkg, "__getattr__")
+        assert callable(glm_image_pkg.__getattr__)
+
+        # And __all__
+        assert hasattr(glm_image_pkg, "__all__")
+
+    def test_module_has_proper_structure(self):
+        """Test that the module has proper Python module structure."""
+        import vllm_omni.model_executor.models.glm_image as glm_image_pkg
+
+        # Should be a module
+        assert hasattr(glm_image_pkg, "__name__")
+        assert glm_image_pkg.__name__ == "vllm_omni.model_executor.models.glm_image"
+
+        # Should have __file__ (even if it's a package)
+        assert hasattr(glm_image_pkg, "__file__")
diff --git a/tests/model_executor/models/omnivoice/__init__.py b/tests/model_executor/models/omnivoice/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/model_executor/models/omnivoice/test_cuda_graph_generator.py b/tests/model_executor/models/omnivoice/test_cuda_graph_generator.py
new file mode 100644
index 000000000..dc6126713
--- /dev/null
+++ b/tests/model_executor/models/omnivoice/test_cuda_graph_generator.py
@@ -0,0 +1,263 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Tests for OmniVoice CUDA Graph generator wrapper numerical equivalence.
+
+Verifies that _OmniVoiceCUDAGraphForward produces results equivalent to
+eager mode across three scenarios:
+  - Exact-size inputs (no padding) → bit-identical
+  - Padded inputs (padded to nearest bucket) → correct slicing, exact match
+    for position-independent models
+  - Oversized inputs (fallback to eager) → bit-identical
+
+Uses a lightweight synthetic generator to keep warmup fast and avoid
+loading actual model weights.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+DEVICE = torch.device("cuda:0")
+NUM_CB = 8
+VOCAB = 1025
+HIDDEN = 64
+CAPTURE_SIZES = [32, 64, 128]
+
+
+# ---------------------------------------------------------------------------
+# Synthetic generator: matches _OmniVoiceCUDAGraphForward's interface
+# ---------------------------------------------------------------------------
+
+
+class _SyntheticGenerator:
+    """Position-independent Embedding+Linear replacing OmniVoiceGenerator's 28-layer transformer.
+
+    Each token position is processed independently (no attention), so padded
+    runs produce bit-identical results to non-padded runs on the same positions.
+    This makes the padding/slicing logic easy to verify exactly.
+    """
+
+    class _Cfg:
+        num_audio_codebook = NUM_CB
+
+    def __init__(self, device: torch.device):
+        self.config = self._Cfg()
+        self.text_embedding = nn.Embedding(1000, HIDDEN).to(device).eval()
+        self._linear = nn.Linear(HIDDEN, NUM_CB * VOCAB, bias=False).to(device).eval()
+        self._rope_cos: torch.Tensor | None = None
+        self._rope_sin: torch.Tensor | None = None
+
+    def _ensure_rope(self, seq_len: int, device: torch.device) -> None:
+        if self._rope_cos is None or self._rope_cos.shape[0] < seq_len:
+            max_len = max(seq_len, 512)
+            self._rope_cos = torch.zeros(max_len, 64, device=device)
+            self._rope_sin = torch.zeros(max_len, 64, device=device)
+
+    def _step_forward(
+        self,
+        input_ids: torch.Tensor,
+        audio_mask: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        two_b, _, S = input_ids.shape
+        x = self.text_embedding(input_ids[:, 0, :].clamp(0, 999))  # [two_b, S, H]
+        logits = self._linear(x)  # [two_b, S, 8*1025]
+        return logits.view(two_b, S, NUM_CB, VOCAB).permute(0, 2, 1, 3)  # [two_b, 8, S, 1025]
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _eager(gen: _SyntheticGenerator, ids: torch.Tensor, mask: torch.Tensor, attn: torch.Tensor) -> torch.Tensor:
+    """Run _step_forward with zero RoPE tensors (synthetic gen ignores them)."""
+    S = ids.shape[-1]
+    cos = torch.zeros(S, 64, device=ids.device)
+    sin = torch.zeros(S, 64, device=ids.device)
+    return gen._step_forward(ids, mask, attn, cos, sin)
+
+
+def _make_inputs(seq_len: int, device: torch.device = DEVICE):
+    two_b = 2
+    ids = torch.randint(0, 100, (two_b, NUM_CB, seq_len), dtype=torch.long, device=device)
+    mask = torch.ones(two_b, seq_len, dtype=torch.bool, device=device)
+    attn = torch.ones(two_b, 1, seq_len, seq_len, dtype=torch.bool, device=device)
+    return ids, mask, attn
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(scope="module")
+def gen():
+    torch.manual_seed(42)
+    return _SyntheticGenerator(DEVICE)
+
+
+@pytest.fixture(scope="module")
+def wrapper(gen):
+    from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import (
+        _OmniVoiceCUDAGraphForward,
+    )
+
+    w = _OmniVoiceCUDAGraphForward(gen, capture_sizes=CAPTURE_SIZES)
+    w.warmup(DEVICE)
+    return w
+
+
+# ---------------------------------------------------------------------------
+# 1. Exact-size inputs → bit-identical
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("seq_len", CAPTURE_SIZES)
+def test_exact_size_bit_identical(gen, wrapper, seq_len):
+    """When input exactly matches a captured bucket, output must be bit-identical to eager."""
+    ids, mask, attn = _make_inputs(seq_len)
+    with torch.no_grad():
+        eager_out = _eager(gen, ids, mask, attn)
+        graph_out = wrapper(ids, mask, attn)
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# 2. Padded inputs → correct shape, values match at actual positions
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("seq_len", [1, 15, 33, 60, 100])
+def test_padded_output_shape(gen, wrapper, seq_len):
+    """Graph output must be sliced back to actual seq_len, not the bucket size."""
+    ids, mask, attn = _make_inputs(seq_len)
+    with torch.no_grad():
+        graph_out = wrapper(ids, mask, attn)
+    assert graph_out.shape == (2, NUM_CB, seq_len, VOCAB)
+
+
+@pytest.mark.parametrize("seq_len", [15, 33, 60, 100])
+def test_padded_output_matches_eager(gen, wrapper, seq_len):
+    """Padded graph output must equal eager output at actual positions.
+
+    The synthetic model has no attention across positions, so zero-padding
+    does not affect non-padded positions — exact match is expected.
+    """
+    ids, mask, attn = _make_inputs(seq_len)
+    with torch.no_grad():
+        eager_out = _eager(gen, ids, mask, attn)
+        graph_out = wrapper(ids, mask, attn)
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# 3. Oversized inputs → fallback to eager (lazy capture), bit-identical
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("seq_len", [129, 200, 256])
+def test_fallback_eager_bit_identical(gen, wrapper, seq_len):
+    """Sequences exceeding the largest bucket fall back to lazy capture → bit-identical."""
+    ids, mask, attn = _make_inputs(seq_len)
+    with torch.no_grad():
+        eager_out = _eager(gen, ids, mask, attn)
+        graph_out = wrapper(ids, mask, attn)
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# 4. Determinism
+# ---------------------------------------------------------------------------
+
+
+def test_deterministic_across_calls(wrapper):
+    """Same input must produce identical output on repeated CUDA graph replays."""
+    ids, mask, attn = _make_inputs(32)
+    with torch.no_grad():
+        out1 = wrapper(ids, mask, attn)
+        out2 = wrapper(ids, mask, attn)
+    torch.testing.assert_close(out1, out2, atol=0, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# 5. _find_bucket logic (CPU, no CUDA graph)
+# ---------------------------------------------------------------------------
+
+
+def test_find_bucket_returns_nearest_bucket():
+    """_find_bucket must return the smallest bucket >= seq_len, or None if all are smaller."""
+    from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import (
+        _OmniVoiceCUDAGraphForward,
+    )
+
+    w = _OmniVoiceCUDAGraphForward.__new__(_OmniVoiceCUDAGraphForward)
+    w._capture_sizes = [32, 64, 128]
+    w._graphs = {}
+
+    assert w._find_bucket(1) == 32
+    assert w._find_bucket(32) == 32
+    assert w._find_bucket(33) == 64
+    assert w._find_bucket(64) == 64
+    assert w._find_bucket(128) == 128
+    assert w._find_bucket(129) is None
+    assert w._find_bucket(1000) is None
+
+
+# ---------------------------------------------------------------------------
+# 6. enable_cuda_graph=False produces same tokens as enable_cuda_graph=True
+# ---------------------------------------------------------------------------
+
+
+def test_cuda_graph_disabled_matches_eager_generator():
+    """OmniVoiceGenerator with enable_cuda_graph=False must produce the same
+    _step_forward output as one with enable_cuda_graph=True (after warmup).
+
+    Uses a 2-layer config so warmup completes quickly without loading weights.
+    """
+    from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import OmniVoiceGenerator
+    from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
+
+    cfg = OmniVoiceConfig()
+    cfg.llm_num_hidden_layers = 2
+    cfg.num_hidden_layers = 2
+
+    torch.manual_seed(0)
+    gen_eager = OmniVoiceGenerator(cfg).to(DEVICE).eval()
+    gen_eager.config.enable_cuda_graph = False
+    gen_eager._cuda_graph_fwd = None
+
+    torch.manual_seed(0)
+    gen_graph = OmniVoiceGenerator(cfg).to(DEVICE).eval()
+    gen_graph.config.cuda_graph_capture_sizes = [64]
+    from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import _OmniVoiceCUDAGraphForward
+
+    gen_graph._cuda_graph_fwd = _OmniVoiceCUDAGraphForward(gen_graph, capture_sizes=[64])
+    gen_graph._cuda_graph_fwd.warmup(DEVICE)
+
+    seq_len = 64
+    ids = torch.zeros(2, cfg.num_audio_codebook, seq_len, dtype=torch.long, device=DEVICE)
+    mask = torch.ones(2, seq_len, dtype=torch.bool, device=DEVICE)
+    attn = torch.ones(2, 1, seq_len, seq_len, dtype=torch.bool, device=DEVICE)
+
+    gen_eager._ensure_rope(seq_len, DEVICE)
+    dtype = gen_eager.text_embedding.weight.dtype
+    cos = gen_eager._rope_cos[:seq_len].to(device=DEVICE, dtype=dtype)
+    sin = gen_eager._rope_sin[:seq_len].to(device=DEVICE, dtype=dtype)
+
+    with torch.no_grad():
+        eager_logits = gen_eager._step_forward(ids, mask, attn, cos, sin)
+        graph_logits = gen_graph._cuda_graph_fwd(ids, mask, attn)
+
+    torch.testing.assert_close(graph_logits, eager_logits, atol=0, rtol=0)
diff --git a/tests/model_executor/models/omnivoice/test_decoder.py b/tests/model_executor/models/omnivoice/test_decoder.py
new file mode 100644
index 000000000..db8e7328e
--- /dev/null
+++ b/tests/model_executor/models/omnivoice/test_decoder.py
@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Correctness and regression tests for OmniVoiceDecoder.
+
+Covers three areas:
+1. NaN regression — verifies the fp16 overflow fix:
+     acoustic decoder must operate in float32 even when fc2 weights are fp16.
+     Demonstrates that the pre-fix code path (no .float() after fc2) produces
+     NaN/Inf with large-scale weights, while the fixed path does not.
+2. Output shape / value sanity — correct [B, 1, T*upsample] shape, finite values.
+3. Acoustic decoder float32 invariant — weights must be float32 after build.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.model_executor.models.omnivoice.omnivoice_decoder import (
+    HiggsAudioRVQ,
+    OmniVoiceDecoder,
+)
+from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+DEVICE = torch.device("cuda:0")
+UPSAMPLE = 4  # synthetic acoustic decoder stride (real DAC uses 960)
+T_FRAMES = 25  # number of audio token frames to decode
+
+
+# ---------------------------------------------------------------------------
+# Tiny acoustic decoder: ConvTranspose1d upsample by UPSAMPLE
+# ---------------------------------------------------------------------------
+
+
+class _TinyAcousticDecoder(nn.Module):
+    def __init__(self, in_channels: int = 256, upsample: int = UPSAMPLE):
+        super().__init__()
+        self.conv = nn.ConvTranspose1d(in_channels, 1, kernel_size=upsample, stride=upsample)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(x)  # [B, 1, T*upsample]
+
+
+# ---------------------------------------------------------------------------
+# Builder
+# ---------------------------------------------------------------------------
+
+
+def _build_decoder(
+    *,
+    fc2_dtype: torch.dtype = torch.float32,
+    acoustic_dtype: torch.dtype = torch.float32,
+    weight_scale: float = 1.0,
+) -> OmniVoiceDecoder:
+    """Build an OmniVoiceDecoder with synthetic components (no model checkpoint required).
+
+    Args:
+        fc2_dtype: dtype for the fc2 linear layer weights.
+        acoustic_dtype: dtype for the acoustic decoder weights.
+        weight_scale: constant weight value for overflow testing.
+    """
+    config = OmniVoiceConfig()
+    decoder = OmniVoiceDecoder(config)
+
+    decoder.quantizer = HiggsAudioRVQ(num_quantizers=8, codebook_size=1024, codebook_dim=64, hidden_size=1024).to(
+        DEVICE
+    )
+
+    decoder.fc2 = nn.Linear(1024, 256).to(DEVICE).to(fc2_dtype)
+
+    acoustic = _TinyAcousticDecoder(in_channels=256, upsample=UPSAMPLE).to(DEVICE).to(acoustic_dtype)
+    if weight_scale != 1.0:
+        with torch.no_grad():
+            nn.init.constant_(acoustic.conv.weight, weight_scale)
+            nn.init.constant_(acoustic.conv.bias, weight_scale)
+    decoder.acoustic_decoder = acoustic
+
+    decoder._loaded = True
+    return decoder
+
+
+# ---------------------------------------------------------------------------
+# 1. NaN regression
+# ---------------------------------------------------------------------------
+
+
+def test_no_nan_fp16_fc2_float32_acoustic():
+    """fp16 fc2 + float32 acoustic decoder (the fixed path) must not produce NaN or Inf."""
+    decoder = _build_decoder(fc2_dtype=torch.float16, acoustic_dtype=torch.float32)
+    tokens = torch.randint(0, 1024, (1, 8, T_FRAMES), device=DEVICE)
+
+    audio = decoder(tokens)
+
+    assert not torch.isnan(audio).any(), "NaN in output — fp16 overflow may have reached acoustic decoder"
+    assert not torch.isinf(audio).any(), "Inf in output"
+
+
+def test_fp16_acoustic_decoder_produces_nan_without_fix():
+    """Pre-fix scenario: fp16 acoustic decoder with large weights overflows to NaN/Inf.
+
+    This test documents the original bug and will fail (raising AssertionError)
+    if the underlying hardware / Triton version somehow avoids overflow, but
+    with weights scaled to 300.0 the ConvTranspose1d intermediate activations
+    reliably exceed fp16 max (~65504) in practice.
+    """
+    decoder = _build_decoder(fc2_dtype=torch.float16, acoustic_dtype=torch.float16, weight_scale=300.0)
+    with torch.no_grad():
+        nn.init.constant_(decoder.fc2.weight, 300.0)
+        nn.init.constant_(decoder.fc2.bias, 300.0)
+        for q in decoder.quantizer.quantizers:
+            nn.init.constant_(q.codebook.weight, 300.0)
+
+    tokens = torch.randint(0, 1024, (1, 8, T_FRAMES), device=DEVICE)
+
+    # Simulate pre-fix forward: no .float() after fc2
+    codes = tokens.transpose(0, 1).long()
+    quantized = decoder.quantizer.decode(codes)
+    quantized_fp16 = decoder.fc2(quantized.transpose(1, 2).to(decoder.fc2.weight.dtype)).transpose(
+        1, 2
+    )  # stays fp16, no .float()
+    audio_broken = decoder.acoustic_decoder(quantized_fp16)
+
+    assert torch.isnan(audio_broken).any() or torch.isinf(audio_broken).any(), (
+        "Expected NaN/Inf in pre-fix fp16 path with large weights — test setup may need a higher weight_scale"
+    )
+
+
+def test_acoustic_decoder_receives_float32_input():
+    """The .float() cast after fc2 must ensure acoustic decoder always gets float32 input,
+    regardless of fc2 weight dtype."""
+    decoder = _build_decoder(fc2_dtype=torch.float16, acoustic_dtype=torch.float32)
+
+    captured_dtypes: list[torch.dtype] = []
+    original_fwd = decoder.acoustic_decoder.forward
+
+    def _hook(x: torch.Tensor) -> torch.Tensor:
+        captured_dtypes.append(x.dtype)
+        return original_fwd(x)
+
+    decoder.acoustic_decoder.forward = _hook  # type: ignore[method-assign]
+    tokens = torch.randint(0, 1024, (1, 8, T_FRAMES), device=DEVICE)
+    decoder(tokens)
+
+    assert captured_dtypes, "acoustic_decoder.forward was never called"
+    assert captured_dtypes[0] == torch.float32, (
+        f"Acoustic decoder received {captured_dtypes[0]}, expected torch.float32 — "
+        "the .float() upcast after fc2 may be missing"
+    )
+
+
+# ---------------------------------------------------------------------------
+# 2. Output shape and value sanity
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("batch,t_frames", [(1, 10), (1, 50), (2, 25)])
+def test_output_shape(batch, t_frames):
+    """Decoder must return [B, 1, T*upsample] waveform tensor."""
+    decoder = _build_decoder()
+    tokens = torch.randint(0, 1024, (batch, 8, t_frames), device=DEVICE)
+
+    audio = decoder(tokens)
+
+    assert audio.shape == (batch, 1, t_frames * UPSAMPLE), (
+        f"Expected ({batch}, 1, {t_frames * UPSAMPLE}), got {audio.shape}"
+    )
+
+
+def test_no_nan_float32_baseline():
+    """Float32 fc2 + float32 acoustic decoder must always produce finite output."""
+    decoder = _build_decoder(fc2_dtype=torch.float32, acoustic_dtype=torch.float32)
+    tokens = torch.randint(0, 1024, (1, 8, T_FRAMES), device=DEVICE)
+
+    audio = decoder(tokens)
+
+    assert not torch.isnan(audio).any()
+    assert not torch.isinf(audio).any()
+
+
+def test_output_is_3d():
+    """Output must be 3-D: [B, 1, samples]. OmniVoiceDecoder.forward unsqueezes if needed."""
+    decoder = _build_decoder()
+    tokens = torch.randint(0, 1024, (1, 8, 10), device=DEVICE)
+
+    audio = decoder(tokens)
+
+    assert audio.dim() == 3, f"Expected 3-D output, got {audio.dim()}-D"
+    assert audio.shape[1] == 1, f"Expected channel dim=1, got {audio.shape[1]}"
+
+
+# ---------------------------------------------------------------------------
+# 3. Acoustic decoder float32 invariant
+# ---------------------------------------------------------------------------
+
+
+def test_acoustic_decoder_weights_are_float32():
+    """Acoustic decoder weights must be float32 — any fp16 parameter is a regression indicator.
+
+    In production, load_weights() calls self.acoustic_decoder.float() to enforce this.
+    The synthetic builder mirrors that invariant.
+    """
+    decoder = _build_decoder(fc2_dtype=torch.float16, acoustic_dtype=torch.float32)
+
+    for name, param in decoder.acoustic_decoder.named_parameters():
+        assert param.dtype == torch.float32, (
+            f"acoustic_decoder.{name} is {param.dtype} — must be float32 (see load_weights fix)"
+        )
diff --git a/tests/model_executor/models/omnivoice/test_triton_kernels.py b/tests/model_executor/models/omnivoice/test_triton_kernels.py
new file mode 100644
index 000000000..1598b0f74
--- /dev/null
+++ b/tests/model_executor/models/omnivoice/test_triton_kernels.py
@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Correctness tests for OmniVoice Triton kernels.
+
+Verifies that RMSNorm, SwiGLU, and fused-add+RMSNorm Triton kernels
+produce results numerically equivalent to the PyTorch reference.
+All computations follow the kernel's internal float32 promotion rules.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+DEVICE = torch.device("cuda:0")
+EPS = 1e-6
+
+try:
+    from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import (
+        _TRITON_AVAILABLE,
+        triton_fused_add_rms_norm,
+        triton_rms_norm,
+        triton_swiglu,
+    )
+except ImportError:
+    _TRITON_AVAILABLE = False
+
+triton_available = pytest.mark.skipif(not _TRITON_AVAILABLE, reason="Triton not available on this platform")
+
+
+# ---------------------------------------------------------------------------
+# PyTorch reference implementations (mirror the kernel's internal arithmetic)
+# ---------------------------------------------------------------------------
+
+
+def _ref_rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+    # Kernel promotes x to float32, computes RMS, casts back to x.dtype, then mul weight
+    x_f32 = x.float()
+    ms = x_f32.pow(2).mean(-1, keepdim=True)
+    normed = (x_f32 * torch.rsqrt(ms + eps)).to(x.dtype)
+    return normed * weight.to(normed.dtype)
+
+
+def _ref_swiglu(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
+    # Kernel: gate → float32, silu, cast to up.dtype, mul up
+    gate_f32 = gate.float()
+    silu = gate_f32 * torch.sigmoid(gate_f32)
+    return silu.to(up.dtype) * up
+
+
+def _ref_fused_add_rms_norm(
+    x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float
+) -> tuple[torch.Tensor, torch.Tensor]:
+    # Kernel: s_f32 = x+r in float32, s stored as x.dtype; y = norm(s_f32)*weight
+    s_f32 = x.float() + residual.float()
+    s = s_f32.to(x.dtype)
+    ms = s_f32.pow(2).mean(-1, keepdim=True)
+    normed = (s_f32 * torch.rsqrt(ms + eps)).to(x.dtype)
+    y = normed * weight.to(normed.dtype)
+    return y, s
+
+
+# ---------------------------------------------------------------------------
+# RMSNorm
+# ---------------------------------------------------------------------------
+
+
+@triton_available
+@pytest.mark.parametrize("rows,cols", [(1, 64), (4, 512), (16, 1024), (32, 3072)])
+def test_rms_norm_float32(rows, cols):
+    torch.manual_seed(0)
+    x = torch.randn(rows, cols, device=DEVICE)
+    w = torch.ones(cols, device=DEVICE)
+
+    ref = _ref_rms_norm(x, w, EPS)
+    out = triton_rms_norm(x, w, EPS)
+
+    torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
+
+
+@triton_available
+@pytest.mark.parametrize("shape", [(2, 16, 64), (4, 32, 512)])
+def test_rms_norm_3d_input(shape):
+    """Triton wrapper must handle 3-D [B, S, H] inputs by flattening and restoring shape."""
+    torch.manual_seed(1)
+    hidden = shape[-1]
+    x = torch.randn(*shape, device=DEVICE)
+    w = torch.ones(hidden, device=DEVICE)
+
+    ref = _ref_rms_norm(x, w, EPS)
+    out = triton_rms_norm(x, w, EPS)
+
+    torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
+    assert out.shape == x.shape
+
+
+@triton_available
+def test_rms_norm_no_nan_on_zero_input():
+    """RMSNorm of the zero vector must not produce NaN (eps prevents divide-by-zero)."""
+    x = torch.zeros(4, 256, device=DEVICE)
+    w = torch.ones(256, device=DEVICE)
+
+    out = triton_rms_norm(x, w, EPS)
+
+    assert not torch.isnan(out).any()
+
+
+# ---------------------------------------------------------------------------
+# SwiGLU
+# ---------------------------------------------------------------------------
+
+
+@triton_available
+@pytest.mark.parametrize("rows,cols", [(1, 64), (4, 512), (16, 3072)])
+def test_swiglu_float32(rows, cols):
+    torch.manual_seed(2)
+    gate = torch.randn(rows, cols, device=DEVICE)
+    up = torch.randn(rows, cols, device=DEVICE)
+
+    ref = _ref_swiglu(gate, up)
+    out = triton_swiglu(gate, up)
+
+    torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
+
+
+@triton_available
+def test_swiglu_negative_gate():
+    """silu(x) = x*sigmoid(x) is defined for negative x; result must not be NaN."""
+    gate = torch.full((4, 64), -10.0, device=DEVICE)
+    up = torch.ones(4, 64, device=DEVICE)
+
+    ref = _ref_swiglu(gate, up)
+    out = triton_swiglu(gate, up)
+
+    torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
+    assert not torch.isnan(out).any()
+
+
+@triton_available
+def test_swiglu_output_shape_preserved():
+    gate = torch.randn(3, 8, 512, device=DEVICE)
+    up = torch.randn(3, 8, 512, device=DEVICE)
+    out = triton_swiglu(gate, up)
+    assert out.shape == gate.shape
+
+
+# ---------------------------------------------------------------------------
+# Fused Add + RMSNorm
+# ---------------------------------------------------------------------------
+
+
+@triton_available
+@pytest.mark.parametrize("rows,cols", [(1, 64), (4, 512), (16, 1024)])
+def test_fused_add_rms_norm_float32(rows, cols):
+    torch.manual_seed(3)
+    x = torch.randn(rows, cols, device=DEVICE)
+    residual = torch.randn(rows, cols, device=DEVICE)
+    weight = torch.ones(cols, device=DEVICE)
+
+    ref_y, ref_s = _ref_fused_add_rms_norm(x, residual, weight, EPS)
+    out_y, out_s = triton_fused_add_rms_norm(x, residual, weight, EPS)
+
+    torch.testing.assert_close(out_y, ref_y, atol=1e-5, rtol=1e-5)
+    # Sum output s = x + r must be bit-exact (no numerics involved, just addition)
+    torch.testing.assert_close(out_s, ref_s, atol=0, rtol=0)
+
+
+@triton_available
+def test_fused_add_rms_norm_sum_is_bitexact():
+    """The residual sum s = x + r must be bit-exact between Triton and PyTorch."""
+    torch.manual_seed(4)
+    x = torch.randn(8, 512, device=DEVICE)
+    r = torch.randn(8, 512, device=DEVICE)
+    w = torch.ones(512, device=DEVICE)
+
+    _, ref_s = _ref_fused_add_rms_norm(x, r, w, EPS)
+    _, out_s = triton_fused_add_rms_norm(x, r, w, EPS)
+
+    torch.testing.assert_close(out_s, ref_s, atol=0, rtol=0)
+
+
+@triton_available
+def test_fused_add_rms_norm_returns_two_tensors():
+    """Wrapper must return exactly two tensors: (normed_output, sum)."""
+    x = torch.randn(4, 256, device=DEVICE)
+    r = torch.randn(4, 256, device=DEVICE)
+    w = torch.ones(256, device=DEVICE)
+
+    result = triton_fused_add_rms_norm(x, r, w, EPS)
+
+    assert len(result) == 2
+    y, s = result
+    assert y.shape == x.shape
+    assert s.shape == x.shape
diff --git a/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py b/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py
index e1ee68991..9a7460bf1 100644
--- a/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py
+++ b/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py
@@ -478,19 +478,29 @@ def test_compute_capture_sizes(kwargs, expected_in, not_expected):
     "batch,channels,seq_len",
     [(2, 64, 1000), (1, 32, 1), (1, 32, 7), (1, 32, 128), (1, 32, 1024), (1, 32, 4096)],
 )
-def test_snakebeta_triton_vs_eager(batch, channels, seq_len):
-    """Fused Triton SnakeBeta kernel must match eager PyTorch output."""
+@pytest.mark.parametrize(
+    "dtype,atol,rtol",
+    [
+        (torch.float32, 1e-5, 1e-5),
+        (torch.bfloat16, 1e-2, 1e-2),
+        (torch.float16, 1e-3, 1e-3),
+    ],
+    ids=["fp32", "bf16", "fp16"],
+)
+def test_snakebeta_triton_vs_eager(batch, channels, seq_len, dtype, atol, rtol):
+    """Fused Triton SnakeBeta kernel must match eager PyTorch output across dtypes."""
     from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta
 
     if not SnakeBeta._init_triton():
         pytest.skip("Triton not available")
 
     torch.manual_seed(42)
-    snake = SnakeBeta(in_features=channels).to(DEVICE).eval()
-    x = torch.randn(batch, channels, seq_len, device=DEVICE)
+    snake = SnakeBeta(in_features=channels).to(DEVICE).to(dtype).eval()
+    x = torch.randn(batch, channels, seq_len, device=DEVICE, dtype=dtype)
 
     with torch.no_grad():
         eager_out = snake._eager_forward(x)
         triton_out = snake._triton_forward(x)
 
-    torch.testing.assert_close(triton_out, eager_out, atol=1e-5, rtol=1e-5)
+    assert triton_out.dtype == x.dtype
+    torch.testing.assert_close(triton_out, eager_out, atol=atol, rtol=rtol)
diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py
index fd3fecaa9..7da61fb33 100644
--- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py
+++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py
@@ -1,34 +1,108 @@
+from collections import OrderedDict
 from types import SimpleNamespace
 
 import numpy as np
 import torch
 
+from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
+    NORMALIZED_REF_AUDIO_KEY,
+    PRECOMPUTED_REF_CODE_KEY,
+    PRECOMPUTED_REF_IDS_KEY,
+    PRECOMPUTED_TEXT_IDS_KEY,
+    Qwen3TTSPromptEmbedsBuilder,
+)
 from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
-    _NORMALIZED_REF_AUDIO_KEY,
-    _PRECOMPUTED_REF_CODE_KEY,
-    _PRECOMPUTED_REF_IDS_KEY,
-    _PRECOMPUTED_TEXT_IDS_KEY,
     Qwen3TTSTalkerForConditionalGeneration,
 )
 
 
-def _make_minimal_talker():
+def _make_minimal_talker(
+    tts_pad_embed: torch.Tensor | None = None,
+    *,
+    build_prompt_embeds=None,
+):
+    """Construct a bare ``Qwen3TTSTalkerForConditionalGeneration`` for preprocess tests.
+
+    Skips ``__init__`` and ``load_weights`` (so we don't need real config /
+    checkpoint files) and pre-populates the minimum set of attributes the
+    preprocess / decode paths read: the persistent ``tts_pad`` buffer and a
+    stub :class:`Qwen3TTSPromptEmbedsBuilder` with a configurable
+    ``build_prompt_embeds`` (defaults to raising, since most preprocess
+    tests don't exercise the builder).
+    """
     model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
     model.talker_config = SimpleNamespace(codec_pad_id=7, num_code_groups=16)
+    if tts_pad_embed is None:
+        tts_pad_embed = torch.zeros((1, 4), dtype=torch.bfloat16)
+    model._tts_pad_embed = tts_pad_embed
+
+    def _default_raise(**_kwargs):
+        raise AssertionError("build_prompt_embeds was not stubbed in this test")
+
+    model._prompt_builder = SimpleNamespace(
+        build_prompt_embeds=build_prompt_embeds if build_prompt_embeds is not None else _default_raise,
+    )
     return model
 
 
+def _make_minimal_builder(
+    *,
+    config: SimpleNamespace | None = None,
+    talker_config: SimpleNamespace | None = None,
+    tts_pad_embed: torch.Tensor | None = None,
+):
+    """Construct a bare :class:`Qwen3TTSPromptEmbedsBuilder` instance.
+
+    Bypasses ``__init__`` so we don't need a real model checkpoint; tests
+    populate the attributes they actually exercise.
+    """
+    builder = Qwen3TTSPromptEmbedsBuilder.__new__(Qwen3TTSPromptEmbedsBuilder)
+    builder._config = config or SimpleNamespace(
+        tts_bos_token_id=100,
+        tts_eos_token_id=101,
+        tts_pad_token_id=102,
+    )
+    builder._talker_config = talker_config or SimpleNamespace(
+        codec_nothink_id=10,
+        codec_think_bos_id=11,
+        codec_think_eos_id=12,
+        codec_think_id=13,
+        codec_language_id={},
+        codec_pad_id=7,
+        codec_bos_id=8,
+        num_code_groups=2,
+        spk_is_dialect={},
+    )
+    builder._model_path = ""
+    builder._text_embedding = None
+    builder._text_projection = lambda x: x
+    builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
+    builder._residual_code_embeddings = lambda: []
+    builder._speaker_encoder = None
+    builder._tts_pad_embed_buffer = (
+        tts_pad_embed if tts_pad_embed is not None else torch.zeros((1, 4), dtype=torch.bfloat16)
+    )
+    builder._encode_ref_audio_batch_fn = lambda wavs, sr, *, device: (_ for _ in ()).throw(
+        AssertionError("encode_ref_audio_batch_fn not stubbed")
+    )
+    builder._speaker_cache = None
+    builder._text_tokenizer = None
+    builder._ref_audio_artifact_cache_max_entries = 256
+    builder._ref_audio_artifact_cache = OrderedDict()
+    builder._resampler_cache = OrderedDict()
+    builder._resampler_cache_max = 16
+    return builder
+
+
 def test_single_token_prefill_uses_prefill_path():
-    model = _make_minimal_talker()
     full_prompt_embeds = torch.arange(12, dtype=torch.float32).reshape(3, 4)
     trailing_text = torch.ones((2, 4), dtype=torch.float32)
-    tts_pad = torch.full((1, 4), 0.5, dtype=torch.float32)
     ref_code = torch.arange(32, dtype=torch.long).reshape(2, 16)
 
     def fake_build_prompt_embeds(*, task_type, info_dict):
-        return full_prompt_embeds, trailing_text, tts_pad, 2, ref_code
+        return full_prompt_embeds, trailing_text, 2, ref_code
 
-    model._build_prompt_embeds = fake_build_prompt_embeds
+    model = _make_minimal_talker(build_prompt_embeds=fake_build_prompt_embeds)
 
     input_ids = torch.tensor([123], dtype=torch.long)
     out_ids, out_embeds, update = model.preprocess(
@@ -47,22 +121,22 @@ def test_single_token_prefill_uses_prefill_path():
     assert update["meta"]["talker_text_offset"] == 0
     assert update["meta"]["ref_code_len"] == 2
     assert torch.equal(update["embed"]["prefill"], full_prompt_embeds)
-    assert torch.equal(update["embed"]["tts_pad"], tts_pad)
+    # ``tts_pad`` is no longer round-tripped through ``info_dict`` — it is
+    # precomputed once into ``model._tts_pad_embed`` at load time.
+    assert "tts_pad" not in update["embed"]
     assert torch.equal(update["hidden_states"]["trailing_text"], trailing_text)
     assert torch.equal(update["codes"]["ref"], ref_code)
     assert update["codes"]["audio"].shape == (1, 16)
 
 
 def test_single_token_prefill_can_be_inferred_from_token_progress():
-    model = _make_minimal_talker()
     full_prompt_embeds = torch.arange(8, dtype=torch.float32).reshape(2, 4)
     trailing_text = torch.ones((1, 4), dtype=torch.float32)
-    tts_pad = torch.zeros((1, 4), dtype=torch.float32)
 
     def fake_build_prompt_embeds(*, task_type, info_dict):
-        return full_prompt_embeds, trailing_text, tts_pad, None, None
+        return full_prompt_embeds, trailing_text, None, None
 
-    model._build_prompt_embeds = fake_build_prompt_embeds
+    model = _make_minimal_talker(build_prompt_embeds=fake_build_prompt_embeds)
 
     out_ids, out_embeds, update = model.preprocess(
         input_ids=torch.tensor([123], dtype=torch.long),
@@ -79,7 +153,8 @@ def test_single_token_prefill_can_be_inferred_from_token_progress():
 
 
 def test_decode_advances_trailing_text_by_offset_without_rewriting_tail():
-    model = _make_minimal_talker()
+    tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16)
+    model = _make_minimal_talker(tts_pad_embed=tts_pad)
 
     def fake_embed_input_ids(input_ids):
         return input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4)
@@ -87,7 +162,6 @@ def test_decode_advances_trailing_text_by_offset_without_rewriting_tail():
     model.embed_input_ids = fake_embed_input_ids
     trailing_text = torch.arange(12, dtype=torch.float32).reshape(3, 4)
     last_hidden = torch.full((4,), 2.0, dtype=torch.float32)
-    tts_pad = torch.full((1, 4), -1.0, dtype=torch.float32)
 
     out_ids, out_embeds, update = model.preprocess(
         input_ids=torch.tensor([123], dtype=torch.long),
@@ -95,7 +169,6 @@ def test_decode_advances_trailing_text_by_offset_without_rewriting_tail():
         text=["hello"],
         task_type=["CustomVoice"],
         hidden_states={"trailing_text": trailing_text, "last": last_hidden},
-        embed={"tts_pad": tts_pad},
         meta={"talker_text_offset": 1},
         _omni_is_prefill=False,
         _omni_num_computed_tokens=2,
@@ -112,7 +185,8 @@ def test_decode_advances_trailing_text_by_offset_without_rewriting_tail():
 
 
 def test_decode_advances_trailing_text_offset_across_multiple_steps():
-    model = _make_minimal_talker()
+    tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16)
+    model = _make_minimal_talker(tts_pad_embed=tts_pad)
 
     def fake_embed_input_ids(input_ids):
         return input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4)
@@ -121,7 +195,6 @@ def test_decode_advances_trailing_text_offset_across_multiple_steps():
     trailing_text = torch.arange(8, dtype=torch.float32).reshape(2, 4)
     state_tail = trailing_text
     last_hidden = torch.full((4,), 2.0, dtype=torch.float32)
-    tts_pad = torch.full((1, 4), -1.0, dtype=torch.float32)
     meta = {"talker_text_offset": 0}
     seen_steps = []
 
@@ -132,7 +205,6 @@ def test_decode_advances_trailing_text_offset_across_multiple_steps():
             text=["hello"],
             task_type=["CustomVoice"],
             hidden_states={"trailing_text": state_tail, "last": last_hidden},
-            embed={"tts_pad": tts_pad},
             meta=meta,
             _omni_is_prefill=False,
             _omni_num_computed_tokens=2,
@@ -151,7 +223,8 @@ def test_decode_advances_trailing_text_offset_across_multiple_steps():
 
 
 def test_decode_compacts_long_trailing_text_after_large_offset():
-    model = _make_minimal_talker()
+    tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16)
+    model = _make_minimal_talker(tts_pad_embed=tts_pad)
 
     def fake_embed_input_ids(input_ids):
         return input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4)
@@ -159,7 +232,6 @@ def test_decode_compacts_long_trailing_text_after_large_offset():
     model.embed_input_ids = fake_embed_input_ids
     trailing_text = torch.arange(130 * 4, dtype=torch.float32).reshape(130, 4)
     last_hidden = torch.full((4,), 2.0, dtype=torch.float32)
-    tts_pad = torch.full((1, 4), -1.0, dtype=torch.float32)
 
     _, _, update = model.preprocess(
         input_ids=torch.tensor([123], dtype=torch.long),
@@ -167,7 +239,6 @@ def test_decode_compacts_long_trailing_text_after_large_offset():
         text=["hello"],
         task_type=["CustomVoice"],
         hidden_states={"trailing_text": trailing_text, "last": last_hidden},
-        embed={"tts_pad": tts_pad},
         meta={"talker_text_offset": 64},
         _omni_is_prefill=False,
         _omni_num_computed_tokens=2,
@@ -180,7 +251,8 @@ def test_decode_compacts_long_trailing_text_after_large_offset():
 
 
 def test_decode_batch_preprocess_matches_decode_state_updates():
-    model = _make_minimal_talker()
+    tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16)
+    model = _make_minimal_talker(tts_pad_embed=tts_pad)
 
     def fake_embed_input_ids(input_ids):
         return input_ids.to(torch.float32).reshape(-1, 1, 1).expand(-1, 1, 4)
@@ -190,7 +262,6 @@ def test_decode_batch_preprocess_matches_decode_state_updates():
     trailing_b = torch.arange(8, dtype=torch.float32).reshape(2, 4) + 100
     last_a = torch.full((4,), 2.0, dtype=torch.float32)
     last_b = torch.full((4,), 3.0, dtype=torch.float32)
-    tts_pad = torch.full((1, 4), -1.0, dtype=torch.float32)
 
     out_ids, out_embeds, past_hidden, text_step, updates = model.preprocess_decode_batch(
         input_ids=torch.tensor([101, 202], dtype=torch.long),
@@ -199,14 +270,12 @@ def test_decode_batch_preprocess_matches_decode_state_updates():
                 "text": ["hello"],
                 "task_type": ["Base"],
                 "hidden_states": {"trailing_text": trailing_a, "last": last_a},
-                "embed": {"tts_pad": tts_pad},
                 "meta": {"talker_text_offset": 1},
             },
             {
                 "text": ["world"],
                 "task_type": ["CustomVoice"],
                 "hidden_states": {"trailing_text": trailing_b, "last": last_b},
-                "embed": {"tts_pad": tts_pad},
                 "meta": {"talker_text_offset": 2},
             },
         ],
@@ -225,53 +294,49 @@ def test_decode_batch_preprocess_matches_decode_state_updates():
     assert updates[1]["hidden_states"]["trailing_text"].numel() == 0
 
 
+def _stub_text_embedding(device_param: torch.nn.Parameter):
+    """Build a lambda that emulates ``nn.Embedding`` for the ``_device()`` helper.
+
+    ``Qwen3TTSPromptEmbedsBuilder._device()`` calls
+    ``next(self._text_embedding.parameters()).device``, so the stub needs a
+    ``parameters()`` attribute *and* must be callable as an embedding.
+    """
+    fn = lambda ids: torch.ones((*ids.shape, 4), device=ids.device)  # noqa: E731
+    fn.parameters = lambda: iter([device_param])  # type: ignore[attr-defined]
+    return fn
+
+
 def test_base_voice_clone_normalizes_ref_audio_once_for_ref_code_and_speaker():
-    model = _make_minimal_talker()
+    builder = _make_minimal_builder()
     device_param = torch.nn.Parameter(torch.empty(0))
-    model.parameters = lambda: iter([device_param])
-    model.config = SimpleNamespace(
-        tts_bos_token_id=100,
-        tts_eos_token_id=101,
-        tts_pad_token_id=102,
-    )
-    model.talker_config = SimpleNamespace(
-        codec_nothink_id=10,
-        codec_think_bos_id=11,
-        codec_think_eos_id=12,
-        codec_think_id=13,
-        codec_language_id={},
-        codec_pad_id=7,
-        codec_bos_id=8,
-        num_code_groups=2,
-        spk_is_dialect={},
-    )
+    builder._text_embedding = _stub_text_embedding(device_param)
+    builder._text_projection = lambda embeds: embeds
+    builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
 
     class FakeTokenizer:
         def __call__(self, *_args, **_kwargs):
             return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)}
 
-    model._get_tokenizer = lambda: FakeTokenizer()
-    model.text_embedding = lambda ids: torch.ones((*ids.shape, 4), device=ids.device)
-    model.text_projection = lambda embeds: embeds
-    model.embed_input_ids = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
-    model._generate_icl_prompt = lambda **kwargs: (
+    builder._text_tokenizer = FakeTokenizer()
+
+    builder._generate_icl_prompt = lambda **kwargs: (
         torch.ones((1, 2, 4), device=kwargs["ref_code"].device),
         torch.ones((1, 4), device=kwargs["ref_code"].device),
     )
 
     normalize_calls = []
     ref_audio = np.arange(1024, dtype=np.float32)
-    model._normalize_ref_audio = lambda raw: normalize_calls.append(raw) or (ref_audio, 16000)
+    builder.normalize_ref_audio = lambda raw: normalize_calls.append(raw) or (ref_audio, 16000)
 
     ref_audio_ids = []
-    model._encode_ref_audio_to_code = lambda wav, _sr: (
-        ref_audio_ids.append(id(wav)) or torch.ones((2, 2), dtype=torch.long)
+    builder._encode_ref_audio_batch_fn = lambda wavs, sr, *, device: (
+        [ref_audio_ids.append(id(w)) or torch.ones((2, 2), dtype=torch.long) for w in wavs]
     )
-    model._extract_speaker_embedding = lambda wav, _sr: (
+    builder.extract_speaker_embedding = lambda wav, _sr: (
         ref_audio_ids.append(id(wav)) or torch.ones(4, dtype=torch.bfloat16)
     )
 
-    _prompt, _trailing, _pad, ref_code_len, ref_code = model._build_prompt_embeds(
+    _prompt, _trailing, ref_code_len, ref_code = builder.build_prompt_embeds(
         task_type="Base",
         info_dict={
             "text": ["hello"],
@@ -288,11 +353,11 @@ def test_base_voice_clone_normalizes_ref_audio_once_for_ref_code_and_speaker():
 
 
 def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate():
-    model = _make_minimal_talker()
+    builder = _make_minimal_builder()
     wav1 = np.arange(2048, dtype=np.float32)
     wav2 = np.arange(3072, dtype=np.float32)
     normalize_calls = []
-    model._normalize_ref_audio = lambda raw: (
+    builder.normalize_ref_audio = lambda raw: (
         normalize_calls.append(raw)
         or (
             wav1 if raw == "a.wav" else wav2,
@@ -300,21 +365,16 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate():
         )
     )
 
-    class FakeSpeechTokenizer:
-        def __init__(self):
-            self.calls = []
+    encode_calls = []
 
-        def encode(self, audios, *, sr, return_dict):
-            self.calls.append((audios, sr, return_dict))
-            return SimpleNamespace(
-                audio_codes=[
-                    torch.full((2, 2), 11, dtype=torch.long),
-                    torch.full((3, 2), 22, dtype=torch.long),
-                ]
-            )
+    def _fake_encode_batch(wavs, sr, *, device):
+        encode_calls.append((wavs, sr))
+        return [
+            torch.full((2, 2), 11, dtype=torch.long),
+            torch.full((3, 2), 22, dtype=torch.long),
+        ]
 
-    tok = FakeSpeechTokenizer()
-    model._ensure_speech_tokenizer_loaded = lambda: tok
+    builder._encode_ref_audio_batch_fn = _fake_encode_batch
 
     class FakeTextTokenizer:
         def __init__(self):
@@ -325,7 +385,7 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate():
             return {"input_ids": [[idx + 1, idx + 2, idx + 3] for idx, _ in enumerate(texts)]}
 
     text_tok = FakeTextTokenizer()
-    model._get_tokenizer = lambda: text_tok
+    builder._text_tokenizer = text_tok
     buf = {
         "r1": {
             "task_type": ["Base"],
@@ -343,43 +403,42 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate():
         },
     }
 
-    model.preprocess_batch(
+    builder.preprocess_batch(
         req_ids=["r1", "r2"],
         model_intermediate_buffer=buf,
         device=torch.device("cpu"),
     )
 
     assert normalize_calls == ["a.wav", "b.wav"]
-    assert len(tok.calls) == 1
-    audios, sr, return_dict = tok.calls[0]
-    assert audios[0] is wav1
-    assert audios[1] is wav2
+    assert len(encode_calls) == 1
+    wavs, sr = encode_calls[0]
+    assert wavs[0] is wav1
+    assert wavs[1] is wav2
     assert sr == 16000
-    assert return_dict is True
-    assert torch.equal(buf["r1"]["codes"][_PRECOMPUTED_REF_CODE_KEY], torch.full((2, 2), 11))
-    assert torch.equal(buf["r2"]["codes"][_PRECOMPUTED_REF_CODE_KEY], torch.full((3, 2), 22))
-    assert buf["r1"][_NORMALIZED_REF_AUDIO_KEY][0] is wav1
-    assert buf["r2"][_NORMALIZED_REF_AUDIO_KEY][0] is wav2
+    assert torch.equal(buf["r1"]["codes"][PRECOMPUTED_REF_CODE_KEY], torch.full((2, 2), 11))
+    assert torch.equal(buf["r2"]["codes"][PRECOMPUTED_REF_CODE_KEY], torch.full((3, 2), 22))
+    assert buf["r1"][NORMALIZED_REF_AUDIO_KEY][0] is wav1
+    assert buf["r2"][NORMALIZED_REF_AUDIO_KEY][0] is wav2
     assert len(text_tok.calls) == 2
-    assert torch.equal(buf["r1"][_PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([1, 2, 3]))
-    assert torch.equal(buf["r2"][_PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([2, 3, 4]))
-    assert torch.equal(buf["r1"][_PRECOMPUTED_REF_IDS_KEY], torch.tensor([1, 2, 3]))
-    assert torch.equal(buf["r2"][_PRECOMPUTED_REF_IDS_KEY], torch.tensor([2, 3, 4]))
+    assert torch.equal(buf["r1"][PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([1, 2, 3]))
+    assert torch.equal(buf["r2"][PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([2, 3, 4]))
+    assert torch.equal(buf["r1"][PRECOMPUTED_REF_IDS_KEY], torch.tensor([1, 2, 3]))
+    assert torch.equal(buf["r2"][PRECOMPUTED_REF_IDS_KEY], torch.tensor([2, 3, 4]))
 
 
 def test_base_voice_clone_batch_preprocess_reuses_singleton_normalized_audio_without_speech_tokenizer():
-    model = _make_minimal_talker()
+    builder = _make_minimal_builder()
     wav = np.arange(2048, dtype=np.float32)
-    model._normalize_ref_audio = lambda raw: (wav, 16000)
-    model._ensure_speech_tokenizer_loaded = lambda: (_ for _ in ()).throw(
-        AssertionError("singleton should not load speech tokenizer")
+    builder.normalize_ref_audio = lambda raw: (wav, 16000)
+    builder._encode_ref_audio_batch_fn = lambda wavs, sr, *, device: (_ for _ in ()).throw(
+        AssertionError("singleton should not call encode_ref_audio_batch")
     )
 
     class FakeTextTokenizer:
         def __call__(self, texts, *, padding=False):
             return {"input_ids": [[7, 8, 9] for _ in texts]}
 
-    model._get_tokenizer = lambda: FakeTextTokenizer()
+    builder._text_tokenizer = FakeTextTokenizer()
     buf = {
         "r1": {
             "task_type": ["Base"],
@@ -390,22 +449,27 @@ def test_base_voice_clone_batch_preprocess_reuses_singleton_normalized_audio_wit
         }
     }
 
-    model.preprocess_batch(
+    builder.preprocess_batch(
         req_ids=["r1"],
         model_intermediate_buffer=buf,
         device=torch.device("cpu"),
     )
 
-    assert buf["r1"][_NORMALIZED_REF_AUDIO_KEY][0] is wav
-    assert torch.equal(buf["r1"][_PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([7, 8, 9]))
-    assert torch.equal(buf["r1"][_PRECOMPUTED_REF_IDS_KEY], torch.tensor([7, 8, 9]))
+    assert buf["r1"][NORMALIZED_REF_AUDIO_KEY][0] is wav
+    assert torch.equal(buf["r1"][PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([7, 8, 9]))
+    assert torch.equal(buf["r1"][PRECOMPUTED_REF_IDS_KEY], torch.tensor([7, 8, 9]))
 
 
 def test_base_voice_clone_batch_preprocess_skips_after_initial_prefill_state_exists():
-    model = _make_minimal_talker()
-    model._normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("normalize not expected"))
-    model._get_tokenizer = lambda: (_ for _ in ()).throw(AssertionError("tokenizer not expected"))
-    model._ensure_speech_tokenizer_loaded = lambda: (_ for _ in ()).throw(
+    builder = _make_minimal_builder()
+    builder.normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("normalize not expected"))
+
+    class _RaisingTokenizer:
+        def __call__(self, *_args, **_kwargs):
+            raise AssertionError("tokenizer not expected")
+
+    builder._text_tokenizer = _RaisingTokenizer()
+    builder._encode_ref_audio_batch_fn = lambda *a, **kw: (_ for _ in ()).throw(
         AssertionError("speech tokenizer not expected")
     )
     buf = {
@@ -419,70 +483,53 @@ def test_base_voice_clone_batch_preprocess_skips_after_initial_prefill_state_exi
         }
     }
 
-    model.preprocess_batch(
+    builder.preprocess_batch(
         req_ids=["r1"],
         model_intermediate_buffer=buf,
         device=torch.device("cpu"),
     )
 
-    assert _PRECOMPUTED_TEXT_IDS_KEY not in buf["r1"]
-    assert _NORMALIZED_REF_AUDIO_KEY not in buf["r1"]
+    assert PRECOMPUTED_TEXT_IDS_KEY not in buf["r1"]
+    assert NORMALIZED_REF_AUDIO_KEY not in buf["r1"]
 
 
 def test_base_voice_clone_uses_batched_ref_code_without_serial_encode():
-    model = _make_minimal_talker()
+    builder = _make_minimal_builder()
     device_param = torch.nn.Parameter(torch.empty(0))
-    model.parameters = lambda: iter([device_param])
-    model.config = SimpleNamespace(
-        tts_bos_token_id=100,
-        tts_eos_token_id=101,
-        tts_pad_token_id=102,
-    )
-    model.talker_config = SimpleNamespace(
-        codec_nothink_id=10,
-        codec_think_bos_id=11,
-        codec_think_eos_id=12,
-        codec_think_id=13,
-        codec_language_id={},
-        codec_pad_id=7,
-        codec_bos_id=8,
-        num_code_groups=2,
-        spk_is_dialect={},
-    )
+    builder._text_embedding = _stub_text_embedding(device_param)
+    builder._text_projection = lambda embeds: embeds
+    builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
 
     class FakeTokenizer:
         def __call__(self, *_args, **_kwargs):
             return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)}
 
-    model._get_tokenizer = lambda: FakeTokenizer()
-    model.text_embedding = lambda ids: torch.ones((*ids.shape, 4), device=ids.device)
-    model.text_projection = lambda embeds: embeds
-    model.embed_input_ids = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
-    model._generate_icl_prompt = lambda **kwargs: (
+    builder._text_tokenizer = FakeTokenizer()
+    builder._generate_icl_prompt = lambda **kwargs: (
         torch.ones((1, 2, 4), device=kwargs["ref_code"].device),
         torch.ones((1, 4), device=kwargs["ref_code"].device),
     )
 
     ref_audio = np.arange(2048, dtype=np.float32)
     ref_code = torch.arange(4, dtype=torch.long).reshape(2, 2)
-    model._normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("serial normalize not expected"))
-    model._encode_ref_audio_to_code = lambda _wav, _sr: (_ for _ in ()).throw(
+    builder.normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("serial normalize not expected"))
+    builder.encode_ref_audio_to_code = lambda _wav, _sr: (_ for _ in ()).throw(
         AssertionError("serial encode not expected")
     )
     speaker_wav_ids = []
-    model._extract_speaker_embedding = lambda wav, _sr: (
+    builder.extract_speaker_embedding = lambda wav, _sr: (
         speaker_wav_ids.append(id(wav)) or torch.ones(4, dtype=torch.bfloat16)
     )
 
-    _prompt, _trailing, _pad, ref_code_len, out_ref_code = model._build_prompt_embeds(
+    _prompt, _trailing, ref_code_len, out_ref_code = builder.build_prompt_embeds(
         task_type="Base",
         info_dict={
             "text": ["hello"],
             "ref_audio": ["ref.wav"],
             "ref_ids": torch.arange(8, dtype=torch.long).reshape(1, -1),
             "non_streaming_mode": [False],
-            "codes": {_PRECOMPUTED_REF_CODE_KEY: ref_code},
-            _NORMALIZED_REF_AUDIO_KEY: (ref_audio, 16000),
+            "codes": {PRECOMPUTED_REF_CODE_KEY: ref_code},
+            NORMALIZED_REF_AUDIO_KEY: (ref_audio, 16000),
         },
     )
 
diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_prompt_len.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_prompt_len.py
new file mode 100644
index 000000000..9493a00f2
--- /dev/null
+++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_prompt_len.py
@@ -0,0 +1,129 @@
+"""Regression tests for `estimate_prompt_len_from_additional_information`.
+
+Pins the 2D `voice_clone_prompt.ref_code` shape behaviour. Applying the
+singleton-batch unwrapper `_first(...)` to that value strips its outer
+dimension and reports `len(ref_code) == num_codebooks` instead of
+`num_frames`, which silently truncates `inputs_embeds` downstream in
+`_build_prompt_embeds`.
+"""
+
+import pytest
+
+from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
+    Qwen3TTSPromptEmbedsBuilder,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _fake_tokenize(text: str, **_kwargs):
+    return [0] * (8 + max(1, len(text.split())))
+
+
+def test_estimate_prompt_len_uses_full_ref_code_length() -> None:
+    num_frames = 318
+    num_codebooks = 8
+    info = {
+        "task_type": ["Base"],
+        "text": ["hello"],
+        "ref_text": ["world"],
+        "voice_clone_prompt": [
+            {
+                "ref_spk_embedding": [0.0] * 512,
+                "ref_code": [[0] * num_codebooks for _ in range(num_frames)],
+                "icl_mode": True,
+            }
+        ],
+        "non_streaming_mode": [True],
+        "language": ["English"],
+    }
+
+    est = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
+        additional_information=info,
+        task_type="Base",
+        tokenize_prompt=_fake_tokenize,
+        codec_language_id={"english": 0},
+        spk_is_dialect=None,
+    )
+
+    # codec_lens = 1 + num_frames = 319; plus text-side and codec-prefix
+    # terms ~20. Would be ~30 if `_first` collapses ref_code to its first row.
+    assert est > 100, f"got {est}; expected ~339. Did `_first(ref_code)` collapse the 2D list again?"
+
+
+def test_estimate_prompt_len_handles_1d_ref_code() -> None:
+    num_frames = 50
+    info = {
+        "task_type": ["Base"],
+        "text": ["hello"],
+        "ref_text": ["world"],
+        "voice_clone_prompt": [
+            {
+                "ref_spk_embedding": [0.0] * 512,
+                "ref_code": list(range(num_frames)),
+                "icl_mode": True,
+            }
+        ],
+        "non_streaming_mode": [True],
+        "language": ["English"],
+    }
+
+    est = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
+        additional_information=info,
+        task_type="Base",
+        tokenize_prompt=_fake_tokenize,
+        codec_language_id={"english": 0},
+        spk_is_dialect=None,
+    )
+
+    assert est > 50, f"got {est}; 1D ref_code must contribute its own length"
+
+
+def test_estimate_prompt_len_uses_list_of_int_ref_ids_from_voice_clone_prompt() -> None:
+    """`voice_clone_prompt.ref_ids` as a list-of-int (the natural pre-tokenized shape)
+    must be read as a sequence, not unwrapped to its first element. Without the fix
+    `_first` collapsed it to an int, the `isinstance(list)` / `hasattr("shape")`
+    branches fell through, and `ref_ids_len` defaulted to 0 — silently dropping the
+    reference-token contribution to the prefix length budget.
+    """
+
+    def make_info(num_ref_ids: int) -> dict:
+        return {
+            "task_type": ["Base"],
+            "text": ["hello"],
+            "ref_text": ["world"],
+            "voice_clone_prompt": [
+                {
+                    "ref_spk_embedding": [0.0] * 512,
+                    "ref_code": [[0] * 8 for _ in range(200)],
+                    "ref_ids": list(range(num_ref_ids)),
+                    "icl_mode": True,
+                }
+            ],
+            "non_streaming_mode": [True],
+            "language": ["English"],
+        }
+
+    est_30 = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
+        additional_information=make_info(30),
+        task_type="Base",
+        tokenize_prompt=_fake_tokenize,
+        codec_language_id={"english": 0},
+        spk_is_dialect=None,
+    )
+    est_50 = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
+        additional_information=make_info(50),
+        task_type="Base",
+        tokenize_prompt=_fake_tokenize,
+        codec_language_id={"english": 0},
+        spk_is_dialect=None,
+    )
+
+    # Each additional ref_id contributes exactly one token to the prefix budget
+    # (the estimator strips a fixed 5-token tail from ref_ids regardless of length).
+    # Under the previous bug `_first` collapsed both lists to a single int, so both
+    # estimates would have been identical (delta = 0).
+    assert est_50 - est_30 == 20, (
+        f"got delta={est_50 - est_30}; expected 20 (= 50 - 30 extra ref_ids). "
+        f"Did `_first(ref_ids)` collapse the list to its first int?"
+    )
diff --git a/tests/model_executor/models/test_encoder_quant_config.py b/tests/model_executor/models/test_encoder_quant_config.py
index 47c7c69c7..28a1efa71 100644
--- a/tests/model_executor/models/test_encoder_quant_config.py
+++ b/tests/model_executor/models/test_encoder_quant_config.py
@@ -58,7 +58,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 @pytest.mark.parametrize("method", sorted(PRE_QUANTIZED_METHODS))
 def test_pre_quantized_returns_none(method: str) -> None:
     """visual_quant_config and audio_quant_config must be None for
-    pre-quantized methods (modelopt, modelopt_fp4, modelopt_mxfp8)."""
+    pre-quantized methods (modelopt, modelopt_fp4, modelopt_mxfp8,
+    modelopt_mixed)."""
     mock_config = MagicMock()
     mock_config.get_name.return_value = method
 
@@ -100,5 +101,5 @@ def test_component_config_passed_through() -> None:
 
 def test_pre_quantized_methods_contains_expected() -> None:
     """Guard against accidental removal of a known pre-quantized method."""
-    expected = {"modelopt", "modelopt_fp4", "modelopt_mxfp8"}
+    expected = {"modelopt", "modelopt_fp4", "modelopt_mxfp8", "modelopt_mixed"}
     assert PRE_QUANTIZED_METHODS == expected
diff --git a/tests/model_executor/models/test_glm_tts_prompt_len.py b/tests/model_executor/models/test_glm_tts_prompt_len.py
new file mode 100644
index 000000000..df12078b4
--- /dev/null
+++ b/tests/model_executor/models/test_glm_tts_prompt_len.py
@@ -0,0 +1,351 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.common.nucleus_ras_sampling import ras_sample_one
+from vllm_omni.model_executor.models.glm_tts import glm_tts as glm_tts_module
+from vllm_omni.model_executor.models.glm_tts.glm_tts import GLMTTSForConditionalGeneration
+from vllm_omni.model_executor.models.glm_tts.text_frontend import GLMTTSTextFrontend
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_glm_tts_ras_fallback_matches_official_random_resample() -> None:
+    sampled = ras_sample_one(
+        torch.tensor([0.0, float("-inf")]),
+        decoded_tokens=[0] * 10,
+        top_p=1.0,
+        top_k=1,
+        win_size=10,
+        tau_r=0.1,
+        generator=None,
+    )
+
+    assert sampled == 0
+
+
+def test_glm_tts_text_frontend_applies_official_symbol_replacements() -> None:
+    frontend = GLMTTSTextFrontend()
+
+    assert frontend.pre_replace("3-2 咯，") == "3减2 喽，"
+    assert frontend.post_replace("α②∈A/B~C²") == "阿尔法二属于A每B到C平方"
+
+
+def test_glm_tts_call_hf_processor_builds_prompt_text_text_boa(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    processor = object.__new__(glm_tts_module.GLMTTSMultiModalProcessor)
+    processor.info = SimpleNamespace(
+        ctx=SimpleNamespace(
+            get_hf_config=lambda *_args, **_kwargs: SimpleNamespace(),
+            model_config=SimpleNamespace(model="fake-model"),
+        )
+    )
+    processor.text_frontend = object()
+    processor.special_ids = {"boa": 99}
+    processor.speech_tokenizer = object()
+    processor.processor_device = torch.device("cpu")
+    processor.campplus_session = object()
+
+    monkeypatch.setattr(
+        processor,
+        "_ensure_cached_runtime_components",
+        lambda *_args, **_kwargs: None,
+    )
+    monkeypatch.setattr(
+        processor,
+        "_encode_text",
+        lambda text: torch.tensor([[11, 12]]) if text.endswith(" ") else torch.tensor([[21, 22, 23]]),
+    )
+    monkeypatch.setattr(processor, "_get_audio", lambda _mm_data: (torch.zeros(8), 24000))
+    monkeypatch.setattr(
+        glm_tts_module,
+        "_normalize_glm_tts_processor_text",
+        lambda _frontend, text, add_trailing_space=False: text + (" " if add_trailing_space else ""),
+    )
+    monkeypatch.setattr(glm_tts_module, "extract_prompt_speech_token", lambda *_args, **_kwargs: [7, 8, 9, 10])
+    monkeypatch.setattr(glm_tts_module, "extract_prompt_feat", lambda *_args, **_kwargs: torch.ones(4, 3))
+    monkeypatch.setattr(glm_tts_module, "extract_spk_embedding", lambda *_args, **_kwargs: [0.1, 0.2, 0.3])
+
+    out = processor._call_hf_processor(
+        prompt="target text",
+        mm_data={"audio": object()},
+        mm_kwargs={"prompt_text": "reference text"},
+        tok_kwargs={},
+    )
+
+    assert out["input_ids"].tolist() == [[11, 12, 21, 22, 23, 99]]
+    assert out["prompt_speech_token"].tolist() == [[7, 8, 9, 10]]
+    assert out["prompt_speech_token_len"][0].item() == 4
+    assert out["glm_tts_text_token_len"][0].item() == 3
+
+
+def test_glm_tts_prompt_updates_append_single_audio_suffix() -> None:
+    processor = object.__new__(glm_tts_module.GLMTTSMultiModalProcessor)
+    out_mm_kwargs = {
+        "audio": [{"prompt_speech_token_len": SimpleNamespace(data=[torch.tensor([5], dtype=torch.long)])}]
+    }
+
+    updates = processor._get_prompt_updates(
+        mm_items=object(),
+        hf_processor_mm_kwargs={},
+        out_mm_kwargs=out_mm_kwargs,
+    )
+
+    assert len(updates) == 1
+    assert updates[0].modality == "audio"
+    assert updates[0].insertion(0) == [1, 1, 1, 1, 1]
+
+
+def test_glm_tts_dummy_inputs_forbid_dummy_wav(monkeypatch: pytest.MonkeyPatch) -> None:
+    builder = object.__new__(glm_tts_module.GLMTTSDummyInputsBuilder)
+    monkeypatch.setattr(
+        builder,
+        "_get_dummy_audios",
+        lambda *_args, **_kwargs: pytest.fail("GLM-TTS dummy inputs must not fabricate wav audio"),
+    )
+
+    assert builder.get_dummy_mm_data(seq_len=16, mm_counts={"audio": 1}) == {}
+
+
+def test_glm_tts_processing_info_declares_audio_token_budget() -> None:
+    info = object.__new__(glm_tts_module.GLMTTSMultiModalProcessingInfo)
+
+    assert info.get_mm_max_tokens_per_item(seq_len=4096, mm_counts={"audio": 1}) == {"audio": 1024}
+    assert info.get_mm_max_tokens_per_item(seq_len=512, mm_counts={"audio": 1}) == {"audio": 512}
+    assert info.get_mm_max_tokens_per_item(seq_len=4096, mm_counts={}) == {}
+
+
+def test_glm_tts_text_only_processor_skips_voice_clone(monkeypatch: pytest.MonkeyPatch) -> None:
+    class FakeTokenizer:
+        def encode(self, text: str) -> list[int]:
+            return [11, 12, 13]
+
+    processor = object.__new__(glm_tts_module.GLMTTSMultiModalProcessor)
+    processor.info = SimpleNamespace(
+        ctx=SimpleNamespace(
+            get_hf_config=lambda *_args, **_kwargs: SimpleNamespace(),
+            model_config=SimpleNamespace(model="fake-model", tokenizer=None, trust_remote_code=True),
+        )
+    )
+
+    monkeypatch.setattr(glm_tts_module, "resolve_glm_tts_model_dir", lambda *_args, **_kwargs: "/fake-root")
+    monkeypatch.setattr(glm_tts_module, "load_glm_tts_tokenizer", lambda *_args, **_kwargs: FakeTokenizer())
+    monkeypatch.setattr(glm_tts_module, "get_glm_tts_special_token_ids", lambda _tokenizer: {"boa": 99})
+    monkeypatch.setattr(glm_tts_module, "GLMTTSTextFrontend", lambda: object())
+    monkeypatch.setattr(
+        glm_tts_module,
+        "_normalize_glm_tts_processor_text",
+        lambda _frontend, text, add_trailing_space=False: text,
+    )
+    monkeypatch.setattr(
+        glm_tts_module,
+        "load_voice_clone_frontend",
+        lambda *_args, **_kwargs: pytest.fail("text-only dummy path should not load voice clone frontend"),
+    )
+
+    out = processor._call_hf_processor(prompt="target text", mm_data={}, mm_kwargs={}, tok_kwargs={})
+
+    assert out["input_ids"].tolist() == [[11, 12, 13, 99]]
+    assert out["prompt_speech_token"].shape == (1, 1024)
+    assert out["prompt_speech_token_len"][0].item() == 1024
+    assert out["glm_tts_prompt_text_token_len"][0].item() == 0
+    assert out["glm_tts_text_token_len"][0].item() == 3
+
+    fields = processor._get_mm_fields_config(out, {})
+    assert {key: value.modality for key, value in fields.items()} == {
+        "prompt_speech_token": "audio",
+        "prompt_speech_token_len": "audio",
+        "glm_tts_prompt_text_token_len": "audio",
+        "glm_tts_text_token_len": "audio",
+    }
+
+    updates = processor._get_prompt_updates(
+        mm_items=object(),
+        hf_processor_mm_kwargs={},
+        out_mm_kwargs={
+            "audio": [{"prompt_speech_token_len": SimpleNamespace(data=[out["prompt_speech_token_len"][0]])}]
+        },
+    )
+    assert len(updates) == 1
+    assert len(updates[0].insertion(0)) == 1024
+
+
+def test_glm_tts_embed_replaces_multimodal_positions_per_request() -> None:
+    class FakeModel(torch.nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.marker = torch.nn.Parameter(torch.zeros(()))
+
+        def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
+            return input_ids.to(torch.float32).reshape(-1, 1)
+
+    model = object.__new__(GLMTTSForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.model_stage = "glm_tts"
+    model.model = FakeModel()
+
+    input_ids = torch.tensor([101, 102, 1, 1, 201, 1, 1, 1])
+    is_multimodal = torch.tensor([False, False, True, True, False, True, True, True])
+    mm_embeds = [
+        torch.tensor([[1001.0], [1002.0]]),
+        torch.tensor([[2001.0], [2002.0], [2003.0]]),
+    ]
+
+    out = model.embed_input_ids(input_ids, multimodal_embeddings=mm_embeds, is_multimodal=is_multimodal)
+
+    assert out.squeeze(-1).tolist() == [
+        101.0,
+        102.0,
+        1001.0,
+        1002.0,
+        201.0,
+        2001.0,
+        2002.0,
+        2003.0,
+    ]
+
+
+def test_glm_tts_embed_multimodal_keeps_all_batched_requests() -> None:
+    class FakeModel(torch.nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.marker = torch.nn.Parameter(torch.zeros(()))
+
+        def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
+            return input_ids.to(torch.float32).reshape(*input_ids.shape, 1)
+
+    model = object.__new__(GLMTTSForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.model_stage = "glm_tts"
+    model.model = FakeModel()
+    model._ats = 1000
+
+    out = model.embed_multimodal(
+        prompt_speech_token=[
+            torch.tensor([[1, 2]]),
+            torch.tensor([[3, 4, 5]]),
+        ]
+    )
+
+    assert out is not None
+    assert len(out) == 2
+    assert out[0].squeeze(-1).tolist() == [1001.0, 1002.0]
+    assert out[1].squeeze(-1).tolist() == [1003.0, 1004.0, 1005.0]
+
+
+def test_glm_tts_embed_multimodal_returns_one_item_for_single_request() -> None:
+    class FakeModel(torch.nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.marker = torch.nn.Parameter(torch.zeros(()))
+
+        def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
+            return input_ids.to(torch.float32).reshape(*input_ids.shape, 1)
+
+    model = object.__new__(GLMTTSForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.model_stage = "glm_tts"
+    model.model = FakeModel()
+    model._ats = 1000
+
+    out = model.embed_multimodal(prompt_speech_token=torch.tensor([[1, 2, 3]]))
+
+    assert out is not None
+    assert len(out) == 1
+    assert out[0].squeeze(-1).tolist() == [1001.0, 1002.0, 1003.0]
+
+
+def test_glm_tts_embed_multimodal_treats_scalar_tensor_list_as_one_item() -> None:
+    class FakeModel(torch.nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.marker = torch.nn.Parameter(torch.zeros(()))
+
+        def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
+            return input_ids.to(torch.float32).reshape(*input_ids.shape, 1)
+
+    model = object.__new__(GLMTTSForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.model_stage = "glm_tts"
+    model.model = FakeModel()
+    model._ats = 1000
+
+    out = model.embed_multimodal(prompt_speech_token=[torch.tensor(1), torch.tensor(2)])
+
+    assert out is not None
+    assert len(out) == 1
+    assert out[0].squeeze(-1).tolist() == [1001.0, 1002.0]
+
+
+def test_glm_tts_preprocess_treats_one_token_text_tail_as_prefill() -> None:
+    model = object.__new__(GLMTTSForConditionalGeneration)
+    model._pad = 0
+    model._ats = 1000
+    model._ate = 2000
+    model._eoa = 3000
+    model._model_dtype = lambda: torch.float32
+
+    input_ids, input_embeds, info = model.preprocess(
+        torch.tensor([123]),
+        torch.tensor([[1.0, 2.0]]),
+        glm_tts_mm_prefill_done=True,
+    )
+
+    assert input_ids.tolist() == [0]
+    assert input_embeds.tolist() == [[1.0, 2.0]]
+    assert info["glm_tts_mm_prefill_done"] is True
+    assert info["speech_tokens"].tolist() == [[-1]]
+
+
+def test_glm_tts_preprocess_initializes_generation_bounds_before_sampling() -> None:
+    model = object.__new__(GLMTTSForConditionalGeneration)
+    model._pad = 0
+    model._ats = 1000
+    model._ate = 2000
+    model._eoa = 3000
+    model._model_dtype = lambda: torch.float32
+    model.config = SimpleNamespace(min_token_text_ratio=2.0, max_token_text_ratio=20.0)
+
+    _, _, info = model.preprocess(
+        torch.tensor([11, 12]),
+        torch.tensor([[1.0], [2.0]]),
+        glm_tts_text_token_len=[torch.tensor([27])],
+    )
+
+    assert info["glm_tts_text_token_len"][0].item() == 27
+
+
+def test_glm_tts_make_omni_output_emits_prompt_speech_len_once() -> None:
+    model = object.__new__(GLMTTSForConditionalGeneration)
+
+    out = model.make_omni_output(
+        torch.randn(2, 4),
+        model_intermediate_buffer=[{}],
+        prompt_speech_token_len=[torch.tensor([9])],
+        glm_tts_text_token_len=[torch.tensor([5])],
+    )
+
+    assert out.multimodal_outputs is not None
+    assert out.multimodal_outputs["prompt_speech_token_len"][0].item() == 9
+    assert out.multimodal_outputs["glm_tts_text_token_len"][0].item() == 5
+
+
+def test_glm_tts_postprocess_caches_prompt_speech_token_len() -> None:
+    model = object.__new__(GLMTTSForConditionalGeneration)
+    model.config = SimpleNamespace(min_token_text_ratio=2.0, max_token_text_ratio=20.0)
+
+    update = model.postprocess(
+        torch.randn(2, 4),
+        multimodal_outputs={
+            "prompt_speech_token_len": [torch.tensor([13])],
+            "glm_tts_text_token_len": [torch.tensor([6])],
+        },
+    )
+
+    assert update["prompt_speech_token_len"].item() == 13
+    assert update["glm_tts_text_token_len"].item() == 6
diff --git a/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py b/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py
index d54533fd0..6debabb0e 100644
--- a/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py
+++ b/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py
@@ -6,13 +6,19 @@ from types import SimpleNamespace
 
 import torch
 
-from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import talker2code2wav_async_chunk, text2flow
+from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
+    talker2code2wav_async_chunk,
+    text2flow,
+    text2flow_full_payload,
+    text2flow_token_only,
+)
 
 
-def _source_output(request_id: str, prompt_ids: list[int], out_ids: list[int], mm: dict):
+def _source_output(request_id: str, prompt_ids: list[int], out_ids: list[int], mm: dict, finished: bool = True):
     return SimpleNamespace(
         request_id=request_id,
         prompt_token_ids=prompt_ids,
+        finished=finished,
         outputs=[SimpleNamespace(token_ids=out_ids, cumulative_token_ids=out_ids, multimodal_output=mm)],
     )
 
@@ -45,8 +51,8 @@ def _transfer_manager(
 
 def test_text2flow_supports_batched_source_outputs():
     source_outputs = [
-        _source_output("req-0", [10, 11], [1, 2, 3], {"speech_token": torch.tensor([[1, 2]])}),
-        _source_output("req-1", [20, 21], [4, 5], {"speech_token": torch.tensor([[3, 4]])}),
+        _source_output("req-0", [10, 11], [1, 2, 3], {"speech_token": torch.tensor([[8, 9]])}),
+        _source_output("req-1", [20, 21], [4, 5], {"speech_token": torch.tensor([[6, 7]])}),
     ]
 
     outputs = text2flow(source_outputs=source_outputs, prompt=None)
@@ -58,6 +64,65 @@ def test_text2flow_supports_batched_source_outputs():
     assert outputs[1]["additional_information"]["ids"]["prompt"] == [20, 21]
 
 
+def test_text2flow_strips_reference_speech_prefix_from_cumulative_ids():
+    source_outputs = [
+        _source_output("req-0", [10, 11], [8, 9, 1, 2, 3], {"speech_token": torch.tensor([[8, 9]])}),
+    ]
+
+    outputs = text2flow(source_outputs=source_outputs, prompt=None)
+
+    assert outputs[0]["prompt_token_ids"] == [1, 2, 3]
+
+
+def test_text2flow_token_only_strips_reference_speech_prefix_from_cumulative_ids():
+    source_outputs = [
+        _source_output(
+            "req-strip",
+            [10, 11],
+            [4, 5, 1, 2, 3],
+            {"embed": {"speech_token": torch.tensor([[4, 5]])}},
+        )
+    ]
+
+    outputs = text2flow_token_only(source_outputs=source_outputs, prompt=None)
+
+    assert len(outputs) == 1
+    assert outputs[0]["prompt_token_ids"] == [1, 2, 3]
+    assert outputs[0]["additional_information"]["ids"]["prompt"] == [10, 11]
+
+
+def test_text2flow_token_only_marks_prompt_trim_for_stop_token_completion():
+    source_outputs = [
+        _source_output(
+            "req-stop",
+            [10, 11],
+            [4, 5, 1, 2, 6562],
+            {"embed": {"speech_token": torch.tensor([[4, 5]])}},
+        )
+    ]
+
+    outputs = text2flow_token_only(source_outputs=source_outputs, prompt=None)
+
+    assert outputs[0]["prompt_token_ids"] == [1, 2, 6562]
+    assert outputs[0]["additional_information"]["meta"]["talker_prefill_offset"] == 2
+
+
+def test_text2flow_full_payload_does_not_send_codec_ids():
+    payload = text2flow_full_payload(
+        None,
+        {
+            "embed.speech_token": torch.tensor([[1, 2]], dtype=torch.long),
+            "codes.audio": torch.tensor([7, 8, 9], dtype=torch.long),
+        },
+        SimpleNamespace(),
+    )
+
+    assert payload is not None
+    assert "codes" not in payload
+    assert "next_stage_prompt_len" not in payload["meta"]
+    assert torch.equal(payload["embed"]["speech_token"], torch.tensor([[1, 2]], dtype=torch.long))
+
+
 def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():
     transfer_manager = _transfer_manager()
     request = SimpleNamespace(
diff --git a/tests/model_executor/stage_input_processors/test_fish_speech_async_chunk.py b/tests/model_executor/stage_input_processors/test_fish_speech_async_chunk.py
new file mode 100644
index 000000000..132615b14
--- /dev/null
+++ b/tests/model_executor/stage_input_processors/test_fish_speech_async_chunk.py
@@ -0,0 +1,229 @@
+from collections import defaultdict
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.stage_input_processors.fish_speech import (
+    slow_ar_to_dac_decoder_async_chunk,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_transfer_manager(
+    *,
+    tensor_payload: bool = False,
+    chunk_frames: int = 2,
+    left_context: int = 1,
+    initial_chunk_frames: int = 0,
+    single_initial_chunk: bool = False,
+    backlog_chunk_frames: int = 0,
+    backlog_load_threshold: float = 0.75,
+    max_num_seqs: int = 1,
+):
+    return SimpleNamespace(
+        code_prompt_token_ids=defaultdict(list),
+        put_req_chunk=defaultdict(int),
+        scheduler_max_num_seqs=max_num_seqs,
+        connector=SimpleNamespace(
+            config={
+                "extra": {
+                    "codec_chunk_frames": chunk_frames,
+                    "codec_left_context_frames": left_context,
+                    "initial_codec_chunk_frames": initial_chunk_frames,
+                    "fish_speech_tensor_codes": tensor_payload,
+                    "fish_speech_single_initial_chunk": single_initial_chunk,
+                    "fish_speech_backlog_codec_chunk_frames": backlog_chunk_frames,
+                    "fish_speech_backlog_load_threshold": backlog_load_threshold,
+                }
+            }
+        ),
+    )
+
+
+def _make_request(req_id="req", *, finished=False):
+    return SimpleNamespace(
+        external_req_id=req_id,
+        additional_information=None,
+        is_finished=lambda: finished,
+    )
+
+
+def _seed_frames(transfer_manager, req_id: str, n_frames: int, n_codebooks: int = 3):
+    frame = torch.arange(1, n_codebooks + 1, dtype=torch.long)
+    transfer_manager.code_prompt_token_ids[req_id] = [frame + i * n_codebooks for i in range(n_frames)]
+
+
+def _call_seeded(transfer_manager, req_id: str = "req", *, n_frames: int, finished: bool = False):
+    _seed_frames(transfer_manager, req_id, n_frames)
+    return slow_ar_to_dac_decoder_async_chunk(
+        transfer_manager,
+        {},
+        _make_request(req_id, finished=finished),
+        is_finished=finished,
+    )
+
+
+def test_slow_ar_to_dac_async_chunk_can_emit_tensor_codes():
+    transfer_manager = _make_transfer_manager(tensor_payload=True)
+    request = _make_request()
+
+    assert (
+        slow_ar_to_dac_decoder_async_chunk(
+            transfer_manager,
+            {"audio_codes": torch.tensor([[1, 2, 3]])},
+            request,
+        )
+        is None
+    )
+    payload = slow_ar_to_dac_decoder_async_chunk(
+        transfer_manager,
+        {"audio_codes": torch.tensor([[4, 5, 6]])},
+        request,
+    )
+
+    assert payload is not None
+    audio_codes = payload.codes.audio
+    assert isinstance(audio_codes, torch.Tensor)
+    assert audio_codes.dtype == torch.long
+    assert torch.equal(audio_codes, torch.tensor([[1, 4], [2, 5], [3, 6]], dtype=torch.long))
+    assert payload.meta.left_context_size == 0
+
+
+def test_slow_ar_to_dac_async_chunk_uses_valid_mask_for_zero_codes():
+    transfer_manager = _make_transfer_manager(tensor_payload=True, chunk_frames=1, left_context=0)
+    request = _make_request()
+
+    assert (
+        slow_ar_to_dac_decoder_async_chunk(
+            transfer_manager,
+            {
+                "audio_codes": torch.tensor([[0, 0, 0]]),
+                "audio_code_valid": torch.tensor([False]),
+            },
+            request,
+        )
+        is None
+    )
+    payload = slow_ar_to_dac_decoder_async_chunk(
+        transfer_manager,
+        {
+            "audio_codes": torch.tensor([[4, 5, 6]]),
+            "audio_code_valid": torch.tensor([True]),
+        },
+        request,
+    )
+
+    assert payload is not None
+    assert torch.equal(payload.codes.audio, torch.tensor([[4], [5], [6]], dtype=torch.long))
+
+
+def test_slow_ar_to_dac_async_chunk_keeps_flat_codes_by_default():
+    transfer_manager = _make_transfer_manager(tensor_payload=False)
+    request = _make_request()
+
+    slow_ar_to_dac_decoder_async_chunk(
+        transfer_manager,
+        {"audio_codes": torch.tensor([[1, 2]])},
+        request,
+    )
+    payload = slow_ar_to_dac_decoder_async_chunk(
+        transfer_manager,
+        {"audio_codes": torch.tensor([[3, 4]])},
+        request,
+    )
+
+    assert payload is not None
+    assert torch.equal(payload.codes.audio, torch.tensor([1, 3, 2, 4], dtype=torch.long))
+
+
+def test_qwen3_style_single_initial_chunk_only_emits_first_chunk_early():
+    transfer_manager = _make_transfer_manager(
+        chunk_frames=5,
+        left_context=5,
+        initial_chunk_frames=2,
+        single_initial_chunk=True,
+    )
+
+    first = _call_seeded(transfer_manager, n_frames=2)
+    assert first is not None
+    assert first.meta.left_context_size == 0
+    assert first.codes.audio.numel() == 3 * 2
+
+    # After the first small chunk, return to normal chunk boundaries:
+    # initial_coverage=2, so the next emit is at 2 + 5 = 7 frames.
+    assert _call_seeded(transfer_manager, n_frames=5) is None
+
+    second = _call_seeded(transfer_manager, n_frames=7)
+    assert second is not None
+    assert second.meta.left_context_size == 2
+    assert second.codes.audio.numel() == 3 * 7
+
+
+def test_backlog_chunk_size_uses_larger_post_initial_boundary():
+    transfer_manager = _make_transfer_manager(
+        chunk_frames=5,
+        left_context=5,
+        initial_chunk_frames=2,
+        single_initial_chunk=True,
+        backlog_chunk_frames=8,
+        backlog_load_threshold=0.5,
+        max_num_seqs=4,
+    )
+    for idx in range(2):
+        _seed_frames(transfer_manager, f"other-{idx}", 10)
+
+    first = _call_seeded(transfer_manager, "req", n_frames=2)
+    assert first is not None
+
+    # Under backlog, steady chunk is 8 instead of 5. The next emit is
+    # initial_coverage=2 plus 8 frames, so n=7 must still be held.
+    assert _call_seeded(transfer_manager, "req", n_frames=7) is None
+    second = _call_seeded(transfer_manager, "req", n_frames=10)
+    assert second is not None
+    assert second.meta.left_context_size == 2
+    assert second.codes.audio.numel() == 3 * 10
+
+
+def test_backlog_chunk_size_keeps_multi_initial_boundary_aligned():
+    transfer_manager = _make_transfer_manager(
+        chunk_frames=25,
+        left_context=25,
+        initial_chunk_frames=10,
+        single_initial_chunk=False,
+        backlog_chunk_frames=50,
+        backlog_load_threshold=0.5,
+        max_num_seqs=4,
+    )
+    for idx in range(2):
+        _seed_frames(transfer_manager, f"other-{idx}", 50)
+
+    first = _call_seeded(transfer_manager, "req", n_frames=10)
+    assert first is not None
+    transfer_manager.put_req_chunk["req"] += 1
+    second = _call_seeded(transfer_manager, "req", n_frames=20)
+    assert second is not None
+    transfer_manager.put_req_chunk["req"] += 1
+
+    assert _call_seeded(transfer_manager, "req", n_frames=50) is None
+    third = _call_seeded(transfer_manager, "req", n_frames=70)
+    assert third is not None
+    assert third.meta.left_context_size == 20
+    assert third.codes.audio.numel() == 3 * 70
+
+
+def test_invalid_backlog_config_fails_loudly():
+    transfer_manager = _make_transfer_manager(
+        chunk_frames=5,
+        left_context=5,
+        backlog_chunk_frames="invalid",
+    )
+    request = _make_request()
+
+    with pytest.raises(ValueError, match="fish_speech_backlog_codec_chunk_frames"):
+        slow_ar_to_dac_decoder_async_chunk(
+            transfer_manager,
+            {"audio_codes": torch.tensor([[1, 2, 3]])},
+            request,
+        )
diff --git a/tests/model_executor/stage_input_processors/test_glm_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_glm_tts_async_chunk.py
new file mode 100644
index 000000000..19bdaf7fe
--- /dev/null
+++ b/tests/model_executor/stage_input_processors/test_glm_tts_async_chunk.py
@@ -0,0 +1,278 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.engine.serialization import serialize_additional_information
+from vllm_omni.model_executor.models.glm_tts.glm_tts_dit_wrapper import (
+    split_request_ids,
+    valid_speech_tokens,
+)
+from vllm_omni.model_executor.stage_input_processors.glm_tts import (
+    ar_to_dit,
+    ar_to_dit_async_chunk,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _Request:
+    external_req_id = "req-1"
+    additional_information = None
+
+    def is_finished(self) -> bool:
+        return False
+
+
+def _transfer_manager() -> SimpleNamespace:
+    connector = SimpleNamespace(
+        config={
+            "extra": {
+                "codec_chunk_frames": 25,
+                "codec_left_context_frames": 25,
+            }
+        }
+    )
+    return SimpleNamespace(connector=connector)
+
+
+def _progressive_transfer_manager() -> SimpleNamespace:
+    connector = SimpleNamespace(
+        config={
+            "extra": {
+                "codec_chunk_frames": [25, 50, 200],
+                "codec_left_context_frames": 25,
+            }
+        }
+    )
+    return SimpleNamespace(connector=connector)
+
+
+def _source_output(mm: dict, request_id: str = "req-0") -> SimpleNamespace:
+    return SimpleNamespace(
+        request_id=request_id,
+        outputs=[
+            SimpleNamespace(
+                multimodal_output=mm,
+            )
+        ],
+    )
+
+
+def _audio_codes(payload) -> list[int]:
+    assert payload.codes is not None
+    assert payload.codes.audio is not None
+    return payload.codes.audio.tolist()
+
+
+def test_ar_to_dit_filters_invalid_tokens_and_preserves_conditioning() -> None:
+    prompt_token = torch.tensor([[1, 2, 3]])
+    prompt_feat = torch.randn(1, 4, 80)
+    embedding = torch.randn(1, 192)
+    source_outputs = [
+        _source_output(
+            {
+                "speech_tokens": torch.tensor([[-1, 0, 7, -1, 32767]]),
+                "prompt_token": prompt_token,
+                "prompt_feat": prompt_feat,
+                "embedding": embedding,
+            }
+        )
+    ]
+
+    outputs = ar_to_dit(source_outputs)
+
+    assert len(outputs) == 1
+    assert outputs[0]["prompt_token_ids"] == [0, 7, 32767]
+    additional_info = outputs[0]["additional_information"]
+    assert additional_info["speech_tokens"] == [0, 7, 32767]
+    # Voice clone data stored in nested embed dict (matches async_chunk path)
+    embed = additional_info["embed"]
+    assert embed["speech_token"] is prompt_token
+    assert embed["speech_feat"] is prompt_feat
+    assert embed["embedding"] is embedding
+    # Meta dict enables block-causal attention in DiT wrapper
+    meta = additional_info["meta"]
+    assert meta["req_id"] == ["req-0"]
+    assert meta["left_context_size"] == 0
+    assert meta["stream_finished"] is True
+    private = additional_info["kv_metadata"]["glm_tts"]
+    assert private["block_pattern"] == [3]  # len([0, 7, 32767])
+    assert private["chunk_sizes_history"] == [3]
+
+
+def test_dit_explicit_empty_speech_tokens_override_scheduler_placeholder() -> None:
+    placeholder = torch.tensor([0], dtype=torch.long)
+
+    token = valid_speech_tokens([], device=placeholder.device, fallback=placeholder)
+
+    assert token.numel() == 0
+
+
+def test_ar_to_dit_async_chunk_restores_serialized_conditioning_payload() -> None:
+    transfer = _transfer_manager()
+    request = _Request()
+    prompt_token = torch.tensor([[1, 2, 3]])
+    prompt_feat = torch.randn(1, 4, 80)
+    embedding = torch.randn(1, 192)
+    request.additional_information = serialize_additional_information(
+        {
+            "prompt_speech_token": prompt_token,
+            "prompt_feat": prompt_feat,
+            "embedding": embedding,
+        }
+    )
+
+    payload = ar_to_dit_async_chunk(
+        transfer,
+        {"speech_tokens": torch.tensor([[0]])},
+        request,
+        is_finished=False,
+    )
+
+    assert payload is None
+    state = transfer.request_payload["req-1"]["_glm_tts_async_state"]
+    prompt_payload = state["prompt_payload"]
+    assert torch.equal(prompt_payload["prompt_speech_token"], prompt_token)
+    assert torch.equal(prompt_payload["prompt_feat"], prompt_feat)
+    assert torch.equal(prompt_payload["embedding"], embedding)
+
+
+def test_split_request_ids_uses_seq_token_counts() -> None:
+    ids = torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.long)
+
+    split = split_request_ids(ids, seq_token_counts=[2, 3, 1])
+
+    assert [part.tolist() for part in split] == [[1, 2], [3, 4, 5], [6]]
+
+
+def test_ar_to_dit_async_chunk_emits_cumulative_prefixes() -> None:
+    transfer = _transfer_manager()
+    request = _Request()
+    payloads = []
+
+    for token in range(75):
+        payload = ar_to_dit_async_chunk(
+            transfer,
+            {"speech_tokens": torch.tensor([[token]])},
+            request,
+            is_finished=False,
+        )
+        if payload is not None:
+            payloads.append(payload)
+
+    assert len(payloads) == 3
+    assert _audio_codes(payloads[0]) == list(range(25))
+    assert payloads[0].meta.left_context_size == 0
+    assert payloads[0].kv_metadata["glm_tts"]["block_pattern"] == [25]
+    assert _audio_codes(payloads[1]) == list(range(50))
+    assert payloads[1].meta.left_context_size == 25
+    assert payloads[1].kv_metadata["glm_tts"]["block_pattern"] == [25]
+    assert _audio_codes(payloads[2]) == list(range(75))
+    assert payloads[2].meta.left_context_size == 50
+    assert payloads[2].kv_metadata["glm_tts"]["block_pattern"] == [25]
+
+
+def test_ar_to_dit_async_chunk_keeps_progressive_block_pattern_fixed() -> None:
+    transfer = _progressive_transfer_manager()
+    request = _Request()
+    payloads = []
+
+    for token in range(30):
+        payload = ar_to_dit_async_chunk(
+            transfer,
+            {"speech_tokens": torch.tensor([[token]])},
+            request,
+            is_finished=False,
+        )
+        if payload is not None:
+            payloads.append(payload)
+
+    final_payload = ar_to_dit_async_chunk(
+        transfer,
+        None,
+        request,
+        is_finished=True,
+    )
+
+    assert len(payloads) == 1
+    assert payloads[0].kv_metadata["glm_tts"]["chunk_sizes_history"] == [25]
+    assert payloads[0].kv_metadata["glm_tts"]["block_pattern"] == [25, 50, 200]
+    assert final_payload is not None
+    assert final_payload.kv_metadata["glm_tts"]["chunk_sizes_history"] == [25, 5]
+    assert final_payload.kv_metadata["glm_tts"]["block_pattern"] == [25, 50, 200]
+
+
+def test_ar_to_dit_async_chunk_terminal_boundary_is_metadata_only() -> None:
+    transfer = _transfer_manager()
+    request = _Request()
+    payloads = []
+
+    for token in range(50):
+        payload = ar_to_dit_async_chunk(
+            transfer,
+            {"speech_tokens": torch.tensor([[token]])},
+            request,
+            is_finished=False,
+        )
+        if payload is not None:
+            payloads.append(payload)
+
+    final_payload = ar_to_dit_async_chunk(
+        transfer,
+        None,
+        request,
+        is_finished=True,
+    )
+
+    assert len(payloads) == 2
+    assert final_payload is not None
+    assert _audio_codes(final_payload) == []
+    assert final_payload.meta.left_context_size == 50
+    assert final_payload.meta.req_id == ["req-1"]
+    assert bool(final_payload.meta.stream_finished.item())
+
+
+def test_ar_to_dit_async_chunk_finished_flushes_partial_unsent_tail() -> None:
+    transfer = _transfer_manager()
+    request = _Request()
+    payloads = []
+
+    for token in range(30):
+        payload = ar_to_dit_async_chunk(
+            transfer,
+            {"speech_tokens": torch.tensor([[token]])},
+            request,
+            is_finished=False,
+        )
+        if payload is not None:
+            payloads.append(payload)
+
+    final_payload = ar_to_dit_async_chunk(
+        transfer,
+        None,
+        request,
+        is_finished=True,
+    )
+
+    assert len(payloads) == 1
+    assert _audio_codes(payloads[0]) == list(range(25))
+    assert payloads[0].meta.left_context_size == 0
+    assert final_payload is not None
+    assert _audio_codes(final_payload) == list(range(30))
+    assert final_payload.meta.left_context_size == 25
+    assert final_payload.meta.req_id == ["req-1"]
+    assert bool(final_payload.meta.finished.item())
+    assert bool(final_payload.meta.stream_finished.item())
+    assert (
+        ar_to_dit_async_chunk(
+            transfer,
+            None,
+            request,
+            is_finished=True,
+        )
+        is None
+    )
diff --git a/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py b/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py
new file mode 100644
index 000000000..0cb61a972
--- /dev/null
+++ b/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py
@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Light coverage for qwen2_5_omni.thinker2talker_full_payload.
+
+Covers the finish-reason-aware stop-row trim contract: when the request
+status is FINISHED_STOPPED, the builder must drop one row from the
+accumulated hidden states (vLLM v1 appends the sampled stop token to
+output_token_ids before check_stop, so the trailing hidden-state row
+corresponds to the stop emission and must not reach the talker).
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import (
+    thinker2talker_full_payload,
+)
+
+
+def _make_request(
+    prompt_token_ids,
+    output_token_ids,
+    status_name: str | None = "FINISHED_STOPPED",
+):
+    status = SimpleNamespace(name=status_name) if status_name else None
+    return SimpleNamespace(
+        request_id="r1",
+        prompt_token_ids=prompt_token_ids,
+        output_token_ids=output_token_ids,
+        all_token_ids=list(prompt_token_ids) + list(output_token_ids),
+        status=status,
+        sampling_params=None,
+    )
+
+
+def test_finished_stopped_trims_one_decode_row():
+    """FINISHED_STOPPED: drop trailing hidden-state row so talker does not
+    consume the stop-emission row.
+    """
+    prompt = [1, 2, 3]
+    output = [10, 11, 12]
+    request = _make_request(prompt, output, status_name="FINISHED_STOPPED")
+    # 6 prompt+output rows + 1 stop-emission row = 7 hidden rows total.
+    hidden = torch.arange(7 * 4, dtype=torch.float32).reshape(7, 4)
+    pooling = {"hidden": hidden}
+
+    payload = thinker2talker_full_payload(transfer_manager=None, pooling_output=pooling, request=request)
+
+    assert payload is not None
+    # ids.output had one trailing stop row dropped: 3 - 1 = 2 remaining.
+    assert payload["ids"]["output"] == output[:-1]
+    # embed.prefill must cover only the prompt rows.
+    assert payload["embed"]["prefill"].shape[0] == len(prompt)
+    # hidden_states.output covers the decode rows minus the dropped stop row.
+    assert payload["hidden_states"]["output"].shape[0] == len(output) - 1
+
+
+def test_finished_length_capped_keeps_all_rows():
+    """FINISHED_LENGTH_CAPPED: no row drop; hidden_states.output covers
+    all decode rows.
+    """
+    prompt = [1, 2, 3]
+    output = [10, 11, 12]
+    request = _make_request(prompt, output, status_name="FINISHED_LENGTH_CAPPED")
+    hidden = torch.arange(6 * 4, dtype=torch.float32).reshape(6, 4)
+    pooling = {"hidden": hidden}
+
+    payload = thinker2talker_full_payload(transfer_manager=None, pooling_output=pooling, request=request)
+
+    assert payload is not None
+    assert payload["ids"]["output"] == output
+    assert payload["embed"]["prefill"].shape[0] == len(prompt)
+    assert payload["hidden_states"]["output"].shape[0] == len(output)
+
+
+def test_missing_hidden_returns_none():
+    """Defensive: pooling_output without "hidden" returns None."""
+    request = _make_request([1, 2], [3], status_name="FINISHED_STOPPED")
+    assert thinker2talker_full_payload(transfer_manager=None, pooling_output={}, request=request) is None
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py
index f11a4654e..4b2bdf3fe 100644
--- a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py
+++ b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py
@@ -20,7 +20,7 @@ def _streaming_context() -> SimpleNamespace:
 
 
 def test_get_streaming_talker_tokens_first_segment(_streaming_context: SimpleNamespace) -> None:
-    inc_p, inc_o, merged, thinker_in = q3._get_streaming_talker_tokens(
+    inc_p, inc_o = q3._get_streaming_talker_tokens(
         "r1",
         [1, 2],
         [10, 11],
@@ -28,13 +28,11 @@ def test_get_streaming_talker_tokens_first_segment(_streaming_context: SimpleNam
     )
     assert inc_p == [1, 2]
     assert inc_o == [10, 11]
-    assert merged == [1, 2, 10, 11]
-    assert thinker_in == [1, 2]
 
 
 def test_get_streaming_talker_tokens_second_segment_accumulates(_streaming_context: SimpleNamespace) -> None:
     q3._get_streaming_talker_tokens("r2", [1, 2], [10, 11], streaming_context=_streaming_context)
-    inc_p, inc_o, merged, thinker_in = q3._get_streaming_talker_tokens(
+    inc_p, inc_o = q3._get_streaming_talker_tokens(
         "r2",
         [1, 2, 3, 4],
         [10, 11, 12, 13],
@@ -42,14 +40,12 @@ def test_get_streaming_talker_tokens_second_segment_accumulates(_streaming_conte
     )
     assert inc_p == [3, 4]
     assert inc_o == [12, 13]
-    assert merged == [1, 2, 10, 3, 4, 12, 13]
-    assert thinker_in == [1, 2, 10, 3, 4]
 
 
 def test_get_streaming_talker_tokens_new_prompt_len_snapshot_truncates(
     _streaming_context: SimpleNamespace,
 ) -> None:
-    inc_p, inc_o, merged, thinker_in = q3._get_streaming_talker_tokens(
+    inc_p, inc_o = q3._get_streaming_talker_tokens(
         "r3",
         [1, 2, 3, 4, 5, 6],
         [10],
@@ -58,8 +54,6 @@ def test_get_streaming_talker_tokens_new_prompt_len_snapshot_truncates(
     )
     assert inc_p == [1, 2, 3, 4]
     assert inc_o == [10]
-    assert merged == [1, 2, 3, 4, 10]
-    assert thinker_in == [1, 2, 3, 4]
 
 
 def test_get_streaming_talker_tokens_clear_state(_streaming_context: SimpleNamespace) -> None:
@@ -82,6 +76,72 @@ def test_get_streaming_codec_delta_len_increments_and_finishes(_streaming_contex
     assert state.talker2code2wav_last_seq_len == 0
 
 
+def test_streaming_input_prefill_chunk_is_cached() -> None:
+    transfer_manager = SimpleNamespace(_pending_streaming_prefills={})
+    request = SimpleNamespace(
+        external_req_id="rt-1",
+        output_token_ids=[100],
+        all_token_ids=[151644, 872, 100],
+        prompt_token_ids=[151644, 872],
+        resumable=True,
+        additional_information=None,
+    )
+    thinker_emb = torch.ones(2, 3)
+    thinker_hid = torch.full((2, 3), 2.0)
+
+    payload = q3._construct_thinker2talker_streaming_input_async_chunk(
+        False,
+        request,
+        thinker_emb,
+        thinker_hid,
+        transfer_manager,
+    )
+
+    assert payload is None
+    cached = transfer_manager._pending_streaming_prefills["rt-1"]
+    assert torch.equal(cached["embed"]["prefill"], thinker_emb)
+    assert torch.equal(cached["hidden_states"]["output"], thinker_hid)
+    assert cached["ids"]["all"] == request.all_token_ids
+    assert cached["ids"]["prompt"] == request.prompt_token_ids
+
+
+def test_streaming_input_prefill_flushes_with_next_decode_chunk() -> None:
+    transfer_manager = SimpleNamespace(
+        _pending_streaming_prefills={
+            "rt-2": {
+                "embed": {"prefill": torch.ones(2, 3)},
+                "hidden_states": {"output": torch.full((2, 3), 2.0)},
+                "ids": {"all": [151644, 872, 100], "prompt": [151644, 872]},
+            }
+        }
+    )
+    request = SimpleNamespace(
+        external_req_id="rt-2",
+        output_token_ids=[101],
+        all_token_ids=[151644, 872, 100, 101],
+        prompt_token_ids=[151644, 872],
+        resumable=True,
+        additional_information=None,
+    )
+    thinker_emb = torch.full((1, 3), 3.0)
+    thinker_hid = torch.full((1, 3), 4.0)
+
+    payload = q3._construct_thinker2talker_streaming_input_async_chunk(
+        False,
+        request,
+        thinker_emb,
+        thinker_hid,
+        transfer_manager,
+    )
+
+    assert payload is not None
+    assert payload.embed.prefill.shape == (3, 3)
+    assert payload.hidden_states.output.shape == (3, 3)
+    assert payload.ids.all == [151644, 872, 100]
+    assert payload.ids.prompt == [151644, 872]
+    assert "rt-2" not in transfer_manager._pending_streaming_prefills
+
+
 def test_talker2code2wav_full_payload_filters_by_output_token_ids() -> None:
     request = SimpleNamespace(
         request_id="codec",
@@ -168,6 +228,7 @@ def test_talker2code2wav_full_payload_keeps_all_zero_codec_rows() -> None:
 
 
 def test_thinker2talker_full_payload_packs_complete_tensors() -> None:
+    """Full-payload path drops the terminal thinker row before talker prefill."""
     request = SimpleNamespace(
         request_id="thinker",
         prompt_token_ids=[151644, 872],
@@ -187,3 +248,741 @@ def test_thinker2talker_full_payload_packs_complete_tensors() -> None:
     assert payload["embed"]["prefill"].device.type == "cpu"
     assert payload["hidden_states"]["output"].device.type == "cpu"
     assert payload["next_stage_prompt_len"] > 0
+    assert payload["embed"]["prefill"].shape[0] == 2
+    assert payload["hidden_states"]["output"].shape[0] == 2
+
+
+def test_accumulator_replaces_keys_in_replace_set() -> None:
+    """REPLACE-key semantics: subsequent emissions of the same key replace, not append."""
+    from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
+
+    class _StubMixin(OmniConnectorModelRunnerMixin):
+        def __init__(self):
+            self._pending_full_payload_send = {}
+            self._full_payload_replace_keys_cached = frozenset({"model_outputs"})
+
+    stub = _StubMixin()
+    stub.accumulate_full_payload_output(
+        "req1",
+        {
+            "model_outputs": torch.tensor([[1.0, 2.0]]),
+            "hidden_states.output": torch.tensor([[10.0]]),
+        },
+        request=None,
+    )
+    stub.accumulate_full_payload_output(
+        "req1",
+        {
+            "model_outputs": torch.tensor([[3.0, 4.0]]),
+            "hidden_states.output": torch.tensor([[20.0]]),
+        },
+        request=None,
+    )
+    output, _ = stub._materialize_full_payload_entry(stub._pending_full_payload_send["req1"])
+    # model_outputs REPLACED (second value only):
+    assert torch.equal(output["model_outputs"], torch.tensor([[3.0, 4.0]]))
+    # hidden_states.output CONCATENATED:
+    assert torch.equal(output["hidden_states.output"], torch.tensor([[10.0], [20.0]]))
+
+
+def test_accumulator_concat_default_when_no_replace_keys() -> None:
+    """Default semantics: 2-D+ tensors concat across emissions when not in replace_keys."""
+    from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
+
+    class _StubMixin(OmniConnectorModelRunnerMixin):
+        def __init__(self):
+            self._pending_full_payload_send = {}
+            self._full_payload_replace_keys_cached = frozenset()
+
+    stub = _StubMixin()
+    stub.accumulate_full_payload_output(
+        "req1",
+        {"embed.prefill": torch.tensor([[1.0]])},
+        request=None,
+    )
+    stub.accumulate_full_payload_output(
+        "req1",
+        {"embed.prefill": torch.tensor([[2.0]])},
+        request=None,
+    )
+    output, _ = stub._materialize_full_payload_entry(stub._pending_full_payload_send["req1"])
+    assert torch.equal(output["embed.prefill"], torch.tensor([[1.0], [2.0]]))
+
+
+def test_covo_audio_llm2code2wav_token_only_smoke() -> None:
+    """Smoke: covo_audio token-only builder returns placeholder prompts sized to audio_codes count."""
+    # source_outputs is a list of objects with .outputs[0].token_ids
+    from vllm_omni.model_executor.models.covo_audio.config_covo_audio import COVO_AUDIO_TOKEN_INDEX
+    from vllm_omni.model_executor.stage_input_processors.covo_audio import (
+        llm2code2wav_token_only,
+    )
+
+    class _Out:
+        def __init__(self, tids):
+            self.token_ids = tids
+
+    class _Wrapper:
+        def __init__(self, tids):
+            self.outputs = [_Out(tids)]
+
+    # 3 codec tokens + 2 non-codec
+    src = [_Wrapper([COVO_AUDIO_TOKEN_INDEX + 0, COVO_AUDIO_TOKEN_INDEX + 1, COVO_AUDIO_TOKEN_INDEX + 2, 100, 200])]
+    out = llm2code2wav_token_only(src)
+    assert len(out) == 1
+    assert len(out[0]["prompt_token_ids"]) == 3
+    assert out[0]["additional_information"] is None
+
+
+def test_covo_audio_llm2code2wav_full_payload_smoke() -> None:
+    """Smoke: covo_audio producer-side payload builder returns audio_codes + finished."""
+    from types import SimpleNamespace
+
+    from vllm_omni.model_executor.models.covo_audio.config_covo_audio import COVO_AUDIO_TOKEN_INDEX
+    from vllm_omni.model_executor.stage_input_processors.covo_audio import (
+        llm2code2wav_full_payload,
+    )
+
+    req = SimpleNamespace(
+        output_token_ids=[COVO_AUDIO_TOKEN_INDEX + 5, COVO_AUDIO_TOKEN_INDEX + 6, 99],
+    )
+    payload = llm2code2wav_full_payload(None, {}, req)
+    assert payload is not None
+    assert payload["codes"]["audio"] == [5, 6]
+    assert payload["meta"]["finished"].item() is True
+
+
+def test_dynin_omni_token_only_smoke() -> None:
+    """Smoke: dynin_omni token-only builders return placeholders."""
+    from vllm_omni.model_executor.stage_input_processors.dynin_omni import (
+        token2text_to_token2image_token_only,
+    )
+
+    class _Out:
+        def __init__(self, tids, mm=None):
+            self.token_ids = tids
+            self.multimodal_output = mm
+
+    class _Wrapper:
+        def __init__(self, tids, mm=None):
+            self.outputs = [_Out(tids, mm)]
+            self.request_id = "r0"
+
+    class _Stage:
+        def __init__(self, outs):
+            self.engine_outputs = outs
+
+    src = [_Wrapper([10, 11, 12])]
+    out = token2text_to_token2image_token_only([_Stage(src)], [0])
+    assert len(out) == 1
+    assert len(out[0]["prompt_token_ids"]) == 3
+    assert out[0]["additional_information"] is None
+
+
+def test_dynin_omni_full_payload_smoke() -> None:
+    """Smoke: dynin_omni producer-side payload builder returns nested OmniPayload + carries metadata."""
+    from types import SimpleNamespace
+
+    from vllm_omni.model_executor.stage_input_processors.dynin_omni import (
+        token2text_to_token2image_full_payload,
+    )
+
+    pooling = {"token_ids": [1, 2, 3]}
+    req = SimpleNamespace(output_token_ids=[], additional_information={"speaker": ["alice"]})
+    payload = token2text_to_token2image_full_payload(None, pooling, req)
+    assert payload is not None
+    assert payload["codes"]["audio"] == [1, 2, 3]
+    assert payload["meta"]["finished"].item() is True
+    # additional_information is normalized + carried forward (speaker stays list-wrapped).
+    assert payload.get("speaker") == ["alice"]
+
+
+def test_qwen2_5_omni_talker2code2wav_token_only_smoke() -> None:
+    """Smoke: qwen2_5_omni talker→code2wav token_only marker + boundary strip."""
+    from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import (
+        TALKER_CODEC_END_TOKEN_ID,
+        TALKER_CODEC_START_TOKEN_ID,
+        talker2code2wav_token_only,
+    )
+
+    class _Out:
+        def __init__(self, tids):
+            self.cumulative_token_ids = tids
+
+    class _Wrap:
+        def __init__(self, tids):
+            self.outputs = [_Out(tids)]
+
+    # 3 inner codes wrapped by START + END
+    src = [_Wrap([TALKER_CODEC_START_TOKEN_ID, 10, 11, 12, TALKER_CODEC_END_TOKEN_ID])]
+    out = talker2code2wav_token_only(src)
+    assert len(out) == 1
+    assert len(out[0]["prompt_token_ids"]) == 3
+    assert out[0]["additional_information"] is None
+
+
+def test_qwen2_5_omni_talker2code2wav_full_payload_smoke() -> None:
+    """Smoke: qwen2_5_omni producer-side payload builder strips boundaries."""
+    from types import SimpleNamespace
+
+    from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import (
+        TALKER_CODEC_END_TOKEN_ID,
+        TALKER_CODEC_START_TOKEN_ID,
+        talker2code2wav_full_payload,
+    )
+
+    req = SimpleNamespace(
+        output_token_ids=[TALKER_CODEC_START_TOKEN_ID, 5, 6, 7, TALKER_CODEC_END_TOKEN_ID],
+    )
+    payload = talker2code2wav_full_payload(None, {}, req)
+    assert payload is not None
+    assert payload["codes"]["audio"] == [5, 6, 7]
+    assert payload["meta"]["finished"].item() is True
+
+
+def test_qwen2_5_omni_talker2code2wav_filters_control_tokens_and_placeholders() -> None:
+    """Qwen2.5 code2wav receives codec ids only, not talker prompt/control ids."""
+    from types import SimpleNamespace
+
+    from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import (
+        TALKER_CODEC_END_TOKEN_ID,
+        TALKER_CODEC_PAD_TOKEN_ID,
+        TALKER_CODEC_START_TOKEN_ID,
+        talker2code2wav_full_payload,
+        talker2code2wav_token_only,
+    )
+
+    class _Out:
+        def __init__(self, tids):
+            self.cumulative_token_ids = tids
+
+    class _Wrap:
+        def __init__(self, tids):
+            self.outputs = [_Out(tids)]
+
+    raw_ids = [
+        TALKER_CODEC_START_TOKEN_ID,
+        TALKER_CODEC_PAD_TOKEN_ID,
+        5,
+        6,
+        TALKER_CODEC_END_TOKEN_ID,
+        -1,
+        -1,
+    ]
+
+    token_only = talker2code2wav_token_only([_Wrap(raw_ids)])
+    assert len(token_only) == 1
+    assert len(token_only[0]["prompt_token_ids"]) == 4
+
+    payload = talker2code2wav_full_payload(None, {}, SimpleNamespace(output_token_ids=raw_ids))
+    assert payload is not None
+    assert payload["codes"]["audio"] == [5, 6, 6, 6]
+    assert payload["meta"]["finished"].item() is True
+
+
+def test_mimo_audio_llm2code2wav_token_only_smoke() -> None:
+    """Smoke: mimo_audio token-only builder sizes prompt."""
+    import torch
+
+    from vllm_omni.model_executor.stage_input_processors.mimo_audio import (
+        llm2code2wav_token_only,
+    )
+
+    class _Out:
+        def __init__(self, mm):
+            self.multimodal_output = mm
+
+    class _Wrap:
+        def __init__(self, mm):
+            self.outputs = [_Out(mm)]
+
+    # 3 batch rows of [1, 8, 4]: prepend_and_flatten_colmajor → 3*1*4*9 = 108
+    codes = torch.arange(96, dtype=torch.long).reshape(3, 1, 8, 4)
+    codes = codes.clamp(min=1)  # ensure nonzero so zero-row filter doesn't drop them
+    src = [_Wrap({"codes": {"audio": codes}})]
+    out = llm2code2wav_token_only(src)
+    assert len(out) == 1
+    assert len(out[0]["prompt_token_ids"]) == 108
+    assert out[0]["additional_information"] is None
+
+
+def test_mimo_audio_llm2code2wav_full_payload_smoke() -> None:
+    """Smoke: mimo_audio producer-side payload builder reads flat codes.audio + flattens."""
+    from types import SimpleNamespace
+
+    import torch
+
+    from vllm_omni.model_executor.stage_input_processors.mimo_audio import (
+        TALKER_CODEC_PAD_TOKEN_ID,
+        llm2code2wav_full_payload,
+    )
+
+    # Simulate accumulator output: 2 steps of [1, 1, 8, 4] CONCAT'd → [2, 1, 8, 4]
+    audio = torch.arange(2 * 1 * 8 * 4, dtype=torch.long).reshape(2, 1, 8, 4)
+    audio = audio.clamp(min=1)  # avoid zero-row drop
+    pooling_output = {"codes.audio": audio}
+    req = SimpleNamespace(output_token_ids=[])
+    payload = llm2code2wav_full_payload(None, pooling_output, req)
+    assert payload is not None
+    assert "codes" in payload and "audio" in payload["codes"]
+    # Flattened length = numel + B*4 (per-batch pad_vec prepended by prepend_and_flatten_colmajor)
+    batch_size = int(audio.shape[0])
+    assert len(payload["codes"]["audio"]) == audio.numel() + batch_size * 4
+    # prepend_and_flatten_colmajor: PAD appears at column start in col-major flatten.
+    # For shape [B=2, 1, 9, 4], each column has 1 PAD then 8 codec vals → PAD at indices 0, 9, 18, 27.
+    out = payload["codes"]["audio"]
+    assert out[0] == TALKER_CODEC_PAD_TOKEN_ID
+    assert out[9] == TALKER_CODEC_PAD_TOKEN_ID
+    assert payload["meta"]["finished"].item() is True
+
+
+def test_mimo_audio_full_payload_nested_fallback() -> None:
+    """Back-compat: full_payload still works if runtime returns nested codes.audio."""
+    from types import SimpleNamespace
+
+    import torch
+
+    from vllm_omni.model_executor.stage_input_processors.mimo_audio import (
+        llm2code2wav_full_payload,
+    )
+
+    audio = torch.arange(1 * 1 * 8 * 4, dtype=torch.long).reshape(1, 1, 8, 4)
+    audio = audio.clamp(min=1)
+    pooling_output = {"codes": {"audio": audio}}  # nested, not flat
+    req = SimpleNamespace(output_token_ids=[])
+    payload = llm2code2wav_full_payload(None, pooling_output, req)
+    assert payload is not None
+    assert len(payload["codes"]["audio"]) == audio.numel() + int(audio.shape[0]) * 4
+
+
+def test_qwen3_tts_talker2code2wav_token_only_smoke() -> None:
+    """Smoke: qwen3_tts token-only sizes placeholder."""
+    import torch
+
+    from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
+        talker2code2wav_token_only,
+    )
+
+    class _Out:
+        def __init__(self, mm, tids):
+            self.multimodal_output = mm
+            self.cumulative_token_ids = tids
+
+    class _Wrap:
+        def __init__(self, mm, tids):
+            self.outputs = [_Out(mm, tids)]
+            self.finished = True
+
+    # 3 valid codec frames Q=16; non-zero & under codebook size
+    audio = torch.arange(3 * 16, dtype=torch.long).reshape(3, 16) + 1
+    mm = {"codes": {"audio": audio}}
+    src = [_Wrap(mm, list(range(10)))]  # seq_len = 9; 3 < 9, no trim
+    out = talker2code2wav_token_only(src)
+    assert len(out) == 1
+    # Codebook-major flat: 16 * 3 = 48
+    assert len(out[0]["prompt_token_ids"]) == 48
+
+
+def test_qwen3_tts_talker2code2wav_full_payload_smoke() -> None:
+    """Smoke: qwen3_tts full_payload reads flat codes.audio + flattens codebook-major."""
+    from types import SimpleNamespace
+
+    import torch
+
+    from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
+        talker2code2wav_full_payload,
+    )
+
+    # 3 valid codec frames [3, 16] CONCAT'd from per-step emits via flatten
+    audio = torch.arange(3 * 16, dtype=torch.long).reshape(3, 16) + 1
+    pooling_output = {"codes.audio": audio}
+    req = SimpleNamespace(output_token_ids=list(range(10)))  # seq_len = 9
+    payload = talker2code2wav_full_payload(None, pooling_output, req)
+    assert payload is not None
+    assert "codes" in payload and "audio" in payload["codes"]
+    # codebook-major: shape [3, 16] -> [16, 3] -> flatten = 48 entries
+    assert isinstance(payload["codes"]["audio"], torch.Tensor)
+    assert payload["codes"]["audio"].shape == (48,)
+    expected = audio.transpose(0, 1).reshape(-1)
+    assert torch.equal(payload["codes"]["audio"], expected)
+    assert payload["meta"]["finished"].item() is True
+
+
+def test_qwen3_tts_full_payload_with_ref_code() -> None:
+    """Exact: ref_code is prepended (not appended) to audio, ref_code_len trims
+    ref, and the flatten is codebook-major.  Protects against ref-append-position
+    regressions, ref_code_len-not-applied bugs, and flatten-order regressions."""
+    from types import SimpleNamespace
+
+    import torch
+
+    from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
+        talker2code2wav_full_payload,
+    )
+
+    # Audio: 3 frames [3, 16] (no filter drops these — all positive, in-range).
+    audio = torch.arange(3 * 16, dtype=torch.long).reshape(3, 16) + 1
+    # Ref code: 2 frames [2, 16] (already 2-D), distinct value range so we can
+    # detect the prepend ordering.
+    ref = torch.arange(2 * 16, dtype=torch.long).reshape(2, 16) + 100
+    pooling_output = {
+        "codes.audio": audio,
+        "codes.ref": [ref],
+        "meta.ref_code_len": torch.tensor([2], dtype=torch.int32),
+    }
+    req = SimpleNamespace(output_token_ids=list(range(10)))  # seq_len = 9 > 3, no audio crop
+    payload = talker2code2wav_full_payload(None, pooling_output, req)
+    assert payload is not None
+
+    # Exact expected: ref (prepended) + audio (no crop since seq_len > rows), then
+    # transpose [5, 16] -> [16, 5] and flatten row-major (codebook-major).
+    expected = torch.cat([ref, audio], dim=0).transpose(0, 1).reshape(-1)
+    assert torch.equal(payload["codes"]["audio"], expected), (
+        f"codec flatten mismatch -- got first 8 = {payload['codes']['audio'][:8].tolist()}, "
+        f"expected first 8 = {expected[:8].tolist()}"
+    )
+    assert payload["codes"]["audio"].shape == (80,)  # 16 quantizers * (2 ref + 3 audio) frames
+
+    # Sanity guards: first codebook-major column = [ref[0,0], ref[1,0], audio[0,0], ...],
+    # so the prepend order must put 100 before 1.
+    first_col = payload["codes"]["audio"][:5].tolist()
+    assert first_col == [100, 116, 1, 17, 33], (
+        f"first column wrong: {first_col} -- ref likely appended instead of prepended"
+    )
+
+
+def test_qwen3_tts_full_payload_nested_fallback() -> None:
+    """Back-compat: full_payload works if pooler returns un-flattened nested dict."""
+    from types import SimpleNamespace
+
+    import torch
+
+    from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
+        talker2code2wav_full_payload,
+    )
+
+    audio = torch.arange(2 * 16, dtype=torch.long).reshape(2, 16) + 1
+    pooling_output = {"codes": {"audio": audio}}  # nested, not flat
+    req = SimpleNamespace(output_token_ids=list(range(10)))
+    payload = talker2code2wav_full_payload(None, pooling_output, req)
+    assert payload is not None
+    assert isinstance(payload["codes"]["audio"], torch.Tensor)
+    assert payload["codes"]["audio"].shape == (32,)  # 16 * 2
+
+
+def test_qwen3_tts_code2wav_prefers_connector_tensor_payload() -> None:
+    """Code2Wav should consume connector codec tensor instead of placeholder zeros."""
+    import torch
+
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav import (
+        _codec_ids_from_payload_or_input,
+    )
+
+    placeholder = torch.zeros(6, dtype=torch.long)
+    codec = torch.arange(12, dtype=torch.long)
+
+    out = _codec_ids_from_payload_or_input(
+        placeholder,
+        {"codes": {"audio": codec}},
+    )
+
+    assert torch.equal(out, codec)
+
+
+def test_qwen3_tts_code2wav_accepts_legacy_list_payload() -> None:
+    """Back-compat: old list full-payloads still override placeholder tokens."""
+    import torch
+
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav import (
+        _codec_ids_from_payload_or_input,
+    )
+
+    placeholder = torch.zeros(6, dtype=torch.long)
+
+    out = _codec_ids_from_payload_or_input(
+        placeholder,
+        {"codes": {"audio": [1, 2, 3, 4]}},
+    )
+
+    assert torch.equal(out, torch.tensor([1, 2, 3, 4], dtype=torch.long))
+
+
+def test_qwen3_tts_code2wav_forward_decodes_connector_payload() -> None:
+    """Forward should decode real connector codes, not token-only placeholders."""
+    from collections import Counter
+
+    import torch
+
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav import (
+        Qwen3TTSCode2Wav,
+    )
+
+    class _Decoder:
+        def __init__(self):
+            self.last_codes = None
+
+        def chunked_decode(self, codes, **_kwargs):
+            self.last_codes = codes.detach().clone()
+            return codes.sum(dim=1).to(torch.float32)
+
+    decoder = _Decoder()
+    model = Qwen3TTSCode2Wav.__new__(Qwen3TTSCode2Wav)
+    torch.nn.Module.__init__(model)
+    model.decoder = decoder
+    model._num_quantizers = 2
+    model._total_upsample = 1
+    model._output_sample_rate = 24000
+    model._decode_chunk_frames = 300
+    model._decode_left_context_frames = 25
+    model._decode_batch_bucket_frames = []
+    model._decode_batch_max_size = 0
+    model._decode_variable_chunk_batch_min_frames = 326
+    model._logged_codec_stats = True
+    model._logged_malformed_codec_lengths = set()
+    model._batch_stats_enabled = False
+    model._batch_stats_log_every = 0
+    model._batch_stats_forwards = 0
+    model._batch_stats_groups = 0
+    model._batch_stats_requests = 0
+    model._batch_stats_padded_frames = 0
+    model._batch_stats_decoded_frames = 0
+    model._batch_stats_actual_frames = Counter()
+    model._batch_stats_bucket_groups = Counter()
+
+    payload_codes = torch.tensor([1, 3, 2, 4], dtype=torch.long)
+    out = model.forward(
+        input_ids=torch.zeros(4, dtype=torch.long),
+        runtime_additional_information=[{"codes": {"audio": payload_codes}, "meta": {}}],
+    )
+
+    assert decoder.last_codes is not None
+    assert torch.equal(decoder.last_codes, torch.tensor([[[1, 3], [2, 4]]], dtype=torch.long))
+    assert torch.equal(out.multimodal_outputs["model_outputs"][0], torch.tensor([3.0, 7.0]))
+
+
+def test_qwen3_tts_codec_filter_and_crop_edge_cases() -> None:
+    """Regression gate for codec filter + seq_len crop on both token_only and full_payload.
+
+    Mixes valid / all-zero / negative / >=_CODEBOOK_SIZE rows.  Asserts:
+    - Token-only placeholder length matches Q * (#kept rows after crop).
+    - Full-payload codes.audio matches the exact codebook-major flatten
+      of the kept-and-cropped rows.
+
+    Protects against future cleanup reverting the codex P2 #3 (negative
+    codec filter) or the _CODEBOOK_SIZE upper bound.
+    """
+    from types import SimpleNamespace
+
+    import torch
+
+    from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
+        _CODEBOOK_SIZE,
+        talker2code2wav_full_payload,
+        talker2code2wav_token_only,
+    )
+
+    Q = 4  # simulated num_quantizers (default is 16; small here for readability)
+    # 7 rows: valid / all-zero / negative / out-of-range / boundary-valid / valid / valid.
+    audio_rows = [
+        [10, 20, 30, 40],  # row 0: valid -> KEEP
+        [0, 0, 0, 0],  # row 1: all-zero -> DROP
+        [50, -1, 60, 70],  # row 2: negative -> DROP
+        [100, _CODEBOOK_SIZE, 110, 120],  # row 3: >= 2048 -> DROP
+        [200, _CODEBOOK_SIZE - 1, 210, 220],  # row 4: boundary 2047 -> KEEP
+        [300, 310, 320, 330],  # row 5: valid -> KEEP
+        [400, 410, 420, 430],  # row 6: valid -> KEEP
+    ]
+    audio = torch.tensor(audio_rows, dtype=torch.long)
+    kept = [audio_rows[i] for i in (0, 4, 5, 6)]  # 4 rows after filter
+
+    # === token_only path ===
+    # cumulative_token_ids of length 4 -> seq_len = 3 -> crop kept[-3:] = rows {4, 5, 6}
+    class _Out:
+        def __init__(self, ctids, mm):
+            self.cumulative_token_ids = ctids
+            self.multimodal_output = mm
+
+    class _Wrap:
+        def __init__(self, ctids, mm):
+            self.outputs = [_Out(ctids, mm)]
+            self.finished = True
+
+    mm = {"codes": {"audio": audio}, "meta": {}}
+    src = [_Wrap(ctids=[1, 2, 3, 4], mm=mm)]
+    out = talker2code2wav_token_only(src, prompt=None)
+    assert len(out) == 1
+    # No ref_code -> ref_frames = 0; expected prompt_len = Q * (#kept-after-crop) = 4 * 3 = 12
+    assert len(out[0]["prompt_token_ids"]) == Q * 3
+
+    # === full_payload path ===
+    pooling_output = {"codes.audio": audio}
+    req = SimpleNamespace(output_token_ids=[1, 2, 3, 4])  # seq_len = 3
+    payload = talker2code2wav_full_payload(None, pooling_output, req)
+    assert payload is not None
+    # After filter + crop, kept rows = [row4, row5, row6] = [[200,2047,210,220],[300,310,320,330],[400,410,420,430]]
+    # Codebook-major flatten: transpose [3, Q] -> [Q, 3] -> reshape(-1)
+    cropped = torch.tensor(kept[-3:], dtype=torch.long)
+    expected = cropped.transpose(0, 1).reshape(-1)
+    assert torch.equal(payload["codes"]["audio"], expected)
+    # Sanity: confirm the boundary-valid 2047 survived (codex P2 #3 regression guard).
+    assert _CODEBOOK_SIZE - 1 in payload["codes"]["audio"].tolist()
+    # Sanity: confirm no negative or >=_CODEBOOK_SIZE codec id leaked through.
+    assert bool(((payload["codes"]["audio"] >= 0) & (payload["codes"]["audio"] < _CODEBOOK_SIZE)).all())
+
+
+def test_cosyvoice3_text2flow_token_only_smoke() -> None:
+    """Smoke: cosyvoice3 token-only carries ids.prompt only."""
+    from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
+        text2flow_token_only,
+    )
+
+    class _Out:
+        def __init__(self, tids):
+            self.cumulative_token_ids = tids
+            self.multimodal_output = {}
+
+    class _Wrap:
+        def __init__(self, output_tids, prompt_tids):
+            self.outputs = [_Out(output_tids)]
+            self.prompt_token_ids = prompt_tids
+            self.finished = True
+
+    # multimodal_output has embed.* + we expect token_only to preserve it.
+    import torch
+
+    embed = {"speech_token": torch.zeros(2, 4)}
+    src = [_Wrap(output_tids=[10, 20, 30], prompt_tids=[1, 2, 3, 4])]
+    src[0].outputs[0].multimodal_output = {"embed": embed}
+    out = text2flow_token_only(src)
+    assert len(out) == 1
+    # prompt_token_ids is the talker's cumulative_token_ids (real codec tokens, not zeros).
+    assert out[0]["prompt_token_ids"] == [10, 20, 30]
+    # additional_information carries ids.prompt PLUS the original multimodal_output (embed.* still inline).
+    # Heavy embed.* removal pending the model_intermediate_buffer plumbing on the code2wav side.
+    assert out[0]["additional_information"]["ids"]["prompt"] == [1, 2, 3, 4]
+    assert "embed" in out[0]["additional_information"]
+
+
+def test_cosyvoice3_text2flow_full_payload_smoke() -> None:
+    """Smoke: cosyvoice3 producer-side reads flat embed.* keys."""
+    from types import SimpleNamespace
+
+    import torch
+
+    from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
+        text2flow_full_payload,
+    )
+
+    speech_token = torch.randn(4, 8)
+    speech_feat = torch.randn(4, 16)
+    embedding = torch.randn(1, 32)
+    pooling_output = {
+        "embed.speech_token": speech_token,
+        "embed.speech_feat": speech_feat,
+        "embed.embedding": embedding,
+    }
+    req = SimpleNamespace(external_req_id="r-1")
+    payload = text2flow_full_payload(None, pooling_output, req)
+    assert payload is not None
+    assert "embed" in payload
+    assert torch.equal(payload["embed"]["speech_token"], speech_token)
+    assert torch.equal(payload["embed"]["speech_feat"], speech_feat)
+    assert torch.equal(payload["embed"]["embedding"], embedding)
+    assert payload["meta"]["finished"].item() is True
+
+
+def test_cosyvoice3_text2flow_full_payload_nested_fallback() -> None:
+    """Back-compat: full_payload works if pooler returns un-flattened nested embed."""
+    from types import SimpleNamespace
+
+    import torch
+
+    from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
+        text2flow_full_payload,
+    )
+
+    speech_token = torch.randn(3, 8)
+    pooling_output = {"embed": {"speech_token": speech_token}}  # nested, not flat
+    req = SimpleNamespace(external_req_id="r-2")
+    payload = text2flow_full_payload(None, pooling_output, req)
+    assert payload is not None
+    assert "speech_token" in payload["embed"]
+    assert torch.equal(payload["embed"]["speech_token"], speech_token)
+
+
+def test_cosyvoice3_full_payload_replace_keys_present() -> None:
+    """Confirm _FULL_PAYLOAD_REPLACE_KEYS lists the three embed.* keys."""
+    from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
+        _FULL_PAYLOAD_REPLACE_KEYS,
+    )
+
+    assert _FULL_PAYLOAD_REPLACE_KEYS == frozenset({"embed.speech_token", "embed.speech_feat", "embed.embedding"})
+
+
+def test_ming_flash_omni_thinker2talker_token_only_smoke() -> None:
+    """Smoke: ming_flash_omni token-only carries voice metadata."""
+    from vllm_omni.model_executor.stage_input_processors.ming_flash_omni import (
+        thinker2talker_token_only,
+    )
+
+    class _Out:
+        def __init__(self, text):
+            self.text = text
+
+    class _Wrap:
+        def __init__(self, text):
+            self.outputs = [_Out(text)]
+
+    class _Prompt:
+        def __init__(self, info):
+            self.additional_information = info
+
+    src = [_Wrap("hello world")]
+    prompt = _Prompt({"voice_name": "ZH_FEMALE", "prompt_text": "ref text"})
+    out = thinker2talker_token_only(src, prompt=prompt)
+    assert len(out) == 1
+    assert out[0]["prompt_token_ids"] == [0]  # talker self-tokenizes; dummy id
+    info = out[0]["additional_information"]
+    assert info["text"] == "hello world"
+    assert info["voice_name"] == "ZH_FEMALE"
+    assert info["prompt_text"] == "ref text"
+    assert info["ming_task"] == "omni"
+
+
+def test_qwen2_5_omni_thinker2talker_token_only_smoke() -> None:
+    """Smoke: qwen2_5_omni thinker token-only allocates prompt slots; bulk payload ships via connector."""
+    from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import (
+        TALKER_CODEC_END_TOKEN_ID,
+        TALKER_CODEC_PAD_TOKEN_ID,
+        TALKER_CODEC_START_TOKEN_ID,
+        thinker2talker_token_only,
+    )
+
+    class _Wrap:
+        def __init__(self, prompt_tids, rid):
+            self.outputs = [object()]
+            self.prompt_token_ids = prompt_tids
+            self.request_id = rid
+
+    class _Prompt(dict):
+        pass
+
+    src = [_Wrap(prompt_tids=[1, 2, 3, 4, 5], rid="r-1")]
+    prompt = [_Prompt(multi_modal_data=None)]
+    out = thinker2talker_token_only(src, prompt=prompt)
+    assert len(out) == 1
+    expected_prompt_len = 1 + len([1, 2, 3, 4, 5]) + 1
+    assert len(out[0]["prompt_token_ids"]) == expected_prompt_len
+    assert out[0]["prompt_token_ids"][0] == TALKER_CODEC_START_TOKEN_ID
+    assert out[0]["prompt_token_ids"][-1] == TALKER_CODEC_END_TOKEN_ID
+    assert all(t == TALKER_CODEC_PAD_TOKEN_ID for t in out[0]["prompt_token_ids"][1:-1])
+    assert out[0]["additional_information"] is None
+
+
+def test_qwen2_5_omni_thinker2talker_full_payload_noop() -> None:
+    """thinker2talker_full_payload returns None when pooling_output lacks the "hidden" key (defensive)."""
+    from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import (
+        thinker2talker_full_payload,
+    )
+
+    payload = thinker2talker_full_payload(None, {"any": "thing"}, None)
+    assert payload is None
diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py
index 05c5ce8b1..e96f79f89 100644
--- a/tests/test_config_factory.py
+++ b/tests/test_config_factory.py
@@ -1580,6 +1580,74 @@ class TestMingFlashOmniPipeline:
         assert len(stages) == 1
         assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniForConditionalGeneration"
 
+    def test_image_pipeline_registered(self):
+        p = _PIPELINE_REGISTRY.get("ming_flash_omni_image")
+        assert p is not None
+        assert p.model_arch == "MingFlashOmniForConditionalGeneration"
+        assert len(p.stages) == 2
+        assert p.validate() == []
+
+    def test_image_thinker_stage(self):
+        s = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(0)
+        assert s.model_stage == "thinker"
+        assert s.execution_type == StageExecutionType.LLM_AR
+        assert s.input_sources == ()
+        assert s.final_output is False
+        assert s.owns_tokenizer is True
+        assert s.requires_multimodal_data is True
+        # Image variant exports hidden states for the diffusion stage.
+        assert s.engine_output_type == "latent"
+        assert s.hf_config_name == "thinker_config"
+        assert s.sampling_constraints["detokenize"] is False
+        assert s.prompt_expand_func is not None
+
+    def test_image_dit_stage(self):
+        s = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(1)
+        assert s.model_stage == "dit"
+        assert s.execution_type == StageExecutionType.DIFFUSION
+        assert s.input_sources == (0,)
+        assert s.final_output is True
+        assert s.final_output_type == "image"
+        assert s.hf_config_name == "image_gen_config"
+        assert s.model_arch == "MingImagePipeline"
+        assert s.custom_process_input_func is not None
+
+    def test_image_processor_wiring_resolves(self):
+        """The prompt_expand_func and custom_process_input_func strings must point to real callables."""
+        thinker = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(0)
+        dit = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(1)
+        for ref in (thinker.prompt_expand_func, dit.custom_process_input_func):
+            module_path, _, attr = ref.rpartition(".")
+            module = importlib.import_module(module_path)
+            assert callable(getattr(module, attr))
+
+    def test_image_yaml_loads_and_merges(self):
+        """deploy/ming_flash_omni_image.yaml parses and routes to the image pipeline."""
+        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_image.yaml"
+        if not deploy_path.exists():
+            pytest.skip("ming_flash_omni_image deploy yaml not found")
+
+        deploy = load_deploy_config(deploy_path)
+        assert len(deploy.stages) == 2
+        assert deploy.async_chunk is False
+        assert deploy.pipeline == "ming_flash_omni_image"
+        assert deploy.connectors is not None
+        assert "shared_memory_connector" in deploy.connectors
+
+        pipeline = _PIPELINE_REGISTRY["ming_flash_omni_image"]
+        stages = merge_pipeline_deploy(pipeline, deploy)
+        assert len(stages) == 2
+        # Stage 0 thinker: AR worker that emits latents.
+        assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniForConditionalGeneration"
+        assert stages[0].yaml_engine_args["engine_output_type"] == "latent"
+        assert stages[0].yaml_extras["default_sampling_params"]["detokenize"] is False
+        assert stages[0].yaml_extras["prompt_expand_func"] is not None
+        # Stage 1 dit: diffusion stage with MingImagePipeline.
+        assert stages[1].yaml_engine_args["model_arch"] == "MingImagePipeline"
+        assert stages[1].custom_process_input_func is not None
+        assert stages[1].final_output is True
+        assert stages[1].final_output_type == "image"
+
 
 class TestBaseConfigInheritance:
     """Test deploy YAML base_config inheritance."""
@@ -1660,7 +1728,7 @@ class TestPlatformOverrides:
         assert deploy.stages[0].tensor_parallel_size == 2
         assert deploy.stages[0].devices == "0,1"
         # Stage 2 unaffected fields stay at base
-        assert deploy.stages[2].enforce_eager is True
+        assert deploy.stages[2].enforce_eager is False
 
     def test_xpu_overrides(self):
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
@@ -1841,12 +1909,19 @@ class TestSentinelDefaultPrecedence:
         )
         assert async_stages[1].custom_process_input_func is None
 
-        # async_chunk=False → stage 0 has no streaming processor, stage 1's
-        # batch-end processor wires up.
+        # async_chunk=False → stage 0 ships the bulk codec via the
+        # worker-connector full-payload producer; stage 1 wires the
+        # ``_token_only`` placeholder so the orchestrator emits no
+        # legacy ``additional_information``-shaped input (PR3 sync-
+        # via-connector data plane).
         sync_stages = merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=False))
-        assert "custom_process_next_stage_input_func" not in sync_stages[0].yaml_engine_args
+        assert (
+            sync_stages[0]
+            .yaml_engine_args["custom_process_next_stage_input_func"]
+            .endswith("talker2code2wav_full_payload")
+        )
         assert sync_stages[1].custom_process_input_func is not None
-        assert sync_stages[1].custom_process_input_func.endswith("talker2code2wav")
+        assert sync_stages[1].custom_process_input_func.endswith("talker2code2wav_token_only")
 
     def test_async_chunk_dispatches_qwen3_omni_processors(self):
         import runpy
@@ -1883,6 +1958,44 @@ class TestSentinelDefaultPrecedence:
             .endswith("talker2code2wav_full_payload")
         )
 
+    def test_ming_flash_omni_topology(self):
+        """Guard ming_flash_omni's PR3 cleanup: stage 0 has no full-payload
+        producer hook (the connector path was removed as fake -- arch is not
+        in ``_FULL_PAYLOAD_INPUT_STAGES``), and stage 1 still wires the
+        legacy ``thinker2talker`` (custom_process_input_func) plus the
+        ``thinker2talker_token_only`` placeholder (sync_process_input_func).
+        Merge under either async_chunk mode must not re-introduce a
+        stage-0 full-payload hook."""
+        from vllm_omni.config.stage_config import DeployConfig, merge_pipeline_deploy
+
+        pipeline = _PIPELINE_REGISTRY["ming_flash_omni"]
+
+        stage0, stage1 = pipeline.stages
+        assert stage0.custom_process_next_stage_input_func is None, (
+            "ming_flash_omni stage 0 must not declare a full-payload producer "
+            "(connector path is not active for this arch)."
+        )
+        assert stage1.custom_process_input_func is not None
+        assert stage1.custom_process_input_func.endswith("thinker2talker")
+        assert stage1.sync_process_input_func is not None
+        assert stage1.sync_process_input_func.endswith("thinker2talker_token_only")
+
+        # async_chunk=True must now be rejected: removing the fake hook means
+        # there is no next-stage input processor for the validator to accept.
+        # (Positive consequence -- users can't accidentally enable async_chunk
+        # on an arch that doesn't actually support it.)
+        import pytest as _pytest
+
+        with _pytest.raises(ValueError, match="async_chunk=True"):
+            merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=True))
+
+        # async_chunk=False merges cleanly and stage-0 yaml_engine_args carries
+        # no spurious full-payload hook.
+        merged = merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=False))
+        assert "custom_process_next_stage_input_func" not in merged[0].yaml_engine_args, (
+            "stage-0 full-payload hook unexpectedly re-appeared in yaml_engine_args"
+        )
+
 
 class TestSamplingConstraintsPrecedence:
     """Test that pipeline sampling_constraints override deploy defaults."""
diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py
index b834d8733..ae6ecc57c 100644
--- a/tests/worker/test_omni_gpu_model_runner.py
+++ b/tests/worker/test_omni_gpu_model_runner.py
@@ -161,11 +161,11 @@ def _make_runner_for_mimo(req_id="r_mimo"):
 
 
 def test_talker_mtp_forward_cpu_updates_inputs_and_info(monkeypatch):
-    # Patch the module-level `set_forward_context` symbol used inside
-    # OmniGPUModelRunner._talker_mtp_forward.
+    # `_talker_mtp_forward` calls `current_omni_platform.set_forward_context`,
+    # which would otherwise dispatch to the real device implementation.
     import vllm_omni.worker.gpu_model_runner as mod  # Must be the same module that defines OmniGPUModelRunner
 
-    monkeypatch.setattr(mod, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
 
     runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
 
@@ -199,7 +199,7 @@ def test_talker_mtp_forward_cpu_updates_inputs_and_info(monkeypatch):
 def test_talker_mtp_forward_cpu_empty_batch_noop(monkeypatch):
     import vllm_omni.worker.gpu_model_runner as mod
 
-    monkeypatch.setattr(mod, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
 
     runner = _make_runner(req_ids=("r1",), hidden_size=4)
 
@@ -215,7 +215,7 @@ def test_talker_mtp_forward_cpu_empty_batch_noop(monkeypatch):
 def test_talker_mtp_forward_ignores_default_sampling_seed_without_request_marker(monkeypatch):
     import vllm_omni.worker.gpu_model_runner as mod
 
-    monkeypatch.setattr(mod, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
 
     runner = _make_runner(req_ids=("r1",), hidden_size=4)
     runner.requests["r1"].sampling_params = SimpleNamespace(seed=42)
@@ -237,7 +237,7 @@ def test_talker_mtp_forward_ignores_default_sampling_seed_without_request_marker
 def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker(monkeypatch):
     import vllm_omni.worker.gpu_model_runner as mod
 
-    monkeypatch.setattr(mod, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
 
     runner = _make_runner(req_ids=("r1",), hidden_size=4)
     runner.requests["r1"].sampling_params = SimpleNamespace(
@@ -281,7 +281,7 @@ def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker
 def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch):
     import vllm_omni.worker.gpu_model_runner as mod
 
-    monkeypatch.setattr(mod, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
 
     runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
     runner.requests["r1"].sampling_params = SimpleNamespace(
@@ -322,7 +322,7 @@ def test_update_intermediate_buffer_writes_to_buffer_and_setattr(monkeypatch):
     (forward path) and mirrors to additional_information_cpu setattr (backward compat)."""
     import vllm_omni.worker.gpu_model_runner as mod
 
-    monkeypatch.setattr(mod, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
 
     runner = _make_runner(req_ids=("r1",), hidden_size=4)
 
@@ -400,16 +400,23 @@ def _make_full_payload_accumulation_runner(
     model_arch="Qwen3OmniMoeForConditionalGeneration",
     model_stage="talker",
     async_chunk=False,
+    final_output=False,
+    custom_process_next_stage_input_func="module.full_payload",
 ):
     runner = object.__new__(OmniConnectorModelRunnerMixin)
     runner.model_config = SimpleNamespace(
         model_arch=model_arch,
         model_stage=model_stage,
         async_chunk=async_chunk,
+        final_output=final_output,
+        custom_process_next_stage_input_func=custom_process_next_stage_input_func,
     )
     runner._custom_process_func = object()
     runner._pending_full_payload_send = {}
     runner._stage_id = 1
+    # Non-None sentinel: the gate short-circuits to False when no connector
+    # is configured at all (terminal stages in pipelines with no connector).
+    runner._omni_connector = object()
     return runner
 
 
@@ -425,7 +432,7 @@ def test_accumulate_full_payload_output_preserves_aligned_all_zero_qwen3_omni_co
 
 
 def test_accumulate_full_payload_output_keeps_misaligned_all_zero_qwen3_omni_codec_rows():
-    # After removing the sender-side zero filter, the accumulator keeps every
+    # After removing the sender-side zero filter, the full-payload accumulator keeps every
     # codec row including misaligned all-zero rows. The downstream consumer
     # (_extract_qwen3_full_payload_codec_rows) is the authoritative crop and
     # filters by output_token_ids.
@@ -472,18 +479,45 @@ def test_accumulate_full_payload_output_keeps_all_zero_qwen3_omni_prefill_placeh
 
 
 def test_full_payload_output_accumulation_hook_matrix():
+    """Producer-side gate: fires iff an explicit next-stage payload hook is loaded.
+
+    A derived `*_full_payload` helper from `custom_process_input_func` is not
+    enough: terminal/input-only consumer stages must not enqueue orphan
+    downstream payloads.
+    """
+    # Thinker / talker producer stages: explicit next-stage payload hook -> gate fires.
     assert _make_full_payload_accumulation_runner(model_stage="thinker")._should_accumulate_full_payload_output()
     assert _make_full_payload_accumulation_runner(model_stage="talker")._should_accumulate_full_payload_output()
-    assert not _make_full_payload_accumulation_runner(model_stage="code2wav")._should_accumulate_full_payload_output()
+
+    # Terminal stage: even if _load_custom_func derived a builder from
+    # custom_process_input_func, final output stages are not producers.
+    runner = _make_full_payload_accumulation_runner(model_stage="code2wav", final_output=True)
+    assert not runner._should_accumulate_full_payload_output()
+
+    # Input-only consumer stage without an explicit producer hook must not
+    # accumulate/send just because a same-module *_full_payload helper exists.
+    runner = _make_full_payload_accumulation_runner(
+        model_stage="token2audio",
+        custom_process_next_stage_input_func=None,
+    )
+    assert not runner._should_accumulate_full_payload_output()
+
+    # async_chunk mode -> gate off.
     assert not _make_full_payload_accumulation_runner(
         model_stage="talker", async_chunk=True
     )._should_accumulate_full_payload_output()
-    assert not _make_full_payload_accumulation_runner(
-        model_arch="Qwen3TTSForConditionalGeneration"
-    )._should_accumulate_full_payload_output()
-    assert not _make_full_payload_accumulation_runner(
-        model_arch="Qwen2_5OmniForConditionalGeneration"
-    )._should_accumulate_full_payload_output()
+
+    # Non-qwen3 arches: gate is arch-agnostic, but if the fixture's arch
+    # does not configure a connector payload builder, its runtime
+    # `_custom_process_func` is None.  Emulate that.
+    runner = _make_full_payload_accumulation_runner(model_arch="Qwen3TTSForConditionalGeneration")
+    runner._custom_process_func = None
+    runner._should_accumulate_full_payload_output_cached = None
+    assert not runner._should_accumulate_full_payload_output()
+    runner = _make_full_payload_accumulation_runner(model_arch="Qwen2_5OmniForConditionalGeneration")
+    runner._custom_process_func = None
+    runner._should_accumulate_full_payload_output_cached = None
+    assert not runner._should_accumulate_full_payload_output()
 
 
 def test_sync_local_stage_payloads_retains_payload_until_request_is_active():
diff --git a/tools/nightly/generate_nightly_perf_excel.py b/tools/nightly/generate_nightly_perf_excel.py
index e76dc4346..54c4b3f34 100644
--- a/tools/nightly/generate_nightly_perf_excel.py
+++ b/tools/nightly/generate_nightly_perf_excel.py
@@ -18,11 +18,13 @@ from openpyxl import Workbook
 from openpyxl.styles import PatternFill
 from openpyxl.utils import get_column_letter
 
-from benchmarks.diffusion.backends import normalize_endpoint
-
 LOGGER = logging.getLogger(__name__)
 
 GREY_BLOCK_FILL = PatternFill(start_color="D3D3D3", fill_type="solid")
+LEGACY_BACKEND_ENDPOINT_ALIASES = {
+    "vllm-omni": "/v1/chat/completions",
+    "openai": "/v1/images/generations",
+}
 
 # Diffusion sheet columns (Qwen-Image diffusion benchmark).
 # Per-stage latency metrics. Unpack from stage_durations_mean/p50/p99 dicts
@@ -145,6 +147,20 @@ def _diffusion_group_key(record: dict[str, Any]) -> tuple[Any, ...]:
     return (record.get("test_name") or "",)
 
 
+def normalize_endpoint(value: str) -> str:
+    """Normalize legacy backend aliases and endpoint paths for report rows."""
+    endpoint = str(value).strip()
+    if not endpoint:
+        raise ValueError("endpoint must not be empty.")
+    endpoint = LEGACY_BACKEND_ENDPOINT_ALIASES.get(
+        endpoint,
+        LEGACY_BACKEND_ENDPOINT_ALIASES.get(endpoint.lstrip("/"), endpoint),
+    )
+    if not endpoint.startswith("/"):
+        endpoint = f"/{endpoint}"
+    return endpoint
+
+
 def _load_summary_columns(script_dir: str) -> list[str]:
     """Load summary column names from a file next to this script; fallback to default if missing."""
     path = os.path.join(script_dir, _COLUMNS_FILENAME)
diff --git a/vllm_omni/attention/__init__.py b/vllm_omni/attention/__init__.py
new file mode 100644
index 000000000..208f01a7c
--- /dev/null
+++ b/vllm_omni/attention/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm_omni/attention/fish_kvcache_attn.py b/vllm_omni/attention/fish_kvcache_attn.py
new file mode 100644
index 000000000..9e552e762
--- /dev/null
+++ b/vllm_omni/attention/fish_kvcache_attn.py
@@ -0,0 +1,191 @@
+from __future__ import annotations
+
+import os
+import threading
+from typing import Any
+
+import torch
+
+FISH_KVCACHE_LONG_SPLIT_TOKENS = 1024
+FISH_KVCACHE_SMALL_PATH_MAX_SEQ_LEN = 1024
+
+_FISH_KVCACHE_ATTN_ENV = "VLLM_OMNI_FISH_KVCACHE_ATTN"
+_ENABLED_VALUES = frozenset({"1", "true", "yes", "on", "required"})
+_DISABLED_VALUES = frozenset({"0", "false", "no", "off", "disabled", "disable"})
+_REQUIRED_VALUES = frozenset({"required"})
+_WORKSPACE_CACHE: dict[tuple[Any, ...], tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {}
+_WORKSPACE_CACHE_LOCK = threading.Lock()
+
+
+def _triton_backend() -> Any:
+    from vllm_omni.attention import fish_kvcache_triton
+
+    return fish_kvcache_triton
+
+
+def is_available() -> bool:
+    return _triton_backend().is_available()
+
+
+def load_error() -> Exception | None:
+    return _triton_backend().load_error()
+
+
+def is_fish_kvcache_attn_enabled() -> bool:
+    value = os.environ.get(_FISH_KVCACHE_ATTN_ENV)
+    if value is None:
+        return True
+    value = value.lower()
+    if value in _DISABLED_VALUES:
+        return False
+    return value in _ENABLED_VALUES or value == ""
+
+
+def is_fish_kvcache_attn_required() -> bool:
+    return os.environ.get(_FISH_KVCACHE_ATTN_ENV, "").lower() in _REQUIRED_VALUES
+
+
+def _is_sliding_window_disabled(sliding_window: Any) -> bool:
+    if sliding_window is None:
+        return True
+    if isinstance(sliding_window, (list, tuple)) and len(sliding_window) == 2:
+        return int(sliding_window[0]) == -1 and int(sliding_window[1]) == -1
+    return False
+
+
+def can_use_fish_kvcache_attn(
+    *,
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    block_table: torch.Tensor | None,
+    seq_lens: torch.Tensor,
+    max_query_len: int,
+    max_seq_len: int,
+    dcp_world_size: int,
+    use_cascade: bool,
+    alibi_slopes: Any,
+    sliding_window: Any,
+    output_scale: torch.Tensor | None = None,
+    output_block_scale: torch.Tensor | None = None,
+) -> bool:
+    if not is_fish_kvcache_attn_enabled():
+        return False
+    if not is_available():
+        return False
+    if max_query_len != 1 or use_cascade or dcp_world_size != 1:
+        return False
+    if block_table is None or alibi_slopes is not None or not _is_sliding_window_disabled(sliding_window):
+        return False
+    if output_scale is not None or output_block_scale is not None:
+        return False
+    if query.dim() != 3 or key_cache.dim() != 4 or value_cache.dim() != 4:
+        return False
+    if block_table.dim() != 2 or seq_lens.dim() != 1:
+        return False
+    if block_table.shape[0] != query.shape[0] or seq_lens.shape[0] != query.shape[0]:
+        return False
+    if query.shape[-1] != 128 or key_cache.shape[-1] != 128:
+        return False
+    if key_cache.shape[1] != 16:
+        return False
+    if query.dtype not in (torch.float16, torch.bfloat16):
+        return False
+    if key_cache.dtype != query.dtype or value_cache.dtype != query.dtype:
+        return False
+    if block_table.dtype != torch.int32 or seq_lens.dtype != torch.int32:
+        return False
+    if max_seq_len <= 0:
+        return False
+    if max_seq_len > block_table.shape[1] * key_cache.shape[1]:
+        return False
+    if not (
+        query.is_contiguous()
+        and key_cache.is_contiguous()
+        and value_cache.is_contiguous()
+        and block_table.is_contiguous()
+        and seq_lens.is_contiguous()
+    ):
+        return False
+    return True
+
+
+def _is_cuda_graph_capturing() -> bool:
+    return bool(torch.cuda.is_current_stream_capturing())
+
+
+def _raise_workspace_capture_miss() -> None:
+    raise RuntimeError("Fish kvcache attention workspace was not prewarmed before CUDA graph capture")
+
+
+def _get_decode_workspace(
+    query: torch.Tensor,
+    max_seq_len: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    batch_size, num_q_heads, head_dim = query.shape
+    if int(max_seq_len) <= FISH_KVCACHE_SMALL_PATH_MAX_SEQ_LEN:
+        key = (query.device.type, query.device.index, "empty")
+        partial_shapes = ((0,), (0,), (0,))
+    else:
+        num_splits = (int(max_seq_len) + FISH_KVCACHE_LONG_SPLIT_TOKENS - 1) // FISH_KVCACHE_LONG_SPLIT_TOKENS
+        key = (
+            query.device.type,
+            query.device.index,
+            int(batch_size),
+            int(num_q_heads),
+            int(head_dim),
+            int(num_splits),
+        )
+        partial_shapes = (
+            (num_splits, batch_size, num_q_heads),
+            (num_splits, batch_size, num_q_heads),
+            (num_splits, batch_size, num_q_heads, head_dim),
+        )
+    with _WORKSPACE_CACHE_LOCK:
+        workspace = _WORKSPACE_CACHE.get(key)
+        if workspace is None:
+            if _is_cuda_graph_capturing():
+                _raise_workspace_capture_miss()
+            workspace = (
+                torch.empty(partial_shapes[0], device=query.device, dtype=torch.float32),
+                torch.empty(partial_shapes[1], device=query.device, dtype=torch.float32),
+                torch.empty(partial_shapes[2], device=query.device, dtype=torch.float32),
+            )
+            _WORKSPACE_CACHE[key] = workspace
+        return workspace
+
+
+def prewarm_fish_kvcache_attn_workspace(
+    query: torch.Tensor,
+    max_seq_len: int,
+) -> None:
+    _get_decode_workspace(query, int(max_seq_len))
+
+
+def fish_decode_kvcache_attn(
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    seq_lens: torch.Tensor,
+    out: torch.Tensor,
+    *,
+    scale: float,
+    max_seq_len: int,
+) -> torch.Tensor:
+    partial_m, partial_l, partial_acc = _get_decode_workspace(query, int(max_seq_len))
+    return _triton_backend().fish_decode_kvcache_attn_triton(
+        query,
+        key_cache,
+        value_cache,
+        block_table,
+        seq_lens,
+        out,
+        scale=float(scale),
+        max_seq_len=int(max_seq_len),
+        small_path_max_seq_len=FISH_KVCACHE_SMALL_PATH_MAX_SEQ_LEN,
+        long_split_tokens=FISH_KVCACHE_LONG_SPLIT_TOKENS,
+        partial_m=partial_m,
+        partial_l=partial_l,
+        partial_acc=partial_acc,
+    )
diff --git a/vllm_omni/attention/fish_kvcache_backend.py b/vllm_omni/attention/fish_kvcache_backend.py
new file mode 100644
index 000000000..3f3e9bdc6
--- /dev/null
+++ b/vllm_omni/attention/fish_kvcache_backend.py
@@ -0,0 +1,406 @@
+from __future__ import annotations
+
+import types
+from collections import Counter
+from collections.abc import Callable
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.attention.fish_kvcache_attn import (
+    FISH_KVCACHE_SMALL_PATH_MAX_SEQ_LEN,
+    can_use_fish_kvcache_attn,
+    fish_decode_kvcache_attn,
+    is_available,
+    is_fish_kvcache_attn_enabled,
+    is_fish_kvcache_attn_required,
+    load_error,
+    prewarm_fish_kvcache_attn_workspace,
+)
+
+logger = init_logger(__name__)
+
+_FISH_SLOW_AR_ARCH = "FishSpeechSlowARForConditionalGeneration"
+_FIRST_SMALL_HIT_LOGGED = False
+_FIRST_LONG_HIT_LOGGED = False
+_FIRST_MISS_LOGGED = False
+_SMALL_HIT_COUNT = 0
+_LONG_HIT_COUNT = 0
+_FALLBACK_COUNTS: Counter[str] = Counter()
+
+
+def _fish_kvcache_enabled() -> bool:
+    return is_fish_kvcache_attn_enabled()
+
+
+def is_fish_kvcache_attn_active_for_model(model_config: Any) -> bool:
+    return getattr(model_config, "model_arch", None) == _FISH_SLOW_AR_ARCH and is_fish_kvcache_attn_enabled()
+
+
+def reset_fish_kvcache_attn_stats() -> None:
+    global _LONG_HIT_COUNT, _SMALL_HIT_COUNT
+    _SMALL_HIT_COUNT = 0
+    _LONG_HIT_COUNT = 0
+    _FALLBACK_COUNTS.clear()
+
+
+def get_fish_kvcache_attn_stats() -> dict[str, Any]:
+    return {
+        "small_hit_count": _SMALL_HIT_COUNT,
+        "long_hit_count": _LONG_HIT_COUNT,
+        "fallback_count_by_reason": dict(_FALLBACK_COUNTS),
+    }
+
+
+def _record_fallback(reason: str) -> None:
+    _FALLBACK_COUNTS[reason] += 1
+
+
+def _small_decode_bucket(max_seq_len: int) -> int:
+    if max_seq_len <= 64:
+        return 64
+    if max_seq_len <= 128:
+        return 128
+    if max_seq_len <= 256:
+        return 256
+    if max_seq_len <= 512:
+        return 512
+    return FISH_KVCACHE_SMALL_PATH_MAX_SEQ_LEN
+
+
+def prewarm_fish_kvcache_attn_capture_workspaces(
+    *,
+    model_config: Any,
+    device: torch.device,
+    dtype: torch.dtype,
+    capture_sizes: list[int] | tuple[int, ...],
+) -> int:
+    """Preallocate Fish attention workspaces used during CUDA graph capture."""
+    if not _fish_kvcache_enabled() or not is_available():
+        return 0
+    if not is_fish_kvcache_attn_active_for_model(model_config):
+        return 0
+
+    text_config = model_config.hf_config.text_config
+    num_heads = int(text_config.num_attention_heads)
+    head_dim = int(text_config.head_dim)
+    max_seq_len = int(model_config.max_model_len)
+    sizes = sorted({int(size) for size in capture_sizes if int(size) > 0})
+    if not sizes:
+        return 0
+
+    # The small path uses an empty workspace keyed only by device. Prewarm it
+    # once, then prewarm every CUDA graph capture batch size for the long path.
+    small_query = torch.empty((1, num_heads, head_dim), device=device, dtype=dtype)
+    prewarm_fish_kvcache_attn_workspace(small_query, 1)
+    prewarmed = 1
+
+    for batch_size in sizes:
+        query = torch.empty((batch_size, num_heads, head_dim), device=device, dtype=dtype)
+        prewarm_fish_kvcache_attn_workspace(query, max_seq_len)
+        prewarmed += 1
+
+    logger.info(
+        "Prewarmed Fish kvcache attention workspaces for capture sizes %s (num_heads=%d head_dim=%d max_seq_len=%d)",
+        sizes,
+        num_heads,
+        head_dim,
+        max_seq_len,
+    )
+    return prewarmed
+
+
+def _attach_seq_lens_cpu_upper_bound(attn_metadata: Any, seq_lens_cpu_upper_bound: torch.Tensor) -> None:
+    if attn_metadata is None:
+        return
+    if isinstance(attn_metadata, dict):
+        for value in attn_metadata.values():
+            _attach_seq_lens_cpu_upper_bound(value, seq_lens_cpu_upper_bound)
+        return
+    if isinstance(attn_metadata, (list, tuple)):
+        for value in attn_metadata:
+            _attach_seq_lens_cpu_upper_bound(value, seq_lens_cpu_upper_bound)
+        return
+    setattr(attn_metadata, "seq_lens_cpu_upper_bound", seq_lens_cpu_upper_bound)
+
+
+def _graph_capture_seq_len_upper_bound(
+    *,
+    model_config: Any,
+    max_query_len: int,
+    for_cudagraph_capture: bool,
+) -> int | None:
+    if not for_cudagraph_capture or int(max_query_len) != 1:
+        return None
+    max_model_len = int(getattr(model_config, "max_model_len", 0))
+    return max(max_model_len, FISH_KVCACHE_SMALL_PATH_MAX_SEQ_LEN)
+
+
+def maybe_attach_fish_kvcache_seq_lens_upper_bound(
+    *,
+    model_config: Any,
+    attn_metadata: Any,
+    input_batch: Any,
+    optimistic_seq_lens_cpu: torch.Tensor,
+    num_reqs: int,
+    num_reqs_padded: int,
+    max_query_len: int,
+    pad_attn: bool,
+    for_cudagraph_capture: bool = False,
+    num_scheduled_tokens_np: Any = None,
+) -> None:
+    """Attach CPU seq-len upper bounds for the Fish-only Triton decode fastpath."""
+    if not is_fish_kvcache_attn_active_for_model(model_config) or int(max_query_len) != 1:
+        return
+
+    upper_bound_rows = int(num_reqs_padded if pad_attn else num_reqs)
+    capture_upper_bound = _graph_capture_seq_len_upper_bound(
+        model_config=model_config,
+        max_query_len=max_query_len,
+        for_cudagraph_capture=for_cudagraph_capture,
+    )
+    if capture_upper_bound is not None:
+        optimistic_seq_lens_cpu[:upper_bound_rows].fill_(capture_upper_bound)
+    else:
+        if num_scheduled_tokens_np is None:
+            return
+        seq_lens_cpu_upper_bound = torch.as_tensor(
+            input_batch.num_computed_tokens_cpu[:num_reqs] + num_scheduled_tokens_np,
+            dtype=torch.int32,
+            device="cpu",
+        )
+        optimistic_seq_lens_cpu[:num_reqs] = seq_lens_cpu_upper_bound
+        if upper_bound_rows > num_reqs:
+            optimistic_seq_lens_cpu[num_reqs:upper_bound_rows].fill_(1)
+    _attach_seq_lens_cpu_upper_bound(
+        attn_metadata,
+        optimistic_seq_lens_cpu[:upper_bound_rows],
+    )
+
+
+def _dispatch_max_seq_len(attn_metadata: Any, active_batch: int, seq_lens: Any) -> int | None:
+    max_seq_len = int(attn_metadata.max_seq_len)
+    if int(attn_metadata.max_query_len) != 1:
+        return max_seq_len
+
+    upper_bound = getattr(attn_metadata, "seq_lens_cpu_upper_bound", None)
+    if upper_bound is None:
+        if is_fish_kvcache_attn_required():
+            raise RuntimeError("Fish kvcache attention requires seq_lens_cpu_upper_bound for decode dispatch")
+        return None
+    if not isinstance(upper_bound, torch.Tensor) or upper_bound.device.type != "cpu":
+        if is_fish_kvcache_attn_required():
+            raise RuntimeError("Fish kvcache attention seq_lens_cpu_upper_bound must be a CPU tensor")
+        return None
+    if int(upper_bound.shape[0]) < active_batch:
+        if is_fish_kvcache_attn_required():
+            raise RuntimeError(
+                "Fish kvcache attention seq_lens_cpu_upper_bound does not cover the active decode batch: "
+                f"upper_bound_rows={int(upper_bound.shape[0])}, active_batch={active_batch}"
+            )
+        return None
+    upper_bound = upper_bound[:active_batch]
+    real_upper_bound = max(1, int(upper_bound.max().item()))
+    if is_fish_kvcache_attn_required():
+        real_max_seq_len = int(seq_lens.max().item())
+        if real_upper_bound < real_max_seq_len:
+            raise RuntimeError(
+                "Fish kvcache attention seq_lens_cpu_upper_bound "
+                f"underestimates real seq_lens: upper_bound={real_upper_bound}, "
+                f"real={real_max_seq_len}"
+            )
+    if real_upper_bound <= FISH_KVCACHE_SMALL_PATH_MAX_SEQ_LEN:
+        return _small_decode_bucket(real_upper_bound)
+    return real_upper_bound
+
+
+def _forward_with_fish_kvcache(
+    impl: Any,
+    original_forward: Callable[..., Any],
+    layer: Any,
+    query: Any,
+    key: Any,
+    value: Any,
+    kv_cache: Any,
+    attn_metadata: Any,
+    output: Any,
+    output_scale: Any = None,
+    output_block_scale: Any = None,
+) -> Any:
+    global _FIRST_LONG_HIT_LOGGED, _FIRST_MISS_LOGGED, _FIRST_SMALL_HIT_LOGGED
+    global _LONG_HIT_COUNT, _SMALL_HIT_COUNT
+
+    if attn_metadata is not None and not attn_metadata.use_cascade:
+        num_actual_tokens = attn_metadata.num_actual_tokens
+        key_cache, value_cache = kv_cache.unbind(0)
+        q = query[:num_actual_tokens]
+        out = output[:num_actual_tokens]
+        active_batch = int(q.shape[0])
+        block_table = attn_metadata.block_table[:active_batch] if attn_metadata.block_table is not None else None
+        seq_lens = attn_metadata.seq_lens[:active_batch]
+        dispatch_max_seq_len = _dispatch_max_seq_len(attn_metadata, active_batch, seq_lens)
+        can_use = dispatch_max_seq_len is not None and can_use_fish_kvcache_attn(
+            query=q,
+            key_cache=key_cache,
+            value_cache=value_cache,
+            block_table=block_table,
+            seq_lens=seq_lens,
+            max_query_len=attn_metadata.max_query_len,
+            max_seq_len=dispatch_max_seq_len,
+            dcp_world_size=impl.dcp_world_size,
+            use_cascade=attn_metadata.use_cascade,
+            alibi_slopes=impl.alibi_slopes,
+            sliding_window=impl.sliding_window,
+            output_scale=output_scale,
+            output_block_scale=output_block_scale,
+        )
+        if can_use:
+            is_long_path = dispatch_max_seq_len > FISH_KVCACHE_SMALL_PATH_MAX_SEQ_LEN
+            if is_long_path:
+                _LONG_HIT_COUNT += 1
+            else:
+                _SMALL_HIT_COUNT += 1
+            should_log_hit = (is_long_path and not _FIRST_LONG_HIT_LOGGED) or (
+                not is_long_path and not _FIRST_SMALL_HIT_LOGGED
+            )
+            if should_log_hit:
+                if is_long_path:
+                    _FIRST_LONG_HIT_LOGGED = True
+                    path_name = "long"
+                else:
+                    _FIRST_SMALL_HIT_LOGGED = True
+                    path_name = "small"
+                logger.info(
+                    "Fish decode-only kvcache attention fast path hit: "
+                    "query_shape=%s key_cache_shape=%s block_table_shape=%s "
+                    "seq_lens_shape=%s max_seq_len=%s path=%s",
+                    tuple(q.shape),
+                    tuple(key_cache.shape),
+                    tuple(attn_metadata.block_table.shape),
+                    tuple(attn_metadata.seq_lens.shape),
+                    dispatch_max_seq_len,
+                    path_name,
+                )
+            fish_decode_kvcache_attn(
+                q,
+                key_cache,
+                value_cache,
+                block_table,
+                seq_lens,
+                out,
+                scale=float(impl.scale),
+                max_seq_len=dispatch_max_seq_len,
+            )
+            return output
+
+        is_decode_request = int(attn_metadata.max_query_len) == 1
+        if is_decode_request and dispatch_max_seq_len is None:
+            _record_fallback("missing_cpu_upper_bound")
+        elif is_decode_request:
+            _record_fallback("guard_miss")
+        else:
+            _record_fallback("non_decode")
+        if is_decode_request and is_fish_kvcache_attn_required():
+            raise RuntimeError(
+                "Fish kvcache attention is required but guard rejected the request: "
+                f"query_shape={tuple(q.shape)} query_dtype={q.dtype} "
+                f"max_query_len={attn_metadata.max_query_len} max_seq_len={dispatch_max_seq_len} "
+                f"dcp_world_size={impl.dcp_world_size} use_cascade={attn_metadata.use_cascade}"
+            )
+
+        if not _FIRST_MISS_LOGGED:
+            _FIRST_MISS_LOGGED = True
+            logger.info(
+                "Fish decode-only kvcache attention fast path miss: "
+                "query_shape=%s query_dtype=%s key_cache_shape=%s key_dtype=%s "
+                "block_table_shape=%s block_dtype=%s seq_lens_shape=%s seq_dtype=%s "
+                "max_query_len=%s max_seq_len=%s dcp_world_size=%s use_cascade=%s",
+                tuple(q.shape),
+                q.dtype,
+                tuple(key_cache.shape),
+                key_cache.dtype,
+                tuple(block_table.shape) if block_table is not None else None,
+                getattr(block_table, "dtype", None),
+                tuple(seq_lens.shape),
+                seq_lens.dtype,
+                attn_metadata.max_query_len,
+                dispatch_max_seq_len,
+                impl.dcp_world_size,
+                attn_metadata.use_cascade,
+            )
+
+    return original_forward(
+        layer,
+        query,
+        key,
+        value,
+        kv_cache,
+        attn_metadata,
+        output,
+        output_scale=output_scale,
+        output_block_scale=output_block_scale,
+    )
+
+
+def install_fish_kvcache_attn_backend(model: Any) -> int:
+    """Install the Fish kvcache fast path on this Fish SlowAR model only."""
+    if not _fish_kvcache_enabled():
+        return 0
+    if not is_available():
+        _record_fallback("implementation_unavailable")
+        if is_fish_kvcache_attn_required():
+            raise RuntimeError(
+                f"VLLM_OMNI_FISH_KVCACHE_ATTN=required but Fish kvcache attention is unavailable: {load_error()!r}"
+            )
+        logger.warning(
+            "Fish kvcache attention is enabled but unavailable: %r",
+            load_error(),
+        )
+        return 0
+
+    installed = 0
+    for layer in model.layers:
+        impl = layer.self_attn.attn.impl
+        if getattr(impl, "_fish_kvcache_attn_installed", False):
+            continue
+
+        original_forward = impl.forward
+
+        def fish_forward(
+            this_impl: Any,
+            layer: Any,
+            query: Any,
+            key: Any,
+            value: Any,
+            kv_cache: Any,
+            attn_metadata: Any,
+            output: Any,
+            output_scale: Any = None,
+            output_block_scale: Any = None,
+            *,
+            _original_forward: Callable[..., Any] = original_forward,
+        ) -> Any:
+            return _forward_with_fish_kvcache(
+                this_impl,
+                _original_forward,
+                layer,
+                query,
+                key,
+                value,
+                kv_cache,
+                attn_metadata,
+                output,
+                output_scale=output_scale,
+                output_block_scale=output_block_scale,
+            )
+
+        impl.forward = types.MethodType(fish_forward, impl)
+        impl._fish_kvcache_attn_installed = True
+        installed += 1
+
+    if installed:
+        logger.info("Installed Fish decode-only kvcache attention backend on %d attention layers", installed)
+    elif is_fish_kvcache_attn_required():
+        raise RuntimeError("Fish kvcache attention is required but installed 0 attention layers")
+    return installed
diff --git a/vllm_omni/attention/fish_kvcache_triton.py b/vllm_omni/attention/fish_kvcache_triton.py
new file mode 100644
index 000000000..35d218903
--- /dev/null
+++ b/vllm_omni/attention/fish_kvcache_triton.py
@@ -0,0 +1,359 @@
+# ruff: noqa: N803
+
+from __future__ import annotations
+
+import torch
+
+try:
+    import triton
+    import triton.language as tl
+except Exception as exc:  # pragma: no cover - depends on optional Triton install
+    triton = None
+    tl = None
+    _LOAD_ERROR: Exception | None = exc
+else:
+    _LOAD_ERROR = None
+
+_BLOCK_SIZE = 16
+_HEAD_DIM = 128
+_BLOCK_N = 128
+
+
+def is_available() -> bool:
+    return _LOAD_ERROR is None and triton is not None and tl is not None
+
+
+def load_error() -> Exception | None:
+    return _LOAD_ERROR
+
+
+if is_available():
+
+    @triton.jit
+    def _small_decode_kernel(
+        Q,
+        K,
+        V,
+        BLOCK_TABLE,
+        SEQ_LENS,
+        OUT,
+        SCALE: tl.constexpr,
+        MAX_SEQ_LEN: tl.constexpr,
+        MAX_BLOCKS_PER_SEQ: tl.constexpr,
+        NUM_Q_HEADS: tl.constexpr,
+        NUM_KV_HEADS: tl.constexpr,
+        KV_GROUP: tl.constexpr,
+        BLOCK_H: tl.constexpr,
+        BLOCK_D: tl.constexpr,
+        BLOCK_N: tl.constexpr,
+        GUARD_MAX_SEQ_LEN: tl.constexpr,
+    ):
+        batch_id = tl.program_id(0)
+        kv_head = tl.program_id(1)
+
+        offs_h = tl.arange(0, BLOCK_H)
+        offs_d = tl.arange(0, BLOCK_D)
+        offs_n = tl.arange(0, BLOCK_N)
+        q_heads = kv_head * KV_GROUP + offs_h
+        mask_h = offs_h < KV_GROUP
+
+        q_offsets = batch_id * NUM_Q_HEADS * BLOCK_D + q_heads[:, None] * BLOCK_D + offs_d[None, :]
+        q = tl.load(Q + q_offsets, mask=mask_h[:, None], other=0.0)
+
+        seq_len = tl.load(SEQ_LENS + batch_id)
+        if GUARD_MAX_SEQ_LEN and seq_len > GUARD_MAX_SEQ_LEN:
+            return
+
+        m_i = tl.full((BLOCK_H,), -float("inf"), tl.float32)
+        l_i = tl.zeros((BLOCK_H,), tl.float32)
+        acc = tl.zeros((BLOCK_H, BLOCK_D), tl.float32)
+
+        start_n = 0
+        while start_n < seq_len:
+            cur_n = start_n + offs_n
+            mask_n = cur_n < seq_len
+            logical_block = cur_n // 16
+            block_offset = cur_n - logical_block * 16
+            physical_block = tl.load(
+                BLOCK_TABLE + batch_id * MAX_BLOCKS_PER_SEQ + logical_block,
+                mask=mask_n,
+                other=0,
+            )
+
+            kv_offsets = (
+                (physical_block[:, None] * 16 + block_offset[:, None]) * NUM_KV_HEADS + kv_head
+            ) * BLOCK_D + offs_d[None, :]
+            k = tl.load(K + kv_offsets, mask=mask_n[:, None], other=0.0)
+            qk = tl.dot(q, tl.trans(k)) * SCALE
+            qk = tl.where(mask_h[:, None] & mask_n[None, :], qk, -float("inf"))
+
+            v = tl.load(V + kv_offsets, mask=mask_n[:, None], other=0.0)
+            m_new = tl.maximum(m_i, tl.max(qk, axis=1))
+            p = tl.exp(qk - m_new[:, None])
+            alpha = tl.exp(m_i - m_new)
+            acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v)
+            l_i = l_i * alpha + tl.sum(p, axis=1)
+            m_i = m_new
+            start_n += BLOCK_N
+
+        denom = tl.where(l_i > 0.0, l_i, 1.0)
+        result = tl.where(l_i[:, None] > 0.0, acc / denom[:, None], 0.0)
+        out_offsets = batch_id * NUM_Q_HEADS * BLOCK_D + q_heads[:, None] * BLOCK_D + offs_d[None, :]
+        tl.store(OUT + out_offsets, result, mask=mask_h[:, None])
+
+    @triton.jit
+    def _split_partial_decode_kernel(
+        Q,
+        K,
+        V,
+        BLOCK_TABLE,
+        SEQ_LENS,
+        PARTIAL_M,
+        PARTIAL_L,
+        PARTIAL_ACC,
+        SCALE: tl.constexpr,
+        MAX_BLOCKS_PER_SEQ: tl.constexpr,
+        NUM_Q_HEADS: tl.constexpr,
+        NUM_KV_HEADS: tl.constexpr,
+        KV_GROUP: tl.constexpr,
+        BLOCK_H: tl.constexpr,
+        BLOCK_D: tl.constexpr,
+        BLOCK_N: tl.constexpr,
+        SPLIT_TOKENS: tl.constexpr,
+        BATCH_SIZE: tl.constexpr,
+        MIN_SEQ_LEN: tl.constexpr,
+    ):
+        batch_id = tl.program_id(0)
+        kv_head = tl.program_id(1)
+        split_id = tl.program_id(2)
+
+        offs_h = tl.arange(0, BLOCK_H)
+        offs_d = tl.arange(0, BLOCK_D)
+        offs_n = tl.arange(0, BLOCK_N)
+        q_heads = kv_head * KV_GROUP + offs_h
+        mask_h = offs_h < KV_GROUP
+
+        partial_head_offsets = (split_id * BATCH_SIZE + batch_id) * NUM_Q_HEADS + q_heads
+        partial_acc_offsets = partial_head_offsets[:, None] * BLOCK_D + offs_d[None, :]
+
+        seq_len = tl.load(SEQ_LENS + batch_id)
+        if MIN_SEQ_LEN and seq_len < MIN_SEQ_LEN:
+            return
+
+        m_i = tl.full((BLOCK_H,), -float("inf"), tl.float32)
+        l_i = tl.zeros((BLOCK_H,), tl.float32)
+        acc = tl.zeros((BLOCK_H, BLOCK_D), tl.float32)
+
+        split_begin = split_id * SPLIT_TOKENS
+        if split_begin >= seq_len:
+            tl.store(PARTIAL_M + partial_head_offsets, m_i, mask=mask_h)
+            tl.store(PARTIAL_L + partial_head_offsets, l_i, mask=mask_h)
+            tl.store(PARTIAL_ACC + partial_acc_offsets, acc, mask=mask_h[:, None])
+            return
+
+        q_offsets = batch_id * NUM_Q_HEADS * BLOCK_D + q_heads[:, None] * BLOCK_D + offs_d[None, :]
+        q = tl.load(Q + q_offsets, mask=mask_h[:, None], other=0.0)
+
+        for split_offset in tl.range(0, SPLIT_TOKENS, BLOCK_N):
+            cur_n = split_begin + split_offset + offs_n
+            mask_n = cur_n < seq_len
+            logical_block = cur_n // 16
+            block_offset = cur_n - logical_block * 16
+            physical_block = tl.load(
+                BLOCK_TABLE + batch_id * MAX_BLOCKS_PER_SEQ + logical_block,
+                mask=mask_n,
+                other=0,
+            )
+
+            kv_offsets = (
+                (physical_block[:, None] * 16 + block_offset[:, None]) * NUM_KV_HEADS + kv_head
+            ) * BLOCK_D + offs_d[None, :]
+            k = tl.load(K + kv_offsets, mask=mask_n[:, None], other=0.0)
+            qk = tl.dot(q, tl.trans(k)) * SCALE
+            qk = tl.where(mask_h[:, None] & mask_n[None, :], qk, -float("inf"))
+
+            v = tl.load(V + kv_offsets, mask=mask_n[:, None], other=0.0)
+            m_new = tl.maximum(m_i, tl.max(qk, axis=1))
+            p = tl.exp(qk - m_new[:, None])
+            alpha = tl.exp(m_i - m_new)
+            acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v)
+            l_i = l_i * alpha + tl.sum(p, axis=1)
+            m_i = m_new
+
+        tl.store(PARTIAL_M + partial_head_offsets, m_i, mask=mask_h)
+        tl.store(PARTIAL_L + partial_head_offsets, l_i, mask=mask_h)
+        tl.store(PARTIAL_ACC + partial_acc_offsets, acc, mask=mask_h[:, None])
+
+    @triton.jit
+    def _split_combine_decode_kernel(
+        SEQ_LENS,
+        PARTIAL_M,
+        PARTIAL_L,
+        PARTIAL_ACC,
+        OUT,
+        NUM_Q_HEADS: tl.constexpr,
+        NUM_KV_HEADS: tl.constexpr,
+        KV_GROUP: tl.constexpr,
+        BLOCK_H: tl.constexpr,
+        BLOCK_D: tl.constexpr,
+        BATCH_SIZE: tl.constexpr,
+        NUM_SPLITS: tl.constexpr,
+        MIN_SEQ_LEN: tl.constexpr,
+    ):
+        batch_id = tl.program_id(0)
+        kv_head = tl.program_id(1)
+
+        seq_len = tl.load(SEQ_LENS + batch_id)
+        if MIN_SEQ_LEN and seq_len < MIN_SEQ_LEN:
+            return
+
+        offs_h = tl.arange(0, BLOCK_H)
+        offs_d = tl.arange(0, BLOCK_D)
+        q_heads = kv_head * KV_GROUP + offs_h
+        mask_h = offs_h < KV_GROUP
+
+        global_m = tl.full((BLOCK_H,), -float("inf"), tl.float32)
+        for split_id in tl.range(0, NUM_SPLITS):
+            partial_head_offsets = (split_id * BATCH_SIZE + batch_id) * NUM_Q_HEADS + q_heads
+            split_m = tl.load(PARTIAL_M + partial_head_offsets, mask=mask_h, other=-float("inf"))
+            global_m = tl.maximum(global_m, split_m)
+
+        global_l = tl.zeros((BLOCK_H,), tl.float32)
+        acc = tl.zeros((BLOCK_H, BLOCK_D), tl.float32)
+        for split_id in tl.range(0, NUM_SPLITS):
+            partial_head_offsets = (split_id * BATCH_SIZE + batch_id) * NUM_Q_HEADS + q_heads
+            split_m = tl.load(PARTIAL_M + partial_head_offsets, mask=mask_h, other=-float("inf"))
+            split_l = tl.load(PARTIAL_L + partial_head_offsets, mask=mask_h, other=0.0)
+            weight = tl.exp(split_m - global_m)
+            partial_acc_offsets = partial_head_offsets[:, None] * BLOCK_D + offs_d[None, :]
+            split_acc = tl.load(PARTIAL_ACC + partial_acc_offsets, mask=mask_h[:, None], other=0.0)
+            global_l += split_l * weight
+            acc += split_acc * weight[:, None]
+
+        denom = tl.where(global_l > 0.0, global_l, 1.0)
+        result = tl.where(global_l[:, None] > 0.0, acc / denom[:, None], 0.0)
+        out_offsets = batch_id * NUM_Q_HEADS * BLOCK_D + q_heads[:, None] * BLOCK_D + offs_d[None, :]
+        tl.store(OUT + out_offsets, result, mask=mask_h[:, None])
+
+
+def fish_decode_kvcache_attn_triton(
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    seq_lens: torch.Tensor,
+    out: torch.Tensor,
+    *,
+    scale: float,
+    max_seq_len: int,
+    small_path_max_seq_len: int,
+    long_split_tokens: int,
+    partial_m: torch.Tensor,
+    partial_l: torch.Tensor,
+    partial_acc: torch.Tensor,
+) -> torch.Tensor:
+    if not is_available():
+        raise RuntimeError(f"Fish Triton attention is unavailable: {_LOAD_ERROR!r}")
+    batch_size, num_q_heads, head_dim = query.shape
+    block_size = key_cache.shape[1]
+    num_kv_heads = key_cache.shape[2]
+    if head_dim != _HEAD_DIM or block_size != _BLOCK_SIZE:
+        raise RuntimeError("Fish Triton attention only supports head_dim=128 and block_size=16")
+    if num_q_heads % num_kv_heads != 0:
+        raise RuntimeError("num_q_heads must be divisible by num_kv_heads")
+
+    kv_group = num_q_heads // num_kv_heads
+    block_h = triton.next_power_of_2(kv_group)
+    max_blocks_per_seq = block_table.shape[1]
+
+    if max_seq_len <= small_path_max_seq_len:
+        _small_decode_kernel[(batch_size, num_kv_heads)](
+            query,
+            key_cache,
+            value_cache,
+            block_table,
+            seq_lens,
+            out,
+            SCALE=float(scale),
+            MAX_SEQ_LEN=int(max_seq_len),
+            MAX_BLOCKS_PER_SEQ=max_blocks_per_seq,
+            NUM_Q_HEADS=num_q_heads,
+            NUM_KV_HEADS=num_kv_heads,
+            KV_GROUP=kv_group,
+            BLOCK_H=block_h,
+            BLOCK_D=head_dim,
+            BLOCK_N=_BLOCK_N,
+            GUARD_MAX_SEQ_LEN=0,
+            num_warps=4,
+            num_stages=3,
+        )
+        return out
+
+    _small_decode_kernel[(batch_size, num_kv_heads)](
+        query,
+        key_cache,
+        value_cache,
+        block_table,
+        seq_lens,
+        out,
+        SCALE=float(scale),
+        MAX_SEQ_LEN=small_path_max_seq_len,
+        MAX_BLOCKS_PER_SEQ=max_blocks_per_seq,
+        NUM_Q_HEADS=num_q_heads,
+        NUM_KV_HEADS=num_kv_heads,
+        KV_GROUP=kv_group,
+        BLOCK_H=block_h,
+        BLOCK_D=head_dim,
+        BLOCK_N=_BLOCK_N,
+        GUARD_MAX_SEQ_LEN=small_path_max_seq_len,
+        num_warps=4,
+        num_stages=3,
+    )
+    num_splits = triton.cdiv(int(max_seq_len), long_split_tokens)
+    expected_partial_acc_shape = (num_splits, batch_size, num_q_heads, head_dim)
+    if tuple(partial_acc.shape) != expected_partial_acc_shape:
+        raise RuntimeError(
+            "Fish Triton long attention workspace shape mismatch: "
+            f"expected partial_acc={expected_partial_acc_shape}, got {tuple(partial_acc.shape)}"
+        )
+    _split_partial_decode_kernel[(batch_size, num_kv_heads, num_splits)](
+        query,
+        key_cache,
+        value_cache,
+        block_table,
+        seq_lens,
+        partial_m,
+        partial_l,
+        partial_acc,
+        SCALE=float(scale),
+        MAX_BLOCKS_PER_SEQ=max_blocks_per_seq,
+        NUM_Q_HEADS=num_q_heads,
+        NUM_KV_HEADS=num_kv_heads,
+        KV_GROUP=kv_group,
+        BLOCK_H=block_h,
+        BLOCK_D=head_dim,
+        BLOCK_N=_BLOCK_N,
+        SPLIT_TOKENS=long_split_tokens,
+        BATCH_SIZE=batch_size,
+        MIN_SEQ_LEN=small_path_max_seq_len + 1,
+        num_warps=4,
+        num_stages=3,
+    )
+    _split_combine_decode_kernel[(batch_size, num_kv_heads)](
+        seq_lens,
+        partial_m,
+        partial_l,
+        partial_acc,
+        out,
+        NUM_Q_HEADS=num_q_heads,
+        NUM_KV_HEADS=num_kv_heads,
+        KV_GROUP=kv_group,
+        BLOCK_H=block_h,
+        BLOCK_D=head_dim,
+        BATCH_SIZE=batch_size,
+        NUM_SPLITS=num_splits,
+        MIN_SEQ_LEN=small_path_max_seq_len + 1,
+        num_warps=4,
+        num_stages=3,
+    )
+    return out
diff --git a/vllm_omni/benchmarks/data_modules/random_multi_modal_dataset.py b/vllm_omni/benchmarks/data_modules/random_multi_modal_dataset.py
index 14ba86cc4..aa597f12a 100644
--- a/vllm_omni/benchmarks/data_modules/random_multi_modal_dataset.py
+++ b/vllm_omni/benchmarks/data_modules/random_multi_modal_dataset.py
@@ -123,10 +123,6 @@ class OmniRandomMultiModalDataset(RandomMultiModalDataset):
                 "23",
                 "-movflags",
                 "+faststart",
-                "-pix_fmt",
-                "yuv420p",
-                "-vf",
-                f"scale={width}:{height}",
             ],
         }
 
diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py
index 00df0dff8..938dd74e7 100644
--- a/vllm_omni/config/pipeline_registry.py
+++ b/vllm_omni/config/pipeline_registry.py
@@ -102,6 +102,10 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
         "vllm_omni.model_executor.models.voxtral_tts.pipeline",
         "VOXTRAL_TTS_PIPELINE",
     ),
+    "glm_tts": (
+        "vllm_omni.model_executor.models.glm_tts.pipeline",
+        "GLM_TTS_PIPELINE",
+    ),
     "fish_qwen3_omni": (
         "vllm_omni.model_executor.models.fish_speech.pipeline",
         "FISH_SPEECH_PIPELINE",
@@ -118,8 +122,16 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
         "vllm_omni.model_executor.models.ming_flash_omni.pipeline",
         "MING_FLASH_OMNI_THINKER_ONLY_PIPELINE",
     ),
+    "ming_flash_omni_image": (
+        "vllm_omni.model_executor.models.ming_flash_omni.pipeline",
+        "MING_FLASH_OMNI_IMAGE_PIPELINE",
+    ),
     "moss_tts_nano": (
         "vllm_omni.model_executor.models.moss_tts_nano.pipeline",
         "MOSS_TTS_NANO_PIPELINE",
     ),
+    "higgs_audio_v2": (
+        "vllm_omni.model_executor.models.higgs_audio_v2.pipeline",
+        "HIGGS_AUDIO_V2_PIPELINE",
+    ),
 }
diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py
index 4188606d7..f6a88546f 100644
--- a/vllm_omni/config/stage_config.py
+++ b/vllm_omni/config/stage_config.py
@@ -10,7 +10,7 @@ import warnings
 from dataclasses import asdict, dataclass, field, fields
 from enum import Enum
 from pathlib import Path
-from typing import Any
+from typing import Any, NamedTuple
 
 from vllm.logger import init_logger
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
@@ -440,6 +440,7 @@ class StageDeployConfig:
     stage_id: int
     devices: str | None = None
     num_replicas: int = 1
+    env: dict[str, Any] | None = None
 
     # Inter-stage connector wiring and request defaults.
     output_connectors: dict[str, str] | None = None
@@ -522,6 +523,7 @@ _STAGE_RESERVED_KEYS = frozenset(
         "stage_id",
         "devices",
         "num_replicas",
+        "env",
         "output_connectors",
         "input_connectors",
         "default_sampling_params",
@@ -542,6 +544,7 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig:
     runtime_cfg = dict(stage_data.get("runtime", {}))
     devices = runtime_cfg.get("devices", stage_data.get("devices"))
     num_replicas = runtime_cfg.get("num_replicas", stage_data.get("num_replicas", 1))
+    env = runtime_cfg.get("env", stage_data.get("env"))
 
     if "engine_args" in stage_data:
         for k, v in stage_data["engine_args"].items():
@@ -556,6 +559,7 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig:
         "stage_id": stage_data["stage_id"],
         "devices": devices,
         "num_replicas": int(num_replicas),
+        "env": env,
     }
     for name, f in _STAGE_DEPLOY_FIELDS.items():
         if name in flat_args:
@@ -695,8 +699,14 @@ def load_deploy_config(path: str | Path) -> DeployConfig:
     return DeployConfig(**kwargs)
 
 
-def _extract_platform_overrides(ps: dict[str, Any]) -> tuple[dict[str, Any], str | None]:
-    """Return ``(overrides, devices)`` from a platform stage entry.
+class PlatformOverrides(NamedTuple):
+    overrides: dict[str, Any]
+    devices: str | None
+    env: dict[str, Any] | None
+
+
+def _extract_platform_overrides(ps: dict[str, Any]) -> PlatformOverrides:
+    """Return overrides, devices, and env from a platform stage entry.
 
     Handles both the nested layout (``engine_args:`` / ``runtime.devices``) and
     the flat layout. ``devices`` is ``None`` when no override is set.
@@ -706,9 +716,9 @@ def _extract_platform_overrides(ps: dict[str, Any]) -> tuple[dict[str, Any], str
         runtime_cfg = ps.get("runtime", {})
         if "num_replicas" in runtime_cfg:
             overrides["num_replicas"] = runtime_cfg["num_replicas"]
-        return overrides, runtime_cfg.get("devices")
-    overrides = {k: v for k, v in ps.items() if k not in ("stage_id", "devices")}
-    return overrides, ps.get("devices")
+        return PlatformOverrides(overrides, runtime_cfg.get("devices"), runtime_cfg.get("env"))
+    overrides = {k: v for k, v in ps.items() if k not in ("stage_id", "devices", "env")}
+    return PlatformOverrides(overrides, ps.get("devices"), ps.get("env"))
 
 
 def _apply_platform_overrides(
@@ -733,10 +743,21 @@ def _apply_platform_overrides(
         base = base_by_id.get(ps["stage_id"])
         if base is None:
             continue
-        overrides, devices = _extract_platform_overrides(ps)
-        if devices is not None:
-            base.devices = devices
-        for key, val in overrides.items():
+        po = _extract_platform_overrides(ps)
+        if po.devices is not None:
+            base.devices = po.devices
+        if po.env is not None:
+            if isinstance(base.env, dict) and isinstance(po.env, dict):
+                base.env = {**base.env, **po.env}
+            else:
+                logger.warning(
+                    "Stage %s env override replaces base env entirely (base type=%s, override type=%s)",
+                    ps["stage_id"],
+                    type(base.env).__name__,
+                    type(po.env).__name__,
+                )
+                base.env = po.env
+        for key, val in po.overrides.items():
             if hasattr(base, key):
                 # Deep-merge dict-valued fields listed in _DEEP_MERGE_KEYS so
                 # platform overlays don't silently clobber sibling keys (e.g.
@@ -924,6 +945,8 @@ def merge_pipeline_deploy(
             if ds.devices is not None:
                 runtime["devices"] = ds.devices
             runtime["num_replicas"] = ds.num_replicas
+            if ds.env is not None:
+                runtime["env"] = ds.env
         runtime["requires_multimodal_data"] = ps.requires_multimodal_data
 
         result.append(
diff --git a/vllm_omni/core/prefix_cache.py b/vllm_omni/core/prefix_cache.py
index 69e7346c4..f80074173 100644
--- a/vllm_omni/core/prefix_cache.py
+++ b/vllm_omni/core/prefix_cache.py
@@ -47,6 +47,7 @@ class OmniTensorPrefixCache:
         self.mm_outputs_cache = {}
         self.mm_cache_keys = set()
         self._new_req_cache_hit_ids: set[str] = set()
+        self._deferred_mm_outputs: dict[str, dict[str, list[tuple[int, torch.Tensor]]]] = {}
 
     def maybe_init_missing_mm_cache_keys(self, multimodal_outputs: dict, seq_len: int):
         """Given multimodal outputs from executing the model, dynamically
@@ -62,7 +63,17 @@ class OmniTensorPrefixCache:
         determined by the warmup.
         """
         for key, val in multimodal_outputs.items():
-            if isinstance(val, torch.Tensor) and val.shape[0] == seq_len and key not in self.mm_cache_keys:
+            # Only cache per-token feature tensors: 2D+ with first dim == seq_len.
+            # A 1D tensor of shape (seq_len,) is a broadcast scalar (per-request
+            # metadata such as ref_code_len / codec_streaming), not per-token data;
+            # caching it by slot causes a shape mismatch when a later request has a
+            # different scheduled seq length.
+            if (
+                isinstance(val, torch.Tensor)
+                and val.ndim >= 2
+                and val.shape[0] == seq_len
+                and key not in self.mm_cache_keys
+            ):
                 feat_dim = val.shape[-1]
                 self.mm_outputs_cache[key] = self._get_cache_tensor(
                     dtype=val.dtype,
@@ -89,6 +100,10 @@ class OmniTensorPrefixCache:
         """Clears the cache hit IDs to prepare for a new engine step."""
         self._new_req_cache_hit_ids.clear()
 
+    def has_prefix_cached_new_req_ids(self) -> bool:
+        """Return True when this step contains a newly scheduled prefix hit."""
+        return bool(self._new_req_cache_hit_ids)
+
     @staticmethod
     def _coerce_to_cpu_tensor(maybe_gpu_tensor: torch.Tensor) -> torch.Tensor:
         """Convert GPU tensors -> contiguous CPU tensors if needed."""
@@ -101,6 +116,7 @@ class OmniTensorPrefixCache:
         num_tokens_unpadded: int,
         slot_mapping: torch.Tensor,
         num_tokens_padded: int | None = None,
+        skip_mm_cache_keys: set[str] | None = None,
     ):
         """Updates the hidden cache state for the provided hidden states and multimodal outputs.
 
@@ -110,10 +126,12 @@ class OmniTensorPrefixCache:
             num_tokens_unpadded: Number of tokens without padding
             slot_mapping: Slot mapping for the input sequence
             num_tokens_padded: Total number of tokens including padding
+            skip_mm_cache_keys: Multimodal keys whose CPU cache writes are deferred
         """
         unpadded_slot_mapping = slot_mapping[:num_tokens_unpadded]
         if num_tokens_padded is None:
             num_tokens_padded = num_tokens_unpadded
+        skip_mm_cache_keys = skip_mm_cache_keys or set()
 
         if hidden_states is not None:
             # Slice to unpadded portion before caching
@@ -136,6 +154,8 @@ class OmniTensorPrefixCache:
 
             for mm_out_key, mm_cache in self.mm_outputs_cache.items():
                 if mm_out_key in multimodal_outputs:
+                    if mm_out_key in skip_mm_cache_keys:
+                        continue
                     # Slice to unpadded portion before caching
                     mm_state = multimodal_outputs[mm_out_key][:num_tokens_unpadded]
                     mm_state = OmniTensorPrefixCache._coerce_to_cpu_tensor(mm_state)
@@ -143,6 +163,165 @@ class OmniTensorPrefixCache:
                     flat_cache[unpadded_slot_mapping] = mm_state
             logger.debug("Writing to mm output cache for %s tokens", num_tokens_unpadded)
 
+    def stage_deferred_mm_outputs(
+        self,
+        query_start_loc: torch.Tensor,
+        input_batch: InputBatch,
+        multimodal_outputs: dict[str, torch.Tensor] | None,
+        num_scheduled_tokens: dict[str, int],
+        deferred_mm_cache_keys: set[str],
+    ) -> None:
+        """Keep GPU multimodal chunks until a request finishes.
+
+        The normal prefix-cache write path copies every cached multimodal tensor
+        to CPU on every step.  For model outputs that are only needed by future
+        full-block prefix hits, we can keep detached GPU chunks and materialize
+        the CPU cache once when the request completes.
+        """
+        if not multimodal_outputs or not deferred_mm_cache_keys:
+            return
+
+        for mm_key in deferred_mm_cache_keys:
+            mm_state = multimodal_outputs.get(mm_key)
+            if not isinstance(mm_state, torch.Tensor) or mm_state.ndim < 2:
+                continue
+            for req_id in input_batch.req_ids:
+                req_idx = input_batch.req_id_to_index[req_id]
+                sched = int(num_scheduled_tokens.get(req_id, 0))
+                if sched <= 0:
+                    continue
+                start = int(query_start_loc[req_idx])
+                end = start + sched
+                if start >= int(mm_state.shape[0]):
+                    continue
+                end = min(end, int(mm_state.shape[0]))
+                if end <= start:
+                    continue
+
+                computed_end = int(input_batch.num_computed_tokens_cpu[req_idx])
+                token_start = max(0, computed_end - sched)
+                chunk = mm_state[start:end].detach()
+                self._deferred_mm_outputs.setdefault(req_id, {}).setdefault(mm_key, []).append((token_start, chunk))
+
+    def commit_deferred_mm_outputs(
+        self,
+        finished_req_ids: set[str] | list[str],
+        input_batch: InputBatch,
+    ) -> None:
+        """Write deferred multimodal chunks into the CPU prefix cache.
+
+        This must run before finished requests are removed from ``input_batch``,
+        because the block table is needed to map logical token positions to KV
+        cache slots.
+        """
+        if not finished_req_ids or not self._deferred_mm_outputs:
+            return
+
+        for req_id in tuple(finished_req_ids):
+            per_key = self._deferred_mm_outputs.pop(req_id, None)
+            if not per_key:
+                continue
+            req_idx = input_batch.req_id_to_index.get(req_id)
+            if req_idx is None:
+                continue
+            for mm_key, chunks in per_key.items():
+                self._commit_deferred_mm_key(req_idx, input_batch, mm_key, chunks)
+
+    def discard_deferred_mm_outputs(self, req_id: str) -> None:
+        """Drop deferred chunks for requests that leave without a cache commit."""
+        self._deferred_mm_outputs.pop(req_id, None)
+
+    def _commit_deferred_mm_key(
+        self,
+        req_idx: int,
+        input_batch: InputBatch,
+        mm_key: str,
+        chunks: list[tuple[int, torch.Tensor]],
+    ) -> None:
+        if not chunks:
+            return
+
+        chunks = sorted(chunks, key=lambda item: item[0])
+        first_start = chunks[0][0]
+        expected_start = first_start
+        tensors: list[torch.Tensor] = []
+        for token_start, tensor in chunks:
+            if tensor.numel() == 0:
+                continue
+            if token_start != expected_start:
+                if tensors:
+                    self._write_deferred_mm_tensor(req_idx, input_batch, mm_key, tensors, first_start)
+                first_start = token_start
+                expected_start = token_start
+                tensors = []
+            tensors.append(tensor)
+            expected_start += int(tensor.shape[0])
+
+        if tensors:
+            self._write_deferred_mm_tensor(req_idx, input_batch, mm_key, tensors, first_start)
+
+    def _write_deferred_mm_tensor(
+        self,
+        req_idx: int,
+        input_batch: InputBatch,
+        mm_key: str,
+        tensors: list[torch.Tensor],
+        token_start: int,
+    ) -> None:
+        mm_state = tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=0)
+        if mm_state.ndim < 2 or mm_state.shape[0] == 0:
+            return
+
+        if mm_key not in self.mm_cache_keys:
+            self.mm_outputs_cache[mm_key] = self._get_cache_tensor(
+                dtype=mm_state.dtype,
+                hidden_size=int(mm_state.shape[-1]),
+            )
+            self.mm_cache_keys.add(mm_key)
+            logger.info(
+                "Initializing deferred multimodal output cache of size %s for key: %s",
+                list(self.mm_outputs_cache[mm_key].shape),
+                mm_key,
+            )
+
+        slots = self._get_slot_ids_for_token_range(
+            req_idx=req_idx,
+            input_batch=input_batch,
+            token_start=token_start,
+            num_tokens=int(mm_state.shape[0]),
+        )
+        if slots.numel() == 0:
+            return
+
+        mm_state_cpu = OmniTensorPrefixCache._coerce_to_cpu_tensor(mm_state[: slots.numel()])
+        mm_cache = self.mm_outputs_cache[mm_key]
+        flat_cache = mm_cache.view(-1, mm_cache.shape[-1])
+        flat_cache[slots] = mm_state_cpu
+
+    def _get_slot_ids_for_token_range(
+        self,
+        req_idx: int,
+        input_batch: InputBatch,
+        token_start: int,
+        num_tokens: int,
+    ) -> torch.Tensor:
+        if num_tokens <= 0:
+            return torch.empty((0,), dtype=torch.long)
+
+        block_table = input_batch.block_table[0].block_table.cpu
+        token_positions = torch.arange(token_start, token_start + num_tokens, dtype=torch.long)
+        block_offsets = token_positions // self.block_size
+        max_blocks = int(block_table.shape[1])
+        valid = block_offsets < max_blocks
+        if not bool(valid.all()):
+            token_positions = token_positions[valid]
+            block_offsets = block_offsets[valid]
+        if token_positions.numel() == 0:
+            return torch.empty((0,), dtype=torch.long)
+
+        block_ids = block_table[req_idx, block_offsets].to(torch.long)
+        return block_ids * self.block_size + (token_positions % self.block_size)
+
     def _coerce_to_payload_dict(
         self,
         element: object,
@@ -151,17 +330,25 @@ class OmniTensorPrefixCache:
         num_scheduled_tokens: dict[str, int],
     ) -> dict[str, object]:
         """Build the multimodal passthrough data per request for
-        the object under consideration. This is identical to the case
-        for no prefix cache when we tensor does have a first dimension
-        matching the seq len.
+        the object under consideration. This mirrors the no-prefix-cache
+        path: a tensor whose first dimension matches the total scheduled
+        token count is sliced per request, so 1D (seq_len,) metadata that
+        is intentionally not cached (e.g. ref_code_len, codec_streaming)
+        is still split per request instead of leaking the whole batch.
         """
+        total_scheduled_tokens = sum(int(num_scheduled_tokens[r]) for r in input_batch.req_ids)
         elem_dict = {}
         for req_id in input_batch.req_ids:
             req_idx = input_batch.req_id_to_index[req_id]
             start = query_start_loc[req_idx]
             end = start + num_scheduled_tokens[req_id]
             elem_dict[req_id] = to_payload_element(
-                element, req_idx, start=start, end=end, pass_lists_through=True, seq_len=None
+                element,
+                req_idx,
+                start=start,
+                end=end,
+                pass_lists_through=True,
+                seq_len=total_scheduled_tokens,
             )
         return elem_dict
 
diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py
index 61ee24547..5c0c489d1 100644
--- a/vllm_omni/core/sched/omni_ar_scheduler.py
+++ b/vllm_omni/core/sched/omni_ar_scheduler.py
@@ -13,7 +13,7 @@ from vllm.v1.core.sched.async_scheduler import AsyncScheduler as AsyncVLLMSchedu
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.core.sched.utils import remove_all
-from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs
+from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
 from vllm.v1.metrics.perf import PerfStats
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
@@ -22,9 +22,8 @@ from vllm.v1.spec_decode.metrics import SpecDecodingStats
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 from vllm_omni.core.sched.omni_scheduling_coordinator import (
     OmniSchedulingCoordinator,
-    uses_qwen3_omni_full_payload_input_coordinator,
+    uses_full_payload_input_coordinator,
 )
-from vllm_omni.core.sched.output import OmniSchedulerOutput
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
 )
@@ -81,7 +80,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         if getattr(model_config, "async_chunk", False):
             self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config)
         self.input_coordinator: OmniSchedulingCoordinator | None = None
-        if uses_qwen3_omni_full_payload_input_coordinator(model_config):
+        if uses_full_payload_input_coordinator(model_config):
             self.input_coordinator = OmniSchedulingCoordinator(
                 scheduler_max_num_seqs=self.vllm_config.scheduler_config.max_num_seqs,
                 stage_id=getattr(model_config, "stage_id", 0),
@@ -211,18 +210,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             for req in list(queue):
                 if getattr(req, "status", None) == RequestStatus.FINISHED_ABORTED:
                     queue.remove(req)
-        connector_output = self._latest_omni_connector_output
-        self._latest_omni_connector_output = None
-        if self.input_coordinator:
-            if connector_output and connector_output.request_metadata:
-                self.input_coordinator.update_request_metadata(
-                    self.requests, connector_output.request_metadata, model_mode="ar"
-                )
-            self.input_coordinator.process_pending_full_payload_inputs(
-                self.waiting,
-                self.running,
-                connector_output.stage_recv_req_ids if connector_output else set(),
-            )
+        self._consume_pending_connector_output(model_mode="ar")
+        self._process_pending_input_timeouts()
 
         if self.chunk_transfer_adapter:
             self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running)
@@ -232,7 +221,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         finally:
             if self.chunk_transfer_adapter:
                 # Add request waiting for chunk to the waiting and running queue
-                self.chunk_transfer_adapter.restore_queues(self.waiting, self.running)
+                self.chunk_transfer_adapter.restore_queues(
+                    self.waiting,
+                    self.running,
+                    scheduler_requests=self.requests,
+                )
             if self.input_coordinator:
                 self.input_coordinator.restore_queues(self.waiting, self.running)
         try:
@@ -274,13 +267,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             finished_reqs = {}
 
         # Wrap in omni scheduler output to carry transfer metadata.
-        base_fields = SchedulerOutput.__dataclass_fields__.keys()
-        base_data = {name: getattr(scheduler_output, name) for name in base_fields}
-        input_regs = self.input_coordinator.pending_input_registrations if self.input_coordinator else []
-        return OmniSchedulerOutput(
-            **base_data,
+        return self._wrap_omni_scheduler_output(
+            scheduler_output,
             finished_requests_needing_kv_transfer=finished_reqs,
-            pending_input_registrations=input_regs,
         )
 
     def update_from_output(
@@ -437,8 +426,19 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # Capture finish_reason BEFORE _handle_stopped_request, which may
                 # reset the status to WAITING for streaming requests that continue.
                 finish_reason = request.get_finished_reason()
-                is_segment_finished = request.is_finished() and request.resumable
+                is_segment_finished = True
                 finished = self._handle_stopped_request(request)
+                if not finished:
+                    # for streaming input request only
+                    if self.chunk_transfer_adapter:
+                        if self.vllm_config.model_config.stage_id != 0:
+                            # Downstream async-chunk stages receive real payloads from the
+                            # connector. This update only resumes polling for the next segment.
+                            self.chunk_transfer_adapter.segment_finished_requests.discard(request.request_id)
+                    request.discard_latest_async_tokens = True
+                    request.num_output_placeholders = 0
+                    request.spec_token_ids = []
+                    request._output_token_ids.clear()
                 if finished:
                     kv_transfer_params = self._free_request(request)
                 if status_before_stop == RequestStatus.RUNNING:
@@ -482,7 +482,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     )
                 )
                 if self.chunk_transfer_adapter is not None:
-                    self.chunk_transfer_adapter.save_async(pooler_output, request)
+                    self.chunk_transfer_adapter.save_async(pooler_output, request, is_segment_finished)
             else:
                 # Invariant: EngineCore returns no partial prefill outputs.
                 assert not prompt_logprobs_tensors
@@ -577,15 +577,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 engine_core_outputs[0] = eco = EngineCoreOutputs()
             eco.scheduler_stats = stats
 
-        omni_output = getattr(model_runner_output, "omni_connector_output", None)
-        if omni_output is not None:
-            self._latest_omni_connector_output = omni_output
-            if self.input_coordinator and omni_output.request_metadata:
-                self.input_coordinator.update_request_metadata(
-                    self.requests,
-                    omni_output.request_metadata,
-                    model_mode="ar",
-                )
+        self._capture_omni_connector_output(model_runner_output)
 
         # Free blocks that were held for transfer (kv_ready and
         # active_kv_transfers updates already done before the per-request loop).
@@ -641,11 +633,36 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         """
         req_id = session.request_id
         self._new_prompt_len_snapshot[req_id] = len(update.prompt_token_ids)
-        if self.vllm_config.model_config.stage_id != 0:
-            self._replace_session_with_streaming_update(session, update)
+        outstanding_async_tokens = getattr(session, "num_output_placeholders", 0)
+        if outstanding_async_tokens > 0:
+            # Async scheduling may already have sampled the previous
+            # segment's next token. Drop that late token instead of
+            # appending it to the new streaming segment.
+            session.discard_latest_async_tokens = True
+            session.num_computed_tokens -= session.num_output_placeholders
+            session.num_output_placeholders = 0
+            session.spec_token_ids = []
+        if self.chunk_transfer_adapter:
+            self.chunk_transfer_adapter.requests_num_chunks_sent.pop(session.external_req_id, None)
+            if self.vllm_config.model_config.stage_id != 0:
+                # Downstream async-chunk stages receive real payloads from the
+                # connector. This update only resumes polling for the next segment.
+                self.chunk_transfer_adapter.segment_finished_requests.discard(session.request_id)
+                # Do not replace prompt/additional_information here; the next
+                # upstream chunk will populate them in chunk transfer adapter.
+                session.arrival_time = update.arrival_time
+                session.sampling_params = update.sampling_params
+                if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ:
+                    self.num_waiting_for_streaming_input -= 1
+                session.status = RequestStatus.WAITING
+                if session in self.skipped_waiting:
+                    self.skipped_waiting.remove_requests((session,))
+                    self._enqueue_waiting_request(session)
 
-        else:
-            super()._update_request_as_session(session, update)
+                if self.log_stats:
+                    session.record_event(EngineCoreEventType.QUEUED)
+                return
+        super()._update_request_as_session(session, update)
 
     def _free_request(self, request: Request, delay_free_blocks: bool = False) -> dict[str, Any] | None:
         # TODO(wzliu)! for offline mode, we should not end process until all data is transferred
@@ -664,70 +681,73 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         if self.finished_req_ids_dict is not None:
             self.finished_req_ids_dict[request.client_index].add(request_id)
 
-        # 2. Omni Specific: Check if we need to transfer KV
-        if self._should_transfer_kv_for_request(request_id):
-            already_triggered = request_id in self.transfer_triggered_requests
-            is_active = request_id in self.active_kv_transfers
+        # Mirror the generation scheduler's try/finally pattern so the
+        # input_coordinator entry is always pruned along every return path,
+        # including the early returns for in-flight / waiting KV transfers
+        # below. _free_input_coordinator_request is a no-op when the
+        # coordinator is None, so the unconditional finally is safe.
+        try:
+            # 2. Omni Specific: Check if we need to transfer KV
+            if self._should_transfer_kv_for_request(request_id):
+                already_triggered = request_id in self.transfer_triggered_requests
+                is_active = request_id in self.active_kv_transfers
 
-            if already_triggered:
-                if is_active:
-                    # It triggered but hasn't finished yet. We MUST wait.
-                    logger.debug(f"[Omni] Request {request_id} finished but transfer is still ACTIVE. Waiting.")
-                    self.waiting_for_transfer_free.add(request_id)
-                    if self.input_coordinator is not None:
-                        self._free_input_coordinator_request(request_id)
-                    kv_xfer_params = None
-                    return kv_xfer_params
-                elif request_id in self.waiting_for_transfer_free:
-                    # Blocks held until KV extraction completes in a future step.
-                    if self.input_coordinator is not None:
-                        self._free_input_coordinator_request(request_id)
-                    return None
+                if already_triggered:
+                    if is_active:
+                        # It triggered but hasn't finished yet. We MUST wait.
+                        logger.debug(f"[Omni] Request {request_id} finished but transfer is still ACTIVE. Waiting.")
+                        self.waiting_for_transfer_free.add(request_id)
+                        kv_xfer_params = None
+                        return kv_xfer_params
+                    elif request_id in self.waiting_for_transfer_free:
+                        # Blocks held until KV extraction completes in a future step.
+                        return None
+                    else:
+                        logger.debug(
+                            f"[Omni] Request {request_id} finished and transfer no longer ACTIVE (extracted/acked). "
+                            "Freeing immediately."
+                        )
                 else:
-                    logger.debug(
-                        f"[Omni] Request {request_id} finished and transfer no longer ACTIVE (extracted/acked). "
-                        "Freeing immediately."
-                    )
-            else:
-                self.waiting_for_transfer_free.add(request_id)
-                confirmed_computed = self._get_confirmed_num_computed_tokens(request)
-                self._mark_request_for_kv_transfer(request_id, confirmed_computed)
-                # Return KV transfer metadata so it propagates to RequestOutput
-                if request_id in self.requests_needing_kv_transfer:
-                    transfer_data = self.requests_needing_kv_transfer[request_id]
-                    kv_xfer_params = {
-                        "past_key_values": transfer_data["block_ids"],
-                        "kv_metadata": {"seq_len": transfer_data["seq_len"], "block_ids": transfer_data["block_ids"]},
-                    }
-                    # Also update request.additional_information for good measure
-                    add_info = getattr(request, "additional_information", None)
-                    # If additional_information is an AdditionalInformationPayload-like object,
-                    # unpack it into a plain dict.
-                    if (
-                        add_info is not None
-                        and hasattr(add_info, "entries")
-                        and isinstance(getattr(add_info, "entries"), dict)
-                    ):
-                        request.additional_information = deserialize_additional_information(add_info)
-                        add_info = request.additional_information
-                    if add_info is None:
-                        request.additional_information = {}
-                        add_info = request.additional_information
-                    if isinstance(add_info, dict):
-                        add_info.update(kv_xfer_params)
+                    self.waiting_for_transfer_free.add(request_id)
+                    confirmed_computed = self._get_confirmed_num_computed_tokens(request)
+                    self._mark_request_for_kv_transfer(request_id, confirmed_computed)
+                    # Return KV transfer metadata so it propagates to RequestOutput
+                    if request_id in self.requests_needing_kv_transfer:
+                        transfer_data = self.requests_needing_kv_transfer[request_id]
+                        kv_xfer_params = {
+                            "past_key_values": transfer_data["block_ids"],
+                            "kv_metadata": {
+                                "seq_len": transfer_data["seq_len"],
+                                "block_ids": transfer_data["block_ids"],
+                            },
+                        }
+                        # Also update request.additional_information for good measure
+                        add_info = getattr(request, "additional_information", None)
+                        # If additional_information is an AdditionalInformationPayload-like object,
+                        # unpack it into a plain dict.
+                        if (
+                            add_info is not None
+                            and hasattr(add_info, "entries")
+                            and isinstance(getattr(add_info, "entries"), dict)
+                        ):
+                            request.additional_information = deserialize_additional_information(add_info)
+                            add_info = request.additional_information
+                        if add_info is None:
+                            request.additional_information = {}
+                            add_info = request.additional_information
+                        if isinstance(add_info, dict):
+                            add_info.update(kv_xfer_params)
 
-                if self.input_coordinator is not None:
-                    self._free_input_coordinator_request(request_id)
-                return kv_xfer_params
+                    return kv_xfer_params
 
-        # 3. Standard Freeing
-        delay_free_blocks |= connector_delay_free_blocks
-        if self.input_coordinator is not None:
+            # 3. Standard Freeing
+            delay_free_blocks |= connector_delay_free_blocks
+            if not delay_free_blocks:
+                self._free_blocks(request)
+
+            return kv_xfer_params
+        finally:
             self._free_input_coordinator_request(request_id)
-        if not delay_free_blocks:
-            self._free_blocks(request)
-
-        return kv_xfer_params
 
     def _free_blocks(self, request: Request):
         # Helper to match base class structure if not directly available
diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py
index 1218d416b..dae3be8ae 100644
--- a/vllm_omni/core/sched/omni_generation_scheduler.py
+++ b/vllm_omni/core/sched/omni_generation_scheduler.py
@@ -26,9 +26,9 @@ from vllm.v1.spec_decode.metrics import SpecDecodingStats
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 from vllm_omni.core.sched.omni_scheduling_coordinator import (
     OmniSchedulingCoordinator,
-    uses_qwen3_omni_full_payload_input_coordinator,
+    uses_full_payload_input_coordinator,
 )
-from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData, OmniSchedulerOutput
+from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
 )
@@ -46,7 +46,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config)
         self._pending_finish_reqs: list[Request] = []
         self.input_coordinator: OmniSchedulingCoordinator | None = None
-        if uses_qwen3_omni_full_payload_input_coordinator(model_config):
+        if uses_full_payload_input_coordinator(model_config):
             self.input_coordinator = OmniSchedulingCoordinator(
                 scheduler_max_num_seqs=self.vllm_config.scheduler_config.max_num_seqs,
                 stage_id=getattr(model_config, "stage_id", 0),
@@ -82,18 +82,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         # Temporary queue: preserve waiting order, do not disturb non-diffusion requests
         skipped_waiting_requests = create_request_queue(self.policy)
         req_index = 0
-        connector_output = self._latest_omni_connector_output
-        self._latest_omni_connector_output = None
-        if self.input_coordinator:
-            if connector_output and connector_output.request_metadata:
-                self.input_coordinator.update_request_metadata(
-                    self.requests, connector_output.request_metadata, model_mode="generation"
-                )
-            self.input_coordinator.process_pending_full_payload_inputs(
-                self.waiting,
-                self.running,
-                connector_output.stage_recv_req_ids if connector_output else set(),
-            )
+        self._consume_pending_connector_output(model_mode="generation")
+        self._process_pending_input_timeouts()
         if self.chunk_transfer_adapter:
             self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running)
 
@@ -112,10 +102,14 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             num_computed_tokens = request.num_computed_tokens
             required_tokens = len(request.prompt_token_ids) - num_computed_tokens
+            if not self.scheduler_config.enable_chunked_prefill and required_tokens > token_budget:
+                # If chunked_prefill is disabled,
+                # we can stop the scheduling here.
+                break
+            # async_chunk: don't schedule placeholder tokens when no new chunk is available.
             if required_tokens <= 0:
-                if (
-                    self.chunk_transfer_adapter is not None
-                    and request.request_id in self.chunk_transfer_adapter.finished_requests
+                if self.chunk_transfer_adapter is not None and self.chunk_transfer_adapter.is_done_receiving_chunks(
+                    request.request_id
                 ):
                     self._pending_finish_reqs.append(request)
                 req_index += 1
@@ -165,7 +159,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             # async_chunk: wait for the first upstream chunk (don't start with placeholders).
             if self.chunk_transfer_adapter is not None and len(request.prompt_token_ids) == 0:
-                if request.request_id in self.chunk_transfer_adapter.finished_requests:
+                if self.chunk_transfer_adapter.is_done_receiving_chunks(request.request_id):
                     self.waiting.pop_request()
                     self._pending_finish_reqs.append(request)
                     continue
@@ -213,19 +207,16 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             if self.chunk_transfer_adapter:
                 # Don't fall back: base scheduler doesn't handle async_chunk
                 # requests with empty prompt_token_ids.
-                self.chunk_transfer_adapter.restore_queues(self.waiting, self.running)
+                self.chunk_transfer_adapter.restore_queues(
+                    self.waiting,
+                    self.running,
+                    scheduler_requests=self.requests,
+                )
             else:
                 res = super().schedule()
                 if self.input_coordinator:
                     self.input_coordinator.restore_queues(self.waiting, self.running)
-                base_fields = SchedulerOutput.__dataclass_fields__.keys()
-                base_data = {name: getattr(res, name) for name in base_fields}
-                return OmniSchedulerOutput(
-                    **base_data,
-                    pending_input_registrations=(
-                        self.input_coordinator.pending_input_registrations if self.input_coordinator else []
-                    ),
-                )
+                return self._wrap_omni_scheduler_output(res)
 
         # Compute common prefix blocks (aligned with v1)
         num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups)
@@ -345,18 +336,15 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             # otherwise they are permanently orphaned in the adapter's
             # internal deques and never scheduled again.
             if self.chunk_transfer_adapter:
-                self.chunk_transfer_adapter.restore_queues(self.waiting, self.running)
+                self.chunk_transfer_adapter.restore_queues(
+                    self.waiting,
+                    self.running,
+                    scheduler_requests=self.requests,
+                )
             if self.input_coordinator:
                 self.input_coordinator.restore_queues(self.waiting, self.running)
 
-        base_fields = SchedulerOutput.__dataclass_fields__.keys()
-        base_data = {name: getattr(scheduler_output, name) for name in base_fields}
-        return OmniSchedulerOutput(
-            **base_data,
-            pending_input_registrations=(
-                self.input_coordinator.pending_input_registrations if self.input_coordinator else []
-            ),
-        )
+        return self._wrap_omni_scheduler_output(scheduler_output)
 
     def finish_requests(self, request_ids, finished_status: RequestStatus) -> list[tuple[str, int]]:
         """Handles the finish signal from outside the scheduler.
@@ -500,7 +488,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 or (self.chunk_transfer_adapter is None and request.num_computed_tokens >= request.num_prompt_tokens)
                 or (
                     self.chunk_transfer_adapter is not None
-                    and request.request_id in self.chunk_transfer_adapter.finished_requests
+                    and self.chunk_transfer_adapter.is_done_receiving_chunks(request.request_id)
                     and request.num_computed_tokens >= len(request.prompt_token_ids)
                 )
             ):
@@ -518,6 +506,10 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                     routed_experts = model_runner_output.routed_experts_dict[req_id]
                 finish_reason = request.get_finished_reason()
                 finished = self._handle_stopped_request(request)
+                if not finished:
+                    # for streaming input request only
+                    if self.chunk_transfer_adapter:
+                        self.chunk_transfer_adapter.segment_finished_requests.discard(req_id)
                 if finished:
                     kv_transfer_params = self._free_request(request)
                     if self.chunk_transfer_adapter is not None:
@@ -670,15 +662,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 engine_core_outputs[0] = eco = EngineCoreOutputs()
             eco.scheduler_stats = stats
 
-        omni_output = getattr(model_runner_output, "omni_connector_output", None)
-        if omni_output is not None:
-            self._latest_omni_connector_output = omni_output
-            if self.input_coordinator and omni_output.request_metadata:
-                self.input_coordinator.update_request_metadata(
-                    self.requests,
-                    omni_output.request_metadata,
-                    model_mode="generation",
-                )
+        self._capture_omni_connector_output(model_runner_output)
 
         return engine_core_outputs
 
@@ -688,4 +672,26 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         Do not expend prompt id using update.
         """
-        self._replace_session_with_streaming_update(session, update)
+        if self.chunk_transfer_adapter:
+            self.chunk_transfer_adapter.segment_finished_requests.discard(session.request_id)
+        session._output_token_ids.clear()
+        session._all_token_ids.clear()
+        new_prompt = update.prompt_token_ids or ()
+        session._all_token_ids.extend(new_prompt)
+        session.num_computed_tokens = 0
+        session.prompt_token_ids = update.prompt_token_ids or ()
+        session.additional_information = update.additional_information or None
+        # Update block hashes for the new tokens.
+        session.update_block_hashes()
+        session.num_prompt_tokens = len(session.prompt_token_ids)
+        session.arrival_time = update.arrival_time
+        session.sampling_params = update.sampling_params
+        if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ:
+            self.num_waiting_for_streaming_input -= 1
+        session.status = RequestStatus.WAITING
+        if session in self.skipped_waiting:
+            self.skipped_waiting.remove_requests((session,))
+            self._enqueue_waiting_request(session)
+
+        if self.log_stats:
+            session.record_event(EngineCoreEventType.QUEUED)
diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py
index 606739e90..c4b1e2429 100644
--- a/vllm_omni/core/sched/omni_scheduler_mixin.py
+++ b/vllm_omni/core/sched/omni_scheduler_mixin.py
@@ -1,7 +1,35 @@
 from __future__ import annotations
 
-from vllm.v1.engine import EngineCoreEventType
-from vllm.v1.request import Request, RequestStatus, StreamingUpdate
+import os
+from typing import Any
+
+from vllm.logger import init_logger
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.request import RequestStatus
+
+from vllm_omni.core.sched.output import OmniChunkRecvHandle, OmniSchedulerOutput
+
+logger = init_logger(__name__)
+
+# Upper bound on how long a request may sit in full-payload-input wait
+# (the state ``OmniSchedulingCoordinator`` records via ``_waiting_since``)
+# before the scheduler force-fails it.  Defends against stuck consumer-side
+# requests when the producer drops a full-payload, send fails, or recv
+# never arrives.  Override per-deployment via
+# VLLM_OMNI_INPUT_WAIT_TIMEOUT_S; set <=0 to disable the safety net.
+#
+# Scope: this constant only covers the full-payload coordinator path
+# (``input_coordinator``).  The async-chunk path uses
+# ``chunk_transfer_adapter`` and is not affected by this constant.
+_INPUT_WAIT_TIMEOUT_RAW = os.environ.get("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "300")
+try:
+    DEFAULT_INPUT_WAIT_TIMEOUT_S: float = float(_INPUT_WAIT_TIMEOUT_RAW)
+except ValueError:
+    logger.warning(
+        "Invalid VLLM_OMNI_INPUT_WAIT_TIMEOUT_S=%r; falling back to 300 seconds.",
+        _INPUT_WAIT_TIMEOUT_RAW,
+    )
+    DEFAULT_INPUT_WAIT_TIMEOUT_S = 300.0
 
 
 class OmniSchedulerMixin:
@@ -13,27 +41,107 @@ class OmniSchedulerMixin:
         if input_coordinator is not None:
             input_coordinator.free_finished_request(request_id)
 
-    def _replace_session_with_streaming_update(
-        self,
-        session: Request,
-        update: StreamingUpdate,
-    ) -> None:
-        """For streaming input: Replace an existing streaming session payload with the latest update."""
-        session._output_token_ids.clear()
-        session._all_token_ids.clear()
-        new_prompt = update.prompt_token_ids or ()
-        session._all_token_ids.extend(new_prompt)
-        session.num_computed_tokens = 0
-        session.prompt_token_ids = update.prompt_token_ids or ()
-        session.additional_information = update.additional_information or None
-        # Update block hashes for the new tokens.
-        session.update_block_hashes()
-        session.num_prompt_tokens = len(session.prompt_token_ids)
-        session.arrival_time = update.arrival_time
-        session.sampling_params = update.sampling_params
-        if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ:
-            self.num_waiting_for_streaming_input -= 1
-        session.status = RequestStatus.WAITING
+    # ------------------------------------------------------------------ #
+    #  Shared scheduler/output helpers (lift the AR / generation duplicates)
+    # ------------------------------------------------------------------ #
 
-        if self.log_stats:
-            session.record_event(EngineCoreEventType.QUEUED)
+    def _consume_pending_connector_output(self, model_mode: str) -> None:
+        """Drain ``self._latest_omni_connector_output`` into the coordinator.
+
+        Called at the top of every ``schedule()`` cycle.  Identical between
+        AR and generation schedulers except for the ``model_mode`` argument
+        forwarded to ``update_request_metadata``.
+        """
+        connector_output = getattr(self, "_latest_omni_connector_output", None)
+        self._latest_omni_connector_output = None
+        input_coordinator = getattr(self, "input_coordinator", None)
+        if input_coordinator is None:
+            return
+        if connector_output and connector_output.request_metadata:
+            input_coordinator.update_request_metadata(
+                self.requests, connector_output.request_metadata, model_mode=model_mode
+            )
+        input_coordinator.process_pending_full_payload_inputs(
+            self.waiting,
+            self.running,
+            connector_output.stage_recv_req_ids if connector_output else set(),
+        )
+
+    def _process_pending_input_timeouts(self) -> None:
+        """Force-fail requests waiting on the full-payload coordinator too long.
+
+        Called at the top of every ``schedule()`` cycle, right after
+        ``_consume_pending_connector_output``.  Without this hook, a request
+        whose producer dropped a payload would sit in the
+        full-payload-input wait state indefinitely (the runner mixin
+        protects ``_pending_load_reqs`` from prune sweeps).
+
+        Reads ``_waiting_since`` timestamps maintained by the input
+        coordinator and delegates to the base scheduler's
+        ``finish_requests`` to mark expired requests FINISHED_ERROR.
+        Disabled when ``DEFAULT_INPUT_WAIT_TIMEOUT_S`` is <= 0.
+
+        Scope: only covers ``input_coordinator`` (full-payload path).
+        Async-chunk requests park in ``chunk_transfer_adapter`` instead
+        and are not handled here -- if a similar safety net is needed
+        for the chunk path, it belongs in the chunk adapter.
+        """
+        if DEFAULT_INPUT_WAIT_TIMEOUT_S <= 0:
+            return
+        input_coordinator = getattr(self, "input_coordinator", None)
+        if input_coordinator is None:
+            return
+        timed_out_ids = input_coordinator.collect_timed_out_request_ids(timeout_s=DEFAULT_INPUT_WAIT_TIMEOUT_S)
+        if not timed_out_ids:
+            return
+        present_ids = {req_id for req_id in timed_out_ids if req_id in self.requests}
+        if not present_ids:
+            return
+        logger.warning(
+            "Marking %d request(s) as FINISHED_ERROR after waiting > %.0fs for connector input: %s",
+            len(present_ids),
+            DEFAULT_INPUT_WAIT_TIMEOUT_S,
+            sorted(present_ids),
+        )
+        self.finish_requests(present_ids, RequestStatus.FINISHED_ERROR)
+
+    def _capture_omni_connector_output(self, model_runner_output: Any) -> None:
+        """Stash the model runner's omni_connector_output for next schedule().
+
+        Called at the tail of every ``update_from_output()`` -- identical
+        between AR and generation schedulers.  Only stashes the output;
+        applying the metadata is the responsibility of
+        ``_consume_pending_connector_output()`` at the start of the next
+        ``schedule()`` cycle.  Applying it twice (once here, once on
+        consume) is unsafe under ``update_request_metadata`` in
+        generation mode, which resets ``prompt_token_ids`` /
+        ``_output_token_ids`` / ``num_computed_tokens`` and would
+        clobber any progress between the two calls.
+        """
+        omni_output = getattr(model_runner_output, "omni_connector_output", None)
+        if omni_output is None:
+            return
+        self._latest_omni_connector_output = omni_output
+
+    def _wrap_omni_scheduler_output(
+        self,
+        base: SchedulerOutput,
+        *,
+        finished_requests_needing_kv_transfer: dict | None = None,
+        pending_input_registrations: list[OmniChunkRecvHandle] | None = None,
+    ) -> OmniSchedulerOutput:
+        """Wrap a base ``SchedulerOutput`` in ``OmniSchedulerOutput``.
+
+        Pulls each base ``SchedulerOutput`` dataclass field via ``getattr``
+        and forwards optional omni-specific fields.  Lifted from 4 separate
+        copy-pastes between AR (1) and generation (3) schedulers.
+        """
+        base_data = {name: getattr(base, name) for name in SchedulerOutput.__dataclass_fields__}
+        input_coordinator = getattr(self, "input_coordinator", None)
+        if pending_input_registrations is None:
+            pending_input_registrations = input_coordinator.pending_input_registrations if input_coordinator else []
+        return OmniSchedulerOutput(
+            **base_data,
+            finished_requests_needing_kv_transfer=finished_requests_needing_kv_transfer or {},
+            pending_input_registrations=pending_input_registrations,
+        )
diff --git a/vllm_omni/core/sched/omni_scheduling_coordinator.py b/vllm_omni/core/sched/omni_scheduling_coordinator.py
index 6c32ed4cd..4056fd938 100644
--- a/vllm_omni/core/sched/omni_scheduling_coordinator.py
+++ b/vllm_omni/core/sched/omni_scheduling_coordinator.py
@@ -19,16 +19,64 @@ from typing import Any
 from vllm.logger import init_logger
 from vllm.v1.request import Request, RequestStatus
 
+from vllm_omni.core.sched.output import OmniChunkRecvHandle
+
 logger = init_logger(__name__)
 
 
-def uses_qwen3_omni_full_payload_input_coordinator(model_config: Any) -> bool:
-    return (
-        getattr(model_config, "stage_id", 0) > 0
-        and not getattr(model_config, "async_chunk", False)
-        and getattr(model_config, "model_arch", None) == "Qwen3OmniMoeForConditionalGeneration"
-        and getattr(model_config, "model_stage", None) in {"talker", "code2wav"}
+# (arch, model_stage) pairs that route their full_payload stage input via
+# the worker connector and therefore need the scheduler-side coordinator to
+# park requests in WAITING_FOR_INPUT until the recv side delivers.  This set
+# must stay aligned with the arch scope of `init_omni_connectors` in
+# gpu_ar_model_runner.py and gpu_generation_model_runner.py.  Adding a stage
+# here without also wiring its worker connector init produces a permanent
+# Stage 1 hang (gate parks the request, no transport ever releases it).
+#
+_FULL_PAYLOAD_INPUT_STAGES: frozenset[tuple[str, str]] = frozenset(
+    {
+        ("Qwen3OmniMoeForConditionalGeneration", "talker"),
+        ("Qwen3OmniMoeForConditionalGeneration", "code2wav"),
+        # qwen2_5_omni thinker->talker uses the real full-payload
+        # producer builder (text_hidden_states routed via
+        # pooler_output["hidden"] -> accumulator -> connector).  Both
+        # stages of qwen2_5_omni are enabled.
+        ("Qwen2_5OmniForConditionalGeneration", "talker"),
+        ("Qwen2_5OmniForConditionalGeneration", "code2wav"),
+        # covo_audio: fused_thinker_talker (Stage 0) -> code2wav (Stage 1).
+        ("CovoAudioForConditionalGeneration", "code2wav"),
+        # mimo_audio: fused_thinker_talker (Stage 0) -> code2wav (Stage 1).
+        ("MiMoAudioModel", "code2wav"),
+        # qwen3_tts: Qwen3TTSTalkerForConditionalGeneration (Stage 0)
+        # -> Qwen3TTSCode2Wav (Stage 1).  Stage 1 is the consumer.
+        ("Qwen3TTSCode2Wav", "code2wav"),
+        # cosyvoice3: cosyvoice3_talker (Stage 0) -> cosyvoice3_code2wav (Stage 1).
+        ("CosyVoice3Model", "cosyvoice3_code2wav"),
+        # dynin: token2text (Stage 0) -> token2image (Stage 1) ->
+        # token2audio (Stage 2).  Producer wires via
+        # custom_process_next_stage_input_func: *_full_payload in deploy yaml.
+        ("DyninOmniForConditionalGeneration", "token2image"),
+        ("DyninOmniForConditionalGeneration", "token2audio"),
+    }
+)
+
+
+def uses_full_payload_input_coordinator(model_config: Any) -> bool:
+    """Returns True iff this stage parks pending requests in
+    WAITING_FOR_INPUT awaiting a full_payload delivery on the worker connector.
+
+    Gated by (model_arch, model_stage) — see _FULL_PAYLOAD_INPUT_STAGES for the
+    rationale on why this is a whitelist instead of a marker-driven structural
+    gate.
+    """
+    if getattr(model_config, "stage_id", 0) <= 0:
+        return False
+    if getattr(model_config, "async_chunk", False):
+        return False
+    key = (
+        getattr(model_config, "model_arch", None),
+        getattr(model_config, "model_stage", None),
     )
+    return key in _FULL_PAYLOAD_INPUT_STAGES
 
 
 class OmniSchedulingCoordinator:
@@ -59,7 +107,11 @@ class OmniSchedulingCoordinator:
 
         # Requests waiting for full_payload stage input (WAITING_FOR_INPUT).
         self._waiting_for_input: deque[Any] = deque()
-        self.pending_input_registrations: list[Any] = []
+        # Per-cycle list of minimal handles to ship to the model runner so it
+        # can call register_chunk_recv().  Typed concretely (not list[Any]) so
+        # the surrounding OmniSchedulerOutput stays msgspec-friendly across
+        # default, PD-disagg, and multi-node executor IPC paths.
+        self.pending_input_registrations: list[OmniChunkRecvHandle] = []
 
         # Monotonic timestamp recording when each request first entered
         # WAITING_FOR_CHUNK or WAITING_FOR_INPUT.  Used by
@@ -166,7 +218,12 @@ class OmniSchedulingCoordinator:
                     self._waiting_since.setdefault(request.request_id, time.monotonic())
                     to_remove.append(request)
                     self._waiting_for_input.append(request)
-                    self.pending_input_registrations.append(request)
+                    self.pending_input_registrations.append(
+                        OmniChunkRecvHandle(
+                            request_id=request.request_id,
+                            external_req_id=getattr(request, "external_req_id", None),
+                        )
+                    )
                 elif request.status == RequestStatus.WAITING_FOR_INPUT:
                     if request.request_id in stage_recv_req_ids:
                         request.status = RequestStatus.WAITING
@@ -174,7 +231,12 @@ class OmniSchedulingCoordinator:
                     else:
                         to_remove.append(request)
                         self._waiting_for_input.append(request)
-                        self.pending_input_registrations.append(request)
+                        self.pending_input_registrations.append(
+                            OmniChunkRecvHandle(
+                                request_id=request.request_id,
+                                external_req_id=getattr(request, "external_req_id", None),
+                            )
+                        )
             if to_remove:
                 # Use the bulk-remove helper: one O(N) sweep instead of N
                 # repeated O(N) removes from a list-backed queue.
diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py
index 800881d9f..29cd87299 100644
--- a/vllm_omni/core/sched/output.py
+++ b/vllm_omni/core/sched/output.py
@@ -1,5 +1,4 @@
 from dataclasses import dataclass, field
-from typing import Any
 
 from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
 from vllm.v1.request import Request
@@ -72,9 +71,26 @@ class OmniCachedRequestData(CachedRequestData):
     additional_information: dict[str, dict | None]
 
 
+@dataclass
+class OmniChunkRecvHandle:
+    """Minimal identifier carried from scheduler to runner for chunk-recv
+    registration.
+
+    The runner's ``register_chunk_recv`` only consumes ``request_id`` and
+    ``external_req_id`` from each pending request, so we ship just those
+    two fields instead of the full Request object.  Concrete typing
+    keeps msgspec serialization deterministic across IPC (default,
+    PD-disagg, multi-node executor variants) and avoids the
+    ``list[Any]`` fallback path.
+    """
+
+    request_id: str
+    external_req_id: str | None = None
+
+
 @dataclass
 class OmniSchedulerOutput(SchedulerOutput):
     """Scheduler output with omni-specific transfer metadata."""
 
     finished_requests_needing_kv_transfer: dict[str, dict] = field(default_factory=dict)
-    pending_input_registrations: list[Any] = field(default_factory=list)
+    pending_input_registrations: list[OmniChunkRecvHandle] = field(default_factory=list)
diff --git a/vllm_omni/data_entry_keys.py b/vllm_omni/data_entry_keys.py
index b239ff5ed..09316b55e 100644
--- a/vllm_omni/data_entry_keys.py
+++ b/vllm_omni/data_entry_keys.py
@@ -59,6 +59,7 @@ class Ids(TypedDict, total=False):
 
 class OmniPayloadMeta(TypedDict, total=False):
     finished: torch.Tensor
+    is_segment_finished: torch.Tensor
     stream_finished: torch.Tensor
     req_id: list[str]
     left_context_size: int
@@ -142,6 +143,7 @@ class IdsStruct(_StructBase):
 
 class MetaStruct(_StructBase):
     finished: torch.Tensor | None = None
+    is_segment_finished: torch.Tensor | None = None
     stream_finished: torch.Tensor | None = None
     req_id: list[str] | None = None
     left_context_size: int | None = None
diff --git a/vllm_omni/deploy/dynin_omni_ci.yaml b/vllm_omni/deploy/dynin_omni_ci.yaml
index 525b7d888..2ddc281e7 100644
--- a/vllm_omni/deploy/dynin_omni_ci.yaml
+++ b/vllm_omni/deploy/dynin_omni_ci.yaml
@@ -14,6 +14,7 @@ stage_args:
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       engine_output_type: latent
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image_full_payload
       trust_remote_code: true
       gpu_memory_utilization: 0.5
       enforce_eager: true
@@ -36,6 +37,7 @@ stage_args:
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       engine_output_type: latent
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2image_to_token2audio_full_payload
       trust_remote_code: true
       gpu_memory_utilization: 0.2
       enforce_eager: true
diff --git a/vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml b/vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml
new file mode 100644
index 000000000..9fec34d69
--- /dev/null
+++ b/vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml
@@ -0,0 +1,72 @@
+# Fish Speech S2 Pro 2-GPU deploy: Stage0 AR on GPU 0, Stage1 DAC on GPU 1.
+#
+# Opt-in profile validated on H100x2. Compared with the single-GPU default
+# (``fish_qwen3_omni.yaml``), it splits AR and DAC across devices to remove
+# AR/DAC contention, enlarges Stage0 AR batching, and lets Stage1 consume
+# multiple DAC chunks in parallel.
+#
+# Select via:
+#   CUDA_VISIBLE_DEVICES=0,1 vllm serve fishaudio/s2-pro --omni \
+#       --deploy-config vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml
+#
+# Or programmatically:
+#   Omni(model="fishaudio/s2-pro",
+#        deploy_config_path="vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml")
+
+async_chunk: true
+enable_chunked_prefill: false
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      codec_streaming: true
+      # Tightened from 10ms to 1ms: chunks become available sooner under
+      # cross-GPU pipelining, so connector idle waiting dominates if left high.
+      connector_get_sleep_s: 0.001
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      # ~21 Hz codec; 25 frames ≈ 1.16s of audio.
+      codec_chunk_frames: 25
+      codec_left_context_frames: 25
+      initial_codec_chunk_frames: 4
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 64
+    gpu_memory_utilization: 0.75
+    enforce_eager: false
+    async_scheduling: false
+    # vLLM scheduler requires max_num_batched_tokens >= max_model_len when
+    # chunked prefill is off (SlowAR decode loop isn't chunked-prefill-safe).
+    max_num_batched_tokens: 16384
+    max_model_len: 16384
+    devices: "0"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.8
+      top_k: 30
+      top_p: 0.9
+      max_tokens: 2048
+      seed: 42
+      repetition_penalty: 1.0
+
+  - stage_id: 1
+    max_num_seqs: 64
+    gpu_memory_utilization: 0.2
+    enforce_eager: true
+    async_scheduling: false
+    max_num_batched_tokens: 16384
+    max_model_len: 16384
+    devices: "1"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      repetition_penalty: 1.0
diff --git a/vllm_omni/deploy/fish_qwen3_omni_high_concurrency_dual_gpu.yaml b/vllm_omni/deploy/fish_qwen3_omni_high_concurrency_dual_gpu.yaml
new file mode 100644
index 000000000..a65609c6c
--- /dev/null
+++ b/vllm_omni/deploy/fish_qwen3_omni_high_concurrency_dual_gpu.yaml
@@ -0,0 +1,68 @@
+# Fish Speech S2 Pro high-concurrency dual-GPU deploy config.
+#
+# Stage0 SlowAR runs on GPU0 and Stage1 DAC runs on GPU1. Use this config
+# when the deployment can trade an extra GPU for lower Stage0/Stage1
+# interference between the AR and DAC stages.
+async_chunk: true
+enable_chunked_prefill: false
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      codec_streaming: true
+      connector_get_sleep_s: 0.001
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      fish_speech_tensor_codes: true
+      fish_speech_single_initial_chunk: true
+      fish_speech_backlog_codec_chunk_frames: 50
+      fish_speech_backlog_load_threshold: 0.75
+      codec_chunk_frames: 25
+      codec_left_context_frames: 25
+      initial_codec_chunk_frames: 4
+      fish_speech_dac_dtype: fp16
+      fish_speech_dac_max_padded_frames: 256
+      fish_speech_dac_max_batch: 0
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 64
+    gpu_memory_utilization: 0.85
+    enforce_eager: false
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 16384
+    max_model_len: 16384
+    devices: "0"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.8
+      top_k: 30
+      top_p: 0.9
+      max_tokens: 2048
+      seed: 42
+      repetition_penalty: 1.0
+
+  - stage_id: 1
+    max_num_seqs: 64
+    gpu_memory_utilization: 0.25
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 16384
+    max_model_len: 16384
+    devices: "1"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      repetition_penalty: 1.0
diff --git a/vllm_omni/deploy/fish_qwen3_omni_high_concurrency_single_gpu.yaml b/vllm_omni/deploy/fish_qwen3_omni_high_concurrency_single_gpu.yaml
new file mode 100644
index 000000000..d771e4163
--- /dev/null
+++ b/vllm_omni/deploy/fish_qwen3_omni_high_concurrency_single_gpu.yaml
@@ -0,0 +1,68 @@
+# Fish Speech S2 Pro high-concurrency single-GPU deploy config.
+#
+# Stage0 SlowAR and Stage1 DAC share GPU0. This config keeps the optimized
+# Fish runtime path explicit instead of changing the default
+# fish_qwen3_omni.yaml deployment.
+async_chunk: true
+enable_chunked_prefill: false
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      codec_streaming: true
+      connector_get_sleep_s: 0.001
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      fish_speech_tensor_codes: true
+      fish_speech_single_initial_chunk: true
+      fish_speech_backlog_codec_chunk_frames: 50
+      fish_speech_backlog_load_threshold: 0.75
+      codec_chunk_frames: 25
+      codec_left_context_frames: 25
+      initial_codec_chunk_frames: 4
+      fish_speech_dac_dtype: fp16
+      fish_speech_dac_max_padded_frames: 256
+      fish_speech_dac_max_batch: 0
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 64
+    gpu_memory_utilization: 0.72
+    enforce_eager: false
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 16384
+    max_model_len: 16384
+    devices: "0"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.8
+      top_k: 30
+      top_p: 0.9
+      max_tokens: 2048
+      seed: 42
+      repetition_penalty: 1.0
+
+  - stage_id: 1
+    max_num_seqs: 64
+    gpu_memory_utilization: 0.18
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 16384
+    max_model_len: 16384
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      repetition_penalty: 1.0
diff --git a/vllm_omni/deploy/glm_image.yaml b/vllm_omni/deploy/glm_image.yaml
index 099df1b15..470bf2cd4 100644
--- a/vllm_omni/deploy/glm_image.yaml
+++ b/vllm_omni/deploy/glm_image.yaml
@@ -46,3 +46,40 @@ stages:
       guidance_scale: 1.5
       height: 1024
       width: 1024
+
+platforms:
+  npu:
+    stages:
+      - stage_id: 0
+        devices: "0,1,2,3"
+        env:
+          # When multiple business processes share a single NPU simultaneously,
+          # port conflict issues may occur. You can refer to the link below to set different parameters for different processes,
+          # replacing the default parameter configuration.
+          # https://www.hiascend.com/document/detail/zh/canncommercial/850/maintenref/envvar/envref_07_0143.html
+          HCCL_IF_BASE_PORT: 23000
+          HCCL_HOST_SOCKET_PORT_RANGE: "23000-23199"
+          HCCL_NPU_SOCKET_PORT_RANGE: "23200-23399"
+        async_scheduling: true
+        enforce_eager: false
+        distributed_executor_backend: "mp"
+        enable_prefix_caching: true
+        tensor_parallel_size: 4
+        sequence_parallel_size: 1
+        compilation_config:
+          cudagraph_capture_sizes: [1, 2]
+          cudagraph_mode: "FULL_DECODE_ONLY"
+        worker_cls: vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker
+        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      - stage_id: 1
+        devices: "0,1,2,3"
+        env:
+          HCCL_IF_BASE_PORT: 24000
+          HCCL_HOST_SOCKET_PORT_RANGE: "24000-24199"
+          HCCL_NPU_SOCKET_PORT_RANGE: "24200-24399"
+        distributed_executor_backend: "mp"
+        parallel_config:
+          tensor_parallel_size: 1
+          sequence_parallel_size: 4
+          ulysses_degree: 4
+          ring_degree: 1
diff --git a/vllm_omni/deploy/glm_tts.yaml b/vllm_omni/deploy/glm_tts.yaml
new file mode 100644
index 000000000..1600b6152
--- /dev/null
+++ b/vllm_omni/deploy/glm_tts.yaml
@@ -0,0 +1,83 @@
+# GLM-TTS deploy: 2-stage AR -> DiT flow-matching.
+#
+# Default mode is async-chunk streaming through SharedMemoryConnector.
+# Stage 0 (AR LLM) streams speech token chunks to Stage 1 (DiT + vocoder)
+# via the shared-memory connector. Stage 1 runs as LLM_GENERATION type
+# to support chunk_transfer_adapter consumption.
+#
+# Stage 1 keeps vLLM eager scheduling but enables the GLM-TTS bucketed DiT
+# CUDA graph wrapper via hf_overrides.use_dit_cuda_graphs. Graph capture falls
+# back to eager if a dynamic shape exceeds the captured buckets.
+#
+# Verified on 1x A40 48GB (both stages share GPU 0).
+# Memory footprint ~16.6 GiB peak; fits 24GB cards with the settings below.
+async_chunk: true
+trust_remote_code: true
+
+connectors:
+  glm_tts_shm:
+    name: SharedMemoryConnector
+    extra:
+      codec_streaming: true
+      connector_get_sleep_s: 0.002
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      # Progressive chunk sizes matching original GLM-TTS defaults.
+      codec_chunk_frames: [25, 50, 200]
+      codec_left_context_frames: 25
+      crossfade_sec: 0.1
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.55
+    # BF16 for AR decode: memory-bandwidth bound on single-batch decode,
+    # bf16 halves HBM reads per step (~1.7x throughput on L4/A40).
+    dtype: bfloat16
+    enforce_eager: false
+    async_scheduling: true
+    enable_prefix_caching: false
+    max_num_batched_tokens: 512
+    max_model_len: 2048
+    devices: "0"
+    output_connectors:
+      to_stage_1: glm_tts_shm
+    default_sampling_params:
+      temperature: 1.0
+      top_k: 25
+      top_p: 0.8
+      max_tokens: 2048
+      min_tokens: 1
+      seed: 42
+      detokenize: false
+    hf_overrides:
+      sample_method: ras
+      ras_top_p: 0.8
+      ras_top_k: 25
+      ras_win_size: 10
+      ras_tau_r: 0.1
+      min_token_text_ratio: 2.0
+      max_token_text_ratio: 20.0
+    skip_mm_profiling: true
+
+  - stage_id: 1
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.25
+    enforce_eager: true
+    max_model_len: 8192
+    max_num_batched_tokens: 2048
+    # Force float32 for the flow-matching DiT stage.  The Euler ODE solver
+    # (10 steps × 2 forward passes each) accumulates rounding errors in half
+    # precision that manifest as audible hiss on H20/H100.  float32 matches
+    # the reference implementation and eliminates the noise floor.
+    dtype: float32
+    devices: "0"
+    enable_prefix_caching: false
+    input_connectors:
+      from_stage_0: glm_tts_shm
+    default_sampling_params:
+      max_tokens: 2048
+    hf_overrides:
+      use_dit_cuda_graphs: true
+    disable_hybrid_kv_cache_manager: true
+    skip_mm_profiling: true
diff --git a/vllm_omni/deploy/higgs_audio_v2.yaml b/vllm_omni/deploy/higgs_audio_v2.yaml
new file mode 100644
index 000000000..07e25d5c1
--- /dev/null
+++ b/vllm_omni/deploy/higgs_audio_v2.yaml
@@ -0,0 +1,83 @@
+# Higgs-audio v2 deploy: Stage 0 (talker) -> Stage 1 (code2wav) via shared memory.
+#
+# Phase-1 scope is plain text -> 24 kHz speech only; voice cloning, multi-speaker,
+# ChatML rich messages, and music/sound-event prompts are rejected by the
+# request validator with explicit 4xx (see
+# vllm_omni/entrypoints/openai/serving_speech.py and
+# vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py).
+#
+# Streaming is wired through the stage-input processor's async-chunk
+# callback, the shared-memory connector ``codec_streaming`` flag, and the
+# WAV/PCM streaming convention in serving_speech.py. To run a strictly
+# non-streaming smoke test, set ``codec_streaming: false`` in the connector
+# extras below.
+async_chunk: false
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      codec_streaming: true
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 5000
+      connector_get_max_wait: 500
+      codec_chunk_frames: 25
+      codec_left_context_frames: 25
+      initial_codec_chunk_frames: 1
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 4
+    gpu_memory_utilization: 0.6
+    # Disable cudagraph capture: the live audio-sampler dispatch is still
+    # iterating, and the host's flash-attn-hopper graph path tripped a CUDA
+    # driver vs runtime mismatch during capture on this image.
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 2048
+    # max_model_len covers prompt + max_new_tokens. Up to ~1500 codec
+    # frames per request fit comfortably under 4096.
+    max_model_len: 4096
+    devices: "0"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      # Upstream serve_engine defaults (temperature=1.0, top_p=0.95,
+      # top_k=50). Greedy (temperature=0) collapses the audio codebook
+      # argmax to a fixed value after a few steps; sampling with diversity
+      # is required for natural-sounding speech.
+      temperature: 1.0
+      top_p: 0.95
+      top_k: 50
+      # Cap Stage-0 to 1024 audio frames (~40 s at 25 fps). The audio-eos
+      # natural stop (128012, forced by the talker's _apply_audio_mode_bias
+      # at ramp-completion) typically fires earlier and trims this.
+      max_tokens: 1024
+      seed: 42
+      repetition_penalty: 1.0
+
+  - stage_id: 1
+    max_num_seqs: 4
+    gpu_memory_utilization: 0.25
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    # Prefill length is num_codebooks * num_frames; 8 * 8192 = 65536 covers
+    # the upper bound at the per-request 250-frame Stage-0 cap above plus
+    # headroom for batching multiple requests.
+    max_num_batched_tokens: 65536
+    max_model_len: 65536
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      repetition_penalty: 1.0
diff --git a/vllm_omni/deploy/hunyuan_image3.yaml b/vllm_omni/deploy/hunyuan_image_3_moe.yaml
similarity index 100%
rename from vllm_omni/deploy/hunyuan_image3.yaml
rename to vllm_omni/deploy/hunyuan_image_3_moe.yaml
diff --git a/vllm_omni/deploy/lance.yaml b/vllm_omni/deploy/lance.yaml
deleted file mode 100644
index bf8c6906f..000000000
--- a/vllm_omni/deploy/lance.yaml
+++ /dev/null
@@ -1,29 +0,0 @@
-# Lance-3B (ByteDance) deploy — single-stage, CUDA. Target: 1x NVIDIA A100 80GB.
-#
-# Lance is BAGEL-lineage (Qwen2-MoT unified AR+diffusion). The single-stage
-# DiT stage is self-contained: it owns the Qwen2-MoT LLM, the Qwen2.5-VL ViT,
-# the Wan2.2 VAE and the tokenizer (text2img / image_edit / x2t_image).
-#
-# Lance-3B BF16 footprint: ~7GB LLM + Qwen2.5-VL ViT + Wan2.2 VAE — comfortably
-# within one A100 80GB. The video checkpoint (Lance_3B_Video) is a follow-up.
-#
-# Select via:
-#   vllm serve bytedance-research/Lance --omni \
-#       --deploy-config vllm_omni/deploy/lance.yaml
-# or programmatically:
-#   Omni(model="bytedance-research/Lance",
-#        deploy_config_path="vllm_omni/deploy/lance.yaml")
-
-pipeline: lance
-async_chunk: false
-
-stages:
-  - stage_id: 0
-    max_num_batched_tokens: 32768
-    max_num_seqs: 1
-    enforce_eager: true
-    trust_remote_code: true
-    enable_prefix_caching: false
-    devices: "0"
-    default_sampling_params:
-      seed: 42
diff --git a/vllm_omni/deploy/ming_flash_omni_image.yaml b/vllm_omni/deploy/ming_flash_omni_image.yaml
new file mode 100644
index 000000000..9f7c6a34d
--- /dev/null
+++ b/vllm_omni/deploy/ming_flash_omni_image.yaml
@@ -0,0 +1,63 @@
+# Ming-flash-omni-2.0 deploy: thinker + image-generation (text-to-image / img2img).
+#
+# Stage 0: Thinker — multimodal understanding -> final-layer hidden states.
+#          TP=4 across GPUs 0-3.
+# Stage 1: DiT (MingImagePipeline) — condition encoder + FlowMatchEuler
+#          denoise + VAE decode. Diffusion worker, runs eager on GPU 4.
+#
+# Cross-stage transport: shared_memory_connector (intra-node).
+#
+# Fields omitted from a stage fall back to StageDeployConfig dataclass
+# defaults (see vllm_omni/config/stage_config.py).
+#   * Stage 0 max_num_seqs=2 leaves room for the CFG-text companion that
+#     `expand_cfg_prompts` (pipeline-level) spawns when a request carries
+#     a negative prompt — positive and negative hidden states both reach
+#     the diffusion stage.
+#   * Image-gen warmup adds latency on top of the thinker;
+#     pass `--stage-init-timeout <seconds>` if the default 300s is tight.
+pipeline: ming_flash_omni_image
+async_chunk: false
+
+trust_remote_code: true
+distributed_executor_backend: mp
+enable_prefix_caching: false
+
+stages:
+  - stage_id: 0
+    # max_num_seqs=2: parent + 1 CFG-text companion (negative prompt).
+    max_num_seqs: 2
+    gpu_memory_utilization: 0.85
+    enforce_eager: true
+    tensor_parallel_size: 4
+    devices: "0,1,2,3"
+    output_connectors:
+      to_stage_1: shared_memory_connector
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 1  # image-gen request: prefill only, no AR decode
+      seed: 42
+
+  - stage_id: 1
+    enforce_eager: true
+    devices: "4"
+    input_connectors:
+      from_stage_0: shared_memory_connector
+    default_sampling_params:
+      seed: 42
+
+connectors:
+  shared_memory_connector:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536  # 64KB threshold
+
+edges:
+  - from: 0
+    to: 1
+    window_size: -1
+    max_inflight: 1
+
+# Ming-flash-omni-2.0 has only been manually verified on CUDA devices
+platforms: {}
diff --git a/vllm_omni/deploy/qwen3_omni_moe.yaml b/vllm_omni/deploy/qwen3_omni_moe.yaml
index 445437c0f..53b7234e6 100644
--- a/vllm_omni/deploy/qwen3_omni_moe.yaml
+++ b/vllm_omni/deploy/qwen3_omni_moe.yaml
@@ -53,12 +53,13 @@ stages:
       repetition_penalty: 1.05
 
   - stage_id: 2
-    max_num_batched_tokens: 51200
+    max_num_batched_tokens: 65536
     max_num_seqs: 64
     gpu_memory_utilization: 0.1
-    enforce_eager: true
+    enforce_eager: false
     trust_remote_code: true
     enable_prefix_caching: false
+    enable_chunked_prefill: false
     async_scheduling: false
     devices: "1"
     input_connectors:
diff --git a/vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml b/vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml
index 791e29a8d..390698e94 100644
--- a/vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml
+++ b/vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml
@@ -23,7 +23,11 @@ connectors:
 
 stages:
   - stage_id: 0
+    max_num_batched_tokens: 32768
+    max_num_seqs: 64
     gpu_memory_utilization: 0.9
+    trust_remote_code: true
+    enable_prefix_caching: false
     devices: "0"
     default_sampling_params:
       temperature: 0.4
@@ -34,7 +38,11 @@ stages:
       repetition_penalty: 1.05
 
   - stage_id: 1
+    max_num_batched_tokens: 32768
+    max_num_seqs: 64
     gpu_memory_utilization: 0.6
+    trust_remote_code: true
+    enable_prefix_caching: false
     devices: "1,2"
     num_replicas: 2
     input_connectors:
@@ -47,11 +55,14 @@ stages:
       repetition_penalty: 1.05
 
   - stage_id: 2
+    max_num_batched_tokens: 65536
+    max_num_seqs: 64
     gpu_memory_utilization: 0.1
-    max_num_seqs: 1
     enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
     async_scheduling: false
-    max_num_batched_tokens: 51200
+    enable_chunked_prefill: false
     devices: "1,2"
     num_replicas: 2
     input_connectors:
@@ -70,20 +81,16 @@ platforms:
       - stage_id: 0
         gpu_memory_utilization: 0.6
         tensor_parallel_size: 2
+        max_num_batched_tokens: 8192
         devices: "0,1"
       - stage_id: 1
         gpu_memory_utilization: 0.6
-        enforce_eager: true
+        max_num_batched_tokens: 8192
         devices: "2"
       - stage_id: 2
         gpu_memory_utilization: 0.3
         devices: "2"
 
-  rocm:
-    stages:
-      - stage_id: 0
-        enforce_eager: true
-
   xpu:
     stages:
       - stage_id: 0
diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml
index 80c316849..55bac5833 100644
--- a/vllm_omni/deploy/qwen3_tts.yaml
+++ b/vllm_omni/deploy/qwen3_tts.yaml
@@ -9,8 +9,8 @@
 #   * Stage 1 (code2wav) sets max_num_batched_tokens=65536 for correctness:
 #     codec prefill length (Q * num_frames) exceeds the 32k default.
 #   * Stage 0 does not set enforce_eager — talker runs cudagraph by default.
-#     Stage 1 sets true because its codec generation loop is not
-#     cudagraph-compatible. NPU platform flips stage 0 to true where
+#     Stage 1 defaults to false so Code2Wav can use its inner cudagraph;
+#     set true to disable that path in eager mode. NPU platform flips stage 0 to true where
 #     cudagraph is not yet verified.
 async_chunk: true
 
@@ -28,13 +28,23 @@ connectors:
       codec_left_context_frames: 72
       # Emit only the first audio chunk early, then return to codec_chunk_frames.
       initial_codec_chunk_frames: 1
+      # Common Stage1 decode buckets:
+          #   no-ref first/steady chunks: 25 / 97 frames
+          #   Base ref-context first/steady chunks: 73 / 169 frames
+          #   decoder internal non-streaming chunks: 325 frames
+      decode_cudagraph_capture_sizes: [25, 73, 97, 169, 325]
 
 stages:
   - stage_id: 0
     max_num_seqs: 10
     gpu_memory_utilization: 0.3
     trust_remote_code: true
-    enable_prefix_caching: false
+    # Prefix caching on by default. The HS opt-out (talker only consumes the
+    # last token's hidden state) and deferred mm-output writes added in this
+    # PR make per-step bookkeeping negligible, and any text repetition across
+    # requests (same voice or shared ref_audio) is a net throughput win.
+    # Disable explicitly if memory headroom is tight on a single-GPU deploy.
+    enable_prefix_caching: true
     async_scheduling: true
     max_num_batched_tokens: 512
     max_model_len: 4096
@@ -54,9 +64,10 @@ stages:
       top_p: 1.0
 
   - stage_id: 1
+    # Tuned for Base voice clone; CustomVoice / VoiceDesign are TTFA-optimal at 1.
     max_num_seqs: 10
     gpu_memory_utilization: 0.3
-    enforce_eager: true
+    enforce_eager: false
     trust_remote_code: true
     enable_prefix_caching: false
     async_scheduling: true
diff --git a/vllm_omni/diffusion/attention/backends/cudnn_attn.py b/vllm_omni/diffusion/attention/backends/cudnn_attn.py
new file mode 100644
index 000000000..f27fe1870
--- /dev/null
+++ b/vllm_omni/diffusion/attention/backends/cudnn_attn.py
@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from torch.nn.attention import SDPBackend, sdpa_kernel
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+from vllm_omni.diffusion.attention.backends.sdpa import _maybe_reshape_attn_mask
+
+logger = init_logger(__name__)
+
+
+class CuDNNAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @classmethod
+    def supports_attention_mask(cls) -> bool:
+        return True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        # cuDNN 9.5+ FMHA on Blackwell supports any head_dim divisible by 8
+        # up to 256 for BF16/FP16. Empty list = "accept any"; the sdpa_kernel
+        # fallback chain handles shapes cuDNN can't take.
+        return []
+
+    @staticmethod
+    def get_name() -> str:
+        return "CUDNN_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> type["CuDNNAttentionImpl"]:
+        return CuDNNAttentionImpl
+
+
+class CuDNNAttentionImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    def forward_cuda(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        attention_mask = None
+        if attn_metadata:
+            attention_mask = _maybe_reshape_attn_mask(query, key, attn_metadata.attn_mask, mask_mode="broadcast_k")
+
+        query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
+        # Pin cuDNN exclusively. A priority list like [CUDNN, FLASH, MATH] hits a
+        # PyTorch SDPA dispatch quirk: when FLASH rejects a non-None attn_mask,
+        # cuDNN gets runtime-disabled in the same call and the dispatcher falls
+        # through to MATH even though cuDNN alone handles the mask fine
+        # (~11 ms vs ~215 ms for MATH on sm_120 HV-1.5 shapes).
+        #
+        # Fall back to the default SDPA dispatcher if cuDNN rejects the shape,
+        # e.g. under torch.compile where Dynamo sees a symbolic head_dim and
+        # cuDNN's kernel selection fails (observed in LTX-2 audio attention).
+        # The unpinned dispatcher then picks EFFICIENT/MATH instead of raising.
+        try:
+            with sdpa_kernel([SDPBackend.CUDNN_ATTENTION]):
+                output = torch.nn.functional.scaled_dot_product_attention(
+                    query,
+                    key,
+                    value,
+                    attn_mask=attention_mask,
+                    dropout_p=0.0,
+                    is_causal=self.causal,
+                    scale=self.softmax_scale,
+                )
+        except RuntimeError as e:
+            if "No available kernel" not in str(e):
+                raise
+            logger.warning_once(
+                "cuDNN SDPA rejected this shape; falling back to default SDPA dispatcher. "
+                "Set DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA for the full dispatcher path on every call."
+            )
+            output = torch.nn.functional.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                attn_mask=attention_mask,
+                dropout_p=0.0,
+                is_causal=self.causal,
+                scale=self.softmax_scale,
+            )
+        return output.permute(0, 2, 1, 3)
diff --git a/vllm_omni/diffusion/attention/backends/flashinfer_attn.py b/vllm_omni/diffusion/attention/backends/flashinfer_attn.py
new file mode 100644
index 000000000..1d2fa893b
--- /dev/null
+++ b/vllm_omni/diffusion/attention/backends/flashinfer_attn.py
@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+
+logger = init_logger(__name__)
+
+try:
+    from flashinfer.prefill import single_prefill_with_kv_cache
+
+    HAS_FLASHINFER = True
+except Exception as e:
+    HAS_FLASHINFER = False
+    logger.warning(
+        "FlashInfer is unavailable; FLASHINFER_ATTN backend will not work. Reason: %s",
+        e,
+    )
+
+
+class FlashInferAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @classmethod
+    def supports_attention_mask(cls) -> bool:
+        # FlashInfer single_prefill_with_kv_cache accepts a ``custom_mask`` kwarg
+        # (2D boolean, ``True`` = keep) for non-causal attention. See reviewer
+        # comment on #3079 and the flashinfer docs:
+        #   https://docs.flashinfer.ai/generated/flashinfer.prefill.single_prefill_with_kv_cache.html
+        return True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        # FlashInfer dense prefill is well-tested for these head_dims on
+        # Ampere/Hopper/Blackwell. Covers the dominant diffusion DiT shapes
+        # (SD3 = 64, Flux/HV/Wan = 128, joint-attn = 256).
+        return [64, 128, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "FLASHINFER_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> type["FlashInferAttentionImpl"]:
+        return FlashInferAttentionImpl
+
+
+class FlashInferAttentionImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    @staticmethod
+    def _pack_mask_for_flashinfer(
+        attn_mask: torch.Tensor, batch_size: int, qo_len: int, kv_len: int
+    ) -> torch.Tensor | None:
+        """Convert a diffusion-style attn_mask into the boolean form
+        FlashInfer's ``custom_mask`` expects (``True`` = keep).
+
+        Returns either ``(qo_len, kv_len)`` (shared across the batch) or
+        ``(batch_size, qo_len, kv_len)`` (per-sample), or ``None`` when the
+        mask is all-keep (elide). Only boolean masks are handled here;
+        additive/float masks raise ``ValueError`` so the caller falls back to
+        SDPA, which applies them with the correct softmax semantics. Shape
+        mismatches also raise ``ValueError``.
+        """
+        mask = attn_mask
+        if mask.dtype != torch.bool:
+            # Additive masks (0 / -inf / -1e4 / finfo.min) cannot be faithfully
+            # reduced to a boolean keep-mask here; SDPA handles them correctly.
+            raise ValueError(
+                f"non-boolean attn_mask (dtype={mask.dtype}); FlashInfer custom_mask "
+                "is boolean-only — deferring to SDPA"
+            )
+        # Diffusion masks arrive as (qo,kv), (1,1,kv), (B,1,1,kv), (B,1,qo,kv)
+        # or (B,H,qo,kv). The mask is identical across heads, so collapse the
+        # head dim, but keep a real batch dim — indexing mask[0] would reuse
+        # sample 0's padding for every sample (wrong under CFG / mixed lengths).
+        if mask.dim() == 4:
+            mask = mask[:, 0]  # (B, qo|1, kv)
+        if mask.dim() == 3 and mask.shape[0] == 1:
+            mask = mask[0]  # (qo|1, kv) — shared across the batch
+        try:
+            if mask.dim() >= 3:
+                mask = mask.broadcast_to((batch_size, qo_len, kv_len))
+            else:
+                mask = mask.broadcast_to((qo_len, kv_len))
+        except RuntimeError as e:
+            raise ValueError(
+                f"attn_mask shape {tuple(attn_mask.shape)} cannot broadcast to "
+                f"(batch={batch_size}, qo_len={qo_len}, kv_len={kv_len})"
+            ) from e
+        if mask.all():
+            return None
+        # ``broadcast_to`` returns a non-contiguous view; materialize for the
+        # kernel, which reads from GPU memory directly.
+        return mask.contiguous()
+
+    def _sdpa_fallback(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+    ) -> torch.Tensor:
+        from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
+
+        impl = SDPAImpl(
+            num_heads=query.shape[2],
+            head_size=query.shape[3],
+            softmax_scale=self.softmax_scale,
+            causal=self.causal,
+        )
+        return impl.forward_cuda(query, key, value, attn_metadata)
+
+    def forward_cuda(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        if not HAS_FLASHINFER:
+            raise ImportError(
+                "FLASHINFER_ATTN backend requires flashinfer. "
+                "Install it or set DIFFUSION_ATTENTION_BACKEND to another backend."
+            )
+
+        # Try the custom_mask path; if the mask can't be packed into the
+        # (qo_len, kv_len) layout FlashInfer expects, fall back to SDPA
+        # instead of risking an illegal-memory-access crash in the kernel.
+        # Input layout is (B, S, H, D); FlashInfer dense prefill takes (S, H, D).
+        batch_size = query.shape[0]
+
+        custom_mask = None
+        if attn_metadata is not None and attn_metadata.attn_mask is not None:
+            try:
+                custom_mask = self._pack_mask_for_flashinfer(
+                    attn_metadata.attn_mask,
+                    batch_size=batch_size,
+                    qo_len=query.shape[1],
+                    kv_len=key.shape[1],
+                )
+            except ValueError as e:
+                logger.debug("Falling back to SDPA for mask path: %s", e)
+                return self._sdpa_fallback(query, key, value, attn_metadata)
+            # FlashInfer cannot combine causal masking with a custom_mask; rather
+            # than silently dropping the explicit mask (diverging from SDPA), let
+            # SDPA handle the causal+mask case correctly.
+            if custom_mask is not None and self.causal:
+                logger.debug("causal=True with explicit attn_mask; deferring to SDPA")
+                return self._sdpa_fallback(query, key, value, attn_metadata)
+
+        outputs = []
+        for b in range(batch_size):
+            kwargs: dict = {
+                "sm_scale": self.softmax_scale,
+                "causal": self.causal,
+                "return_lse": False,
+            }
+            if custom_mask is not None:
+                kwargs["custom_mask"] = custom_mask if custom_mask.dim() == 2 else custom_mask[b]
+            out = single_prefill_with_kv_cache(query[b], key[b], value[b], **kwargs)
+            outputs.append(out)
+
+        if batch_size == 1:
+            return outputs[0].unsqueeze(0)
+        return torch.stack(outputs, dim=0)
diff --git a/vllm_omni/diffusion/attention/backends/registry.py b/vllm_omni/diffusion/attention/backends/registry.py
index a77223d16..8c22cc2cd 100644
--- a/vllm_omni/diffusion/attention/backends/registry.py
+++ b/vllm_omni/diffusion/attention/backends/registry.py
@@ -59,6 +59,9 @@ class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta):
     FLASH_ATTN = "vllm_omni.diffusion.attention.backends.flash_attn.FlashAttentionBackend"
     TORCH_SDPA = "vllm_omni.diffusion.attention.backends.sdpa.SDPABackend"
     SAGE_ATTN = "vllm_omni.diffusion.attention.backends.sage_attn.SageAttentionBackend"
+    SAGE_ATTN_3 = "vllm_omni.diffusion.attention.backends.sage_attn3.SageAttention3Backend"
+    CUDNN_ATTN = "vllm_omni.diffusion.attention.backends.cudnn_attn.CuDNNAttentionBackend"
+    FLASHINFER_ATTN = "vllm_omni.diffusion.attention.backends.flashinfer_attn.FlashInferAttentionBackend"
 
     def get_path(self, include_classname: bool = True) -> str:
         """Get the class path for this backend (respects overrides).
diff --git a/vllm_omni/diffusion/attention/backends/sage_attn3.py b/vllm_omni/diffusion/attention/backends/sage_attn3.py
new file mode 100644
index 000000000..51ac3a48a
--- /dev/null
+++ b/vllm_omni/diffusion/attention/backends/sage_attn3.py
@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+import torch.nn.functional as F
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+
+logger = init_logger(__name__)
+
+try:
+    from sageattn3 import sageattn3_blackwell  # noqa: F401
+except ImportError:
+    logger.warning(
+        "SageAttention3Backend is not available. Install `sageattn3` from "
+        "https://github.com/thu-ml/SageAttention/tree/main/sageattention3_blackwell"
+    )
+    raise ImportError
+
+
+# Wrapping sageattn3_blackwell as a torch.library custom op keeps it opaque to
+# torch.compile. Otherwise Dynamo graph-breaks on the raw pybind11 kernel and
+# Inductor fails scheduling with KeyError: 'op5'. The hasattr guard keeps this
+# idempotent across test re-imports that pop the module from sys.modules.
+if not hasattr(torch.ops.vllm_omni, "sageattn3_blackwell"):
+
+    @torch.library.custom_op("vllm_omni::sageattn3_blackwell", mutates_args=())
+    def _sageattn3_blackwell_op(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        is_causal: bool,
+    ) -> torch.Tensor:
+        from sageattn3 import sageattn3_blackwell as _kernel
+
+        return _kernel(query, key, value, is_causal=is_causal)
+
+    @_sageattn3_blackwell_op.register_fake
+    def _(query, key, value, is_causal):
+        return torch.empty_like(query)
+
+
+_sageattn3_blackwell_op = torch.ops.vllm_omni.sageattn3_blackwell
+
+
+class SageAttention3Backend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "SAGE_ATTN_3"
+
+    @staticmethod
+    def get_impl_cls() -> type["SageAttention3Impl"]:
+        return SageAttention3Impl
+
+
+class SageAttention3Impl(AttentionImpl):
+    _warned_gqa_fallback_global: bool = False
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.dropout = extra_impl_args.get("dropout_p", 0.0)
+
+    def forward_cuda(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        query = query.transpose(1, 2).contiguous()
+        key = key.transpose(1, 2).contiguous()
+        value = value.transpose(1, 2).contiguous()
+
+        if key.shape[1] != query.shape[1]:
+            if query.shape[1] % key.shape[1] != 0:
+                raise ValueError(
+                    "GQA/MQA requires query heads to be a multiple of KV heads, "
+                    f"got q_heads={query.shape[1]} and kv_heads={key.shape[1]}"
+                )
+            if not type(self)._warned_gqa_fallback_global:
+                logger.warning("SageAttention3 does not support GQA/MQA (Hq != Hkv); falling back to torch SDPA.")
+                type(self)._warned_gqa_fallback_global = True
+            output = F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                is_causal=self.causal,
+                dropout_p=self.dropout,
+                scale=self.softmax_scale,
+                enable_gqa=True,
+            )
+        else:
+            output = _sageattn3_blackwell_op(query, key, value, self.causal)
+
+        return output.transpose(1, 2).contiguous()
diff --git a/vllm_omni/diffusion/attention/backends/sdpa.py b/vllm_omni/diffusion/attention/backends/sdpa.py
index ab71e753b..f8d607e48 100644
--- a/vllm_omni/diffusion/attention/backends/sdpa.py
+++ b/vllm_omni/diffusion/attention/backends/sdpa.py
@@ -30,10 +30,6 @@ def _maybe_reshape_attn_mask(
       - broadcast_k: [batch_size, 1, 1, seq_len_k]
       - full_qk: [batch_size, 1, seq_len_q, seq_len_k]
     """
-    # Skip Attention Mask if all values are 1, `None` mask can speedup the computation
-    if attn_mask is not None and torch.all(attn_mask != 0):
-        attn_mask = None
-
     # Reshape Attention Mask
     # 2D [batch_size, seq_len_k] mask only.
     if (
@@ -89,6 +85,7 @@ class SDPAImpl(AttentionImpl):
     ) -> None:
         self.causal = causal
         self.softmax_scale = softmax_scale
+        self.requires_gqa = num_heads != num_kv_heads
         if backend_kwargs:
             logger.warning("SDPAImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
 
@@ -115,6 +112,7 @@ class SDPAImpl(AttentionImpl):
             dropout_p=0.0,
             is_causal=self.causal,
             scale=self.softmax_scale,
+            enable_gqa=self.requires_gqa,
         )
         out = output.permute(0, 2, 1, 3)
         return out
diff --git a/vllm_omni/diffusion/attention/parallel/ring.py b/vllm_omni/diffusion/attention/parallel/ring.py
index 94cd0d205..745b25ccf 100644
--- a/vllm_omni/diffusion/attention/parallel/ring.py
+++ b/vllm_omni/diffusion/attention/parallel/ring.py
@@ -10,7 +10,7 @@ import torch
 from vllm.logger import init_logger
 
 # import torch.distributed as dist # Not used directly here, but good practice if needed
-from vllm_omni.diffusion.attention.backends.ring.ring_globals import HAS_FA3, HAS_FLASH_ATTN
+from vllm_omni.diffusion.attention.backends.ring.ring_globals import HAS_AITER, HAS_FA3, HAS_FLASH_ATTN
 from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType
 from vllm_omni.diffusion.attention.parallel.base import (
     ParallelAttentionContext,
@@ -110,13 +110,12 @@ class RingParallelAttention:
         if backend_pref is not None:
             backend_pref = backend_pref.lower()
 
-        # Determine attention type with fallback chain: FA3 -> FA2 -> SDPA
         # FP32 is not supported by Flash Attention, force SDPA
         if query.dtype == torch.float32:
             backend_pref = "sdpa"
-        elif not HAS_FA3 and not HAS_FLASH_ATTN:
+        elif not HAS_FA3 and not HAS_FLASH_ATTN and not HAS_AITER:
             if backend_pref != "sdpa":
-                logger.warning_once("Flash Attention (FA2/FA3) is not available! Force enabling SDPA.")
+                logger.warning_once("Flash Attention (FA2/FA3/AITER) is not available! Force enabling SDPA.")
             backend_pref = "sdpa"
 
         # Extract joint tensors
@@ -147,7 +146,13 @@ class RingParallelAttention:
         from vllm_omni.diffusion.attention.backends.ring_flash_attn import ring_flash_attn_func
 
         # Prefer FA3 over FA2 for better performance (FA3 supports Ampere/Ada/Hopper)
-        attn_type = AttnType.FA3 if HAS_FA3 else AttnType.FA
+        # On ROCm, use AITER
+        if HAS_FA3:
+            attn_type = AttnType.FA3
+        elif HAS_AITER:
+            attn_type = AttnType.AITER
+        else:
+            attn_type = AttnType.FA
 
         return ring_flash_attn_func(
             query,
diff --git a/vllm_omni/diffusion/cache/__init__.py b/vllm_omni/diffusion/cache/__init__.py
index a5968f612..2d87a089a 100644
--- a/vllm_omni/diffusion/cache/__init__.py
+++ b/vllm_omni/diffusion/cache/__init__.py
@@ -5,12 +5,19 @@ Cache module for diffusion model inference acceleration.
 
 This module provides a unified cache backend system for different caching strategies:
 - TeaCache: Timestep Embedding Aware Cache for adaptive transformer caching
+- MagCache: Magnitude-based Cache for adaptive transformer caching
 - cache-dit: DBCache, SCM, and TaylorSeer caching strategies
 
 Cache backends are instantiated directly via their constructors and configured via OmniDiffusionConfig.
 """
 
 from vllm_omni.diffusion.cache.base import CacheBackend
+from vllm_omni.diffusion.cache.prompt_embed_cache import (
+    PromptEmbedCache,
+    install_prompt_embed_cache,
+    resolve_prompt_embed_cache_config,
+    uninstall_prompt_embed_cache,
+)
 from vllm_omni.diffusion.cache.teacache import (
     CacheContext,
     TeaCacheConfig,
@@ -20,8 +27,12 @@ from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend
 
 __all__ = [
     "CacheBackend",
-    "TeaCacheConfig",
     "CacheContext",
     "TeaCacheBackend",
+    "TeaCacheConfig",
     "apply_teacache_hook",
+    "PromptEmbedCache",
+    "install_prompt_embed_cache",
+    "resolve_prompt_embed_cache_config",
+    "uninstall_prompt_embed_cache",
 ]
diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py
index daf9b3939..005857642 100644
--- a/vllm_omni/diffusion/cache/cache_dit_backend.py
+++ b/vllm_omni/diffusion/cache/cache_dit_backend.py
@@ -18,6 +18,7 @@ from cache_dit import BlockAdapter, DBCacheConfig, ForwardPattern, ParamsModifie
 from cache_dit.caching.block_adapters import FakeDiffusionPipeline
 from cache_dit.caching.cache_adapters.cache_adapter import CachedAdapter
 from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2
+from cache_dit.caching.cache_blocks.pattern_3_4_5 import CachedBlocks_Pattern_3_4_5
 from cache_dit.caching.cache_contexts import BasicCacheConfig
 from cache_dit.caching.cache_contexts.cache_manager import CachedContextManager
 from vllm.logger import init_logger
@@ -1170,6 +1171,162 @@ def enable_cache_for_bagel(pipeline: Any, cache_config: Any) -> Callable[[int],
     return refresh_cache_context
 
 
+class SensenovaCachedBlocks(CachedBlocks_Pattern_3_4_5):
+    """
+    Custom CachedBlocks for SenseNova-U1 that only caches image-token hidden
+    states during denoising.
+    """
+
+    @classmethod
+    def _is_denoising_call(cls, kwargs: dict[str, Any]) -> bool:
+        if kwargs.get("cache_dit_skip", False):
+            return False
+
+        # Prefix/text forwards either omit image_gen_indicators or update the
+        # DynamicCache. Denoising forwards are gen-only and use update_cache=False.
+        if kwargs.get("update_cache", True):
+            return False
+
+        exist_gen = kwargs.get("exist_gen")
+        exist_und = kwargs.get("exist_und")
+        if exist_gen is None or exist_und is None:
+            image_gen_indicators = kwargs.get("image_gen_indicators")
+            if image_gen_indicators is None:
+                return False
+            exist_gen = image_gen_indicators.any().item()
+            exist_und = (~image_gen_indicators).any().item()
+
+        return exist_gen and not exist_und
+
+    @staticmethod
+    def _strip_cache_only_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
+        kwargs = dict(kwargs)
+        kwargs.pop("cache_dit_skip", None)
+        return kwargs
+
+    def forward(self, hidden_states: torch.Tensor, *args, **kwargs):
+        block_kwargs = self._strip_cache_only_kwargs(kwargs)
+        if not self._is_denoising_call(kwargs):
+            hidden_states, new_encoder_hidden_states = self.call_blocks(
+                hidden_states,
+                *args,
+                **block_kwargs,
+            )
+            return self._process_forward_outputs(hidden_states, new_encoder_hidden_states)
+
+        return super().forward(hidden_states, *args, **block_kwargs)
+
+
+class SensenovaCachedAdapter(CachedAdapter):
+    """Custom CachedAdapter for SenseNova-U1 that uses SensenovaCachedBlocks."""
+
+    @classmethod
+    def collect_unified_blocks(
+        cls,
+        block_adapter: BlockAdapter,
+        contexts_kwargs: list[dict],
+    ) -> list[dict[str, torch.nn.ModuleList]]:
+        BlockAdapter.assert_normalized(block_adapter)
+
+        total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = []
+        assert hasattr(block_adapter.pipe, "_context_manager")
+
+        for i in range(len(block_adapter.transformer)):
+            unified_blocks_bind_context = {}
+            for j in range(len(block_adapter.blocks[i])):
+                cache_config: BasicCacheConfig = contexts_kwargs[i * len(block_adapter.blocks[i]) + j]["cache_config"]
+                unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList(
+                    [
+                        SensenovaCachedBlocks(
+                            # 0. Transformer blocks configuration
+                            block_adapter.blocks[i][j],
+                            transformer=block_adapter.transformer[i],
+                            forward_pattern=block_adapter.forward_pattern[i][j],
+                            check_forward_pattern=block_adapter.check_forward_pattern,
+                            check_num_outputs=block_adapter.check_num_outputs,
+                            # 1. Cache/Prune context configuration
+                            cache_prefix=block_adapter.blocks_name[i][j],
+                            cache_context=block_adapter.unique_blocks_name[i][j],
+                            context_manager=block_adapter.pipe._context_manager,
+                            cache_type=cache_config.cache_type,
+                        )
+                    ]
+                )
+
+            total_cached_blocks.append(unified_blocks_bind_context)
+
+        return total_cached_blocks
+
+
+def enable_cache_for_sensenova_u1(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for SenseNova-U1 model.
+
+    Args:
+        pipeline: The SenseNova-U1 pipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
+
+    Returns:
+        A refresh function that can be called to update cache context with new num_inference_steps.
+    """
+    transformer = getattr(getattr(pipeline, "language_model", None), "model", None)
+    if transformer is None or not hasattr(transformer, "layers"):
+        raise ValueError("SenseNovaU1Pipeline cache-dit expects pipeline.language_model.model.layers.")
+
+    # Build DBCacheConfig
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    calibrator_config = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    modifier = ParamsModifier(
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
+    logger.info(
+        f"Enabling cache-dit on SenseNova-U1 decoder layers: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
+
+    SensenovaCachedAdapter.apply(
+        BlockAdapter(
+            transformer=transformer,
+            blocks=transformer.layers,
+            forward_pattern=ForwardPattern.Pattern_3,
+            params_modifiers=[modifier],
+            check_forward_pattern=True,
+            has_separate_cfg=True,
+        ),
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        transformer = pipeline.language_model.model
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
+
+
 def enable_cache_for_flux2(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
     """Enable cache-dit for Flux.2-dev pipeline.
 
@@ -1320,6 +1477,74 @@ def enable_cache_for_glm_image(pipeline: Any, cache_config: Any) -> Callable[[in
     return refresh_cache_context
 
 
+def enable_cache_for_dreamid_omni(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for DreamID-Omni fused pipeline.
+
+    Args:
+        pipeline: The DreamIDOmni pipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
+
+    Returns:
+        A refresh function that can be called with a new ``num_inference_steps``
+        to update the cache context for the pipeline.
+    """
+    transformer = getattr(pipeline, "transformer", None)
+    fused_blocks = getattr(transformer, "fused_blocks", None)
+    if transformer is None or fused_blocks is None:
+        raise ValueError("DreamIDOmniPipeline cache-dit expects pipeline.transformer.fused_blocks.")
+
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    calibrator = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    modifier = ParamsModifier(
+        cache_config=db_cache_config,
+        calibrator_config=calibrator,
+    )
+
+    logger.info(
+        f"Enabling cache-dit on DreamID-Omni fused blocks: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=transformer,
+            blocks=fused_blocks,
+            forward_pattern=ForwardPattern.Pattern_0,
+            params_modifiers=[modifier],
+            has_separate_cfg=True,
+        ),
+        cache_config=db_cache_config,
+    )
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        transformer = pipeline.transformer
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
+
+
 def enable_cache_for_ernie_image(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
     """Enable cache-dit for ERNIE-Image pipeline.
 
@@ -1413,7 +1638,7 @@ def enable_cache_for_helios(pipeline: Any, cache_config: Any) -> Callable[[int],
             params_modifiers=[
                 ParamsModifier(cache_config=db_cache_config, calibrator_config=calibrator_config),
             ],
-            has_separate_cfg=False,
+            has_separate_cfg=True,
         ),
         cache_config=db_cache_config,
         calibrator_config=calibrator_config,
@@ -1519,8 +1744,10 @@ CUSTOM_DIT_ENABLERS.update(
         "LTX23Pipeline": enable_cache_for_ltx2,
         "LTX23ImageToVideoPipeline": enable_cache_for_ltx2,
         "BagelPipeline": enable_cache_for_bagel,
+        "SenseNovaU1Pipeline": enable_cache_for_sensenova_u1,
         "GlmImagePipeline": enable_cache_for_glm_image,
         "Flux2Pipeline": enable_cache_for_flux2,
+        "DreamIDOmniPipeline": enable_cache_for_dreamid_omni,
         "ErnieImagePipeline": enable_cache_for_ernie_image,
         "HunyuanVideo15Pipeline": enable_cache_for_hunyuan_video_15,
         "HunyuanVideo15I2VPipeline": enable_cache_for_hunyuan_video_15,
diff --git a/vllm_omni/diffusion/cache/magcache/__init__.py b/vllm_omni/diffusion/cache/magcache/__init__.py
new file mode 100644
index 000000000..08f0d88fa
--- /dev/null
+++ b/vllm_omni/diffusion/cache/magcache/__init__.py
@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.diffusion.cache.magcache.config import MagCacheConfig
+from vllm_omni.diffusion.cache.magcache.hook import (
+    MagCacheBlockHook,
+    MagCacheHeadHook,
+    MagCacheState,
+    apply_mag_cache_hook,
+)
+from vllm_omni.diffusion.cache.magcache.strategy import (
+    Flux2MagCacheStrategy,
+    FluxMagCacheStrategy,
+    MagCacheStrategy,
+    MagCacheStrategyRegistry,
+    get_strategy,
+    register_strategy,
+)
+
+__all__ = [
+    "Flux2MagCacheStrategy",
+    "FluxMagCacheStrategy",
+    "MagCacheBlockHook",
+    "MagCacheConfig",
+    "MagCacheHeadHook",
+    "MagCacheState",
+    "MagCacheStrategy",
+    "MagCacheStrategyRegistry",
+    "apply_mag_cache_hook",
+    "get_strategy",
+    "register_strategy",
+]
diff --git a/vllm_omni/diffusion/cache/magcache/backend.py b/vllm_omni/diffusion/cache/magcache/backend.py
new file mode 100644
index 000000000..4ad427d3a
--- /dev/null
+++ b/vllm_omni/diffusion/cache/magcache/backend.py
@@ -0,0 +1,177 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+MagCache backend implementation.
+
+This module provides the MagCache backend that implements the CacheBackend
+interface using the hooks-based MagCache system.
+"""
+
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.cache.base import CacheBackend
+from vllm_omni.diffusion.cache.magcache.config import MagCacheConfig
+from vllm_omni.diffusion.cache.magcache.hook import (
+    apply_mag_cache_hook,
+)
+from vllm_omni.diffusion.cache.magcache.strategy import (
+    get_strategy,
+)
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+
+logger = init_logger(__name__)
+
+
+class MagCacheBackend(CacheBackend):
+    """
+    MagCache implementation using hooks.
+
+    MagCache (Magnitude-based Cache) is an adaptive caching technique that
+    speeds up diffusion inference by reusing transformer block computations
+    based on accumulated magnitude error between timesteps.
+
+    The backend applies MagCache hooks to the transformer which intercept the
+    forward pass and implement the caching logic transparently.
+
+    Example:
+        >>> from vllm_omni.diffusion.data import DiffusionCacheConfig
+        >>> from vllm_omni.diffusion.cache.magcache.config import MagCacheConfig
+        >>> from vllm_omni.diffusion.cache.magcache.strategy import FluxMagCacheStrategy
+        >>> cache_config = DiffusionCacheConfig(
+        ...     mag_ratios=FluxMagCacheStrategy.FLUX_MAG_RATIOS,
+        ...     num_inference_steps=28,
+        ...     mag_threshold=0.24,
+        ...     mag_max_skip_steps=5,
+        ...     mag_retention_ratio=0.1,
+        ... )
+        >>> backend = MagCacheBackend(cache_config)
+        >>> backend.enable(pipeline)
+        >>> backend.refresh(pipeline, num_inference_steps=50)
+    """
+
+    def __init__(self, config: DiffusionCacheConfig):
+        super().__init__(config)
+        self._registered = False
+        self._magcache_config: MagCacheConfig | None = None
+        self._transformer_id: int | None = None
+
+    def enable(self, pipeline: Any) -> None:
+        """Enable MagCache on transformer using hooks.
+
+        This creates a MagCacheConfig from the backend's DiffusionCacheConfig
+        and applies the MagCache hook to the transformer.
+
+        Args:
+            pipeline: Diffusion pipeline instance. Extracts transformer and transformer_type:
+                     - transformer: pipeline.transformer
+                     - transformer_type: pipeline.transformer.__class__.__name__
+        """
+        transformer = pipeline.transformer
+        transformer_type = transformer.__class__.__name__
+
+        num_inference_steps = self.config.num_inference_steps or 28
+
+        mag_ratios = self.config.mag_ratios
+        strategy = None
+
+        if mag_ratios is None and not self.config.mag_calibrate:
+            strategy = get_strategy(transformer_type)
+            original_ratios = strategy.mag_ratios
+
+            if len(original_ratios) != num_inference_steps and hasattr(strategy, "nearest_interp"):
+                mag_ratios = strategy.nearest_interp(original_ratios, num_inference_steps)
+                logger.info(
+                    f"MagCache: Interpolated mag_ratios from {len(original_ratios)} to {num_inference_steps} steps"
+                )
+            else:
+                mag_ratios = original_ratios
+                if len(original_ratios) != num_inference_steps:
+                    logger.warning(
+                        f"MagCache: mag_ratios length ({len(original_ratios)}) != "
+                        f"num_inference_steps ({num_inference_steps}), "
+                        f"this may cause unexpected behavior"
+                    )
+
+            logger.info(f"MagCache: Using mag_ratios from {type(strategy).__name__}")
+
+        if mag_ratios is None and not self.config.mag_calibrate:
+            raise ValueError(
+                f"mag_ratios must be provided for MagCache. "
+                f"For {transformer_type}, you need to provide mag_ratios or run in calibrate mode."
+            )
+
+        self._magcache_config = MagCacheConfig(
+            transformer_type=transformer_type,
+            threshold=self.config.mag_threshold,
+            max_skip_steps=self.config.mag_max_skip_steps,
+            retention_ratio=self.config.mag_retention_ratio,
+            num_inference_steps=num_inference_steps,
+            mag_calibrate=self.config.mag_calibrate,
+            mag_ratios=mag_ratios if not self.config.mag_calibrate else None,
+        )
+        self._transformer_id = id(transformer)
+
+        apply_mag_cache_hook(transformer, self._magcache_config, strategy=strategy)
+
+        self._registered = True
+        self.enabled = True
+
+    def refresh(self, pipeline: Any, num_inference_steps: int) -> None:
+        """Refresh MagCache state for new generation.
+
+        Clears all cached residuals and resets counters/accumulators.
+        Should be called before each generation to ensure clean state.
+
+        Args:
+            pipeline: Diffusion pipeline instance. Extracts transformer via pipeline.transformer.
+            num_inference_steps: Number of inference steps for the current generation.
+                                May be used for cache context updates.
+        """
+        transformer = pipeline.transformer
+        current_transformer_id = id(transformer)
+
+        needs_re_register = False
+
+        if self._registered and hasattr(self, "_transformer_id"):
+            if current_transformer_id != self._transformer_id:
+                logger.warning(
+                    f"Transformer was replaced (id changed from {self._transformer_id} "
+                    f"to {current_transformer_id}), re-registering hooks"
+                )
+                needs_re_register = True
+
+        if not self._registered or needs_re_register:
+            self.enable(pipeline)
+            return
+
+        blocks_with_hooks = []
+
+        for name, submodule in transformer.named_children():
+            if not isinstance(submodule, torch.nn.ModuleList):
+                continue
+            for index, block in enumerate(submodule):
+                registry = getattr(block, "_hook_registry", None)
+                if registry is not None and len(registry._hooks) > 0:
+                    blocks_with_hooks.append((f"{name}.{index}", block, registry))
+
+        if not blocks_with_hooks:
+            logger.warning("No hooks found on transformer blocks, re-registering")
+            apply_mag_cache_hook(transformer, self._magcache_config)
+            self._transformer_id = current_transformer_id
+        else:
+            for name, block, registry in blocks_with_hooks:
+                for hook in registry._hooks.values():
+                    if hasattr(hook, "reset_state"):
+                        hook.reset_state(block)
+
+    def is_enabled(self) -> bool:
+        """Check if MagCache is enabled.
+
+        Returns:
+            True if enabled, False otherwise.
+        """
+        return self.enabled
diff --git a/vllm_omni/diffusion/cache/magcache/config.py b/vllm_omni/diffusion/cache/magcache/config.py
new file mode 100644
index 000000000..6ccab8fa5
--- /dev/null
+++ b/vllm_omni/diffusion/cache/magcache/config.py
@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class MagCacheConfig:
+    """
+    Configuration for MagCache applied to transformer models.
+
+    MagCache (Magnitude-based Cache) is an adaptive caching technique that speeds up
+    diffusion model inference by reusing transformer block computations based on
+    magnitude ratios between consecutive timesteps.
+
+    Reference: https://github.com/Zehong-Ma/MagCache
+
+    Args:
+        threshold: Accumulated error threshold. Higher = more aggressive skipping (faster, lower quality).
+            Default: 0.24
+        max_skip_steps: Max consecutive skip steps (K).
+            Default: 5
+        retention_ratio: Fraction of initial steps where skipping is disabled (stability).
+            Default: 0.1
+        num_inference_steps: Total inference steps. Required for retention step calculation.
+            Default: 28
+        mag_ratios: Pre-computed magnitude ratios per step. Calibrate or use strategy defaults.
+            Default: None
+        mag_calibrate: If True, runs without skipping and logs norm_ratios for calibration.
+            Default: False
+        transformer_type: Transformer class name for logging.
+            Default: "FluxTransformer2DModel"
+    """
+
+    threshold: float = 0.24
+    max_skip_steps: int = 5
+    retention_ratio: float = 0.1
+    num_inference_steps: int = 28
+    mag_ratios: torch.Tensor | list[float] | None = None
+    mag_calibrate: bool = False
+    transformer_type: str = "FluxTransformer2DModel"
+
+    def __post_init__(self) -> None:
+        """Validate and set default coefficients."""
+        if self.threshold <= 0:
+            raise ValueError(f"threshold must be positive, got {self.threshold}")
+
+        if self.max_skip_steps <= 0:
+            raise ValueError(f"max_skip_steps must be positive, got {self.max_skip_steps}")
+
+        if not 0 < self.retention_ratio < 1:
+            raise ValueError(f"retention_ratio must be in (0, 1), got {self.retention_ratio}")
+
+        if self.num_inference_steps is None:
+            raise ValueError(
+                "num_inference_steps must be provided for MagCache. "
+                "This is required to determine retention steps and interpolate mag_ratios. "
+                "For Flux models, use num_inference_steps=28."
+            )
+
+        if self.num_inference_steps <= 0:
+            raise ValueError(f"num_inference_steps must be positive, got {self.num_inference_steps}")
+
+        if not self.mag_calibrate and self.mag_ratios is None:
+            raise ValueError(
+                "mag_ratios must be provided for MagCache inference because these ratios "
+                "are model-dependent. To get them for your model:\n"
+                "1. Initialize MagCacheConfig(mag_calibrate=True, ...)\n"
+                "2. Run inference on your model once.\n"
+                "3. Copy the printed ratios array and pass it to mag_ratios in the config.\n"
+                "For Flux models, you can import FLUX_MAG_RATIOS from vllm_omni.diffusion.cache.magcache.strategy."
+            )
+
+        if not self.mag_calibrate and self.mag_ratios is not None:
+            if not torch.is_tensor(self.mag_ratios):
+                self.mag_ratios = torch.tensor(self.mag_ratios)
diff --git a/vllm_omni/diffusion/cache/magcache/hook.py b/vllm_omni/diffusion/cache/magcache/hook.py
new file mode 100644
index 000000000..6a440e020
--- /dev/null
+++ b/vllm_omni/diffusion/cache/magcache/hook.py
@@ -0,0 +1,610 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Hook-based MagCache implementation for vLLM-Omni.
+
+This module implements a diffusers-style hook system for MagCache (Magnitude-based Cache),
+providing adaptive caching for diffusion model inference.
+
+MagCache speeds up inference by skipping transformer block computations when the accumulated
+magnitude error is below a threshold, reusing cached residuals instead.
+
+Based on: https://github.com/Zehong-Ma/MagCache
+Reference: diffusers/src/diffusers/hooks/mag_cache.py
+
+Architecture:
+- MagCacheStrategy: Model-specific strategy for preprocessing/postprocessing
+- MagCacheState: Per-step state tracking residuals and accumulated error
+- MagCacheHeadHook: Decides whether to skip based on accumulated error
+- MagCacheBlockHook: Computes and stores residuals at tail block
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from diffusers.hooks._helpers import TransformerBlockMetadata, TransformerBlockRegistry
+from diffusers.utils.torch_utils import unwrap_module
+
+from vllm_omni.diffusion.cache.magcache.config import MagCacheConfig
+from vllm_omni.diffusion.cache.magcache.state import MagCacheState
+from vllm_omni.diffusion.cache.magcache.strategy import MagCacheStrategy, MagCacheStrategyRegistry
+from vllm_omni.diffusion.hooks.base import HookRegistry, ModelHook, StateManager
+from vllm_omni.logger import init_logger
+
+logger = init_logger(__name__)
+
+_MAG_CACHE_LEADER_BLOCK_HOOK = "mag_cache_leader_block_hook"
+_MAG_CACHE_BLOCK_HOOK = "mag_cache_block_hook"
+
+
+class MagCacheHeadHook(ModelHook):
+    """Head block hook for MagCache - decides whether to skip computation."""
+
+    _HOOK_NAME = "mag_cache_head"
+
+    def __init__(
+        self,
+        state_manager: StateManager,
+        config: MagCacheConfig,
+        strategy: MagCacheStrategy | None = None,
+        is_tail: bool = False,
+    ):
+        super().__init__()
+        self.state_manager = state_manager
+        self.config = config
+        self._strategy = strategy
+        self._metadata = None
+        self._is_tail = is_tail
+
+    def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
+        unwrapped_module = unwrap_module(module)
+        block_class = unwrapped_module.__class__
+
+        try:
+            self._metadata = TransformerBlockRegistry.get(block_class)
+        except ValueError:
+            if self._strategy is not None:
+                metadata = self._strategy.register_block_metadata(block_class)
+                if metadata is not None:
+                    TransformerBlockRegistry.register(model_class=block_class, metadata=metadata)
+                else:
+                    TransformerBlockRegistry.register(
+                        model_class=block_class,
+                        metadata=TransformerBlockMetadata(
+                            return_hidden_states_index=1,
+                            return_encoder_hidden_states_index=0,
+                        ),
+                    )
+            else:
+                TransformerBlockRegistry.register(
+                    model_class=block_class,
+                    metadata=TransformerBlockMetadata(
+                        return_hidden_states_index=1,
+                        return_encoder_hidden_states_index=0,
+                    ),
+                )
+            self._metadata = TransformerBlockRegistry.get(block_class)
+
+        return module
+
+    @torch.compiler.disable
+    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
+        if self.state_manager._current_context is None:
+            self.state_manager.set_context("magcache")
+
+        if hasattr(self._metadata, "hidden_states_argument_name"):
+            arg_name = self._metadata.hidden_states_argument_name
+        else:
+            arg_name = "hidden_states"
+        hidden_states = self._metadata._get_parameter_from_args_kwargs(arg_name, args, kwargs)
+
+        state: MagCacheState = self.state_manager.get_state()
+        state.head_block_input = hidden_states
+
+        if state._is_first_step:
+            state.accumulated_ratio = 1.0
+            state.accumulated_err = 0.0
+            state.accumulated_steps = 0
+            state._is_first_step = False
+
+        should_compute = True
+
+        if self.config.mag_calibrate:
+            should_compute = True
+        else:
+            current_step = state.step_index
+            if current_step >= len(self.config.mag_ratios):
+                current_scale = 1.0
+            else:
+                current_scale = self.config.mag_ratios[current_step]
+
+            retention_step = int(self.config.retention_ratio * self.config.num_inference_steps + 0.5)
+
+            if current_step >= retention_step:
+                state.accumulated_ratio *= current_scale
+                state.accumulated_steps += 1
+                state.accumulated_err += abs(1.0 - state.accumulated_ratio)
+
+                if (
+                    state.previous_residual is not None
+                    and state.accumulated_err <= self.config.threshold
+                    and state.accumulated_steps <= self.config.max_skip_steps
+                ):
+                    should_compute = False
+                else:
+                    state.accumulated_ratio = 1.0
+                    state.accumulated_steps = 0
+                    state.accumulated_err = 0.0
+
+        state.should_compute = should_compute
+
+        if not should_compute:
+            res = state.previous_residual
+
+            if isinstance(res, tuple):
+                res = tuple(r.to(hidden_states.device) for r in res)
+
+                if self._strategy is not None:
+                    original_encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs(
+                        "encoder_hidden_states", args, kwargs
+                    )
+                    if original_encoder_hidden_states.device != res[1].device:
+                        original_encoder_hidden_states = original_encoder_hidden_states.to(res[1].device)
+                    output, enc_output = self._strategy.apply_residual_tuple(
+                        hidden_states, original_encoder_hidden_states, res
+                    )
+                    ret_list = [None] * 2
+                    ret_list[self._metadata.return_hidden_states_index] = output
+                    ret_list[self._metadata.return_encoder_hidden_states_index] = enc_output
+                    return self.log_cache_hit(state, output, ret_list)
+                else:
+                    raise RuntimeError(
+                        f"MagCache residual is tuple but no strategy available for {self._metadata.transformer_type}. "
+                        f"Please register a MagCacheStrategy for this model."
+                    )
+            elif res.device != hidden_states.device:
+                res = res.to(hidden_states.device)
+
+            if self._strategy is not None:
+                output = self._strategy.apply_residual(hidden_states, res)
+            elif res.shape == hidden_states.shape:
+                output = hidden_states + res
+            elif (
+                hidden_states.ndim == 3
+                and res.ndim == 3
+                and hidden_states.shape[0] == res.shape[0]
+                and hidden_states.shape[2] == res.shape[2]
+            ):
+                diff = hidden_states.shape[1] - res.shape[1]
+                if diff > 0:
+                    output = hidden_states.clone()
+                    output[:, diff:, :] = output[:, diff:, :] + res
+                else:
+                    output = hidden_states + res
+            else:
+                output = hidden_states + res
+
+            if self._metadata.return_encoder_hidden_states_index is not None:
+                original_encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs(
+                    "encoder_hidden_states", args, kwargs
+                )
+                max_idx = max(
+                    self._metadata.return_hidden_states_index, self._metadata.return_encoder_hidden_states_index
+                )
+                ret_list = [None] * (max_idx + 1)
+                ret_list[self._metadata.return_hidden_states_index] = output
+                ret_list[self._metadata.return_encoder_hidden_states_index] = original_encoder_hidden_states
+                return self.log_cache_hit(state, output, ret_list)
+            else:
+                return self.log_cache_hit(state, output, None)
+        else:
+            output = self.fn_ref.original_forward(*args, **kwargs)
+            result = self.log_cache_miss(state, output)
+
+        if self._is_tail:
+            if isinstance(output, tuple):
+                out_hidden = output[self._metadata.return_hidden_states_index]
+            else:
+                out_hidden = output
+
+            in_hidden = state.head_block_input
+
+            if in_hidden is not None:
+                if self._strategy is not None:
+                    residual = self._strategy.compute_residual(out_hidden, in_hidden)
+                elif out_hidden.shape == in_hidden.shape:
+                    residual = out_hidden - in_hidden
+                else:
+                    residual = out_hidden
+
+                if self.config.mag_calibrate:
+                    self._perform_calibration_head(state, residual)
+
+                state.previous_residual = residual
+                self._advance_step_head(state)
+
+        return result
+
+    def _perform_calibration_head(
+        self,
+        state: MagCacheState,
+        current_residual: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+    ) -> None:
+        if self._strategy is not None:
+            ratio, std, cos_dis = self._strategy.compute_calibration_metrics(current_residual, state.previous_residual)
+        else:
+            if state.previous_residual is None:
+                ratio, std, cos_dis = 1.0, 0.0, 0.0
+            else:
+                ratio, std, cos_dis = 1.0, 0.0, 0.0
+
+        state.calibration_ratios.append(ratio)
+        state.norm_ratios.append(round(ratio, 5))
+        state.norm_stds.append(round(std, 5))
+        state.cos_dises.append(round(cos_dis, 5))
+
+    def _advance_step_head(self, state: MagCacheState) -> None:
+        state.step_index += 1
+        if state.step_index >= self.config.num_inference_steps:
+            if self.config.mag_calibrate:
+                logger.info("MagCache calibration complete.")
+                logger.info(f"norm_ratios: {state.norm_ratios}")
+                logger.info(f"norm_stds: {state.norm_stds}")
+                logger.info(f"cos_dises: {state.cos_dises}")
+                logger.info("Copy these values to DiffusionCacheConfig(mag_ratios=...) for production use")
+
+            state.step_index = 0
+            state.accumulated_ratio = 1.0
+            state.accumulated_steps = 0
+            state.accumulated_err = 0.0
+            state.previous_residual = None
+            state.calibration_ratios = []
+            state.norm_ratios = []
+            state.norm_stds = []
+            state.cos_dises = []
+            state._is_first_step = True
+
+    def log_cache_hit(self, state: MagCacheState, output, ret):
+        step = state.step_index
+        if state.previous_residual is not None:
+            if isinstance(state.previous_residual, tuple):
+                residual_shape = tuple(r.shape for r in state.previous_residual)
+            else:
+                residual_shape = state.previous_residual.shape
+        else:
+            residual_shape = "None"
+        logger.debug(
+            f"[MagCache][HEAD] STEP={step}: CACHE_HIT (err={state.accumulated_err:.6f}, "
+            f"steps_skipped={state.accumulated_steps}, residual_shape={residual_shape}"
+        )
+        return ret if ret is not None else output
+
+    def log_cache_miss(self, state: MagCacheState, output):
+        step = state.step_index
+        residual_norm = 0.0
+        if state.previous_residual is not None:
+            if isinstance(state.previous_residual, tuple):
+                residual_norm = sum(float(torch.norm(r).item()) for r in state.previous_residual)
+            else:
+                residual_norm = float(torch.norm(state.previous_residual).item())
+        logger.debug(
+            f"[MagCache][HEAD] STEP={step}: CACHE_MISS "
+            f"(err={state.accumulated_err:.6f}, acc_ratio={state.accumulated_ratio:.6f}, "
+            f"residual_norm={residual_norm:.6f}, threshold={self.config.threshold}, "
+            f"max_skip={self.config.max_skip_steps})"
+        )
+        return output
+
+    def reset_state(self, module: torch.nn.Module) -> torch.nn.Module:
+        self.state_manager.reset()
+        return module
+
+
+class MagCacheBlockHook(ModelHook):
+    """Block hook for MagCache - computes residuals at tail block."""
+
+    _HOOK_NAME = "mag_cache_block"
+
+    def __init__(
+        self,
+        state_manager: StateManager,
+        is_tail: bool = False,
+        config: MagCacheConfig | None = None,
+        strategy: MagCacheStrategy | None = None,
+    ):
+        super().__init__()
+        self.state_manager = state_manager
+        self.is_tail = is_tail
+        self.config = config
+        self._strategy = strategy
+        self._metadata = None
+
+    def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
+        unwrapped_module = unwrap_module(module)
+        block_class = unwrapped_module.__class__
+
+        try:
+            self._metadata = TransformerBlockRegistry.get(block_class)
+        except ValueError:
+            if self._strategy is not None:
+                metadata = self._strategy.register_block_metadata(block_class)
+                if metadata is not None:
+                    TransformerBlockRegistry.register(model_class=block_class, metadata=metadata)
+                else:
+                    TransformerBlockRegistry.register(
+                        model_class=block_class,
+                        metadata=TransformerBlockMetadata(
+                            return_hidden_states_index=1,
+                            return_encoder_hidden_states_index=0,
+                        ),
+                    )
+            else:
+                TransformerBlockRegistry.register(
+                    model_class=block_class,
+                    metadata=TransformerBlockMetadata(
+                        return_hidden_states_index=1,
+                        return_encoder_hidden_states_index=0,
+                    ),
+                )
+            self._metadata = TransformerBlockRegistry.get(block_class)
+
+        return module
+
+    def reset_state(self, module: torch.nn.Module) -> torch.nn.Module:
+        self.state_manager.reset()
+        return module
+
+    @torch.compiler.disable
+    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
+        if self.state_manager._current_context is None:
+            self.state_manager.set_context("magcache")
+        state: MagCacheState = self.state_manager.get_state()
+
+        if not state.should_compute:
+            res = state.previous_residual
+            if res is None:
+                res = torch.zeros_like(args[0])
+
+            if hasattr(self._metadata, "hidden_states_argument_name"):
+                arg_name = self._metadata.hidden_states_argument_name
+            else:
+                arg_name = "hidden_states"
+            hidden_states = self._metadata._get_parameter_from_args_kwargs(arg_name, args, kwargs)
+
+            if self.is_tail:
+                self.advance_step(state)
+
+            if self._metadata.return_encoder_hidden_states_index is not None:
+                encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs(
+                    "encoder_hidden_states", args, kwargs
+                )
+
+                if self._strategy is not None:
+                    out_hidden, enc_out = self._strategy.apply_residual_tuple(hidden_states, encoder_hidden_states, res)
+                else:
+                    out_hidden = hidden_states + res
+                    enc_out = encoder_hidden_states
+
+                max_idx = max(
+                    self._metadata.return_hidden_states_index, self._metadata.return_encoder_hidden_states_index
+                )
+                ret_list = [None] * (max_idx + 1)
+                ret_list[self._metadata.return_hidden_states_index] = out_hidden
+                ret_list[self._metadata.return_encoder_hidden_states_index] = enc_out
+                return tuple(ret_list)
+
+            if self._strategy is not None:
+                output = self._strategy.apply_residual(hidden_states, res)
+            else:
+                output = hidden_states + res
+
+            return output
+
+        output = self.fn_ref.original_forward(*args, **kwargs)
+
+        if self.is_tail:
+            if isinstance(output, tuple):
+                out_hidden = output[self._metadata.return_hidden_states_index]
+            else:
+                out_hidden = output
+
+            in_hidden = state.head_block_input
+
+            if in_hidden is None:
+                return output
+
+            if self._strategy is not None:
+                residual = self._strategy.compute_residual(output, in_hidden)
+            elif out_hidden.shape == in_hidden.shape:
+                residual = out_hidden - in_hidden
+            elif out_hidden.ndim == 3 and in_hidden.ndim == 3 and out_hidden.shape[2] == in_hidden.shape[2]:
+                residual = out_hidden - in_hidden
+            else:
+                residual = out_hidden
+
+            if self.config.mag_calibrate:
+                self.perform_calibration(state, residual)
+
+            state.previous_residual = residual
+            self.advance_step(state)
+
+            self.log_residual_computed(state, residual)
+
+        return output
+
+    def log_residual_computed(
+        self,
+        state: MagCacheState,
+        residual: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+    ) -> None:
+        step = state.step_index
+        if residual is None:
+            residual_norm = 0.0
+            residual_shape = "None"
+        elif isinstance(residual, tuple):
+            residual_norm = sum(float(torch.norm(r).item()) for r in residual)
+            residual_shape = tuple(r.shape for r in residual)
+        else:
+            residual_norm = float(torch.norm(residual).item())
+            residual_shape = residual.shape
+        logger.debug(
+            f"[MagCache][TAIL] STEP={step}: RESIDUAL_COMPUTED (norm={residual_norm:.6f}, shape={residual_shape})"
+        )
+
+    def perform_calibration(
+        self,
+        state: MagCacheState,
+        current_residual: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+    ) -> None:
+        if self._strategy is not None:
+            ratio, std, cos_dis = self._strategy.compute_calibration_metrics(current_residual, state.previous_residual)
+        else:
+
+            def _get_norm(residual: torch.Tensor | tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+                if isinstance(residual, tuple):
+                    return sum(torch.linalg.norm(r.float(), dim=-1) for r in residual)
+                return torch.linalg.norm(residual.float(), dim=-1)
+
+            if state.previous_residual is None:
+                ratio, std, cos_dis = 1.0, 0.0, 0.0
+            else:
+                curr_norm = _get_norm(current_residual)
+                prev_norm = _get_norm(state.previous_residual)
+                ratio = (curr_norm / (prev_norm + 1e-8)).mean().item()
+                std = (curr_norm / (prev_norm + 1e-8)).std().item()
+                cos_dis = (
+                    (
+                        1
+                        - F.cosine_similarity(
+                            current_residual.flatten(0, -2) if current_residual.ndim > 2 else current_residual,
+                            state.previous_residual.flatten(0, -2)
+                            if state.previous_residual.ndim > 2
+                            else state.previous_residual,
+                            dim=-1,
+                            eps=1e-8,
+                        )
+                    )
+                    .mean()
+                    .item()
+                )
+
+        state.calibration_ratios.append(ratio)
+        state.norm_ratios.append(round(ratio, 5))
+        state.norm_stds.append(round(std, 5))
+        state.cos_dises.append(round(cos_dis, 5))
+
+    def advance_step(self, state: MagCacheState) -> None:
+        state.step_index += 1
+        if state.step_index >= self.config.num_inference_steps:
+            if self.config.mag_calibrate:
+                logger.info("MagCache calibration complete.")
+                logger.info(f"norm_ratios: {state.norm_ratios}")
+                logger.info(f"norm_stds: {state.norm_stds}")
+                logger.info(f"cos_dises: {state.cos_dises}")
+                logger.info("Copy these values to DiffusionCacheConfig(mag_ratios=...) for production use")
+
+            state.step_index = 0
+            state.accumulated_ratio = 1.0
+            state.accumulated_steps = 0
+            state.accumulated_err = 0.0
+            state.previous_residual = None
+            state.calibration_ratios = []
+            state.norm_ratios = []
+            state.norm_stds = []
+            state.cos_dises = []
+            state._is_first_step = True
+
+
+def apply_mag_cache_hook(
+    module: torch.nn.Module,
+    config: MagCacheConfig,
+    strategy: MagCacheStrategy | None = None,
+) -> None:
+    """Apply MagCache optimization to a transformer module.
+
+    Args:
+        module: Transformer model to optimize (e.g., FluxTransformer2DModel)
+        config: MagCacheConfig specifying caching parameters
+        strategy: Optional strategy to use. If None, will be looked up from registry.
+    """
+    HookRegistry.check_if_exists_or_initialize(module)
+
+    transformer_type = config.transformer_type
+    if strategy is None:
+        strategy = MagCacheStrategyRegistry.get_if_exists(transformer_type)
+
+    if strategy is None:
+        logger.warning(
+            f"MagCache: No strategy found for '{transformer_type}'. "
+            f"Using default behavior. Available strategies: {list(MagCacheStrategyRegistry._registry.keys())}"
+        )
+    else:
+        strategy_name = type(strategy).__name__
+        logger.info(f"MagCache: Applying {strategy_name} for '{transformer_type}'")
+        if hasattr(type(strategy), "register_blocks"):
+            type(strategy).register_blocks()
+
+    state_manager = StateManager(MagCacheState, (), {})
+    remaining_blocks = []
+
+    for name, submodule in module.named_children():
+        if not isinstance(submodule, torch.nn.ModuleList):
+            continue
+        for index, block in enumerate(submodule):
+            remaining_blocks.append((f"{name}.{index}", block))
+
+    if not remaining_blocks:
+        logger.warning("MagCache: No transformer blocks found to apply hooks.")
+        return
+
+    if len(remaining_blocks) == 1:
+        name, block = remaining_blocks[0]
+        logger.info(f"MagCache: Applying Head+Tail Hook to single block '{name}'")
+        _apply_mag_cache_head_hook(block, state_manager, config, strategy, is_tail=True)
+        return
+
+    head_block_name, head_block = remaining_blocks.pop(0)
+    tail_block_name, tail_block = remaining_blocks.pop(-1)
+
+    logger.info(f"MagCache: Applying Head Hook to {head_block_name}")
+    _apply_mag_cache_head_hook(head_block, state_manager, config, strategy)
+
+    for name, block in remaining_blocks:
+        _apply_mag_cache_block_hook(block, state_manager, config, strategy=strategy)
+
+    logger.info(f"MagCache: Applying Tail Hook to {tail_block_name}")
+    _apply_mag_cache_block_hook(tail_block, state_manager, config, is_tail=True, strategy=strategy)
+
+
+def _apply_mag_cache_head_hook(
+    block: torch.nn.Module,
+    state_manager: StateManager,
+    config: MagCacheConfig,
+    strategy: MagCacheStrategy | None = None,
+    is_tail: bool = False,
+) -> None:
+    registry = HookRegistry.check_if_exists_or_initialize(block)
+
+    if registry.get_hook(_MAG_CACHE_LEADER_BLOCK_HOOK) is not None:
+        registry.remove_hook(_MAG_CACHE_LEADER_BLOCK_HOOK)
+
+    hook = MagCacheHeadHook(state_manager, config, strategy, is_tail=is_tail)
+    registry.register_hook(_MAG_CACHE_LEADER_BLOCK_HOOK, hook)
+
+
+def _apply_mag_cache_block_hook(
+    block: torch.nn.Module,
+    state_manager: StateManager,
+    config: MagCacheConfig,
+    is_tail: bool = False,
+    strategy: MagCacheStrategy | None = None,
+) -> None:
+    registry = HookRegistry.check_if_exists_or_initialize(block)
+
+    if registry.get_hook(_MAG_CACHE_BLOCK_HOOK) is not None:
+        registry.remove_hook(_MAG_CACHE_BLOCK_HOOK)
+
+    hook = MagCacheBlockHook(state_manager, is_tail, config, strategy)
+    registry.register_hook(_MAG_CACHE_BLOCK_HOOK, hook)
diff --git a/vllm_omni/diffusion/cache/magcache/state.py b/vllm_omni/diffusion/cache/magcache/state.py
new file mode 100644
index 000000000..220f3709a
--- /dev/null
+++ b/vllm_omni/diffusion/cache/magcache/state.py
@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+MagCache state management.
+
+This module contains the MagCacheState class which manages the internal state
+for MagCache caching logic, including residuals, accumulated metrics, and step tracking.
+"""
+
+import torch
+
+
+class MagCacheState:
+    """State management for MagCache caching logic."""
+
+    def __init__(self) -> None:
+        """Initialize empty MagCache state."""
+        self.previous_residual: torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None = None
+        self.head_block_input: torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None = None
+        self.should_compute: bool = True
+        self.accumulated_ratio: float = 1.0
+        self.accumulated_err: float = 0.0
+        self.accumulated_steps: int = 0
+        self.step_index: int = 0
+        self.calibration_ratios: list[float] = []
+        self.norm_ratios: list[float] = []
+        self.norm_stds: list[float] = []
+        self.cos_dises: list[float] = []
+        self._is_first_step: bool = True
+
+    def reset(self) -> None:
+        """Reset all state variables for a new inference run."""
+        self.previous_residual = None
+        self.head_block_input = None
+        self.should_compute = True
+        self.accumulated_ratio = 1.0
+        self.accumulated_err = 0.0
+        self.accumulated_steps = 0
+        self.step_index = 0
+        self.calibration_ratios = []
+        self.norm_ratios = []
+        self.norm_stds = []
+        self.cos_dises = []
+        self._is_first_step = True
diff --git a/vllm_omni/diffusion/cache/magcache/strategy.py b/vllm_omni/diffusion/cache/magcache/strategy.py
new file mode 100644
index 000000000..7743ff2c7
--- /dev/null
+++ b/vllm_omni/diffusion/cache/magcache/strategy.py
@@ -0,0 +1,488 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+MagCache strategy definitions for different model architectures.
+
+This module provides model-specific strategies for MagCache, allowing easy
+extension to new models by implementing the MagCacheStrategy interface.
+
+Architecture:
+- MagCacheStrategy: Abstract base class defining the strategy interface
+- FluxMagCacheStrategy: Strategy for Flux (dual-stream) models
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+import torch
+from diffusers.hooks._helpers import TransformerBlockMetadata
+
+
+class MagCacheStrategy(ABC):
+    """
+    Abstract base class for MagCache strategies.
+
+    Each model architecture requires a specific strategy to handle:
+    - Residual computation (how to calculate the residual for caching)
+    - Residual application (how to apply cached residual)
+    - Model-specific magnitude ratios
+
+    Implement this class to add support for new model architectures.
+    """
+
+    @property
+    @abstractmethod
+    def mag_ratios(self) -> torch.Tensor:
+        """Return the default mag_ratios tensor for this model.
+
+        This tensor defines caching ratios for each transformer block.
+        Values should be calibrated for the specific model architecture.
+
+        Returns:
+            1D tensor of mag_ratios (one per transformer block).
+        """
+        pass
+
+    def register_block_metadata(self, block_class: type) -> TransformerBlockMetadata | None:
+        """Register model-specific transformer block metadata.
+
+        Override this method to provide custom metadata for transformer blocks
+        that have non-standard output formats (e.g., tuple returns).
+
+        Args:
+            block_class: The transformer block class to register.
+
+        Returns:
+            TransformerBlockMetadata if custom registration is needed, None otherwise.
+        """
+        return None
+
+    def compute_residual(
+        self,
+        output: torch.Tensor,
+        head_input: torch.Tensor,
+    ) -> torch.Tensor:
+        """Compute residual between block output and input.
+
+        Default implementation: output - head_input.
+        Override this method for models with non-standard output formats.
+
+        Args:
+            output: Output from transformer blocks.
+            head_input: Input to the first block.
+
+        Returns:
+            Residual tensor for caching.
+        """
+        return output - head_input
+
+    def apply_residual(
+        self,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+    ) -> torch.Tensor:
+        """Apply cached residual to hidden states.
+
+        Default implementation: add residual to hidden_states.
+        This works for most model architectures.
+
+        Args:
+            hidden_states: Current hidden states.
+            residual: Cached residual to apply.
+
+        Returns:
+            Hidden states with residual added.
+        """
+        return hidden_states + residual
+
+    def apply_residual_tuple(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        residual: tuple[torch.Tensor, torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Apply cached residual tuple to both hidden_states and encoder_hidden_states.
+
+        Default implementation: add residuals separately.
+        Override this method for models with specific residual application logic.
+
+        Args:
+            hidden_states: Current hidden states.
+            encoder_hidden_states: Current encoder hidden states.
+            residual: Tuple of (hidden_states_residual, encoder_hidden_states_residual).
+
+        Returns:
+            Tuple of (hidden_states, encoder_hidden_states) with residuals applied.
+        """
+        h_res, e_res = residual
+        return hidden_states + h_res, encoder_hidden_states + e_res
+
+    def compute_calibration_metrics(
+        self,
+        current_residual: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        previous_residual: torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None,
+    ) -> tuple[float, float, float]:
+        """Compute calibration metrics for mag_ratios generation.
+
+        Default implementation computes norm ratios and cosine dissimilarity.
+        Override this method for models with custom metric computation.
+
+        Args:
+            current_residual: Residual from the current step.
+            previous_residual: Residual from the previous step (None for first step).
+
+        Returns:
+            Tuple of (norm_ratio, norm_std, cos_dis):
+            - norm_ratio: Mean ratio of current to previous residual norms
+            - norm_std: Standard deviation of the norm ratios
+            - cos_dis: Mean cosine dissimilarity (1 - cosine_similarity)
+        """
+        import torch.nn.functional as F
+
+        if previous_residual is None:
+            return 1.0, 0.0, 0.0
+
+        curr_norm = torch.linalg.norm(current_residual.float(), dim=-1)
+        prev_norm = torch.linalg.norm(previous_residual.float(), dim=-1)
+
+        ratio = (curr_norm / (prev_norm + 1e-8)).mean().item()
+        std = (curr_norm / (prev_norm + 1e-8)).std().item()
+        cos_dis = (1 - F.cosine_similarity(current_residual, previous_residual, dim=-1, eps=1e-8)).mean().item()
+
+        return ratio, std, cos_dis
+
+    def get_calibration_metrics_names(self) -> tuple[str, str, str]:
+        """Return the names of calibration metrics for logging.
+
+        Returns:
+            Tuple of metric names in order: (norm_ratio_name, norm_std_name, cos_dis_name)
+        """
+        return ("norm_ratio", "norm_std", "cos_dis")
+
+
+class FluxMagCacheStrategy(MagCacheStrategy):
+    """
+    MagCache strategy for Flux (dual-stream) models.
+
+    Flux architecture:
+    - transformer blocks (dual-stream): image tokens and text tokens
+      processed independently with separate weights
+    - single transformer blocks (single-stream): concatenated sequence
+      (image + text tokens) shares the same group of weights
+    - Final norm_out and proj_out layers
+
+    This strategy provides:
+    - mag_ratios: Pre-computed magnitude ratios for Flux (28 steps)
+    - compute_residual: Handles tuple output format
+    - apply_residual_tuple: Handles decoder residual only
+    """
+
+    FLUX_MAG_RATIOS = torch.tensor(
+        [
+            1.0,
+            1.07313,
+            1.21035,
+            1.04432,
+            1.06818,
+            1.05547,
+            1.0183,
+            1.03405,
+            1.02574,
+            1.03042,
+            1.02739,
+            1.01955,
+            1.01585,
+            1.02439,
+            1.01154,
+            1.01377,
+            1.00994,
+            1.01444,
+            1.00839,
+            1.02269,
+            1.0007,
+            1.00714,
+            1.00484,
+            1.01381,
+            1.00426,
+            0.99764,
+            1.00778,
+            1.00233,
+        ]
+    )
+
+    @property
+    def mag_ratios(self) -> torch.Tensor:
+        """Return default mag_ratios for Flux model."""
+        return self.FLUX_MAG_RATIOS
+
+    def compute_residual(
+        self,
+        output: torch.Tensor,
+        head_input: torch.Tensor,
+    ) -> torch.Tensor:
+        """Compute residual for Flux output format (tuple or single tensor).
+
+        Flux single transformer blocks return a tuple, so we extract
+        the decoder output (index 1) before computing residual.
+        """
+        if isinstance(output, tuple):
+            decoder_output = output[1] if len(output) > 1 else output[0]
+            return decoder_output - head_input
+        else:
+            return output - head_input
+
+    def apply_residual_tuple(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        residual: tuple[torch.Tensor, torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Apply residual tuple for Flux - only add decoder residual.
+
+        Flux has separate image and text processing, so the residual
+        is only applied to the decoder (image) branch.
+        """
+        if isinstance(residual, tuple):
+            decoder_residual = residual[1]
+        else:
+            decoder_residual = residual
+
+        output = hidden_states + decoder_residual
+        enc_output = encoder_hidden_states
+
+        return output, enc_output
+
+    @staticmethod
+    def nearest_interp(src_array: torch.Tensor, target_length: int) -> torch.Tensor:
+        """Interpolate mag_ratios to target length using nearest neighbor."""
+        src_length = len(src_array)
+        if target_length == 1:
+            return src_array[-1:]
+
+        scale = (src_length - 1) / (target_length - 1)
+        grid = torch.arange(target_length, device=src_array.device, dtype=torch.float32)
+        mapped_indices = torch.round(grid * scale).long()
+
+        return src_array[mapped_indices]
+
+
+class Flux2MagCacheStrategy(FluxMagCacheStrategy):
+    """MagCache strategy for Flux2 model.
+
+    Flux2 shares the same dual-stream architecture as Flux, but may have
+    different tensor shapes in some transformer blocks, requiring special
+    handling in residual computation.
+    """
+
+    FLUX2_MAG_RATIOS = torch.tensor(
+        [
+            1.0,
+            0.96528,
+            1.11559,
+            1.0565,
+            1.00425,
+            1.0805,
+            0.98616,
+            1.09289,
+            1.03196,
+            1.06679,
+            1.03941,
+            1.05375,
+            1.03128,
+            1.05349,
+            1.01983,
+            1.05535,
+            1.0662,
+            1.05748,
+            1.00318,
+            1.05222,
+            1.04556,
+            1.0506,
+            1.05058,
+            1.05219,
+            1.02025,
+            1.05052,
+            1.04143,
+            1.0498,
+        ]
+    )
+
+    @property
+    def mag_ratios(self) -> torch.Tensor:
+        """Return default mag_ratios for Flux2 model."""
+        return self.FLUX2_MAG_RATIOS
+
+    def register_block_metadata(self, block_class: type) -> TransformerBlockMetadata | None:
+        """Register Flux2-specific block metadata based on block type.
+
+        Flux2 has two block types with different output formats:
+        - Dual-stream (Flux2TransformerBlock): returns (encoder_hidden_states, hidden_states)
+        - Single-stream (Flux2SingleTransformerBlock): returns single tensor
+        """
+        class_name = block_class.__name__
+        is_single_stream = "Single" in class_name
+
+        if is_single_stream:
+            return TransformerBlockMetadata(
+                return_hidden_states_index=0,
+                return_encoder_hidden_states_index=None,
+            )
+        else:
+            return TransformerBlockMetadata(
+                return_hidden_states_index=1,
+                return_encoder_hidden_states_index=0,
+            )
+
+    def compute_residual(
+        self,
+        output: torch.Tensor,
+        head_input: torch.Tensor,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        """Compute residual for Flux2 with dual-stream support.
+
+        For dual-stream blocks, computes residual for both encoder and decoder branches.
+        For single-stream blocks: if shapes match, computes output - input; otherwise returns output.
+        """
+        if isinstance(output, tuple):
+            enc_output, dec_output = output[0], output[1]
+
+            if isinstance(head_input, tuple):
+                enc_head, dec_head = head_input[0], head_input[1]
+            else:
+                enc_head = head_input[:, : enc_output.shape[1], ...]
+                dec_head = head_input[:, enc_output.shape[1] :, ...]
+
+            enc_residual = enc_output - enc_head
+            dec_residual = dec_output - dec_head
+
+            return (enc_residual, dec_residual)
+        else:
+            if output.shape == head_input.shape:
+                return output - head_input
+            else:
+                return output
+
+    def apply_residual_tuple(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        residual: tuple[torch.Tensor, torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Apply residual tuple for Flux2 with shape mismatch handling.
+
+        For Flux2 dual-stream blocks, handles the case where hidden_states
+        and decoder_residual may have different shapes.
+        """
+        if isinstance(residual, tuple):
+            decoder_residual = residual[1]
+        else:
+            decoder_residual = residual
+
+        if hidden_states.shape == decoder_residual.shape:
+            output = hidden_states + decoder_residual
+        else:
+            output = hidden_states
+        enc_output = encoder_hidden_states
+
+        return output, enc_output
+
+    def apply_residual(
+        self,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+    ) -> torch.Tensor:
+        """Apply residual for Flux2.
+
+        For single-stream blocks: if shapes match, adds residual; otherwise returns input.
+        For dual-stream blocks: applies decoder residual only.
+        """
+        if isinstance(residual, tuple):
+            dec_residual = residual[1]
+            if hidden_states.shape == dec_residual.shape:
+                return hidden_states + dec_residual
+            else:
+                return hidden_states
+        else:
+            if residual.shape == hidden_states.shape:
+                return hidden_states + residual
+            else:
+                return hidden_states
+
+
+class MagCacheStrategyRegistry:
+    """Registry for MagCache strategies by transformer type."""
+
+    _registry: dict[str, MagCacheStrategy] = {}
+
+    @classmethod
+    def register(cls, name: str, strategy: MagCacheStrategy) -> None:
+        """Register a strategy with explicit name.
+
+        Args:
+            name: Transformer model type identifier (e.g., "FluxTransformer2DModel")
+            strategy: MagCacheStrategy instance
+        """
+        cls._registry[name] = strategy
+
+    @classmethod
+    def get(cls, transformer_type: str) -> MagCacheStrategy:
+        """Get strategy for given transformer type."""
+        if transformer_type not in cls._registry:
+            available = list(cls._registry.keys())
+            raise ValueError(f"Unknown model type: '{transformer_type}'. Available types: {available}")
+        return cls._registry[transformer_type]
+
+    @classmethod
+    def get_if_exists(cls, transformer_type: str) -> MagCacheStrategy | None:
+        """Get strategy if exists, None otherwise."""
+        return cls._registry.get(transformer_type)
+
+
+MagCacheStrategyRegistry.register("FluxTransformer2DModel", FluxMagCacheStrategy())
+
+MagCacheStrategyRegistry.register("Flux2Transformer2DModel", Flux2MagCacheStrategy())
+
+
+def register_strategy(
+    transformer_cls_name: str,
+    strategy: MagCacheStrategy,
+) -> None:
+    """Register a MagCache strategy for a model type.
+
+    This allows extending MagCache support to new models without modifying
+    the core MagCache code.
+
+    Args:
+        transformer_cls_name: Transformer model type identifier (class name or type string)
+                               Must match pipeline.transformer.__class__.__name__
+        strategy: MagCacheStrategy instance for this model type
+
+    Example:
+        >>> class MyModelMagCacheStrategy(MagCacheStrategy):
+        ...     @property
+        ...     def mag_ratios(self):
+        ...         return torch.tensor([...])
+        >>> register_strategy("MyModelTransformer", MyModelMagCacheStrategy())
+    """
+    MagCacheStrategyRegistry.register(transformer_cls_name, strategy)
+
+
+def get_strategy(transformer_cls_name: str) -> MagCacheStrategy:
+    """Get strategy function for given transformer class.
+
+    This function looks up the strategy based on the exact transformer_cls_name string,
+    which should match the transformer type in the pipeline (i.e., pipeline.transformer.__class__.__name__).
+
+    Args:
+        transformer_cls_name: Transformer class name (e.g., "FluxTransformer2DModel")
+                              Must exactly match a registered strategy.
+
+    Returns:
+        MagCacheStrategy instance for the model
+
+    Raises:
+        ValueError: If model type not found in registry
+    """
+    return MagCacheStrategyRegistry.get(transformer_cls_name)
diff --git a/vllm_omni/diffusion/cache/prompt_embed_cache.py b/vllm_omni/diffusion/cache/prompt_embed_cache.py
new file mode 100644
index 000000000..f95197a30
--- /dev/null
+++ b/vllm_omni/diffusion/cache/prompt_embed_cache.py
@@ -0,0 +1,337 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Prompt-embedding cache for diffusion pipelines.
+
+Text encoders are the single most expensive per-request preprocessing step in
+most diffusion pipelines, and their inputs are frequently identical across
+requests (e.g. in GRPO-style rollouts the same prompt is submitted many times
+with different seeds to sample different images). Re-running the text encoder
+for each of those requests wastes compute and GPU memory.
+
+This module provides a small LRU cache plus a transparent wrapper for a
+pipeline's ``encode_prompt`` method. Because almost every diffusion pipeline in
+``vllm_omni/diffusion/models`` routes all text-encoder invocations through
+``self.encode_prompt(...)``, wrapping that single method is sufficient to
+cache results model-wide without per-pipeline edits.
+
+Design points:
+    * The wrapper is installed by :class:`DiffusionModelRunner` after the
+      pipeline has loaded, so each runner process owns its own cache.
+    * Cache keys are derived from the bound ``encode_prompt`` arguments. Only
+      inputs we can safely hash (``str`` / ``int`` / ``float`` / ``bool`` /
+      ``None`` / ``bytes`` / ``torch.device`` / ``torch.dtype`` / numpy
+      scalars / nested lists/tuples/dicts of those) participate in the key.
+      If any argument is a tensor, PIL image, or other non-trivial object, we
+      bypass the cache for that call to guarantee correctness.
+    * If a caller passes precomputed ``*_embeds`` into ``encode_prompt`` the
+      wrapper also bypasses the cache because the call is already short-circuit.
+    * Cache values are detached tensors. Downstream pipeline code typically
+      does non-inplace ops (``.repeat``, ``.view``, slicing); we do not clone
+      on hit since outputs are treated as read-only.
+"""
+
+from __future__ import annotations
+
+import functools
+import inspect
+import os
+import threading
+from collections import OrderedDict
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+# Sentinel returned by :func:`_hashable` when an input cannot be hashed
+# safely (e.g. tensor / PIL image / custom object). Any such input disables
+# caching for that specific call.
+_NOT_HASHABLE = object()
+
+# Sentinel returned by :meth:`PromptEmbedCache.get` to distinguish a cache
+# miss from a legitimate cached value of ``None``.
+_CACHE_MISS = object()
+
+# Argument names whose presence of a non-None value means the caller is
+# providing precomputed embeddings. In every pipeline we inspected, passing
+# these makes ``encode_prompt`` a cheap passthrough, so caching it has no
+# benefit and just wastes memory on large tensor keys.
+_PRECOMPUTED_EMBED_ARGS = (
+    "prompt_embeds",
+    "negative_prompt_embeds",
+    "pooled_prompt_embeds",
+    "negative_pooled_prompt_embeds",
+    "prompt_embeds_mask",
+    "negative_prompt_embeds_mask",
+)
+
+
+def _hashable(obj: Any) -> Any:
+    """Convert *obj* into a hashable representation or return ``_NOT_HASHABLE``.
+
+    Only inputs whose textual / scalar identity fully determines the
+    ``encode_prompt`` output are considered safe. Tensors and PIL images (which
+    flow through e.g. image-edit pipelines) deliberately disable caching for
+    that call rather than risk stale results. Common value-like configuration
+    objects such as ``torch.device`` and ``torch.dtype`` are normalized to
+    stable string forms so they can participate in cache keys safely.
+    """
+    if obj is None or isinstance(obj, (str, int, float, bool, bytes)):
+        return obj
+    if isinstance(obj, torch.device):
+        return ("__torch_device__", str(obj))
+    if isinstance(obj, torch.dtype):
+        return ("__torch_dtype__", str(obj))
+    if type(obj).__module__.split(".", 1)[0] == "numpy" and hasattr(obj, "item"):
+        try:
+            return _hashable(obj.item())
+        except (ValueError, TypeError):
+            return _NOT_HASHABLE
+    if isinstance(obj, (list, tuple)):
+        out = []
+        for item in obj:
+            h = _hashable(item)
+            if h is _NOT_HASHABLE:
+                return _NOT_HASHABLE
+            out.append(h)
+        return ("__seq__", tuple(out))
+    if isinstance(obj, dict):
+        items = []
+        for k in sorted(obj.keys(), key=repr):
+            v = _hashable(obj[k])
+            if v is _NOT_HASHABLE:
+                return _NOT_HASHABLE
+            items.append((repr(k), v))
+        return ("__dict__", tuple(items))
+    # torch.Tensor, PIL.Image, np.ndarray, and any other object → not safe.
+    return _NOT_HASHABLE
+
+
+def _detach_output(value: Any) -> Any:
+    """Detach tensors in the return value so the cache never holds grad state."""
+    if isinstance(value, torch.Tensor):
+        return value.detach()
+    if isinstance(value, tuple):
+        detached = tuple(_detach_output(v) for v in value)
+        # Preserve namedtuple types so attribute access keeps working.
+        if hasattr(value, "_fields"):
+            return type(value)(*detached)
+        return detached
+    if isinstance(value, list):
+        return [_detach_output(v) for v in value]
+    if isinstance(value, dict):
+        return {k: _detach_output(v) for k, v in value.items()}
+    return value
+
+
+class PromptEmbedCache:
+    """Thread-safe LRU cache for ``encode_prompt`` outputs.
+
+    The cache stores whatever ``encode_prompt`` returns (tensor, tuple of
+    tensors, ``None``, etc.). Lookup / insertion is O(1) amortized. Eviction
+    is least-recently-used.
+    """
+
+    def __init__(self, max_size: int = 32, enabled: bool = True) -> None:
+        if max_size <= 0:
+            raise ValueError(f"max_size must be positive, got {max_size}")
+        self.max_size = max_size
+        self.enabled = enabled
+        self._store: OrderedDict[Any, Any] = OrderedDict()
+        self._lock = threading.Lock()
+        self.hits = 0
+        self.misses = 0
+        self.bypassed = 0
+
+    def get(self, key: Any) -> Any:
+        """Return the cached value or ``_CACHE_MISS`` if absent.
+
+        Returning a sentinel (rather than ``None``) lets callers cache
+        legitimate ``None`` results from the wrapped function.
+        """
+        with self._lock:
+            if key not in self._store:
+                self.misses += 1
+                return _CACHE_MISS
+            self._store.move_to_end(key)
+            self.hits += 1
+            return self._store[key]
+
+    def put(self, key: Any, value: Any) -> None:
+        with self._lock:
+            self._store[key] = value
+            self._store.move_to_end(key)
+            while len(self._store) > self.max_size:
+                self._store.popitem(last=False)
+
+    def clear(self) -> None:
+        with self._lock:
+            self._store.clear()
+
+    def stats(self) -> dict[str, int]:
+        with self._lock:
+            return {
+                "size": len(self._store),
+                "max_size": self.max_size,
+                "hits": self.hits,
+                "misses": self.misses,
+                "bypassed": self.bypassed,
+            }
+
+
+def _build_key(
+    signature: inspect.Signature,
+    model_tag: str,
+    args: tuple,
+    kwargs: dict,
+) -> Any:
+    """Return a hashable cache key or ``None`` if this call must bypass cache.
+
+    We bind the raw positional / keyword arguments to the method signature so
+    that keyword and positional invocations with the same underlying values
+    collide in the cache.
+    """
+    try:
+        bound = signature.bind_partial(*args, **kwargs)
+    except TypeError:
+        return None
+    bound.apply_defaults()
+
+    # Bypass when the caller supplies precomputed embeddings.
+    for name in _PRECOMPUTED_EMBED_ARGS:
+        if name in bound.arguments and bound.arguments[name] is not None:
+            return None
+
+    items: list[tuple[str, Any]] = []
+    for name in sorted(bound.arguments.keys()):
+        h = _hashable(bound.arguments[name])
+        if h is _NOT_HASHABLE:
+            return None
+        items.append((name, h))
+    return (model_tag, tuple(items))
+
+
+def install_prompt_embed_cache(
+    pipeline: Any,
+    *,
+    max_size: int = 32,
+    enabled: bool = True,
+    model_tag: str | None = None,
+) -> PromptEmbedCache | None:
+    """Wrap ``pipeline.encode_prompt`` so results are cached by argument identity.
+
+    Idempotent: calling twice on the same pipeline is a no-op and returns the
+    existing cache. Returns ``None`` if the pipeline has no ``encode_prompt``
+    method.
+    """
+    existing = getattr(pipeline, "_prompt_embed_cache", None)
+    if existing is not None:
+        return existing
+
+    encode_fn = getattr(pipeline, "encode_prompt", None)
+    if encode_fn is None or not callable(encode_fn):
+        logger.debug(
+            "Prompt-embed cache: pipeline %s has no encode_prompt(); skipping.",
+            type(pipeline).__name__,
+        )
+        return None
+
+    try:
+        signature = inspect.signature(encode_fn)
+    except (TypeError, ValueError) as e:
+        logger.warning(
+            "Prompt-embed cache: cannot introspect encode_prompt on %s (%s); skipping.",
+            type(pipeline).__name__,
+            e,
+        )
+        return None
+
+    cache = PromptEmbedCache(max_size=max_size, enabled=enabled)
+    tag = model_tag or type(pipeline).__name__
+
+    @functools.wraps(encode_fn)
+    def cached_encode_prompt(*args: Any, **kwargs: Any) -> Any:
+        if not cache.enabled:
+            return encode_fn(*args, **kwargs)
+        key = _build_key(signature, tag, args, kwargs)
+        if key is None:
+            with cache._lock:
+                cache.bypassed += 1
+            return encode_fn(*args, **kwargs)
+        hit = cache.get(key)
+        if hit is not _CACHE_MISS:
+            return hit
+        out = encode_fn(*args, **kwargs)
+        cache.put(key, _detach_output(out))
+        return out
+
+    cached_encode_prompt.__wrapped__ = encode_fn  # type: ignore[attr-defined]
+    # Use object.__setattr__ to bypass diffusers ModelMixin attribute checks.
+    try:
+        pipeline.encode_prompt = cached_encode_prompt  # type: ignore[assignment]
+    except Exception:
+        object.__setattr__(pipeline, "encode_prompt", cached_encode_prompt)
+    object.__setattr__(pipeline, "_prompt_embed_cache", cache)
+    object.__setattr__(pipeline, "_prompt_embed_cache_original_fn", encode_fn)
+
+    logger.info(
+        "Prompt-embed cache installed on %s (max_size=%d, enabled=%s).",
+        tag,
+        max_size,
+        enabled,
+    )
+    return cache
+
+
+def uninstall_prompt_embed_cache(pipeline: Any) -> None:
+    """Restore the original ``encode_prompt`` on *pipeline* if wrapped."""
+    original = getattr(pipeline, "_prompt_embed_cache_original_fn", None)
+    if original is None:
+        return
+    try:
+        pipeline.encode_prompt = original  # type: ignore[assignment]
+    except Exception:
+        object.__setattr__(pipeline, "encode_prompt", original)
+    for attr in ("_prompt_embed_cache", "_prompt_embed_cache_original_fn"):
+        if hasattr(pipeline, attr):
+            try:
+                delattr(pipeline, attr)
+            except Exception:
+                object.__setattr__(pipeline, attr, None)
+
+
+def resolve_prompt_embed_cache_config(
+    enable: bool | None = None,
+    max_size: int | None = None,
+) -> tuple[bool, int]:
+    """Combine explicit args with env-var overrides.
+
+    Environment variables (useful for quick enablement in GRPO jobs without
+    touching config files):
+
+        ``OMNI_DIFFUSION_PROMPT_EMBED_CACHE`` (``1``/``0``/``true``/``false``)
+        ``OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE`` (positive int)
+    """
+    env_enable = os.environ.get("OMNI_DIFFUSION_PROMPT_EMBED_CACHE")
+    if env_enable is not None:
+        parsed = env_enable.strip().lower()
+        if parsed in ("1", "true", "yes", "on"):
+            enable = True
+        elif parsed in ("0", "false", "no", "off"):
+            enable = False
+
+    env_size = os.environ.get("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE")
+    if env_size is not None:
+        try:
+            env_size_int = int(env_size)
+            if env_size_int > 0:
+                max_size = env_size_int
+        except ValueError:
+            logger.warning(
+                "Ignoring non-integer OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE=%r.",
+                env_size,
+            )
+
+    return bool(enable), int(max_size or 32)
diff --git a/vllm_omni/diffusion/cache/selector.py b/vllm_omni/diffusion/cache/selector.py
index 7c09bf664..e6c2e70e3 100644
--- a/vllm_omni/diffusion/cache/selector.py
+++ b/vllm_omni/diffusion/cache/selector.py
@@ -2,6 +2,7 @@ from typing import Any
 
 from vllm_omni.diffusion.cache.base import CacheBackend
 from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTBackend
+from vllm_omni.diffusion.cache.magcache.backend import MagCacheBackend
 from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend
 from vllm_omni.diffusion.data import DiffusionCacheConfig
 
@@ -12,14 +13,15 @@ def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBack
     This is a selector function that routes to the appropriate backend implementation.
     - cache_dit: Uses CacheDiTBackend with enable()/refresh() interface
     - tea_cache: Uses TeaCacheBackend with enable()/refresh() interface
+    - mag_cache: Uses MagCacheBackend with enable()/refresh() interface
 
     Args:
-        cache_backend: Cache backend name ("cache_dit", "tea_cache", or None).
+        cache_backend: Cache backend name ("cache_dit", "tea_cache", "mag_cache", or None).
         cache_config: Cache configuration (dict or DiffusionCacheConfig instance).
 
     Returns:
-        Cache backend instance (CacheDiTBackend or TeaCacheBackend) if cache_backend is set,
-        None otherwise.
+        Cache backend instance (CacheDiTBackend, TeaCacheBackend, or MagCacheBackend)
+        if cache_backend is set, None otherwise.
 
     Raises:
         ValueError: If cache_backend is unsupported.
@@ -34,5 +36,9 @@ def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBack
         return CacheDiTBackend(cache_config)
     elif cache_backend == "tea_cache":
         return TeaCacheBackend(cache_config)
+    elif cache_backend == "mag_cache":
+        return MagCacheBackend(cache_config)
     else:
-        raise ValueError(f"Unsupported cache backend: {cache_backend}. Supported: 'cache_dit', 'tea_cache'")
+        raise ValueError(
+            f"Unsupported cache backend: {cache_backend}. Supported: 'cache_dit', 'tea_cache', 'mag_cache'"
+        )
diff --git a/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py b/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py
index 37ea454de..f808b301b 100644
--- a/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py
+++ b/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py
@@ -77,7 +77,7 @@ class DefaultAdapter:
                 os.path.join("transformer", "config.json"),
                 od_config.model,
             )
-            od_config.tf_model_config = TransformerConfig.from_dict(tf_config_dict)
+            od_config.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict))
 
         loader = DiffusersPipelineLoader(LoadConfig(), od_config=od_config)
         # load_model will handle dtypes / device placement, put in .eval() mode
@@ -202,6 +202,7 @@ class TeaCacheCoefficientEstimator:
         self.hook.start_collection()
         req = OmniDiffusionRequest(
             prompts=[prompt],
+            request_id="teacache-coefficient-estimator",
             sampling_params=OmniDiffusionSamplingParams(
                 num_inference_steps=generate_kwargs.get("num_inference_steps", 20),
                 seed=generate_kwargs.get("seed", 42),
diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py
index 5b7ab14e5..1f0c38aee 100644
--- a/vllm_omni/diffusion/cache/teacache/extractors.py
+++ b/vllm_omni/diffusion/cache/teacache/extractors.py
@@ -277,8 +277,10 @@ def extract_qwen_context(
 
     def postprocess(h):
         """Apply Qwen-specific output postprocessing."""
-        t = temb.chunk(2, dim=0)[0] if module.zero_cond_t else temb
-        h = module.norm_out(h, t)
+        if getattr(module, "zero_cond_t", False):
+            h = module.norm_out(h, temb.chunk(2, dim=0)[0])
+        else:
+            h = module.norm_out(h, temb)
         output = module.proj_out(h)
         if not return_dict:
             return (output,)
@@ -305,12 +307,9 @@ def extract_bagel_context(
     packed_vae_position_ids: torch.LongTensor,
     packed_text_ids: torch.LongTensor,
     packed_text_indexes: torch.LongTensor,
-    packed_indexes: torch.LongTensor,
     packed_position_ids: torch.LongTensor,
     packed_seqlens: torch.IntTensor,
-    key_values_lens: torch.IntTensor,
     past_key_values: Any,
-    packed_key_value_indexes: torch.LongTensor,
     **kwargs: Any,
 ) -> CacheContext:
     """
@@ -324,12 +323,9 @@ def extract_bagel_context(
         packed_vae_position_ids: Position IDs for VAE tokens
         packed_text_ids: Text token IDs
         packed_text_indexes: Indexes for text tokens in packed sequence
-        packed_indexes: Global indexes
         packed_position_ids: Global position IDs
         packed_seqlens: Sequence lengths
-        key_values_lens: KV cache lengths
         past_key_values: KV cache
-        packed_key_value_indexes: KV cache indexes
         **kwargs: Additional keyword arguments
 
     Returns:
@@ -373,10 +369,7 @@ def extract_bagel_context(
             packed_query_sequence=packed_sequence,
             query_lens=packed_seqlens,
             packed_query_position_ids=packed_position_ids,
-            packed_query_indexes=packed_indexes,
             past_key_values=past_key_values,
-            key_values_lens=key_values_lens,
-            packed_key_value_indexes=packed_key_value_indexes,
             update_past_key_values=False,
             is_causal=False,
             **extra_inputs,
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index 3ba8b24f2..8e9bdbf32 100644
--- a/vllm_omni/diffusion/data.py
+++ b/vllm_omni/diffusion/data.py
@@ -11,7 +11,7 @@ from typing import TYPE_CHECKING, Any
 import diffusers
 import torch
 from PIL import Image
-from pydantic import model_validator
+from pydantic import Field, model_validator
 from typing_extensions import Self
 from vllm.config.utils import config
 from vllm.logger import init_logger
@@ -281,7 +281,7 @@ class TransformerConfig:
 @dataclass
 class DiffusionCacheConfig:
     """
-    Configuration for cache adapters (TeaCache, cache-dit, etc.).
+    Configuration for cache adapters (TeaCache, cache-dit, MagCache, etc.).
 
     This dataclass provides a unified interface for cache configuration parameters.
     It can be initialized from a dictionary and accessed via attributes.
@@ -291,6 +291,8 @@ class DiffusionCacheConfig:
         - cache-dit: Fn_compute_blocks, Bn_compute_blocks, max_warmup_steps,
                     residual_diff_threshold, enable_taylorseer, taylorseer_order,
                     scm_steps_mask_policy, scm_steps_policy
+        - MagCache: mag_threshold, mag_max_skip_steps, mag_retention_ratio,
+                    mag_ratios, mag_calibrate
 
     Example:
         >>> # From dict (user-facing API) - partial config uses defaults for missing keys
@@ -300,7 +302,7 @@ class DiffusionCacheConfig:
         >>> print(config.Fn_compute_blocks)  # 8 (default)
         >>> # Empty dict uses all defaults
         >>> default_config = DiffusionCacheConfig.from_dict({})
-        >>> print(default_config.rel_l1_thresh)  # 0.2 (default)
+        >>> print(config.rel_l1_thresh)  # 0.2 (default)
     """
 
     # TeaCache parameters [tea_cache only]
@@ -308,6 +310,18 @@ class DiffusionCacheConfig:
     rel_l1_thresh: float = 0.2
     coefficients: list[float] | None = None  # Uses model-specific defaults if None
 
+    # MagCache parameters [mag_cache only]
+    # Default: 0.24 threshold for accumulated magnitude error
+    mag_threshold: float = 0.24
+    # Default: 5 maximum consecutive skip steps (K)
+    mag_max_skip_steps: int = 5
+    # Default: 0.1 fraction of initial steps where skipping is disabled (stability)
+    mag_retention_ratio: float = 0.1
+    # Default: None magnitude ratios (model-specific, required for inference)
+    mag_ratios: list[float] | None = None
+    # Default: False calibration mode (computes mag_ratios on first run)
+    mag_calibrate: bool = False
+
     # cache-dit parameters [cache-dit only]
     # Default: 1 forward compute block (optimized for single-transformer models)
     # Use 1 as default instead of cache-dit's 8, optimized for single-transformer models
@@ -432,6 +446,15 @@ class OmniDiffusionConfig:
     cache_config: DiffusionCacheConfig | dict[str, Any] = field(default_factory=dict)
     enable_cache_dit_summary: bool = False
 
+    # Prompt-embedding cache. When enabled, ``DiffusionModelRunner`` wraps the
+    # pipeline's ``encode_prompt`` so repeated calls with identical prompt
+    # arguments (e.g. GRPO rollouts that sample the same prompt many times
+    # with different seeds) reuse the text-encoder output instead of re-running
+    # it. Safe against inputs that cannot be hashed (tensors, PIL images):
+    # those calls transparently bypass the cache.
+    enable_prompt_embed_cache: bool = False
+    prompt_embed_cache_size: int = 32
+
     # Distributed executor backend
     distributed_executor_backend: str = "mp"
     nccl_port: int | None = None
@@ -593,6 +616,9 @@ class OmniDiffusionConfig:
     # Maximum number of sequences to generate in a batch
     max_num_seqs: int = 1
 
+    # Supplementary model specific parameters
+    extras: dict[str, Any] = Field(default_factory=dict)
+
     @property
     def is_moe(self) -> bool:
         num_experts = self.tf_model_config.get("num_experts", None)
@@ -606,6 +632,25 @@ class OmniDiffusionConfig:
 
         return False
 
+    def _resolve_master_port(self) -> int:
+        """Resolve torch.distributed master port without unnecessary random jitter.
+
+        Precedence:
+        1. ``MASTER_PORT`` environment variable (set by orchestrators for multi-replica launch).
+        2. Explicit ``master_port`` passed at construction time.
+        3. An OS-assigned ephemeral port when neither is provided.
+        """
+        from vllm.utils.network_utils import get_open_port
+
+        from vllm_omni.diffusion import envs
+
+        env_port = envs.MASTER_PORT
+        if env_port is not None:
+            return self.settle_port(env_port, port_inc=37)
+        if self.master_port is not None:
+            return self.settle_port(self.master_port, port_inc=37)
+        return self.settle_port(get_open_port(), port_inc=37)
+
     def settle_port(self, port: int, port_inc: int = 42, max_attempts: int = 100) -> int:
         """
         Find an available port with retry logic.
@@ -642,9 +687,7 @@ class OmniDiffusionConfig:
         )
 
     def __post_init__(self):
-        # TODO: remove hard code
-        initial_master_port = (self.master_port or 30005) + random.randint(0, 100)
-        self.master_port = self.settle_port(initial_master_port, 37)
+        self.master_port = self._resolve_master_port()
 
         if isinstance(self.profiler_config, dict):
             from vllm.config import ProfilerConfig
@@ -745,8 +788,11 @@ class OmniDiffusionConfig:
             return
 
         is_checkpoint_fp8 = bool(getattr(tf_config.quant_config, "is_checkpoint_fp8_serialized", False))
-        should_use_checkpoint_config = self.quantization_config is None or (
-            is_checkpoint_fp8 and self._is_generic_fp8_quant_config(self.quantization_config)
+        is_checkpoint_nvfp4 = bool(getattr(tf_config.quant_config, "is_checkpoint_nvfp4_serialized", False))
+        should_use_checkpoint_config = (
+            self.quantization_config is None
+            or (is_checkpoint_fp8 and self._is_generic_fp8_quant_config(self.quantization_config))
+            or (is_checkpoint_nvfp4 and self._is_generic_nvfp4_quant_config(self.quantization_config))
         )
         if should_use_checkpoint_config:
             self.quantization_config = tf_config.quant_config
@@ -766,6 +812,17 @@ class OmniDiffusionConfig:
             return quant_config.get_name() == "fp8"
         return False
 
+    @staticmethod
+    def _is_generic_nvfp4_quant_config(quant_config: object) -> bool:
+        if isinstance(quant_config, str):
+            return quant_config.lower() in {"fp4", "nvfp4", "modelopt_fp4"}
+        if isinstance(quant_config, Mapping):
+            method = quant_config.get("method", quant_config.get("quant_method"))
+            return isinstance(method, str) and method.lower() in {"fp4", "nvfp4", "modelopt_fp4"}
+        if hasattr(quant_config, "get_name"):
+            return quant_config.get_name() == "modelopt_fp4"
+        return False
+
     def set_tf_model_config(self, tf_config: "TransformerConfig") -> None:
         """Assign `tf_model_config` and propagate quantization if detected.
 
@@ -894,6 +951,17 @@ class OmniDiffusionConfig:
                     self.model_class_name = "SenseNovaU1Pipeline"
                     self.tf_model_config = TransformerConfig()
                     self.update_multimodal_support()
+                elif "BailingMM2NativeForConditionalGeneration" in architectures or model_type in (
+                    "bailingmm_moe_v2_lite",
+                    "ming_flash_omni",
+                    "ming_flash_omni_thinker",
+                ):
+                    # Ming-flash-omni-2.0 — imagegen stage uses the custom
+                    # ``MingImagePipeline`` (ZImage DiT + Qwen2 connector). See
+                    # vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py.
+                    self.model_class_name = "MingImagePipeline"
+                    self.tf_model_config = TransformerConfig()
+                    self.update_multimodal_support()
                 elif model_type == "nextstep":
                     if self.model_class_name is None:
                         self.model_class_name = "NextStep11Pipeline"
@@ -999,6 +1067,28 @@ class DiffusionOutput:
     # memory usage info
     peak_memory_mb: float = 0.0
 
+    # When True, move all tensor fields (including tensors inside
+    # ``custom_output``) to CPU at construction time. Useful when the output
+    # is shipped across process boundaries (e.g. step-execution mode) and the
+    # receiving side must not initialise a stray CUDA context.
+    to_cpu: bool = False
+
+    def __post_init__(self) -> None:
+        if not self.to_cpu:
+            return
+
+        def _maybe_to_cpu(value: Any) -> Any:
+            if isinstance(value, torch.Tensor):
+                return value.detach().cpu()
+            return value
+
+        self.output = _maybe_to_cpu(self.output)
+        self.trajectory_timesteps = _maybe_to_cpu(self.trajectory_timesteps)
+        self.trajectory_latents = _maybe_to_cpu(self.trajectory_latents)
+        self.trajectory_log_probs = _maybe_to_cpu(self.trajectory_log_probs)
+        if self.custom_output:
+            self.custom_output = {k: _maybe_to_cpu(v) for k, v in self.custom_output.items()}
+
 
 class DiffusionRequestAbortedError(RuntimeError):
     """Raised when a diffusion request ends via user-visible abort."""
diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index c13bd3c0c..b4bfecb77 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/vllm_omni/diffusion/diffusion_engine.py
@@ -30,10 +30,10 @@ from vllm_omni.diffusion.registry import (
     get_diffusion_post_process_func,
     get_diffusion_pre_process_func,
 )
-from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
 from vllm_omni.diffusion.sched import RequestScheduler, SchedulerInterface, StepScheduler
 from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus
-from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
+from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput, RunnerOutput
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -83,6 +83,13 @@ def supports_audio_output(model_class_name: str) -> bool:
     return bool(getattr(model_cls, "support_audio_output", False))
 
 
+def get_dummy_run_num_frames(model_class_name: str, supports_audio_input: bool) -> int:
+    model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
+    if model_cls is not None and hasattr(model_cls, "dummy_run_num_frames"):
+        return int(getattr(model_cls, "dummy_run_num_frames"))
+    return 2 if supports_audio_input or supports_audio_output(model_class_name) else 1
+
+
 def get_extra_body_params(model_class_name: str) -> frozenset[str]:
     """Return the set of extra_body keys accepted by a pipeline.
 
@@ -217,7 +224,7 @@ class DiffusionEngine:
             logger.warning("Output is None, returning empty OmniRequestOutput")
             return [
                 OmniRequestOutput.from_diffusion(
-                    request_id=request.request_ids[i] if i < len(request.request_ids) else "",
+                    request_id=request.request_id,
                     images=[],
                     prompt=prompt,
                     metrics={},
@@ -304,7 +311,7 @@ class DiffusionEngine:
         if len(request.prompts) == 1:
             # Single request: return single OmniRequestOutput
             prompt = request.prompts[0]
-            request_id = request.request_ids[0] if request.request_ids else ""
+            request_id = request.request_id
 
             if is_text_output:
                 return [
@@ -368,10 +375,9 @@ class DiffusionEngine:
             # Split images based on num_outputs_per_prompt for each request
             results = []
             output_idx = 0
+            request_id = request.request_id
 
             for i, prompt in enumerate(request.prompts):
-                request_id = request.request_ids[i] if i < len(request.request_ids) else ""
-
                 # Get images for this request
                 num_outputs = request.sampling_params.num_outputs_per_prompt
                 start_idx = output_idx
@@ -467,17 +473,17 @@ class DiffusionEngine:
                 runner_output = self.execute_fn(sched_output)
             except Exception as exc:
                 logger.error(
-                    "Execution failed for diffusion requests %s", sched_output.scheduled_req_ids, exc_info=True
+                    "Execution failed for diffusion requests %s", sched_output.scheduled_request_ids, exc_info=True
                 )
                 runner_output = BatchRunnerOutput.from_list(
                     [
                         RunnerOutput(
-                            req_id=req_id,
+                            request_id=request_id,
                             step_index=None,
                             finished=True,
                             result=DiffusionOutput(error=str(exc)),
                         )
-                        for req_id in sched_output.scheduled_req_ids
+                        for request_id in sched_output.scheduled_request_ids
                     ]
                 )
 
@@ -544,7 +550,7 @@ class DiffusionEngine:
     def _handle_finished_requests(
         self,
         finished_ids: set[str],
-        runner_output: RunnerOutput | None = None,
+        runner_output: BaseRunnerOutput | None = None,
         missing_result_error: str = "Diffusion execution finished without a final output",
     ):
         for rid in finished_ids:
@@ -553,7 +559,7 @@ class DiffusionEngine:
             if fut is None:
                 continue
             if runner_output is not None:
-                _output = runner_output.get_req_output(rid)
+                _output = runner_output.get_request_output(rid)
             else:
                 _output = None
             out = self._finalize_finished_request(rid, _output, missing_result_error)
@@ -579,11 +585,11 @@ class DiffusionEngine:
             if self._closed:
                 raise RuntimeError("DiffusionEngine is closed.")
             fut = self.main_loop.create_future()
-            sched_req_id = self.scheduler.add_request(request)
-            self._out_queue[sched_req_id] = fut
+            request_id = self.scheduler.add_request(request)
+            self._out_queue[request_id] = fut
             self._cv.notify_all()
 
-        return sched_req_id
+        return request_id
 
     async def get_result(self, request_id: str) -> DiffusionOutput:
         fut = self._out_queue.get(request_id)
@@ -599,37 +605,37 @@ class DiffusionEngine:
     async def async_add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
         # No lock needed: add_request is already protected by self._cv, and
         # all executor calls are serialized inside the busy loop.
-        sched_req_id = self.add_request(request)
-        return await self.get_result(sched_req_id)
+        request_id = self.add_request(request)
+        return await self.get_result(request_id)
 
     def add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
         with self._rpc_lock:
             if self._closed:
                 raise RuntimeError("DiffusionEngine is closed.")
-            target_sched_req_id = self.scheduler.add_request(request)
+            target_request_id = self.scheduler.add_request(request)
 
             # keep scheduling and executing until the target request is finished
             while True:
                 self._process_aborts_queue()
                 sched_output = self.scheduler.schedule()
                 if sched_output.is_empty:
-                    if target_sched_req_id in sched_output.finished_req_ids:
-                        return self._finalize_finished_request(target_sched_req_id)
+                    if target_request_id in sched_output.finished_req_ids:
+                        return self._finalize_finished_request(target_request_id)
                     if not self.scheduler.has_requests():
                         raise RuntimeError("Diffusion scheduler has no runnable requests.")
                     continue
 
                 # NOTE: add_req_and_wait_for_response() is synchronous, will be only called
                 # within _dummy_run, only one request will be scheduled
-                sched_req_id = sched_output.scheduled_req_ids[0]
+                request_id = sched_output.scheduled_request_ids[0]
                 try:
                     runner_output = self.execute_fn(sched_output)
                 except EngineDeadError:
                     raise
                 except Exception as exc:
-                    logger.error("Execution failed for diffusion request %s", sched_req_id, exc_info=True)
+                    logger.error("Execution failed for diffusion request %s", request_id, exc_info=True)
                     runner_output = RunnerOutput(
-                        req_id=sched_req_id,
+                        request_id=request_id,
                         step_index=None,
                         finished=True,
                         result=DiffusionOutput(error=str(exc)),
@@ -642,10 +648,10 @@ class DiffusionEngine:
                 # sync func should receive one result
                 if not isinstance(runner_output, RunnerOutput) and not len(runner_output) == 1:
                     raise ValueError("Sync func should receive one result at one time")
-                if target_sched_req_id in finished_req_ids:
-                    runner_output = runner_output.get_req_output(target_sched_req_id)
+                if target_request_id in finished_req_ids:
+                    runner_output = runner_output.get_request_output(target_request_id)
                     return self._finalize_finished_request(
-                        target_sched_req_id,
+                        target_request_id,
                         runner_output=runner_output,
                         missing_result_error="Diffusion execution finished without a final output.",
                     )
@@ -691,13 +697,15 @@ class DiffusionEngine:
             dummy_audio = np.random.randn(audio_sr * 2).astype(np.float32)
             prompt.setdefault("multi_modal_data", {})["audio"] = dummy_audio
 
+        num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input)
         req = OmniDiffusionRequest(
             prompts=[prompt],
-            request_ids=["dummy_req_id"],
+            request_id=DUMMY_DIFFUSION_REQUEST_ID,
             sampling_params=OmniDiffusionSamplingParams(
                 height=height,
                 width=width,
                 num_inference_steps=num_inference_steps,
+                num_frames=num_frames,
                 # Keep warmup path minimal and robust across text encoders.
                 # Some models may fail when warmup implicitly triggers
                 # classifier-free guidance with an empty negative prompt.
@@ -899,34 +907,27 @@ class DiffusionEngine:
     def _abort_requests(self, request_ids: str | Iterable[str]) -> None:
         request_ids = [request_ids] if isinstance(request_ids, str) else list(request_ids)
 
-        sched_req_ids: list[str] = []
         for request_id in dict.fromkeys(request_ids):
-            sched_req_id = self.scheduler.get_sched_req_id(request_id)
-            if sched_req_id is not None:
-                sched_req_ids.append(sched_req_id)
-
-        for sched_req_id in dict.fromkeys(sched_req_ids):
-            if self.scheduler.get_request_state(sched_req_id) is not None:
-                self.scheduler.finish_requests(sched_req_id, DiffusionRequestStatus.FINISHED_ABORTED)
+            if self.scheduler.get_request_state(request_id) is not None:
+                self.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_ABORTED)
 
     def _finalize_finished_request(
         self,
-        sched_req_id: str,
+        request_id: str,
         runner_output: RunnerOutput | None = None,
         missing_result_error: str = "Diffusion scheduler finished target request without execution output.",
     ) -> DiffusionOutput:
-        state = self.scheduler.get_request_state(sched_req_id)
-        popped_state = self.scheduler.pop_request_state(sched_req_id)
+        state = self.scheduler.get_request_state(request_id)
+        popped_state = self.scheduler.pop_request_state(request_id)
         state = state or popped_state
 
         if state is None:
-            raise RuntimeError(f"Diffusion scheduler lost state for request {sched_req_id}.")
+            raise RuntimeError(f"Diffusion scheduler lost state for request {request_id}.")
 
         if state.status == DiffusionRequestStatus.FINISHED_ABORTED:
-            request_id = state.req.request_ids[0] if state.req.request_ids else sched_req_id
             return DiffusionOutput(
                 aborted=True,
-                abort_message=f"Request {request_id} aborted.",
+                abort_message=f"Request {state.req.request_id} aborted.",
             )
 
         if runner_output is not None and runner_output.result is not None:
diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_hunyuan_video_15.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_hunyuan_video_15.py
new file mode 100644
index 000000000..5011f2a45
--- /dev/null
+++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_hunyuan_video_15.py
@@ -0,0 +1,150 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+import torch
+from diffusers import AutoencoderKLHunyuanVideo15
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+    DistributedOperator,
+    DistributedVaeMixin,
+    GridSpec,
+    TileTask,
+)
+
+logger = init_logger(__name__)
+
+
+class DistributedAutoencoderKLHunyuanVideo15(AutoencoderKLHunyuanVideo15, DistributedVaeMixin):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any):
+        model = super().from_pretrained(*args, **kwargs)
+        model.init_distributed()
+        return model
+
+    def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+        _, _, _, height, width = z.shape
+        overlap_height = int(self.tile_latent_min_height * (1 - self.tile_overlap_factor))
+        overlap_width = int(self.tile_latent_min_width * (1 - self.tile_overlap_factor))
+        blend_height = int(self.tile_sample_min_height * self.tile_overlap_factor)
+        blend_width = int(self.tile_sample_min_width * self.tile_overlap_factor)
+        row_limit_height = self.tile_sample_min_height - blend_height
+        row_limit_width = self.tile_sample_min_width - blend_width
+
+        tiletask_list = []
+        for i in range(0, height, overlap_height):
+            for j in range(0, width, overlap_width):
+                tile = z[:, :, :, i : i + self.tile_latent_min_height, j : j + self.tile_latent_min_width]
+                tiletask_list.append(
+                    TileTask(
+                        len(tiletask_list),
+                        (i // overlap_height, j // overlap_width),
+                        tile,
+                        workload=tile.shape[3] * tile.shape[4],
+                    )
+                )
+
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+            tile_spec={
+                "blend_height": blend_height,
+                "blend_width": blend_width,
+                "row_limit_height": row_limit_height,
+                "row_limit_width": row_limit_width,
+            },
+            output_dtype=self.dtype,
+        )
+        return tiletask_list, grid_spec
+
+    def tile_exec(self, task: TileTask) -> torch.Tensor:
+        return self.decoder(task.tensor)
+
+    def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
+        grid_h, grid_w = grid_spec.grid_shape
+        result_rows = []
+        for i in range(grid_h):
+            result_row = []
+            for j in range(grid_w):
+                tile = coord_tensor_map[(i, j)]
+                if i > 0:
+                    tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+                if j > 0:
+                    tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+                result_row.append(
+                    tile[
+                        :,
+                        :,
+                        :,
+                        : grid_spec.tile_spec["row_limit_height"],
+                        : grid_spec.tile_spec["row_limit_width"],
+                    ]
+                )
+            result_rows.append(torch.cat(result_row, dim=-1))
+        return torch.cat(result_rows, dim=-2)
+
+    def encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+        _, _, _, height, width = x.shape
+        overlap_height = int(self.tile_sample_min_height * (1 - self.tile_overlap_factor))
+        overlap_width = int(self.tile_sample_min_width * (1 - self.tile_overlap_factor))
+        blend_height = int(self.tile_latent_min_height * self.tile_overlap_factor)
+        blend_width = int(self.tile_latent_min_width * self.tile_overlap_factor)
+        row_limit_height = self.tile_latent_min_height - blend_height
+        row_limit_width = self.tile_latent_min_width - blend_width
+
+        tiletask_list = []
+        for i in range(0, height, overlap_height):
+            for j in range(0, width, overlap_width):
+                tile = x[:, :, :, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width]
+                tiletask_list.append(
+                    TileTask(
+                        len(tiletask_list),
+                        (i // overlap_height, j // overlap_width),
+                        tile,
+                        workload=tile.shape[3] * tile.shape[4],
+                    )
+                )
+
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+            tile_spec={
+                "blend_height": blend_height,
+                "blend_width": blend_width,
+                "row_limit_height": row_limit_height,
+                "row_limit_width": row_limit_width,
+            },
+            output_dtype=self.dtype,
+        )
+        return tiletask_list, grid_spec
+
+    def encode_tile_exec(self, task: TileTask) -> torch.Tensor:
+        return self.encoder(task.tensor)
+
+    def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+        if not self.is_distributed_enabled():
+            return super().tiled_encode(x)
+
+        logger.debug("Encode running with distributed executor")
+        return self.distributed_executor.execute(
+            x,
+            DistributedOperator(
+                split=self.encode_tile_split,
+                exec=self.encode_tile_exec,
+                merge=self.tile_merge,
+            ),
+            broadcast_result=True,
+        )
+
+    def tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
+        if not self.is_distributed_enabled():
+            return super().tiled_decode(z)
+
+        logger.debug("Decode running with distributed executor")
+        return self.distributed_executor.execute(
+            z,
+            DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
+            broadcast_result=True,
+        )
diff --git a/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py
index 209f65625..8fb4ebeb3 100644
--- a/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py
+++ b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py
@@ -125,7 +125,7 @@ class DistributedVaeExecutor:
 
         # 2. local decode
         assigned = self._balance_tasks(tiletask_list, pp_size)
-        local_tasks = assigned[self.rank] if pp_size <= self.world_size else []
+        local_tasks = assigned[self.rank] if self.rank < pp_size else []
         local_results = [(t.tile_id, operator.exec(t)) for t in local_tasks]
 
         # 3. compute shape per rank
diff --git a/vllm_omni/diffusion/distributed/utils.py b/vllm_omni/diffusion/distributed/utils.py
index 92d00ef4b..851328e62 100644
--- a/vllm_omni/diffusion/distributed/utils.py
+++ b/vllm_omni/diffusion/distributed/utils.py
@@ -10,5 +10,7 @@ from vllm_omni.platforms import current_omni_platform
 
 def get_local_device() -> torch.device:
     """Return the torch device for the current rank based on detected device type."""
+    if current_omni_platform.is_initialized():
+        return current_omni_platform.get_torch_device(current_omni_platform.current_device())
     local_rank = int(os.environ.get("LOCAL_RANK", 0))
     return current_omni_platform.get_torch_device(local_rank)
diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py
index 6727122d4..1bd848226 100644
--- a/vllm_omni/diffusion/executor/multiproc_executor.py
+++ b/vllm_omni/diffusion/executor/multiproc_executor.py
@@ -60,11 +60,11 @@ class BackgroundResources:
             for proc in self.processes:
                 if not proc.is_alive():
                     continue
-                proc.join(30)
+                proc.join(5)
                 if proc.is_alive():
                     logger.warning("Terminating diffusion worker %s after timeout", proc.name)
                     proc.terminate()
-                    proc.join(30)
+                    proc.join(5)
 
 
 class MultiprocDiffusionExecutor(DiffusionExecutor):
@@ -311,7 +311,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             raise RuntimeError(f"Unexpected response type for execute_request: {type(result)!r}")
 
         return RunnerOutput(
-            req_id=new_req.sched_req_id,
+            request_id=new_req.request_id,
             step_index=None,
             finished=True,
             result=result,
@@ -334,9 +334,9 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         # TODO: Remove this fallback; DiffusionOutput cannot faithfully represent
         # failed multi-request step batches.
         if isinstance(result, DiffusionOutput):
-            req_id = scheduler_output.scheduled_req_ids[0] if scheduler_output.scheduled_req_ids else ""
+            request_id = scheduler_output.scheduled_request_ids[0] if scheduler_output.scheduled_request_ids else ""
             return RunnerOutput(
-                req_id=req_id,
+                request_id=request_id,
                 step_index=None,
                 finished=True,
                 result=result,
diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py
index 35c93c799..f2b6927d2 100644
--- a/vllm_omni/diffusion/forward_context.py
+++ b/vllm_omni/diffusion/forward_context.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from contextlib import contextmanager
 from dataclasses import dataclass
+from typing import TYPE_CHECKING
 
 import vllm.ir
 from vllm.config import VllmConfig
@@ -11,6 +12,9 @@ from vllm_omni.diffusion.attention.backends.abstract import (
 )
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 
+if TYPE_CHECKING:
+    import torch
+
 
 @dataclass
 class ForwardContext:
@@ -23,6 +27,8 @@ class ForwardContext:
     attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None
     split_text_embed_in_sp: bool = False
     denoise_step_idx: int | None = None
+    # Per-request reference latent for img2img DiT models (e.g. Ming)
+    ref_latent: torch.Tensor | None = None
     # whether to split the text embed in sequence parallel, if True, the text embed will be split in sequence parallel
 
     # Sequence Parallel padding support
@@ -171,3 +177,14 @@ def set_forward_context_denoise_step_idx(step_idx: int | None) -> None:
     """Set the current diffusion denoise step on the active ForwardContext."""
     if _forward_context is not None:
         _forward_context.denoise_step_idx = step_idx
+
+
+def set_forward_context_ref_latent(ref_latent: torch.Tensor | None) -> None:
+    """Set the per-request reference latent on the active ForwardContext.
+
+    Used by img2img-capable DiT models (e.g. Ming-flash-omni-2.0) so the
+    transformer can read the reference latent from request scope instead of
+    module instance state.
+    """
+    if _forward_context is not None:
+        _forward_context.ref_latent = ref_latent
diff --git a/vllm_omni/diffusion/hooks/base.py b/vllm_omni/diffusion/hooks/base.py
index 517c66158..6863a4594 100644
--- a/vllm_omni/diffusion/hooks/base.py
+++ b/vllm_omni/diffusion/hooks/base.py
@@ -15,6 +15,9 @@ from dataclasses import dataclass
 from typing import Any
 
 import torch.nn as nn
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
 
 
 class BaseState:
@@ -27,17 +30,24 @@ class BaseState:
 class StateManager:
     """Manage per-context hook state instances."""
 
-    def __init__(self, state_cls: Callable[[], BaseState]):
+    def __init__(self, state_cls: Callable[[], BaseState], init_args: tuple = (), init_kwargs: dict | None = None):
         self._state_cls = state_cls
+        self._init_args = init_args
+        self._init_kwargs = init_kwargs or {}
         self._states: dict[str, BaseState] = {}
         self._context: str = "default"
 
+    @property
+    def _current_context(self) -> str | None:
+        """Alias for _context for compatibility with diffusers hook code."""
+        return self._context if self._context != "default" else None
+
     def set_context(self, name: str) -> None:
         self._context = name or "default"
 
     def get_state(self) -> BaseState:
         if self._context not in self._states:
-            self._states[self._context] = self._state_cls()
+            self._states[self._context] = self._state_cls(*self._init_args, **self._init_kwargs)
         return self._states[self._context]
 
     def reset(self) -> None:
@@ -190,6 +200,22 @@ class HookRegistry:
 
         return registry
 
+    @classmethod
+    def check_if_exists_or_initialize(cls, module: nn.Module) -> HookRegistry:
+        """Get existing registry or create a new one for the module.
+
+        This method ensures a HookRegistry exists on the module and returns it.
+        If a registry doesn't exist, it creates one and attaches it to the module.
+        This is equivalent to get_or_create() for compatibility with diffusers API.
+
+        Args:
+            module: The module to get/create a registry for.
+
+        Returns:
+            The HookRegistry for this module.
+        """
+        return cls.get_or_create(module)
+
     def update_sorted_hooks(self):
         """Sort hooks by name, which dictates pre/post process order."""
         sorted_hooks = [self._hooks[k] for k in sorted(self._hooks) if self._hooks[k] != self._new_fwd_impl_hook]
@@ -205,7 +231,23 @@ class HookRegistry:
             name: Unique name for this hook.
             hook: The hook instance to register.
         """
+        if name in self._hooks:
+            logger.warning(f"Hook with name '{name}' already exists. Overwriting existing hook.")
+            self.remove_hook(name)
+
         hook.initialize_hook(self.module)
+
+        if hasattr(hook, "fn_ref"):
+            hook.fn_ref.original_forward = self.module._omni_original_forward
+        else:
+            original_forward = self.module._omni_original_forward  # type: ignore[attr-defined]
+
+            class _FnRef:
+                def __init__(self, orig_forward):
+                    self.original_forward = orig_forward
+
+            hook.fn_ref = _FnRef(original_forward)
+
         self._hooks[name] = hook
         # We can only have one hook that overrides new_forward,
         # since we don't currently have a mechanism for combining them.
@@ -290,3 +332,17 @@ class HookRegistry:
         hook = self._hooks.get(name)
         if hook is not None:
             hook.reset_state(self.module)
+
+    def reset(self) -> None:
+        """Reset all hooks and clear the registry.
+
+        This removes all hooks from the registry and resets each hook's state.
+        Also restores module.forward to its original implementation.
+        """
+        for name, hook in list(self._hooks.items()):
+            hook.reset_state(self.module)
+        self._hooks.clear()
+
+        if hasattr(self.module, "_omni_original_forward"):
+            self.module.forward = self.module._omni_original_forward  # type: ignore[attr-defined]
+            delattr(self.module, "_omni_original_forward")
diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py
index b6f881d1c..d8eb2ac53 100644
--- a/vllm_omni/diffusion/inline_stage_diffusion_client.py
+++ b/vllm_omni/diffusion/inline_stage_diffusion_client.py
@@ -128,7 +128,6 @@ class InlineStageDiffusionClient(StageClientBase):
             request = OmniDiffusionRequest(
                 prompts=[prompt],
                 sampling_params=sampling_params,
-                request_ids=[request_id],
                 request_id=request_id,
                 kv_sender_info=kv_sender_info,
             )
@@ -187,7 +186,6 @@ class InlineStageDiffusionClient(StageClientBase):
             request = OmniDiffusionRequest(
                 prompts=prompts,
                 sampling_params=sampling_params,
-                request_ids=[f"{request_id}-{i}" for i in range(len(prompts))],
                 request_id=request_id,
                 kv_sender_info=kv_sender_info,
             )
diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py
index b5c4f9df9..b4fa1209f 100644
--- a/vllm_omni/diffusion/layers/rope.py
+++ b/vllm_omni/diffusion/layers/rope.py
@@ -94,10 +94,18 @@ class RotaryEmbedding(CustomOp):
         self.interleaved = not is_neox_style
         self.apply_rotary_emb_flash_attn = None
         self.has_mindie = False
+        # ``find_spec("flash_attn")`` is True as long as *any* package publishes
+        # the ``flash_attn`` namespace — including ``flash-attn-4``, which ships
+        # only ``flash_attn.cute`` and no ``flash_attn.ops``. Guard the import
+        # so a partial namespace doesn't crash the RoPE init; the CUDA forward
+        # path uses ``vllm.vllm_flash_attn.layers.rotary`` anyway.
         if find_spec("flash_attn") is not None:
-            from flash_attn.ops.triton.rotary import apply_rotary
+            try:
+                from flash_attn.ops.triton.rotary import apply_rotary
 
-            self.apply_rotary_emb_flash_attn = apply_rotary
+                self.apply_rotary_emb_flash_attn = apply_rotary
+            except ImportError:
+                pass
         if find_spec("mindiesd") is not None:
             self.has_mindie = True
 
diff --git a/vllm_omni/diffusion/model_loader/checkpoint_adapters/__init__.py b/vllm_omni/diffusion/model_loader/checkpoint_adapters/__init__.py
index cc20b2da4..584562193 100644
--- a/vllm_omni/diffusion/model_loader/checkpoint_adapters/__init__.py
+++ b/vllm_omni/diffusion/model_loader/checkpoint_adapters/__init__.py
@@ -2,7 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from torch import nn
 
-from .modelopt import ModelOptFp8CheckpointAdapter
+from .modelopt import (
+    ModelOptFp8CheckpointAdapter,
+    ModelOptMixedPrecisionCheckpointAdapter,
+    ModelOptNvFp4CheckpointAdapter,
+)
 
 
 def get_checkpoint_adapter(
@@ -10,13 +14,19 @@ def get_checkpoint_adapter(
     source: object,
     quant_config: object | None,
     use_safetensors: bool,
-) -> ModelOptFp8CheckpointAdapter | None:
+) -> ModelOptFp8CheckpointAdapter | ModelOptNvFp4CheckpointAdapter | ModelOptMixedPrecisionCheckpointAdapter | None:
     if ModelOptFp8CheckpointAdapter.is_compatible(source, quant_config, use_safetensors):
         return ModelOptFp8CheckpointAdapter(model, source)
+    if ModelOptNvFp4CheckpointAdapter.is_compatible(source, quant_config, use_safetensors):
+        return ModelOptNvFp4CheckpointAdapter(model, source)
+    if ModelOptMixedPrecisionCheckpointAdapter.is_compatible(source, quant_config, use_safetensors):
+        return ModelOptMixedPrecisionCheckpointAdapter(model, source)
     return None
 
 
 __all__ = [
     "ModelOptFp8CheckpointAdapter",
+    "ModelOptMixedPrecisionCheckpointAdapter",
+    "ModelOptNvFp4CheckpointAdapter",
     "get_checkpoint_adapter",
 ]
diff --git a/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py b/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py
index 6102a55e6..04fc8ea32 100644
--- a/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py
+++ b/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py
@@ -11,7 +11,7 @@ from vllm.model_executor.utils import get_packed_modules_mapping
 
 logger = init_logger(__name__)
 
-MODEL_OPT_SCALE_SUFFIXES = (".input_scale", ".weight_scale", ".weight_scale_inv")
+MODEL_OPT_SCALE_SUFFIXES = (".input_scale", ".weight_scale", ".weight_scale_2", ".weight_scale_inv")
 DEFAULT_PACKED_MODULES_MAPPING = {
     "to_qkv": ("to_q", "to_k", "to_v"),
     "add_kv_proj": ("add_q_proj", "add_k_proj", "add_v_proj"),
@@ -249,3 +249,24 @@ class ModelOptFp8CheckpointAdapter:
 
         self._check_pending_weights(state)
         self._log_adaptation_summary(state)
+
+
+class ModelOptNvFp4CheckpointAdapter(ModelOptFp8CheckpointAdapter):
+    @staticmethod
+    def _is_checkpoint_quant_config(quant_config: object | None) -> bool:
+        return (
+            quant_config is not None
+            and hasattr(quant_config, "get_name")
+            and quant_config.get_name() == "modelopt_fp4"
+            and bool(getattr(quant_config, "is_checkpoint_nvfp4_serialized", False))
+        )
+
+
+class ModelOptMixedPrecisionCheckpointAdapter(ModelOptFp8CheckpointAdapter):
+    @staticmethod
+    def _is_checkpoint_quant_config(quant_config: object | None) -> bool:
+        return (
+            quant_config is not None
+            and hasattr(quant_config, "get_name")
+            and quant_config.get_name() == "modelopt_mixed"
+        )
diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py
index f4ebfcb24..c9f81b64d 100644
--- a/vllm_omni/diffusion/model_loader/diffusers_loader.py
+++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py
@@ -55,6 +55,23 @@ TRANSFORMER_WEIGHTS_INDEX = "model.safetensors.index.json"
 INDEX_FILES = [DIFFUSION_MODEL_WEIGHTS_INDEX, TRANSFORMER_WEIGHTS_INDEX]
 
 
+def _resolve_custom_pipeline_cls(custom_pipeline_name: str | type | None) -> type:
+    """Resolve a custom pipeline reference to a class.
+
+    Accepts either a fully qualified name string (resolved via import) or an
+    already-imported class object (returned as-is).
+    """
+    if custom_pipeline_name is None:
+        raise ValueError("custom_pipeline_name is required for load_format='custom_pipeline'")
+    if isinstance(custom_pipeline_name, str):
+        return resolve_obj_by_qualname(custom_pipeline_name)
+    if isinstance(custom_pipeline_name, type):
+        return custom_pipeline_name
+    raise TypeError(
+        f"custom_pipeline_name must be a qualified name string or a class, got {type(custom_pipeline_name).__name__}"
+    )
+
+
 class DiffusersPipelineLoader:
     """Model loader that can load diffusers pipeline components from disk."""
 
@@ -294,17 +311,32 @@ class DiffusersPipelineLoader:
         od_config: OmniDiffusionConfig,
         load_device: str,
         load_format: str | None = "default",
-        custom_pipeline_name: str | None = None,
+        custom_pipeline_name: str | type[nn.Module] | None = None,
         device: torch.device | None = None,
     ) -> nn.Module:
         """Load a model with the given configurations."""
         if load_format is None:
             load_format = "default"
         self.od_config = od_config
-        # CPU offload + FP8: load weights on device for FP8 quantization
+        # CPU offload + quantization: for offline-quantized models (e.g., AutoRound MXFP8),
+        # weights are already quantized in the checkpoint — load directly on CPU.
+        # For online quantization, load on device so quantization can run on accelerator,
+        # then move back to CPU afterward.
+        offload_after_quant = False
         if load_device == "cpu" and od_config.quantization_config is not None:
-            load_device = device.type
-            logger.info(f"Quantization enabled with CPU offload, using {load_device} for weight loading")
+            quant_cfg = od_config.quantization_config
+            is_offline = getattr(quant_cfg, "data_type", None) == "mx_fp" or getattr(
+                quant_cfg, "is_checkpoint_quantized", False
+            )
+            if not is_offline:
+                load_device = device.type
+                offload_after_quant = True
+                logger.info(
+                    "Online quantization with CPU offload, using %s for weight loading (will offload back to CPU)",
+                    load_device,
+                )
+            else:
+                logger.info("Offline-quantized model with CPU offload, loading weights directly on CPU")
 
         target_device = torch.device(load_device)
         with set_default_torch_dtype(od_config.dtype):
@@ -313,19 +345,34 @@ class DiffusersPipelineLoader:
                     od_config, target_device=device, load_format=load_format, custom_pipeline_name=custom_pipeline_name
                 )
             else:
-                with target_device:
-                    if load_format == "default":
-                        model = initialize_model(od_config)
-                    elif load_format == "diffusers":
-                        model = DiffusersAdapterPipeline(od_config=od_config, device=target_device)
-                    elif load_format == "custom_pipeline":
-                        from vllm_omni.diffusion.config import set_current_diffusion_config
+                if load_format == "custom_pipeline":
+                    # NOTE: Custom pipelines call HuggingFace `from_pretrained(...).to(device)`
+                    # internally. If we construct them under `with target_device:` (CUDA),
+                    # safetensors takes a direct-to-GPU fast path that calls `cudaMalloc`
+                    # via the driver API and BYPASSES PyTorch's caching allocator.
+                    # That makes those bytes invisible to CuMemAllocator, so `sleep()`
+                    # cannot offload/unmap them and GPU memory stays pinned.
+                    #
+                    # Fix: build the custom pipeline on CPU first (no default device
+                    # context), then explicitly move it to the target device. The
+                    # subsequent `.to(target_device)` issues `torch.empty(..., device=cuda)`
+                    # + `copy_`, which goes through the caching allocator and is fully
+                    # tracked by CuMemAllocator.
+                    from vllm_omni.diffusion.config import set_current_diffusion_config
 
-                        model_cls = resolve_obj_by_qualname(custom_pipeline_name)
-                        with set_current_diffusion_config(od_config):
-                            model = model_cls(od_config=od_config)
-                    else:
-                        raise ValueError(f"Unknown load_format: {load_format}")
+                    model_cls = _resolve_custom_pipeline_cls(custom_pipeline_name)
+                    with set_current_diffusion_config(od_config):
+                        model = model_cls(od_config=od_config)
+                    if target_device.type != "cpu":
+                        model.to(target_device)
+                else:
+                    with target_device:
+                        if load_format == "default":
+                            model = initialize_model(od_config)
+                        elif load_format == "diffusers":
+                            model = DiffusersAdapterPipeline(od_config=od_config, device=target_device)
+                        else:
+                            raise ValueError(f"Unknown load_format: {load_format}")
                 logger.debug("Loading weights on %s ...", load_device)
                 if load_format == "diffusers":
                     # DiffusersAdapterPipeline.load_weights() calls
@@ -342,6 +389,10 @@ class DiffusersPipelineLoader:
             # This is needed for vLLM's quantization methods that need to transform weights
             self._process_weights_after_loading(model, target_device)
 
+            if offload_after_quant:
+                model.to("cpu")
+                logger.info("Quantization complete, offloaded model back to CPU")
+
         return model.eval()
 
     def _process_weights_after_loading(self, model: nn.Module, target_device: torch.device) -> None:
@@ -563,7 +614,7 @@ class DiffusersPipelineLoader:
         od_config: OmniDiffusionConfig,
         target_device: torch.device,
         load_format: str = "default",
-        custom_pipeline_name: str | None = None,
+        custom_pipeline_name: str | type[nn.Module] | None = None,
     ) -> nn.Module:
         """Load model with HSDP sharding for inference.
 
@@ -591,9 +642,9 @@ class DiffusersPipelineLoader:
         if load_format == "default":
             model = initialize_model(od_config)
         elif load_format == "custom_pipeline":
+            model_cls = _resolve_custom_pipeline_cls(custom_pipeline_name)
             from vllm_omni.diffusion.config import set_current_diffusion_config
 
-            model_cls = resolve_obj_by_qualname(custom_pipeline_name)
             with set_current_diffusion_config(od_config):
                 model = model_cls(od_config=od_config)
         self.load_weights(model)
diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py
index d333aa481..e1bd93f94 100644
--- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py
+++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py
@@ -14,7 +14,6 @@ import numpy as np
 import torch
 import torch.distributed as dist
 from torch import nn
-from torch.nn.attention.flex_attention import flex_attention
 from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
 from transformers.models.qwen2.modeling_qwen2 import (
     Qwen2PreTrainedModel,
@@ -36,7 +35,6 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.transformers_utils.configs.bagel import BagelConfig
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata as DiffusionAttentionMetadata
-from vllm_omni.diffusion.attention.backends.utils.fa import flash_attn_varlen_func
 from vllm_omni.diffusion.attention.layer import Attention as DiffusionAttention
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.distributed.parallel_state import (
@@ -241,11 +239,6 @@ class BagelMLP(nn.Module):
         return x
 
 
-torch._dynamo.config.cache_size_limit = 512
-torch._dynamo.config.accumulated_cache_size_limit = 4096
-flex_attention = torch.compile(flex_attention)
-
-
 class Qwen2MoTConfig(Qwen2Config):
     """Configuration for Qwen2MoT (Mixture of Tokens) model.
 
@@ -414,45 +407,55 @@ class PackedAttentionMoT(nn.Module):
 
         self.rotary_op = RotaryEmbedding(is_neox_style=True)
 
-        # SP (Ulysses / Ring) attention for generation mode denoising
-        sp_size = parallel_config.sequence_parallel_size if parallel_config is not None else 1
-        if sp_size is not None and sp_size > 1:
-            self.sp_attn = DiffusionAttention(
-                num_heads=self.total_num_heads,
-                head_size=self.head_dim,
-                softmax_scale=1.0 / (self.head_dim**0.5),
-                causal=False,
-                num_kv_heads=self.total_num_kv_heads,
-            )
-        else:
-            self.sp_attn = None
+        self.attn_causal = DiffusionAttention(
+            num_heads=self.total_num_heads,
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=True,
+            num_kv_heads=self.total_num_kv_heads,
+        )
+        self.attn_noncausal = DiffusionAttention(
+            num_heads=self.total_num_heads,
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=False,
+            num_kv_heads=self.total_num_kv_heads,
+        )
 
     def _is_sp_active(self) -> bool:
         """Check if SP is active for this attention layer."""
-        if self.sp_attn is None:
-            return False
         if not is_forward_context_available():
             return False
         return get_forward_context().sp_active
 
-    def _forward_sp_gen(
+    def _forward_gen(
         self,
         packed_query_sequence: torch.Tensor,
         packed_query_position_embeddings: torch.Tensor,
         past_key_values: NaiveCache | None,
         packed_vae_token_indexes: torch.Tensor,
         packed_text_indexes: torch.Tensor,
+        update_past_key_values: bool = False,
     ) -> tuple[torch.Tensor, NaiveCache | None]:
-        """SP-aware attention for gen mode denoising.
+        """Forward pass for generation mode.
 
-        Converts packed format to batched (1, S, H, D) and uses the diffusion
-        Attention layer (Ulysses / Ring) with joint mechanism:
-          - Main Q/K/V: VAE tokens (split across SP ranks)
-          - Joint Q: text marker Q (replicated)
-          - Joint K/V: KV cache K/V + text marker K/V (replicated)
+        This path does the following:
+
+        1. Apply qkv projection to the text seq & vae seqs
+        2. Reshape both to 3D & apply RMS norms
+        3. Apply RoPE to text / VAE components independently
+        4. Create the full K/V; Bagel currently manages its own KV cache
+           (NaiveCache) independently since it is a diffusion model
+        5. Apply non-causal attention, while taking sequence parallelism into account
+        6. Apply output projections on the split parts
+        7. Merge back into the packed format
+        8. Update the NaiveCache.
+
+        TODO (Alex): it would be best to remove packing from Bagel to simplify the code.
+        Currently we shouldn't need it in the model, and it would be ideal to handle
+        packing/batching etc in a more model agnostic way.
         """
         packed_query_sequence = packed_query_sequence.to(torch.bfloat16)
-
         packed_text_query_sequence = packed_query_sequence[packed_text_indexes]
         packed_vae_query_sequence = packed_query_sequence[packed_vae_token_indexes]
 
@@ -481,8 +484,6 @@ class PackedAttentionMoT(nn.Module):
         # Apply RoPE - need to build per-token cos/sin for text and vae separately
         # packed_query_position_embeddings are ordered as the packed sequence
         cos_full, sin_full = [x[..., : self.head_dim // 2] for x in packed_query_position_embeddings]
-
-        # Extract cos/sin for text and vae positions
         text_cos = cos_full[packed_text_indexes]
         text_sin = sin_full[packed_text_indexes]
         vae_cos = cos_full[packed_vae_token_indexes]
@@ -504,33 +505,35 @@ class PackedAttentionMoT(nn.Module):
         if past_key_values is not None and past_key_values.key_cache[self.layer_idx] is not None:
             cache_k = past_key_values.key_cache[self.layer_idx]
             cache_v = past_key_values.value_cache[self.layer_idx]
-            joint_k = torch.cat([cache_k, text_k], dim=0).unsqueeze(0)
-            joint_v = torch.cat([cache_v, text_v], dim=0).unsqueeze(0)
+            ctx_k = torch.cat([cache_k, text_k], dim=0)
+            ctx_v = torch.cat([cache_v, text_v], dim=0)
         else:
-            joint_k = text_k.unsqueeze(0)
-            joint_v = text_v.unsqueeze(0)
+            ctx_k = text_k
+            ctx_v = text_v
 
-        # Reshape to batched (1, S, H, D) for diffusion Attention
-        vae_q_4d = vae_q.unsqueeze(0)
-        vae_k_4d = vae_k.unsqueeze(0)
-        vae_v_4d = vae_v.unsqueeze(0)
-        text_q_4d = text_q.unsqueeze(0)
+        # NOTE: we reshape to batched (1, S, H, D) for diffusion Attention
+        # attn_out should be: (1, text_len + local_vae_len, H, D)
+        if self._is_sp_active():
+            # Joint mechanism keeps text+cache replicated across SP ranks
+            attn_out = self.attn_noncausal(
+                vae_q.unsqueeze(0),
+                vae_k.unsqueeze(0),
+                vae_v.unsqueeze(0),
+                DiffusionAttentionMetadata(
+                    joint_query=text_q.unsqueeze(0),
+                    joint_key=ctx_k.unsqueeze(0),
+                    joint_value=ctx_v.unsqueeze(0),
+                    joint_strategy="front",
+                ),
+            )
+        else:
+            q = torch.cat([text_q, vae_q], dim=0).unsqueeze(0)
+            k = torch.cat([ctx_k, vae_k], dim=0).unsqueeze(0)
+            v = torch.cat([ctx_v, vae_v], dim=0).unsqueeze(0)
+            attn_out = self.attn_noncausal(q, k, v)
 
-        # Call SP-aware attention: VAE as main, text+cache as joint
-        attn_out = self.sp_attn(
-            vae_q_4d,
-            vae_k_4d,
-            vae_v_4d,
-            DiffusionAttentionMetadata(
-                joint_query=text_q_4d,
-                joint_key=joint_k,
-                joint_value=joint_v,
-                joint_strategy="front",
-            ),
-        )
-        # attn_out: (1, text_len + local_vae_len, H, D)
         text_len = text_q.shape[0]
-        attn_out = attn_out.squeeze(0)  # (text_len + local_vae_len, H, D)
+        attn_out = attn_out.squeeze(0)
         text_attn = attn_out[:text_len].reshape(text_len, self.q_size)
         vae_attn = attn_out[text_len:].reshape(-1, self.q_size)
 
@@ -544,145 +547,108 @@ class PackedAttentionMoT(nn.Module):
         full_output[packed_text_indexes] = text_out
         full_output[packed_vae_token_indexes] = vae_out
 
+        if update_past_key_values:
+            new_k = torch.cat([ctx_k, vae_k], dim=0)
+            new_v = torch.cat([ctx_v, vae_v], dim=0)
+            past_key_values.key_cache[self.layer_idx] = new_k
+            past_key_values.value_cache[self.layer_idx] = new_v
+
         return full_output, past_key_values
 
+    def _forward_und(
+        self,
+        packed_query_sequence: torch.Tensor,
+        packed_query_position_embeddings: torch.Tensor,
+        past_key_values: NaiveCache | None,
+        is_causal: bool,
+        update_past_key_values: bool = True,
+    ) -> tuple[torch.Tensor, NaiveCache | None]:
+        """Forward pass for understanding mode.
+
+        This path does the following (not hard to read):
+
+        1. Apply qkv projection to the text seq
+        2. Reshape to 3D & apply RMS norms
+        3. Apply RoPE
+        4. Create the full K/V; Bagel currently manages its own KV cache
+           (NaiveCache) independently since it is a diffusion model
+        5. Apply attention based on causality kwarg
+        6. Apply output projection
+        7. Update the NaiveCache.
+        """
+
+        qkv, _ = self.qkv_proj(packed_query_sequence)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q = q.view(-1, self.num_heads, self.head_dim)
+        k = k.view(-1, self.num_kv_heads, self.head_dim)
+        v = v.view(-1, self.num_kv_heads, self.head_dim)
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        cos, sin = [x[..., : self.head_dim // 2] for x in packed_query_position_embeddings]
+        q = self.rotary_op(q.to(cos.dtype).unsqueeze(0), cos, sin).squeeze(0)
+        k = self.rotary_op(k.to(cos.dtype).unsqueeze(0), cos, sin).squeeze(0)
+
+        q = q.to(torch.bfloat16)
+        k = k.to(torch.bfloat16)
+        v = v.to(torch.bfloat16)
+
+        if past_key_values is not None and past_key_values.key_cache[self.layer_idx] is not None:
+            cache_k = past_key_values.key_cache[self.layer_idx]
+            cache_v = past_key_values.value_cache[self.layer_idx]
+            full_k = torch.cat([cache_k, k], dim=0)
+            full_v = torch.cat([cache_v, v], dim=0)
+        else:
+            full_k = k
+            full_v = v
+
+        attn = self.attn_causal if is_causal else self.attn_noncausal
+        attn_out = attn(
+            q.unsqueeze(0),
+            full_k.unsqueeze(0),
+            full_v.unsqueeze(0),
+        )
+
+        attn_out = attn_out.squeeze(0).reshape(-1, self.q_size)
+        attn_out, _ = self.o_proj(attn_out)
+
+        if update_past_key_values:
+            past_key_values.key_cache[self.layer_idx] = full_k
+            past_key_values.value_cache[self.layer_idx] = full_v
+
+        return attn_out, past_key_values
+
     def forward(
         self,
         packed_query_sequence: torch.Tensor,
         query_lens: torch.Tensor,
         packed_query_position_embeddings: torch.Tensor,
-        packed_query_indexes: torch.Tensor,
         past_key_values: NaiveCache | None = None,
-        key_values_lens: torch.Tensor | None = None,
-        packed_key_value_indexes: torch.Tensor | None = None,
         update_past_key_values=True,
         is_causal=True,
         mode="und",
         packed_vae_token_indexes=None,
         packed_text_indexes=None,
     ):
-        # SP path for gen-mode denoising (non-causal, no KV update)
-        if (
-            mode == "gen"
-            and not update_past_key_values
-            and not is_causal
-            and self._is_sp_active()
-            and packed_vae_token_indexes is not None
-            and packed_text_indexes is not None
-        ):
-            return self._forward_sp_gen(
+        if mode == "gen":
+            if is_causal:
+                raise ValueError("Generation model for Bagel requires non-causal attention")
+            return self._forward_gen(
                 packed_query_sequence=packed_query_sequence,
                 packed_query_position_embeddings=packed_query_position_embeddings,
                 past_key_values=past_key_values,
                 packed_vae_token_indexes=packed_vae_token_indexes,
                 packed_text_indexes=packed_text_indexes,
+                update_past_key_values=update_past_key_values,
             )
 
-        if mode == "und":
-            qkv, _ = self.qkv_proj(packed_query_sequence)
-            q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
-            packed_query_states = q.view(-1, self.num_heads, self.head_dim)
-            packed_key_states = k.view(-1, self.num_kv_heads, self.head_dim)
-            packed_value_states = v.view(-1, self.num_kv_heads, self.head_dim)
-            packed_query_states = self.q_norm(packed_query_states)
-            packed_key_states = self.k_norm(packed_key_states)
-        elif mode == "gen":
-            packed_query_sequence = packed_query_sequence.to(torch.bfloat16)
-
-            packed_text_query_sequence = packed_query_sequence[packed_text_indexes]
-            packed_vae_query_sequence = packed_query_sequence[packed_vae_token_indexes]
-
-            # Project text tokens through base qkv
-            text_qkv, _ = self.qkv_proj(packed_text_query_sequence)
-            text_q, text_k, text_v = text_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
-
-            # Project vae tokens through moe_gen qkv
-            vae_qkv, _ = self.qkv_proj_moe_gen(packed_vae_query_sequence)
-            vae_q, vae_k, vae_v = vae_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
-
-            # Merge into packed tensors
-            total_len = packed_query_sequence.shape[0]
-            packed_query_states = packed_query_sequence.new_zeros((total_len, self.q_size))
-            packed_key_states = packed_query_sequence.new_zeros((total_len, self.kv_size))
-            packed_value_states = packed_query_sequence.new_zeros((total_len, self.kv_size))
-
-            packed_query_states[packed_text_indexes] = text_q
-            packed_query_states[packed_vae_token_indexes] = vae_q
-            packed_key_states[packed_text_indexes] = text_k
-            packed_key_states[packed_vae_token_indexes] = vae_k
-            packed_value_states[packed_text_indexes] = text_v
-            packed_value_states[packed_vae_token_indexes] = vae_v
-
-            packed_query_states = packed_query_states.view(-1, self.num_heads, self.head_dim)
-            packed_key_states = packed_key_states.view(-1, self.num_kv_heads, self.head_dim)
-            packed_value_states = packed_value_states.view(-1, self.num_kv_heads, self.head_dim)
-
-            packed_query_states = packed_query_states.to(torch.float32)
-            packed_query_states[packed_text_indexes] = self.q_norm(packed_query_states[packed_text_indexes])
-            packed_query_states[packed_vae_token_indexes] = self.q_norm_moe_gen(
-                packed_query_states[packed_vae_token_indexes]
-            )
-
-            packed_key_states = packed_key_states.to(torch.float32)
-            packed_key_states[packed_text_indexes] = self.k_norm(packed_key_states[packed_text_indexes])
-            packed_key_states[packed_vae_token_indexes] = self.k_norm_moe_gen(
-                packed_key_states[packed_vae_token_indexes]
-            )
-
-        cos, sin = [x[..., : self.head_dim // 2] for x in packed_query_position_embeddings]
-        packed_query_states = self.rotary_op(packed_query_states.to(cos.dtype).unsqueeze(0), cos, sin).squeeze(0)
-        packed_key_states = self.rotary_op(packed_key_states.to(cos.dtype).unsqueeze(0), cos, sin).squeeze(0)
-
-        packed_query_states = packed_query_states.to(torch.bfloat16)
-        packed_key_states = packed_key_states.to(torch.bfloat16)
-        packed_value_states = packed_value_states.to(torch.bfloat16)
-
-        if past_key_values is not None and past_key_values.key_cache[self.layer_idx] is not None:
-            past_key_states = past_key_values.key_cache[self.layer_idx]
-            past_value_states = past_key_values.value_cache[self.layer_idx]
-
-            seqlens = sum(query_lens) + sum(key_values_lens)
-            merged_key_states = past_key_states.new_zeros(size=[seqlens, self.num_kv_heads, self.head_dim])
-            merged_value_states = past_key_states.new_zeros(size=[seqlens, self.num_kv_heads, self.head_dim])
-            merged_key_states[packed_query_indexes] = packed_key_states
-            merged_key_states[packed_key_value_indexes] = past_key_states
-            merged_value_states[packed_query_indexes] = packed_value_states
-            merged_value_states[packed_key_value_indexes] = past_value_states
-            key_values_lens = key_values_lens + query_lens
-        else:
-            merged_key_states = packed_key_states
-            merged_value_states = packed_value_states
-            key_values_lens = query_lens
-
-        cu_seqlens_q = torch.nn.functional.pad(torch.cumsum(query_lens, dim=0), (1, 0))
-        cu_seqlens_k = torch.nn.functional.pad(torch.cumsum(key_values_lens, dim=0), (1, 0))
-
-        packed_attn_output = flash_attn_varlen_func(
-            q=packed_query_states,
-            k=merged_key_states,
-            v=merged_value_states,
-            cu_seqlens_q=cu_seqlens_q.to(torch.int32),
-            cu_seqlens_k=cu_seqlens_k.to(torch.int32),
-            max_seqlen_q=max(query_lens).item(),
-            max_seqlen_k=max(key_values_lens).item(),
-            causal=is_causal,
+        return self._forward_und(
+            packed_query_sequence=packed_query_sequence,
+            packed_query_position_embeddings=packed_query_position_embeddings,
+            past_key_values=past_key_values,
+            is_causal=is_causal,
+            update_past_key_values=update_past_key_values,
         )
-        packed_attn_output = packed_attn_output.reshape(-1, self.q_size)
-        if mode == "und":
-            packed_attn_output, _ = self.o_proj(packed_attn_output)
-        elif mode == "gen":
-            text_out, _ = self.o_proj(packed_attn_output[packed_text_indexes])
-            vae_out, _ = self.o_proj_moe_gen(packed_attn_output[packed_vae_token_indexes])
-            full_output = text_out.new_zeros((packed_attn_output.shape[0], self.hidden_size))
-            full_output[packed_text_indexes] = text_out
-            full_output[packed_vae_token_indexes] = vae_out
-            packed_attn_output = full_output
-
-        if update_past_key_values:
-            past_key_values.key_cache[self.layer_idx] = merged_key_states
-            past_key_values.value_cache[self.layer_idx] = merged_value_states
-
-        return packed_attn_output, past_key_values
 
 
 class Qwen2MoTDecoderLayer(nn.Module):
@@ -728,10 +694,7 @@ class Qwen2MoTDecoderLayer(nn.Module):
         packed_query_sequence: torch.Tensor | None = None,
         query_lens: torch.Tensor = None,
         packed_query_position_embeddings: torch.Tensor = None,
-        packed_query_indexes: torch.Tensor = None,
         past_key_values: NaiveCache | None = None,
-        key_values_lens: torch.Tensor | None = None,
-        packed_key_value_indexes: torch.Tensor | None = None,
         update_past_key_values=True,
         is_causal=True,
         mode="und",
@@ -758,10 +721,7 @@ class Qwen2MoTDecoderLayer(nn.Module):
             packed_query_sequence=packed_query_sequence,
             query_lens=query_lens,
             packed_query_position_embeddings=packed_query_position_embeddings,
-            packed_query_indexes=packed_query_indexes,
             past_key_values=past_key_values,
-            key_values_lens=key_values_lens,
-            packed_key_value_indexes=packed_key_value_indexes,
             update_past_key_values=update_past_key_values,
             is_causal=is_causal,
             mode=mode,
@@ -842,10 +802,7 @@ class Qwen2MoTModel(Qwen2PreTrainedModel):
         packed_query_sequence: torch.Tensor | None = None,
         query_lens: torch.Tensor | None = None,
         packed_query_position_ids: torch.Tensor | None = None,
-        packed_query_indexes: torch.Tensor | None = None,
         past_key_values: NaiveCache | None = None,
-        key_values_lens: torch.Tensor | None = None,
-        packed_key_value_indexes: torch.Tensor | None = None,
         update_past_key_values=True,
         is_causal=True,
         mode="und",
@@ -883,15 +840,16 @@ class Qwen2MoTModel(Qwen2PreTrainedModel):
                 )
 
         for layer_idx, decoder_layer in enumerate(self.layers):
+            # TODO (Alex): Remove encoder_hidden_states as a kwarg; currently we keep it
+            # for compatibility with the current custom CacheDiT adapter, as we need to be
+            # careful to not break the NaiveCache handling when switching from pattern
+            # 0 -> 4.
             packed_query_sequence, past_key_values = decoder_layer(
                 hidden_states=packed_query_sequence,
                 encoder_hidden_states=None,
                 query_lens=query_lens,
                 packed_query_position_embeddings=packed_query_position_embeddings,
-                packed_query_indexes=packed_query_indexes,
                 past_key_values=past_key_values,
-                key_values_lens=key_values_lens,
-                packed_key_value_indexes=packed_key_value_indexes,
                 update_past_key_values=update_past_key_values,
                 is_causal=is_causal,
                 **extra_inputs,
@@ -959,10 +917,7 @@ class Qwen2MoTForCausalLM(Qwen2PreTrainedModel):
         packed_query_sequence: torch.Tensor | None = None,
         query_lens: torch.Tensor | None = None,
         packed_query_position_ids: torch.Tensor | None = None,
-        packed_query_indexes: torch.Tensor | None = None,
         past_key_values: NaiveCache | None = None,
-        key_values_lens: torch.Tensor | None = None,
-        packed_key_value_indexes: torch.Tensor | None = None,
         update_past_key_values=True,
         is_causal=True,
         mode="und",
@@ -975,10 +930,7 @@ class Qwen2MoTForCausalLM(Qwen2PreTrainedModel):
             packed_query_sequence=packed_query_sequence,
             query_lens=query_lens,
             packed_query_position_ids=packed_query_position_ids,
-            packed_query_indexes=packed_query_indexes,
             past_key_values=past_key_values,
-            key_values_lens=key_values_lens,
-            packed_key_value_indexes=packed_key_value_indexes,
             update_past_key_values=update_past_key_values,
             is_causal=is_causal,
             mode=mode,
@@ -1266,32 +1218,21 @@ class Bagel(nn.Module):
         packed_text_ids = list()
         packed_text_position_ids = list()
         text_token_lens = list()
-        packed_text_indexes = list()
-        packed_key_value_indexes = list()
 
-        curr = 0
         newlens, new_rope = list(), list()
         for prompt, curr_kvlen, curr_position_id in zip(prompts, curr_kvlens, curr_rope):
-            packed_key_value_indexes.extend(range(curr, curr + curr_kvlen))
-            curr += curr_kvlen
-
             text_ids = tokenizer.encode(prompt, add_special_tokens=False)
             text_ids = [new_token_ids["bos_token_id"]] + text_ids + [new_token_ids["eos_token_id"]]
             text_token_lens.append(len(text_ids))
             packed_text_ids.extend(text_ids)
             packed_text_position_ids.extend(range(curr_position_id, curr_position_id + len(text_ids)))
-            packed_text_indexes.extend(range(curr, curr + len(text_ids)))
             newlens.append(curr_kvlen + len(text_ids))
             new_rope.append(curr_position_id + len(text_ids))
-            curr += len(text_ids)
 
         generation_input = {
             "text_token_lens": torch.tensor(text_token_lens, dtype=torch.int),
             "packed_text_ids": torch.tensor(packed_text_ids, dtype=torch.long),
             "packed_text_position_ids": torch.tensor(packed_text_position_ids, dtype=torch.long),
-            "packed_text_indexes": torch.tensor(packed_text_indexes, dtype=torch.long),
-            "packed_key_value_indexes": torch.tensor(packed_key_value_indexes, dtype=torch.long),
-            "key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int),
         }
 
         return generation_input, newlens, new_rope
@@ -1302,9 +1243,6 @@ class Bagel(nn.Module):
         packed_text_ids: torch.IntTensor,
         packed_text_position_ids: torch.LongTensor,
         text_token_lens: torch.LongTensor,
-        packed_text_indexes: torch.LongTensor,
-        packed_key_value_indexes: torch.LongTensor,
-        key_values_lens: torch.IntTensor,
     ):
         extra_inputs = {}
         if self.use_moe:
@@ -1314,10 +1252,7 @@ class Bagel(nn.Module):
             packed_text_ids=packed_text_ids,
             query_lens=text_token_lens,
             packed_query_position_ids=packed_text_position_ids,
-            packed_query_indexes=packed_text_indexes,
             past_key_values=past_key_values,
-            packed_key_value_indexes=packed_key_value_indexes,
-            key_values_lens=key_values_lens,
             update_past_key_values=True,
             is_causal=True,
             **extra_inputs,
@@ -1330,20 +1265,14 @@ class Bagel(nn.Module):
         patchified_vae_latent_shapes, packed_vae_position_ids = list(), list()
         packed_vae_token_indexes = list()
         packed_text_ids, packed_text_indexes = list(), list()
-        packed_seqlens, packed_position_ids, packed_indexes = list(), list(), list()
-        packed_key_value_indexes = list()
+        packed_seqlens, packed_position_ids = list(), list()
 
-        _curr = curr = 0
+        _curr = 0
         vae_image_tensors = list()
         newlens, new_rope = list(), list()
         for image, curr_kvlen, curr_position_id in zip(images, curr_kvlens, curr_rope):
-            packed_key_value_indexes.extend(range(curr, curr + curr_kvlen))
-            curr += curr_kvlen
-
             packed_text_ids.append(new_token_ids["start_of_image"])
             packed_text_indexes.append(_curr)
-            packed_indexes.append(curr)
-            curr += 1
             _curr += 1
 
             image_tensor = transforms(image)
@@ -1362,14 +1291,10 @@ class Bagel(nn.Module):
 
             num_img_tokens = w * h
             packed_vae_token_indexes.extend(range(_curr, _curr + num_img_tokens))
-            packed_indexes.extend(range(curr, curr + num_img_tokens))
-            curr += num_img_tokens
             _curr += num_img_tokens
 
             packed_text_ids.append(new_token_ids["end_of_image"])
             packed_text_indexes.append(_curr)
-            packed_indexes.append(curr)
-            curr += 1
             _curr += 1
 
             packed_position_ids.extend([curr_position_id] * (num_img_tokens + 2))
@@ -1393,9 +1318,6 @@ class Bagel(nn.Module):
             "packed_text_indexes": torch.tensor(packed_text_indexes, dtype=torch.long),
             "packed_position_ids": torch.tensor(packed_position_ids, dtype=torch.long),
             "packed_seqlens": torch.tensor(packed_seqlens, dtype=torch.int),
-            "packed_indexes": torch.tensor(packed_indexes, dtype=torch.long),
-            "packed_key_value_indexes": torch.tensor(packed_key_value_indexes, dtype=torch.long),
-            "key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int),
         }
 
         return generation_input, newlens, new_rope
@@ -1413,9 +1335,6 @@ class Bagel(nn.Module):
         packed_text_indexes: torch.LongTensor,
         packed_position_ids: torch.LongTensor,
         packed_seqlens: torch.IntTensor,
-        packed_indexes: torch.LongTensor,
-        key_values_lens: torch.IntTensor,
-        packed_key_value_indexes: torch.Tensor,
     ):
         padded_latent = vae_model.encode(padded_images)
 
@@ -1430,6 +1349,18 @@ class Bagel(nn.Module):
         packed_timestep_embeds = self.time_embedder(packed_timesteps)
         packed_latent = self.vae2llm(packed_latent) + packed_timestep_embeds + packed_pos_embed
 
+        # Mirror forward_cache_update_vit: build a full packed_sequence so MoE gen-mode
+        # indexes (packed_text_indexes / packed_vae_token_indexes) match tensor length.
+        packed_text_embedding = self.language_model.forward(
+            packed_text_ids=packed_text_ids,
+            return_embeddings_only=True,
+        ).packed_query_sequence
+        packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), self.hidden_size))
+        packed_sequence[packed_text_indexes] = packed_text_embedding
+        if packed_latent.dtype != packed_sequence.dtype:
+            packed_latent = packed_latent.to(packed_sequence.dtype)
+        packed_sequence[packed_vae_token_indexes] = packed_latent
+
         extra_inputs = {}
         if self.use_moe:
             extra_inputs = {
@@ -1439,13 +1370,10 @@ class Bagel(nn.Module):
             }
 
         output = self.language_model.forward(
-            packed_text_ids=packed_text_ids,
+            packed_query_sequence=packed_sequence,
             query_lens=packed_seqlens,
             packed_query_position_ids=packed_position_ids,
-            packed_query_indexes=packed_indexes,
             past_key_values=past_key_values,
-            key_values_lens=key_values_lens,
-            packed_key_value_indexes=packed_key_value_indexes,
             update_past_key_values=True,
             is_causal=False,
             **extra_inputs,
@@ -1458,19 +1386,13 @@ class Bagel(nn.Module):
         packed_vit_token_indexes = list()
         vit_token_seqlens, packed_vit_tokens, packed_vit_position_ids = list(), list(), list()
         packed_text_ids, packed_text_indexes = list(), list()
-        packed_seqlens, packed_position_ids, packed_indexes = list(), list(), list()
-        packed_key_value_indexes = list()
+        packed_seqlens, packed_position_ids = list(), list()
 
-        _curr = curr = 0
+        _curr = 0
         newlens, new_rope = list(), list()
         for image, curr_kvlen, curr_position_id in zip(images, curr_kvlens, curr_rope):
-            packed_key_value_indexes.extend(range(curr, curr + curr_kvlen))
-            curr += curr_kvlen
-
             packed_text_ids.append(new_token_ids["start_of_image"])
             packed_text_indexes.append(_curr)
-            packed_indexes.append(curr)
-            curr += 1
             _curr += 1
 
             image_tensor = transforms(image)
@@ -1486,14 +1408,10 @@ class Bagel(nn.Module):
             packed_vit_position_ids.append(vit_position_ids)
             vit_token_seqlens.append(num_img_tokens)
             packed_vit_token_indexes.extend(range(_curr, _curr + num_img_tokens))
-            packed_indexes.extend(range(curr, curr + num_img_tokens))
-            curr += num_img_tokens
             _curr += num_img_tokens
 
             packed_text_ids.append(new_token_ids["end_of_image"])
             packed_text_indexes.append(_curr)
-            packed_indexes.append(curr)
-            curr += 1
             _curr += 1
 
             packed_position_ids.extend([curr_position_id] * (num_img_tokens + 2))
@@ -1510,9 +1428,6 @@ class Bagel(nn.Module):
             "packed_vit_token_indexes": torch.tensor(packed_vit_token_indexes, dtype=torch.long),
             "packed_position_ids": torch.tensor(packed_position_ids, dtype=torch.long),
             "packed_seqlens": torch.tensor(packed_seqlens, dtype=torch.int),
-            "packed_indexes": torch.tensor(packed_indexes, dtype=torch.long),
-            "packed_key_value_indexes": torch.tensor(packed_key_value_indexes, dtype=torch.long),
-            "key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int),
         }
 
         return generation_input, newlens, new_rope
@@ -1528,9 +1443,6 @@ class Bagel(nn.Module):
         vit_token_seqlens: torch.IntTensor,
         packed_position_ids: torch.LongTensor,
         packed_seqlens: torch.IntTensor,
-        packed_indexes: torch.LongTensor,
-        packed_key_value_indexes: torch.LongTensor,
-        key_values_lens: torch.IntTensor,
     ):
         packed_text_embedding = self.language_model.forward(
             packed_text_ids=packed_text_ids,
@@ -1563,10 +1475,7 @@ class Bagel(nn.Module):
             packed_query_sequence=packed_sequence,
             query_lens=packed_seqlens,
             packed_query_position_ids=packed_position_ids,
-            packed_query_indexes=packed_indexes,
             past_key_values=past_key_values,
-            packed_key_value_indexes=packed_key_value_indexes,
-            key_values_lens=key_values_lens,
             update_past_key_values=True,
             is_causal=False,
             **extra_inputs,
@@ -1578,19 +1487,12 @@ class Bagel(nn.Module):
     def prepare_input(self, curr_kvlens, curr_rope, image_sizes, new_token_ids=None):
         packed_text_ids, packed_text_indexes = list(), list()
         packed_vae_position_ids, packed_vae_token_indexes, packed_init_noises = list(), list(), list()
-        packed_position_ids, packed_seqlens, packed_indexes = list(), list(), list()
-        packed_key_value_indexes = list()
+        packed_position_ids, packed_seqlens = list(), list()
 
-        query_curr = curr = 0
+        query_curr = 0
         for (H, W), curr_kvlen, curr_position_id in zip(image_sizes, curr_kvlens, curr_rope):
-            packed_key_value_indexes.extend(range(curr, curr + curr_kvlen))
-            curr += curr_kvlen
-
             packed_text_ids.append(new_token_ids["start_of_image"])
             packed_text_indexes.append(query_curr)
-
-            packed_indexes.append(curr)
-            curr += 1
             query_curr += 1
 
             vae_position_ids = self.get_flattened_position_ids(
@@ -1604,16 +1506,10 @@ class Bagel(nn.Module):
             packed_init_noises.append(torch.randn(num_image_tokens, self.latent_channel * self.latent_patch_size**2))
             packed_vae_token_indexes.extend(range(query_curr, query_curr + num_image_tokens))
             packed_seqlens.append(num_image_tokens + 2)
-
-            packed_indexes.extend(range(curr, curr + num_image_tokens))
-            curr += num_image_tokens
             query_curr += num_image_tokens
 
             packed_text_ids.append(new_token_ids["end_of_image"])
             packed_text_indexes.append(query_curr)
-
-            packed_indexes.append(curr)
-            curr += 1
             query_curr += 1
 
             packed_position_ids.extend([curr_position_id] * (num_image_tokens + 2))
@@ -1627,9 +1523,6 @@ class Bagel(nn.Module):
             "packed_vae_token_indexes": torch.tensor(packed_vae_token_indexes, dtype=torch.long),
             "packed_seqlens": torch.tensor(packed_seqlens, dtype=torch.int),
             "packed_position_ids": torch.tensor(packed_position_ids, dtype=torch.long),
-            "key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int),
-            "packed_indexes": torch.tensor(packed_indexes, dtype=torch.long),
-            "packed_key_value_indexes": torch.tensor(packed_key_value_indexes, dtype=torch.long),
         }
 
         return generation_input
@@ -1638,34 +1531,15 @@ class Bagel(nn.Module):
         return self.prepare_input(curr_kvlens, curr_rope, image_sizes, new_token_ids)
 
     def prepare_vae_latent_cfg(self, curr_kvlens, curr_rope, image_sizes):
-        packed_position_ids, packed_indexes, packed_key_value_indexes = list(), list(), list()
+        packed_position_ids = list()
 
-        query_curr = curr = 0
         for (H, W), curr_kvlen, curr_position_id in zip(image_sizes, curr_kvlens, curr_rope):
-            packed_key_value_indexes.extend(range(curr, curr + curr_kvlen))
-            curr += curr_kvlen
-
-            packed_indexes.append(curr)
-            curr += 1
-            query_curr += 1
-
             h, w = H // self.latent_downsample, W // self.latent_downsample
             num_image_tokens = h * w
-            packed_indexes.extend(range(curr, curr + num_image_tokens))
-            curr += num_image_tokens
-            query_curr += num_image_tokens
-
-            packed_indexes.append(curr)
-            curr += 1
-            query_curr += 1
-
             packed_position_ids.extend([curr_position_id] * (num_image_tokens + 2))
 
         generation_input = {
             "cfg_packed_position_ids": torch.tensor(packed_position_ids, dtype=torch.long),
-            "cfg_key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int),
-            "cfg_packed_query_indexes": torch.tensor(packed_indexes, dtype=torch.long),
-            "cfg_packed_key_value_indexes": torch.tensor(packed_key_value_indexes, dtype=torch.long),
         }
 
         return generation_input
@@ -1692,21 +1566,15 @@ class Bagel(nn.Module):
 
         Ported from the original BAGEL ``Bagel.prepare_start_tokens``.
         """
-        packed_start_tokens, packed_key_value_indexes = list(), list()
+        packed_start_tokens = list()
         packed_query_position_ids = list()
 
-        curr = 0
         for curr_kvlen, curr_position_id in zip(curr_kvlens, curr_rope):
-            packed_key_value_indexes.extend(range(curr, curr + curr_kvlen))
             packed_start_tokens.append(new_token_ids["bos_token_id"])
             packed_query_position_ids.append(curr_position_id)
-            curr += curr_kvlen
-
         generation_input = {
             "packed_start_tokens": torch.tensor(packed_start_tokens, dtype=torch.long),
             "packed_query_position_ids": torch.tensor(packed_query_position_ids, dtype=torch.long),
-            "key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int),
-            "packed_key_value_indexes": torch.tensor(packed_key_value_indexes, dtype=torch.long),
         }
         return generation_input
 
@@ -1714,8 +1582,6 @@ class Bagel(nn.Module):
     def generate_text(
         self,
         past_key_values: NaiveCache,
-        packed_key_value_indexes: torch.LongTensor,
-        key_values_lens: torch.IntTensor,
         packed_start_tokens: torch.LongTensor,
         packed_query_position_ids: torch.LongTensor,
         max_length: int,
@@ -1734,26 +1600,12 @@ class Bagel(nn.Module):
         while step < max_length:
             generated_sequence.append(curr_tokens)
             query_lens = torch.ones_like(curr_tokens)
-            packed_query_indexes = torch.cumsum(key_values_lens, dim=0) + torch.arange(
-                0,
-                len(key_values_lens),
-                device=key_values_lens.device,
-                dtype=key_values_lens.dtype,
-            )
-
-            uppacked = list(packed_key_value_indexes.split(key_values_lens.tolist(), dim=0))
-            for i in range(len(uppacked)):
-                uppacked[i] += i
-            packed_key_value_indexes = torch.cat(uppacked, dim=0)
 
             output = self.language_model(
                 packed_text_ids=curr_tokens,
                 query_lens=query_lens,
                 packed_query_position_ids=packed_query_position_ids,
-                packed_query_indexes=packed_query_indexes,
                 past_key_values=past_key_values,
-                key_values_lens=key_values_lens,
-                packed_key_value_indexes=packed_key_value_indexes,
                 update_past_key_values=True,
                 is_causal=True,
                 mode="und",
@@ -1768,13 +1620,6 @@ class Bagel(nn.Module):
             else:
                 curr_tokens = torch.argmax(pred_logits, dim=-1)
 
-            uppacked = list(packed_key_value_indexes.split(key_values_lens.tolist(), dim=0))
-            for i in range(len(uppacked)):
-                uppacked[i] = torch.cat(
-                    [uppacked[i], torch.tensor([uppacked[i][-1] + 1], device=uppacked[i].device)], dim=0
-                )
-            packed_key_value_indexes = torch.cat(uppacked, dim=0)
-            key_values_lens = key_values_lens + 1
             packed_query_position_ids = packed_query_position_ids + 1
             step += 1
 
@@ -1793,10 +1638,7 @@ class Bagel(nn.Module):
         packed_vae_token_indexes: torch.LongTensor,
         packed_seqlens: torch.IntTensor,
         packed_position_ids: torch.LongTensor,
-        packed_indexes: torch.LongTensor,
         past_key_values: NaiveCache,
-        key_values_lens: torch.IntTensor,
-        packed_key_value_indexes: torch.LongTensor,
         num_timesteps: int = 24,
         timestep_shift: float = 1.0,
         cfg_renorm_min: float = 0.0,
@@ -1804,18 +1646,12 @@ class Bagel(nn.Module):
         cfg_interval: tuple[float, float] = [0, 1],
         # cfg_text
         cfg_text_scale: float = 1.0,
-        cfg_text_packed_query_indexes: torch.LongTensor | None = None,
         cfg_text_packed_position_ids: torch.LongTensor | None = None,
         cfg_text_past_key_values: NaiveCache | None = None,
-        cfg_text_key_values_lens: torch.IntTensor | None = None,
-        cfg_text_packed_key_value_indexes: torch.LongTensor | None = None,
         # cfg_img
         cfg_img_scale: float = 1.0,
-        cfg_img_packed_query_indexes: torch.LongTensor | None = None,
         cfg_img_packed_position_ids: torch.LongTensor | None = None,
         cfg_img_past_key_values: NaiveCache | None = None,
-        cfg_img_key_values_lens: torch.IntTensor | None = None,
-        cfg_img_packed_key_value_indexes: torch.LongTensor | None = None,
         return_trajectory_latents: bool = False,
         scheduler: object | None = None,
         scheduler_kwargs: dict | None = None,
@@ -1868,25 +1704,16 @@ class Bagel(nn.Module):
                 packed_vae_token_indexes=packed_vae_token_indexes,
                 packed_seqlens=packed_seqlens,
                 packed_position_ids=packed_position_ids,
-                packed_indexes=packed_indexes,
                 past_key_values=past_key_values,
-                key_values_lens=key_values_lens,
-                packed_key_value_indexes=packed_key_value_indexes,
                 cfg_renorm_min=cfg_renorm_min,
                 cfg_renorm_type=cfg_renorm_type,
                 cfg_interval=cfg_interval,
                 cfg_text_scale=cfg_text_scale,
-                cfg_text_packed_query_indexes=cfg_text_packed_query_indexes,
                 cfg_text_packed_position_ids=cfg_text_packed_position_ids,
                 cfg_text_past_key_values=cfg_text_past_key_values,
-                cfg_text_key_values_lens=cfg_text_key_values_lens,
-                cfg_text_packed_key_value_indexes=cfg_text_packed_key_value_indexes,
                 cfg_img_scale=cfg_img_scale,
-                cfg_img_packed_query_indexes=cfg_img_packed_query_indexes,
                 cfg_img_packed_position_ids=cfg_img_packed_position_ids,
                 cfg_img_past_key_values=cfg_img_past_key_values,
-                cfg_img_key_values_lens=cfg_img_key_values_lens,
-                cfg_img_packed_key_value_indexes=cfg_img_packed_key_value_indexes,
                 return_trajectory_latents=return_trajectory_latents,
                 scheduler=scheduler,
                 scheduler_kwargs=scheduler_kwargs,
@@ -1921,31 +1748,22 @@ class Bagel(nn.Module):
 
                 v_t = self.forward_single_branch(
                     **common,
-                    packed_indexes=packed_indexes,
                     packed_position_ids=packed_position_ids,
-                    key_values_lens=key_values_lens,
                     past_key_values=past_key_values,
-                    packed_key_value_indexes=packed_key_value_indexes,
                 )
 
                 if cfg_text_scale_ > 1.0:
                     cfg_text_v_t = self.forward_single_branch(
                         **common,
-                        packed_indexes=cfg_text_packed_query_indexes,
                         packed_position_ids=cfg_text_packed_position_ids,
-                        key_values_lens=cfg_text_key_values_lens,
                         past_key_values=cfg_text_past_key_values,
-                        packed_key_value_indexes=cfg_text_packed_key_value_indexes,
                     )
                     cfg_img_v_t = None
                     if cfg_img_scale_ > 1.0:
                         cfg_img_v_t = self.forward_single_branch(
                             **common,
-                            packed_indexes=cfg_img_packed_query_indexes,
                             packed_position_ids=cfg_img_packed_position_ids,
-                            key_values_lens=cfg_img_key_values_lens,
                             past_key_values=cfg_img_past_key_values,
-                            packed_key_value_indexes=cfg_img_packed_key_value_indexes,
                         )
                     v_t = self._combine_cfg(
                         v_t,
@@ -1990,12 +1808,9 @@ class Bagel(nn.Module):
                     packed_vae_position_ids=packed_vae_position_ids,
                     packed_text_ids=packed_text_ids,
                     packed_text_indexes=packed_text_indexes,
-                    packed_indexes=packed_indexes,
                     packed_position_ids=packed_position_ids,
                     packed_seqlens=packed_seqlens,
-                    key_values_lens=key_values_lens,
                     past_key_values=past_key_values,
-                    packed_key_value_indexes=packed_key_value_indexes,
                 )
                 if scheduler is not None:
                     out = scheduler.step(v_t.to(x_t.device), timesteps[i], x_t, dts[i], **_sched_kw)
@@ -2019,48 +1834,25 @@ class Bagel(nn.Module):
             seq_len = int(packed_seqlens.sum())
 
             # Branch 0: main (gen_context), always present
-            branches_qi = [packed_indexes]
-            branches_kvi = [packed_key_value_indexes]
-            branches_kvl = [key_values_lens]
             branches_pid = [packed_position_ids]
             branches_cache = [past_key_values]
 
             # Branch 1: cfg_text (unconditional text), always present when use_cfg_text
-            branches_qi.append(cfg_text_packed_query_indexes)
-            branches_kvi.append(cfg_text_packed_key_value_indexes)
-            branches_kvl.append(cfg_text_key_values_lens)
             branches_pid.append(cfg_text_packed_position_ids)
             branches_cache.append(cfg_text_past_key_values)
 
             # Branch 2: cfg_img (text-only, no image), optional
             if use_cfg_img:
-                branches_qi.append(cfg_img_packed_query_indexes)
-                branches_kvi.append(cfg_img_packed_key_value_indexes)
-                branches_kvl.append(cfg_img_key_values_lens)
                 branches_pid.append(cfg_img_packed_position_ids)
                 branches_cache.append(cfg_img_past_key_values)
 
             num_branches = len(branches_cache)
 
-            # Compute per-branch offsets in the merged KV+Q attention tensor
-            merged_offsets = [0]
-            for b_idx in range(num_branches):
-                merged_offsets.append(merged_offsets[-1] + int(branches_kvl[b_idx].sum()) + seq_len)
-
             cfg_batched = {
                 "num_branches": num_branches,
                 "seq_len": seq_len,
                 "batched_query_lens": packed_seqlens.repeat(num_branches),
-                # Concatenate along the sequence (last) dim so this works for both
-                # scalar position ids ``(S,)`` and 3-D mRoPE ``(3, S)`` layouts.
-                "batched_position_ids": torch.cat(branches_pid, dim=-1),
-                "batched_kv_lens": torch.cat(branches_kvl),
-                "batched_query_indexes": torch.cat(
-                    [qi + merged_offsets[b_idx] for b_idx, qi in enumerate(branches_qi)]
-                ),
-                "batched_kv_indexes": torch.cat(
-                    [kvi + merged_offsets[b_idx] for b_idx, kvi in enumerate(branches_kvi)]
-                ),
+                "batched_position_ids": torch.cat(branches_pid),
                 "batched_text_indexes": torch.cat(
                     [packed_text_indexes + b_idx * seq_len for b_idx in range(num_branches)]
                 ),
@@ -2095,11 +1887,8 @@ class Bagel(nn.Module):
                 packed_text_ids=packed_text_ids,
                 packed_text_indexes=packed_text_indexes,
                 packed_position_ids=packed_position_ids,
-                packed_indexes=packed_indexes,
                 packed_seqlens=packed_seqlens,
-                key_values_lens=key_values_lens,
                 past_key_values=past_key_values,
-                packed_key_value_indexes=packed_key_value_indexes,
                 cfg_renorm_min=cfg_renorm_min,
                 cfg_renorm_type=cfg_renorm_type,
                 cfg_text_scale=cfg_text_scale_,
@@ -2138,25 +1927,16 @@ class Bagel(nn.Module):
         packed_vae_token_indexes: torch.LongTensor,
         packed_seqlens: torch.IntTensor,
         packed_position_ids: torch.LongTensor,
-        packed_indexes: torch.LongTensor,
         past_key_values: NaiveCache,
-        key_values_lens: torch.IntTensor,
-        packed_key_value_indexes: torch.LongTensor,
         cfg_renorm_min: float,
         cfg_renorm_type: str,
         cfg_interval: tuple[float, float],
         cfg_text_scale: float,
-        cfg_text_packed_query_indexes: torch.LongTensor | None,
         cfg_text_packed_position_ids: torch.LongTensor | None,
         cfg_text_past_key_values: NaiveCache | None,
-        cfg_text_key_values_lens: torch.IntTensor | None,
-        cfg_text_packed_key_value_indexes: torch.LongTensor | None,
         cfg_img_scale: float,
-        cfg_img_packed_query_indexes: torch.LongTensor | None,
         cfg_img_packed_position_ids: torch.LongTensor | None,
         cfg_img_past_key_values: NaiveCache | None,
-        cfg_img_key_values_lens: torch.IntTensor | None,
-        cfg_img_packed_key_value_indexes: torch.LongTensor | None,
         return_trajectory_latents: bool = False,
         scheduler: object | None = None,
         scheduler_kwargs: dict | None = None,
@@ -2196,24 +1976,15 @@ class Bagel(nn.Module):
         if cfg_rank == 0:
             # Gen branch: use main inputs directly
             branch_position_ids = packed_position_ids
-            branch_indexes = packed_indexes
             branch_past_key_values = past_key_values
-            branch_key_values_lens = key_values_lens
-            branch_key_value_indexes = packed_key_value_indexes
         elif cfg_rank == 1:
             # Text CFG branch
             branch_position_ids = cfg_text_packed_position_ids
-            branch_indexes = cfg_text_packed_query_indexes
             branch_past_key_values = cfg_text_past_key_values
-            branch_key_values_lens = cfg_text_key_values_lens
-            branch_key_value_indexes = cfg_text_packed_key_value_indexes
         elif cfg_rank == 2:
             # Image CFG branch
             branch_position_ids = cfg_img_packed_position_ids
-            branch_indexes = cfg_img_packed_query_indexes
             branch_past_key_values = cfg_img_past_key_values
-            branch_key_values_lens = cfg_img_key_values_lens
-            branch_key_value_indexes = cfg_img_packed_key_value_indexes
         else:
             raise RuntimeError(f"Unexpected cfg_rank={cfg_rank} for Bagel 3-branch CFG parallel")
 
@@ -2246,12 +2017,9 @@ class Bagel(nn.Module):
                     packed_vae_position_ids=packed_vae_position_ids,
                     packed_text_ids=packed_text_ids,
                     packed_text_indexes=packed_text_indexes,
-                    packed_indexes=branch_indexes,
                     packed_position_ids=branch_position_ids,
                     packed_seqlens=packed_seqlens,
-                    key_values_lens=branch_key_values_lens,
                     past_key_values=branch_past_key_values,
-                    packed_key_value_indexes=branch_key_value_indexes,
                 )
 
                 gathered = cfg_group.all_gather(local_v_t, separate_tensors=True)
@@ -2273,12 +2041,9 @@ class Bagel(nn.Module):
                     packed_vae_position_ids=packed_vae_position_ids,
                     packed_text_ids=packed_text_ids,
                     packed_text_indexes=packed_text_indexes,
-                    packed_indexes=packed_indexes,
                     packed_position_ids=packed_position_ids,
                     packed_seqlens=packed_seqlens,
-                    key_values_lens=key_values_lens,
                     past_key_values=past_key_values,
-                    packed_key_value_indexes=packed_key_value_indexes,
                 )
 
             if scheduler is not None:
@@ -2356,12 +2121,9 @@ class Bagel(nn.Module):
         packed_vae_position_ids: torch.LongTensor,
         packed_text_ids: torch.LongTensor,
         packed_text_indexes: torch.LongTensor,
-        packed_indexes: torch.LongTensor,
         packed_position_ids: torch.LongTensor,
         packed_seqlens: torch.IntTensor,
-        key_values_lens: torch.IntTensor,
         past_key_values: NaiveCache,
-        packed_key_value_indexes: torch.LongTensor,
     ) -> torch.Tensor:
         """Run a single-branch forward pass (no CFG batching).
 
@@ -2405,24 +2167,6 @@ class Bagel(nn.Module):
                 x_t_emb = x_t_emb.to(packed_sequence.dtype)
             packed_sequence[local_vae_indexes] = x_t_emb
 
-            # Build local packed_indexes for KV cache merging.
-            # In the denoising loop packed_indexes is always contiguous
-            # (arange(kv_len, kv_len + total)), so we can safely build
-            # the local slice from scratch.
-            local_total = int(local_seqlens.sum())
-            kv_len = int(key_values_lens.sum())
-            original_total = int(packed_seqlens.sum())
-            assert torch.equal(
-                packed_indexes,
-                torch.arange(kv_len, kv_len + original_total, device=packed_indexes.device, dtype=packed_indexes.dtype),
-            ), "packed_indexes must be contiguous for SP; non-contiguous layout not supported"
-            local_packed_indexes = torch.arange(
-                kv_len,
-                kv_len + local_total,
-                device=packed_indexes.device,
-                dtype=packed_indexes.dtype,
-            )
-
             extra_inputs = {}
             if self.use_moe:
                 extra_inputs["mode"] = "gen"
@@ -2433,10 +2177,7 @@ class Bagel(nn.Module):
                 packed_query_sequence=packed_sequence,
                 query_lens=local_seqlens,
                 packed_query_position_ids=local_position_ids,
-                packed_query_indexes=local_packed_indexes,
                 past_key_values=past_key_values,
-                key_values_lens=key_values_lens,
-                packed_key_value_indexes=packed_key_value_indexes,
                 update_past_key_values=False,
                 is_causal=False,
                 **extra_inputs,
@@ -2472,10 +2213,7 @@ class Bagel(nn.Module):
             packed_query_sequence=packed_sequence,
             query_lens=packed_seqlens,
             packed_query_position_ids=packed_position_ids,
-            packed_query_indexes=packed_indexes,
             past_key_values=past_key_values,
-            key_values_lens=key_values_lens,
-            packed_key_value_indexes=packed_key_value_indexes,
             update_past_key_values=False,
             is_causal=False,
             **extra_inputs,
@@ -2492,12 +2230,9 @@ class Bagel(nn.Module):
         packed_vae_position_ids: torch.LongTensor,
         packed_text_ids: torch.LongTensor,
         packed_text_indexes: torch.LongTensor,
-        packed_indexes: torch.LongTensor,
         packed_position_ids: torch.LongTensor,
         packed_seqlens: torch.IntTensor,
-        key_values_lens: torch.IntTensor,
         past_key_values: NaiveCache,
-        packed_key_value_indexes: torch.LongTensor,
         cfg_renorm_min: float = 0.0,
         cfg_renorm_type: str = "global",
         cfg_text_scale: float = 1.0,
@@ -2543,10 +2278,7 @@ class Bagel(nn.Module):
                 packed_query_sequence=batched_sequence,
                 query_lens=cfg_batched["batched_query_lens"],
                 packed_query_position_ids=cfg_batched["batched_position_ids"],
-                packed_query_indexes=cfg_batched["batched_query_indexes"],
                 past_key_values=cfg_batched["merged_cache"],
-                key_values_lens=cfg_batched["batched_kv_lens"],
-                packed_key_value_indexes=cfg_batched["batched_kv_indexes"],
                 update_past_key_values=False,
                 is_causal=False,
                 **extra_inputs,
@@ -2579,10 +2311,7 @@ class Bagel(nn.Module):
                 packed_query_sequence=packed_sequence,
                 query_lens=packed_seqlens,
                 packed_query_position_ids=packed_position_ids,
-                packed_query_indexes=packed_indexes,
                 past_key_values=past_key_values,
-                key_values_lens=key_values_lens,
-                packed_key_value_indexes=packed_key_value_indexes,
                 update_past_key_values=False,
                 is_causal=False,
                 **extra_inputs,
diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py
index a9bd897a9..137ba3b3c 100644
--- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py
+++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py
@@ -793,13 +793,7 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
                 cfg_renorm_type=gen_params.cfg_renorm_type,
                 **generation_input,
                 cfg_text_packed_position_ids=generation_input_cfg_text["cfg_packed_position_ids"],
-                cfg_text_packed_query_indexes=generation_input_cfg_text["cfg_packed_query_indexes"],
-                cfg_text_key_values_lens=generation_input_cfg_text["cfg_key_values_lens"],
-                cfg_text_packed_key_value_indexes=generation_input_cfg_text["cfg_packed_key_value_indexes"],
                 cfg_img_packed_position_ids=generation_input_cfg_img["cfg_packed_position_ids"],
-                cfg_img_packed_query_indexes=generation_input_cfg_img["cfg_packed_query_indexes"],
-                cfg_img_key_values_lens=generation_input_cfg_img["cfg_key_values_lens"],
-                cfg_img_packed_key_value_indexes=generation_input_cfg_img["cfg_packed_key_value_indexes"],
                 return_trajectory_latents=req.sampling_params.return_trajectory_latents,
                 scheduler=self.scheduler,
                 scheduler_kwargs=self.scheduler_kwargs,
diff --git a/vllm_omni/diffusion/models/dreamid_omni/fusion.py b/vllm_omni/diffusion/models/dreamid_omni/fusion.py
index a7ba745e9..423935427 100644
--- a/vllm_omni/diffusion/models/dreamid_omni/fusion.py
+++ b/vllm_omni/diffusion/models/dreamid_omni/fusion.py
@@ -1,18 +1,28 @@
 import re
+from typing import TYPE_CHECKING
 
 import torch
 import torch.nn as nn
+from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import ColumnParallelLinear
 
 from vllm_omni.diffusion.attention.layer import Attention
 
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
 try:
-    from dreamid_omni.modules.model import WanLayerNorm, WanRMSNorm
+    from dreamid_omni.modules.model import WanLayerNorm
 except ImportError:
     raise ImportError("Failed to import from dependency 'dreamid_omni'.")
 
 from vllm_omni.diffusion.distributed.utils import get_local_device
-from vllm_omni.diffusion.models.dreamid_omni.wan2_2 import WanModel, rope_apply
+from vllm_omni.diffusion.models.dreamid_omni.wan2_2 import (
+    DistributedRMSNorm,
+    WanModel,
+    rope_apply,
+)
 
 logger = init_logger(__name__)
 
@@ -128,8 +138,8 @@ class FusedBlock(nn.Module):
 
     def forward(
         self,
-        vid,
-        audio,
+        hidden_states,
+        encoder_hidden_states,
         attn: Attention,
         vid_e,
         vid_seq_lens,
@@ -148,6 +158,8 @@ class FusedBlock(nn.Module):
         audio_ref_lengths,
         audio_freqs_scaling,
     ):
+        vid = hidden_states
+        audio = encoder_hidden_states
         vid_block = self.vid_block
         audio_block = self.audio_block
 
@@ -241,26 +253,36 @@ class FusedBlock(nn.Module):
 class FusionModel(nn.Module):
     _layerwise_offload_blocks_attrs = ["fused_blocks"]
 
+    packed_modules_mapping = {
+        "to_qkv": ["q", "k", "v"],
+    }
+
     @staticmethod
     def _is_fused_block(name: str, module) -> bool:
         return "fused_blocks" in name and name.split(".")[-1].isdigit()
 
     _hsdp_shard_conditions = [_is_fused_block]
 
-    def __init__(self, video_config=None, audio_config=None):
+    def __init__(
+        self,
+        video_config=None,
+        audio_config=None,
+        *,
+        quant_config: "QuantizationConfig | None" = None,
+    ):
         super().__init__()
         has_video = True
         has_audio = True
         self.device = get_local_device()
         if video_config is not None:
-            self.video_model = WanModel(**video_config)
+            self.video_model = WanModel(quant_config=quant_config, prefix="video_model", **video_config)
         else:
             has_video = False
             self.video_model = None
             logger.warning("No video model is provided!")
 
         if audio_config is not None:
-            self.audio_model = WanModel(**audio_config)
+            self.audio_model = WanModel(quant_config=quant_config, prefix="audio_model", **audio_config)
         else:
             has_audio = False
             self.audio_model = None
@@ -270,11 +292,12 @@ class FusionModel(nn.Module):
             assert len(self.video_model.blocks) == len(self.audio_model.blocks)
             self.num_blocks = len(self.video_model.blocks)
 
-            self.inject_cross_attention_kv_projections()
+            self.inject_cross_attention_kv_projections(quant_config=quant_config)
 
-        self.num_heads = self.video_model.num_heads
+        tp_size = get_tensor_model_parallel_world_size()
+        self.full_num_heads = self.video_model.num_heads
+        self.num_heads = self.full_num_heads // tp_size
         self.head_dim = self.video_model.dim // self.video_model.num_heads
-        # Make a single shared instance to pass in at forward time
         self.attn = Attention(
             num_heads=self.num_heads,
             head_size=self.head_dim,
@@ -313,21 +336,42 @@ class FusionModel(nn.Module):
 
         return super().load_state_dict(state_dict, strict=strict, assign=assign)
 
-    def inject_cross_attention_kv_projections(self):
-        for vid_block in self.video_model.blocks:
-            vid_block.cross_attn.k_fusion = nn.Linear(vid_block.dim, vid_block.dim)
-            vid_block.cross_attn.v_fusion = nn.Linear(vid_block.dim, vid_block.dim)
-            vid_block.cross_attn.pre_attn_norm_fusion = WanLayerNorm(vid_block.dim, elementwise_affine=True)
-            vid_block.cross_attn.norm_k_fusion = (
-                WanRMSNorm(vid_block.dim, eps=1e-6) if vid_block.qk_norm else nn.Identity()
+    def inject_cross_attention_kv_projections(
+        self,
+        quant_config: "QuantizationConfig | None" = None,
+    ):
+        def _make_kv_proj(dim: int, prefix: str) -> ColumnParallelLinear:
+            # K/V projections for the cross-modal (a2v / v2a) attention path.
+            # Inputs are full latents, so we quantize matched to self_attn.
+            return ColumnParallelLinear(
+                dim,
+                dim,
+                bias=True,
+                gather_output=False,
+                return_bias=False,
+                quant_config=quant_config,
+                prefix=prefix,
             )
 
-        for audio_block in self.audio_model.blocks:
-            audio_block.cross_attn.k_fusion = nn.Linear(audio_block.dim, audio_block.dim)
-            audio_block.cross_attn.v_fusion = nn.Linear(audio_block.dim, audio_block.dim)
-            audio_block.cross_attn.pre_attn_norm_fusion = WanLayerNorm(audio_block.dim, elementwise_affine=True)
-            audio_block.cross_attn.norm_k_fusion = (
-                WanRMSNorm(audio_block.dim, eps=1e-6) if audio_block.qk_norm else nn.Identity()
+        tp_size = get_tensor_model_parallel_world_size()
+        for i, vid_block in enumerate(self.video_model.blocks):
+            cross = vid_block.cross_attn
+            block_prefix = f"video_model.blocks.{i}.cross_attn"
+            cross.k_fusion = _make_kv_proj(vid_block.dim, f"{block_prefix}.k_fusion")
+            cross.v_fusion = _make_kv_proj(vid_block.dim, f"{block_prefix}.v_fusion")
+            cross.pre_attn_norm_fusion = WanLayerNorm(vid_block.dim, elementwise_affine=True)
+            cross.norm_k_fusion = (
+                DistributedRMSNorm(vid_block.dim // tp_size, eps=1e-6) if vid_block.qk_norm else nn.Identity()
+            )
+
+        for i, audio_block in enumerate(self.audio_model.blocks):
+            cross = audio_block.cross_attn
+            block_prefix = f"audio_model.blocks.{i}.cross_attn"
+            cross.k_fusion = _make_kv_proj(audio_block.dim, f"{block_prefix}.k_fusion")
+            cross.v_fusion = _make_kv_proj(audio_block.dim, f"{block_prefix}.v_fusion")
+            cross.pre_attn_norm_fusion = WanLayerNorm(audio_block.dim, elementwise_affine=True)
+            cross.norm_k_fusion = (
+                DistributedRMSNorm(audio_block.dim // tp_size, eps=1e-6) if audio_block.qk_norm else nn.Identity()
             )
 
     def _detach_blocks_from_backbones(self) -> None:
diff --git a/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py b/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py
index cc932f8c1..47b905de7 100644
--- a/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py
+++ b/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py
@@ -13,6 +13,11 @@ from PIL import Image, ImageOps
 from torch import nn
 from torchvision.transforms import Compose, Normalize
 from tqdm import tqdm
+from vllm.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
@@ -124,7 +129,8 @@ class DreamIDOmniPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Suppor
         self.text_encoder = self.text_model.model
 
         # Fusion model — weights are loaded later via load_weights()
-        self.model = FusionModel(VIDEO_CONFIG, AUDIO_CONFIG)
+        quant_config = getattr(self.od_config, "quantization_config", None)
+        self.model = FusionModel(VIDEO_CONFIG, AUDIO_CONFIG, quant_config=quant_config)
         self.transformer = self.model
 
         fusion_path = self.od_config.tf_model_config.get("fusion", None)
@@ -236,10 +242,75 @@ class DreamIDOmniPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Suppor
         return ref_vae_latents, ref_audio_lengths
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        prefix = "model."
-        state_dict = {name[len(prefix) :]: tensor for name, tensor in weights if name.startswith(prefix)}
-        self.model.load_state_dict(state_dict, strict=True)
-        return {prefix + k for k in state_dict}
+        pipeline_prefix = "model."
+
+        qkv_stacked_mapping = [
+            (".self_attn.to_qkv", ".self_attn.q", "q"),
+            (".self_attn.to_qkv", ".self_attn.k", "k"),
+            (".self_attn.to_qkv", ".self_attn.v", "v"),
+        ]
+
+        params_dict = dict(self.model.named_parameters())
+        loaded_local: set[str] = set()
+
+        tp_size = get_tensor_model_parallel_world_size()
+        tp_rank = get_tensor_model_parallel_rank() if tp_size > 1 else 0
+
+        def _maybe_shard_weight(weight: torch.Tensor, param: torch.Tensor) -> torch.Tensor:
+            """Slice a checkpoint tensor to this rank's TP shard when the
+            param is a TP-local shape but the checkpoint stores the full one.
+            """
+            if tp_size <= 1 or weight.shape == param.shape:
+                return weight
+            if weight.ndim == 1 and weight.numel() == param.numel() * tp_size:
+                return weight.chunk(tp_size, dim=0)[tp_rank]
+            if weight.ndim == 2:
+                if weight.shape[0] == param.shape[0] * tp_size:
+                    return weight.chunk(tp_size, dim=0)[tp_rank]
+                if weight.shape[1] == param.shape[1] * tp_size:
+                    return weight.chunk(tp_size, dim=1)[tp_rank]
+            return weight
+
+        for full_name, loaded_weight in weights:
+            if not full_name.startswith(pipeline_prefix):
+                continue
+            name = full_name[len(pipeline_prefix) :]
+
+            # Fused-QKV stacking: route checkpoint's separate
+            # - ...self_attn.{q,k,v}: weights into the corresponding
+            # - ...self_attn.to_qkv : via its sharded weight_loader.
+            stacked_hit = False
+            for fused_name, raw_name, shard_id in qkv_stacked_mapping:
+                if raw_name not in name:
+                    continue
+                stacked = name.replace(raw_name, fused_name)
+                if stacked not in params_dict:
+                    continue
+                param = params_dict[stacked]
+                param.weight_loader(param, loaded_weight, shard_id)
+                loaded_local.add(stacked)
+                stacked_hit = True
+                break
+            if stacked_hit:
+                continue
+
+            if name not in params_dict:
+                logger.warning(f"Skipping DreamID-Omni weight {full_name} -- not found in model parameters")
+                continue
+            param = params_dict[name]
+            weight_loader = getattr(param, "weight_loader", None)
+            if weight_loader is not None:
+                # Parallel Linears handle TP sharding internally
+                weight_loader(param, loaded_weight)
+            else:
+                default_weight_loader(param, _maybe_shard_weight(loaded_weight, param))
+            loaded_local.add(name)
+
+        # Detach must happen before `DiffusersPipelineLoader._process_weights_after_loading`
+        # and before any forward pass that triggers layerwise-offload hooks.
+        self.model._detach_blocks_from_backbones()
+
+        return {pipeline_prefix + k for k in loaded_local}
 
     def get_scheduler_time_steps(self, sampling_steps, solver_name="unipc", device=0, shift=5.0):
         torch.manual_seed(4)
@@ -375,6 +446,33 @@ class DreamIDOmniPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Suppor
             )[0].squeeze(0)
         return video_noise, audio_noise
 
+    def encode_prompt(
+        self,
+        prompt: str,
+        video_negative_prompt: str = "",
+        audio_negative_prompt: str = "",
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Encode the positive and negative prompts via ``self.text_model``.
+
+        This is the single text-encoder entrypoint so that the runner-level
+        prompt-embedding cache (see
+        ``vllm_omni/diffusion/cache/prompt_embed_cache.py``) can transparently
+        memoize results when the same prompts are submitted repeatedly (e.g.
+        GRPO rollouts that sample the same prompt with different seeds).
+
+        Returns:
+            (audio_pos, video_pos, video_neg, audio_neg) embeddings cast to
+            ``self.target_dtype``. ``audio_pos`` and ``video_pos`` are the
+            same tensor, matching the original inline behavior.
+        """
+        text_embeddings = self.text_model(
+            [prompt, video_negative_prompt, audio_negative_prompt],
+            device=self.device,
+        )
+        text_embeddings = [emb.to(self.target_dtype) for emb in text_embeddings]
+        # Index 0 is the positive embedding, shared between video and audio branches.
+        return text_embeddings[0], text_embeddings[0], text_embeddings[1], text_embeddings[2]
+
     def forward(
         self,
         request: OmniDiffusionRequest,
@@ -432,12 +530,13 @@ class DreamIDOmniPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Suppor
         )
 
         # 3. text embedding
-        text_embeddings = self.text_model([prompt, video_negative_prompt, audio_negative_prompt], device=self.device)
-        text_embeddings = [emb.to(self.target_dtype) for emb in text_embeddings]
-        text_embeddings_audio_pos = text_embeddings[0]
-        text_embeddings_video_pos = text_embeddings[0]
-        text_embeddings_video_neg = text_embeddings[1]
-        text_embeddings_audio_neg = text_embeddings[2]
+        text_embeddings_audio_pos, text_embeddings_video_pos, text_embeddings_video_neg, text_embeddings_audio_neg = (
+            self.encode_prompt(
+                prompt=prompt,
+                video_negative_prompt=video_negative_prompt,
+                audio_negative_prompt=audio_negative_prompt,
+            )
+        )
 
         video_latent_h, video_latent_w = height // 16, width // 16
 
diff --git a/vllm_omni/diffusion/models/dreamid_omni/wan2_2.py b/vllm_omni/diffusion/models/dreamid_omni/wan2_2.py
index db289daf9..c7e464cc8 100644
--- a/vllm_omni/diffusion/models/dreamid_omni/wan2_2.py
+++ b/vllm_omni/diffusion/models/dreamid_omni/wan2_2.py
@@ -6,6 +6,16 @@ import torch.amp as amp
 import torch.nn as nn
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.models.modeling_utils import ModelMixin
+from vllm.distributed import (
+    get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_reduce,
+)
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
 from vllm_omni.diffusion.attention.layer import Attention
 
@@ -17,7 +27,6 @@ try:
         MLPProj,
         ModulationAdd,
         WanLayerNorm,
-        WanRMSNorm,
         rope_apply,  # diff from  wan2.2
         rope_params,
         sinusoidal_embedding_1d,
@@ -26,42 +35,99 @@ except ImportError:
     raise ImportError("Failed to import from dependency 'dreamid_omni'.")
 
 
+class DistributedRMSNorm(nn.Module):
+    """RMSNorm that computes global RMS across tensor parallel ranks.
+
+    Mirrors `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py::DistributedRMSNorm`
+    """
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        tp_size = get_tensor_model_parallel_world_size()
+        x_dtype = x.dtype
+        x_float = x.float()
+        local_sum_sq = x_float.pow(2).sum(dim=-1, keepdim=True)
+        local_count = x.shape[-1]
+
+        if tp_size > 1:
+            local_sum_sq = tensor_model_parallel_all_reduce(local_sum_sq)
+            local_count = local_count * tp_size
+
+        rms = torch.sqrt(local_sum_sq / local_count + self.eps)
+        return (x_float / rms * self.weight.float()).to(x_dtype)
+
+
 class WanSelfAttention(nn.Module):
     """Optimized self-attention module using vLLM layers."""
 
-    def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        window_size=(-1, -1),
+        qk_norm=True,
+        eps=1e-6,
+        *,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
         assert dim % num_heads == 0
         super().__init__()
         self.dim = dim
-        self.num_heads = num_heads
         self.head_dim = dim // num_heads
         self.window_size = window_size
         self.qk_norm = qk_norm
         self.eps = eps
 
-        # layers
-        self.q = nn.Linear(dim, dim)
-        self.k = nn.Linear(dim, dim)
-        self.v = nn.Linear(dim, dim)
-        self.o = nn.Linear(dim, dim)
-        self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+        self.to_qkv = QKVParallelLinear(
+            hidden_size=dim,
+            head_size=self.head_dim,
+            total_num_heads=num_heads,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_qkv" if prefix else "to_qkv",
+        )
+        self.num_heads = self.to_qkv.num_heads
+        self.num_kv_heads = self.to_qkv.num_kv_heads
+        self.tp_inner_dim = self.num_heads * self.head_dim
+        self.tp_kv_dim = self.num_kv_heads * self.head_dim
+
+        self.o = RowParallelLinear(
+            dim,
+            dim,
+            bias=True,
+            input_is_parallel=True,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o" if prefix else "o",
+        )
+        self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = DistributedRMSNorm(self.tp_kv_dim, eps=eps) if qk_norm else nn.Identity()
         # Unified attention layer
         self.attn = Attention(
             num_heads=self.num_heads,
             head_size=self.head_dim,
-            num_kv_heads=self.num_heads,
+            num_kv_heads=self.num_kv_heads,
             softmax_scale=1.0 / (self.head_dim**0.5),
             causal=False,
+            prefix=prefix,
         )
 
     # query, key, value function
     def qkv_fn(self, x):
-        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+        b, s = x.shape[:2]
+        qkv, _ = self.to_qkv(x)
+        q_size = self.num_heads * self.head_dim
+        kv_size = self.num_kv_heads * self.head_dim
+        q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
 
-        q = self.norm_q(self.q(x)).view(b, s, n, d)
-        k = self.norm_k(self.k(x)).view(b, s, n, d)
-        v = self.v(x).view(b, s, n, d)
+        q = self.norm_q(q).view(b, s, self.num_heads, self.head_dim)
+        k = self.norm_k(k).view(b, s, self.num_kv_heads, self.head_dim)
+        v = v.view(b, s, self.num_kv_heads, self.head_dim)
         return q, k, v
 
     def forward(self, x, seq_lens, grid_sizes, freqs, ref_lengths=None, freqs_scaling=None):
@@ -87,14 +153,97 @@ class WanSelfAttention(nn.Module):
         return x
 
 
-class WanT2VCrossAttention(WanSelfAttention):
-    def qkv_fn(self, x, context):
-        b, n, d = x.size(0), self.num_heads, self.head_dim
+class WanT2VCrossAttention(nn.Module):
+    """Text -> latent cross-attention."""
 
-        # compute query, key, value
-        q = self.norm_q(self.q(x)).view(b, -1, n, d)
-        k = self.norm_k(self.k(context)).view(b, -1, n, d)
-        v = self.v(context).view(b, -1, n, d)
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        window_size=(-1, -1),
+        qk_norm=True,
+        eps=1e-6,
+        *,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        assert dim % num_heads == 0
+        self.dim = dim
+        self.head_dim = dim // num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.eps = eps
+
+        # Separate Q/K/V — Q's input differs from K/V's input (text context)
+        self.q = ColumnParallelLinear(
+            dim,
+            dim,
+            bias=True,
+            gather_output=False,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.q" if prefix else "q",
+        )
+        self.k = ColumnParallelLinear(
+            dim,
+            dim,
+            bias=True,
+            gather_output=False,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.k" if prefix else "k",
+        )
+        self.v = ColumnParallelLinear(
+            dim,
+            dim,
+            bias=True,
+            gather_output=False,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.v" if prefix else "v",
+        )
+
+        tp_size = get_tensor_model_parallel_world_size()
+        self.num_heads = num_heads // tp_size
+        self.num_kv_heads = self.num_heads
+        self.tp_inner_dim = self.num_heads * self.head_dim
+
+        self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
+
+        self.o = RowParallelLinear(
+            dim,
+            dim,
+            bias=True,
+            input_is_parallel=True,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o" if prefix else "o",
+        )
+
+        # Kept for interface parity with WanSelfAttention;
+        # the active cross-attention path in DreamID-Omni uses the shared FusionModel.attn.
+        # Disable kv quant to prevent affecting quality.
+        self.attn = Attention(
+            num_heads=self.num_heads,
+            head_size=self.head_dim,
+            num_kv_heads=self.num_kv_heads,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=False,
+            prefix=prefix,
+            disable_kv_quant=True,
+        )
+
+    def qkv_fn(self, x, context):
+        b = x.size(0)
+
+        q = self.q(x)
+        k = self.k(context)
+        v = self.v(context)
+        q = self.norm_q(q).view(b, -1, self.num_heads, self.head_dim)
+        k = self.norm_k(k).view(b, -1, self.num_kv_heads, self.head_dim)
+        v = v.view(b, -1, self.num_kv_heads, self.head_dim)
 
         return q, k, v
 
@@ -142,6 +291,9 @@ class WanAttentionBlock(nn.Module):
         cross_attn_norm=False,
         eps=1e-6,
         additional_emb_length=None,
+        *,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.dim = dim
@@ -154,7 +306,15 @@ class WanAttentionBlock(nn.Module):
 
         # layers
         self.norm1 = WanLayerNorm(dim, eps)
-        self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, eps)
+        self.self_attn = WanSelfAttention(
+            dim,
+            num_heads,
+            window_size,
+            qk_norm,
+            eps,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn" if prefix else "self_attn",
+        )
         self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
         if cross_attn_type == "i2v_cross_attn":
             assert additional_emb_length is not None, "additional_emb_length should be specified for i2v_cross_attn"
@@ -167,9 +327,32 @@ class WanAttentionBlock(nn.Module):
                 (-1, -1),
                 qk_norm,
                 eps,
+                quant_config=quant_config,
+                prefix=f"{prefix}.cross_attn" if prefix else "cross_attn",
             )
         self.norm2 = WanLayerNorm(dim, eps)
-        self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim))
+        # FFN layout: Stay compatible with upstream DreamID-Omni checkpoints
+        self.ffn = nn.Sequential(
+            ColumnParallelLinear(
+                dim,
+                ffn_dim,
+                bias=True,
+                gather_output=False,
+                return_bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.ffn.0" if prefix else "ffn.0",
+            ),
+            nn.GELU(approximate="tanh"),
+            RowParallelLinear(
+                ffn_dim,
+                dim,
+                bias=True,
+                input_is_parallel=True,
+                return_bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.ffn.2" if prefix else "ffn.2",
+            ),
+        )
 
         # modulation
         self.modulation = ModulationAdd(dim, 6)
@@ -243,6 +426,9 @@ class WanModel(ModelMixin, ConfigMixin):
         cross_attn_norm=True,
         temporal_rope_scaling_factor=1.0,
         eps=1e-6,
+        *,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
     ):
         r"""
         Initialize the diffusion model backbone.
@@ -357,8 +543,10 @@ class WanModel(ModelMixin, ConfigMixin):
                     cross_attn_norm,
                     eps,
                     additional_emb_length,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.blocks.{layer_idx}" if prefix else f"blocks.{layer_idx}",
                 )
-                for _ in range(num_layers)
+                for layer_idx in range(num_layers)
             ]
         )
 
diff --git a/vllm_omni/diffusion/models/ernie_image/pipeline_ernie_image.py b/vllm_omni/diffusion/models/ernie_image/pipeline_ernie_image.py
index 7b5d80571..cd60cc93b 100644
--- a/vllm_omni/diffusion/models/ernie_image/pipeline_ernie_image.py
+++ b/vllm_omni/diffusion/models/ernie_image/pipeline_ernie_image.py
@@ -247,8 +247,10 @@ class ErnieImagePipeline(
 
     @staticmethod
     def _is_warmup_request(req: OmniDiffusionRequest) -> bool:
-        request_ids = getattr(req, "request_ids", None) or ()
-        return len(request_ids) == 1 and request_ids[0] == "dummy_req_id"
+        is_dummy_run = getattr(req, "is_dummy_run", None)
+        if callable(is_dummy_run):
+            return bool(is_dummy_run())
+        return OmniDiffusionRequest.is_dummy_run_request_id(getattr(req, "request_id", None))
 
     @staticmethod
     def _should_apply_pe(req: OmniDiffusionRequest) -> bool:
diff --git a/vllm_omni/diffusion/models/flux2/flux2_transformer.py b/vllm_omni/diffusion/models/flux2/flux2_transformer.py
index ef2cdd025..0c6896c13 100644
--- a/vllm_omni/diffusion/models/flux2/flux2_transformer.py
+++ b/vllm_omni/diffusion/models/flux2/flux2_transformer.py
@@ -25,7 +25,13 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.layers.rope import RotaryEmbedding
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput,
+    SequenceParallelOutput,
+)
+from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
@@ -91,6 +97,7 @@ class Flux2FeedForward(nn.Module):
 class Flux2Attention(nn.Module):
     def __init__(
         self,
+        parallel_config: DiffusionParallelConfig,
         query_dim: int,
         heads: int = 8,
         dim_head: int = 64,
@@ -106,6 +113,7 @@ class Flux2Attention(nn.Module):
         prefix: str = "",
     ):
         super().__init__()
+        self.parallel_config = parallel_config
         self.head_dim = dim_head
         self.inner_dim = out_dim if out_dim is not None else dim_head * heads
         self.query_dim = query_dim
@@ -206,34 +214,82 @@ class Flux2Attention(nn.Module):
             encoder_query = self.norm_added_q(encoder_query)
             encoder_key = self.norm_added_k(encoder_key)
 
-            query = torch.cat([encoder_query, query], dim=1)
-            key = torch.cat([encoder_key, key], dim=1)
-            value = torch.cat([encoder_value, value], dim=1)
-
-        if image_rotary_emb is not None:
-            cos, sin = image_rotary_emb
-            cos = cos.to(query.dtype)
-            sin = sin.to(query.dtype)
-            query = self.rope(query, cos, sin)
-            key = self.rope(key, cos, sin)
-
-        attn_metadata = None
-        if attention_mask is not None:
-            if attention_mask.dim() == 3:
-                attention_mask = attention_mask.unsqueeze(1)
-            attn_metadata = AttentionMetadata(attn_mask=attention_mask)
-
-        hidden_states = self.attn(query, key, value, attn_metadata)
-        hidden_states = hidden_states.flatten(2, 3).to(query.dtype)
-
-        if encoder_hidden_states is not None:
-            context_len = encoder_hidden_states.shape[1]
-            encoder_hidden_states, hidden_states = hidden_states.split_with_sizes(
-                [context_len, hidden_states.shape[1] - context_len],
-                dim=1,
+            sp_size = self.parallel_config.sequence_parallel_size
+            use_sp_joint_attention = (
+                sp_size is not None
+                and sp_size > 1
+                and is_forward_context_available()
+                and not get_forward_context().split_text_embed_in_sp
             )
-            # Contiguous for FP8 quantization in RowParallelLinear
-            encoder_hidden_states = self.to_add_out(encoder_hidden_states.contiguous())
+
+            if use_sp_joint_attention and image_rotary_emb is not None:
+                cos, sin = image_rotary_emb
+                cos = cos.to(query.dtype)
+                sin = sin.to(query.dtype)
+                txt_len = encoder_query.shape[1]
+                txt_cos, img_cos = cos[:txt_len], cos[txt_len:]
+                txt_sin, img_sin = sin[:txt_len], sin[txt_len:]
+
+                query = self.rope(query, img_cos, img_sin)
+                key = self.rope(key, img_cos, img_sin)
+                encoder_query = self.rope(encoder_query, txt_cos, txt_sin)
+                encoder_key = self.rope(encoder_key, txt_cos, txt_sin)
+
+                attn_metadata = AttentionMetadata(
+                    joint_query=encoder_query,
+                    joint_key=encoder_key,
+                    joint_value=encoder_value,
+                    joint_strategy="front",
+                )
+                hidden_states_mask: torch.Tensor | None = kwargs.get("hidden_states_mask", None)
+                encoder_hidden_states_mask: torch.Tensor | None = kwargs.get("encoder_hidden_states_mask", None)
+                if hidden_states_mask is not None:
+                    attn_metadata.attn_mask = hidden_states_mask
+                if encoder_hidden_states_mask is not None:
+                    attn_metadata.joint_attn_mask = encoder_hidden_states_mask
+
+                hidden_states = self.attn(query, key, value, attn_metadata)
+                hidden_states = hidden_states.flatten(2, 3).to(query.dtype)
+
+                txt_len = encoder_hidden_states.shape[1]
+                encoder_hidden_states, hidden_states = hidden_states.split_with_sizes(
+                    [txt_len, hidden_states.shape[1] - txt_len],
+                    dim=1,
+                )
+                encoder_hidden_states = self.to_add_out(encoder_hidden_states.contiguous())
+            else:
+                query = torch.cat([encoder_query, query], dim=1)
+                key = torch.cat([encoder_key, key], dim=1)
+                value = torch.cat([encoder_value, value], dim=1)
+
+                query, key = apply_rope_to_qk(self.rope, query, key, image_rotary_emb)
+
+                attn_metadata = None
+                if attention_mask is not None:
+                    if attention_mask.dim() == 3:
+                        attention_mask = attention_mask.unsqueeze(1)
+                    attn_metadata = AttentionMetadata(attn_mask=attention_mask)
+
+                hidden_states = self.attn(query, key, value, attn_metadata)
+                hidden_states = hidden_states.flatten(2, 3).to(query.dtype)
+
+                context_len = encoder_hidden_states.shape[1]
+                encoder_hidden_states, hidden_states = hidden_states.split_with_sizes(
+                    [context_len, hidden_states.shape[1] - context_len],
+                    dim=1,
+                )
+                encoder_hidden_states = self.to_add_out(encoder_hidden_states.contiguous())
+        else:
+            query, key = apply_rope_to_qk(self.rope, query, key, image_rotary_emb)
+
+            attn_metadata = None
+            if attention_mask is not None:
+                if attention_mask.dim() == 3:
+                    attention_mask = attention_mask.unsqueeze(1)
+                attn_metadata = AttentionMetadata(attn_mask=attention_mask)
+
+            hidden_states = self.attn(query, key, value, attn_metadata)
+            hidden_states = hidden_states.flatten(2, 3).to(query.dtype)
 
         hidden_states = self.to_out[0](hidden_states.contiguous())
         hidden_states = self.to_out[1](hidden_states)
@@ -250,6 +306,7 @@ class Flux2ParallelSelfAttention(nn.Module):
 
     def __init__(
         self,
+        parallel_config: DiffusionParallelConfig,
         query_dim: int,
         heads: int = 8,
         dim_head: int = 64,
@@ -265,6 +322,7 @@ class Flux2ParallelSelfAttention(nn.Module):
         prefix: str = "",
     ):
         super().__init__()
+        self.parallel_config = parallel_config
         self.head_dim = dim_head
         self.inner_dim = out_dim if out_dim is not None else dim_head * heads
         self.query_dim = query_dim
@@ -327,20 +385,59 @@ class Flux2ParallelSelfAttention(nn.Module):
         query = self.norm_q(query)
         key = self.norm_k(key)
 
-        if image_rotary_emb is not None:
+        sp_size = self.parallel_config.sequence_parallel_size
+        text_seq_len = kwargs.get("text_seq_len", None)
+        use_sp_single_stream = (
+            sp_size is not None
+            and sp_size > 1
+            and text_seq_len is not None
+            and is_forward_context_available()
+            and not get_forward_context().split_text_embed_in_sp
+        )
+
+        if use_sp_single_stream and image_rotary_emb is not None:
             cos, sin = image_rotary_emb
             cos = cos.to(query.dtype)
             sin = sin.to(query.dtype)
-            query = self.rope(query, cos, sin)
-            key = self.rope(key, cos, sin)
+            txt_cos, img_cos = cos[:text_seq_len], cos[text_seq_len:]
+            txt_sin, img_sin = sin[:text_seq_len], sin[text_seq_len:]
 
-        attn_metadata = None
-        if attention_mask is not None:
-            if attention_mask.dim() == 3:
-                attention_mask = attention_mask.unsqueeze(1)
-            attn_metadata = AttentionMetadata(attn_mask=attention_mask)
+            img_query = query[:, text_seq_len:]
+            img_key = key[:, text_seq_len:]
+            img_value = value[:, text_seq_len:]
+            text_query = query[:, :text_seq_len]
+            text_key = key[:, :text_seq_len]
+            text_value = value[:, :text_seq_len]
 
-        attn_output = self.attn(query, key, value, attn_metadata)
+            img_query = self.rope(img_query, img_cos, img_sin)
+            img_key = self.rope(img_key, img_cos, img_sin)
+            text_query = self.rope(text_query, txt_cos, txt_sin)
+            text_key = self.rope(text_key, txt_cos, txt_sin)
+
+            attn_metadata = AttentionMetadata(
+                joint_query=text_query,
+                joint_key=text_key,
+                joint_value=text_value,
+                joint_strategy="front",
+            )
+            hidden_states_mask: torch.Tensor | None = kwargs.get("hidden_states_mask", None)
+            encoder_hidden_states_mask: torch.Tensor | None = kwargs.get("encoder_hidden_states_mask", None)
+            if hidden_states_mask is not None:
+                attn_metadata.attn_mask = hidden_states_mask
+            if encoder_hidden_states_mask is not None:
+                attn_metadata.joint_attn_mask = encoder_hidden_states_mask
+
+            attn_output = self.attn(img_query, img_key, img_value, attn_metadata)
+        else:
+            query, key = apply_rope_to_qk(self.rope, query, key, image_rotary_emb)
+
+            attn_metadata = None
+            if attention_mask is not None:
+                if attention_mask.dim() == 3:
+                    attention_mask = attention_mask.unsqueeze(1)
+                attn_metadata = AttentionMetadata(attn_mask=attention_mask)
+
+            attn_output = self.attn(query, key, value, attn_metadata)
         attn_output = attn_output.flatten(2, 3).to(query.dtype)
 
         mlp_hidden_states = self.mlp_act_fn(mlp_hidden_states)
@@ -352,6 +449,7 @@ class Flux2ParallelSelfAttention(nn.Module):
 class Flux2SingleTransformerBlock(nn.Module):
     def __init__(
         self,
+        parallel_config: DiffusionParallelConfig,
         dim: int,
         num_attention_heads: int,
         attention_head_dim: int,
@@ -364,6 +462,7 @@ class Flux2SingleTransformerBlock(nn.Module):
         super().__init__()
         self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
         self.attn = Flux2ParallelSelfAttention(
+            parallel_config=parallel_config,
             query_dim=dim,
             dim_head=attention_head_dim,
             heads=num_attention_heads,
@@ -387,6 +486,13 @@ class Flux2SingleTransformerBlock(nn.Module):
         split_hidden_states: bool = False,
         text_seq_len: int | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        """
+        Forward pass for Flux2SingleTransformerBlock with SP support.
+
+        In SP mode: image hidden_states is chunked (B, img_len/SP, D),
+        text encoder_hidden_states is full (B, txt_len, D).
+        The block concatenates them for joint attention.
+        """
         if encoder_hidden_states is not None:
             text_seq_len = encoder_hidden_states.shape[1]
             hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
@@ -400,6 +506,7 @@ class Flux2SingleTransformerBlock(nn.Module):
         attn_output = self.attn(
             hidden_states=norm_hidden_states,
             image_rotary_emb=image_rotary_emb,
+            text_seq_len=text_seq_len,
             **joint_attention_kwargs,
         )
 
@@ -416,6 +523,7 @@ class Flux2SingleTransformerBlock(nn.Module):
 class Flux2TransformerBlock(nn.Module):
     def __init__(
         self,
+        parallel_config: DiffusionParallelConfig,
         dim: int,
         num_attention_heads: int,
         attention_head_dim: int,
@@ -432,6 +540,7 @@ class Flux2TransformerBlock(nn.Module):
         self.norm1_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
 
         self.attn = Flux2Attention(
+            parallel_config=parallel_config,
             query_dim=dim,
             added_kv_proj_dim=dim,
             dim_head=attention_head_dim,
@@ -543,6 +652,47 @@ class Flux2PosEmbed(nn.Module):
         return freqs_cos, freqs_sin
 
 
+class Flux2RopePrepare(nn.Module):
+    """Prepares RoPE embeddings for sequence parallel.
+
+    This module encapsulates the RoPE computation for Flux.2-dev.
+    For dual-stream attention, text components (outputs 0, 1) are replicated
+    across SP ranks, while image components (outputs 2, 3) are sharded.
+
+    NOTE: The hidden_states projection is handled separately in forward()
+    so that _sp_plan can shard it at the root level.
+    """
+
+    def __init__(self, pos_embed: Flux2PosEmbed):
+        super().__init__()
+        self.pos_embed = pos_embed
+
+    def forward(
+        self,
+        img_ids: torch.Tensor,
+        txt_ids: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Compute RoPE embeddings for text and image sequences.
+
+        Args:
+            img_ids: Image position IDs (img_seq_len, n_axes)
+            txt_ids: Text position IDs (txt_seq_len, n_axes)
+
+        Returns:
+            Tuple of cosine / sine components for text & image
+            in the order: (txt_cos, txt_sin, img_cos, img_sin)
+
+        NOTE: careful about output orders if this is refactored in the
+        future; we need to match the _sp_plan indices, since text
+        components (0 & 1) need to be replicated across SP ranks,
+        while image components (2 & 3) must be sharded.
+        """
+        img_freqs_cos, img_freqs_sin = self.pos_embed(img_ids)
+        txt_freqs_cos, txt_freqs_sin = self.pos_embed(txt_ids)
+        return txt_freqs_cos, txt_freqs_sin, img_freqs_cos, img_freqs_sin
+
+
 class Flux2TimestepGuidanceEmbeddings(nn.Module):
     def __init__(
         self,
@@ -601,9 +751,22 @@ class Flux2Modulation(nn.Module):
 class Flux2Transformer2DModel(nn.Module):
     """
     The Transformer model introduced in Flux 2.
+
+    Supports Sequence Parallelism (Ulysses and Ring) when configured via OmniDiffusionConfig.
     """
 
     _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"]
+    _sp_plan = {
+        "": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True),
+        },
+        "rope_prepare": {
+            2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True),
+            3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True),
+        },
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+    """SP plan: shard hidden_states at root level, shard img_freqs at rope_prepare, gather output at proj_out."""
 
     @staticmethod
     def _is_transformer_block(name: str, module) -> bool:
@@ -628,6 +791,7 @@ class Flux2Transformer2DModel(nn.Module):
         eps: float = 1e-6,
         guidance_embeds: bool = True,
         quant_config: "QuantizationConfig | None" = None,
+        od_config: OmniDiffusionConfig | None = None,
     ):
         super().__init__()
         self.guidance_embeds = guidance_embeds
@@ -652,7 +816,13 @@ class Flux2Transformer2DModel(nn.Module):
             guidance_embeds=guidance_embeds,
         )
 
+        if od_config is not None:
+            self.parallel_config = od_config.parallel_config
+        else:
+            self.parallel_config = DiffusionParallelConfig()
+
         self.pos_embed = Flux2PosEmbed(theta=rope_theta, axes_dim=axes_dims_rope)
+        self.rope_prepare = Flux2RopePrepare(self.pos_embed)
         self.time_guidance_embed = Flux2TimestepGuidanceEmbeddings(
             in_channels=timestep_guidance_channels,
             embedding_dim=self.inner_dim,
@@ -670,6 +840,7 @@ class Flux2Transformer2DModel(nn.Module):
         self.transformer_blocks = nn.ModuleList(
             [
                 Flux2TransformerBlock(
+                    parallel_config=self.parallel_config,
                     dim=self.inner_dim,
                     num_attention_heads=num_attention_heads,
                     attention_head_dim=attention_head_dim,
@@ -686,6 +857,7 @@ class Flux2Transformer2DModel(nn.Module):
         self.single_transformer_blocks = nn.ModuleList(
             [
                 Flux2SingleTransformerBlock(
+                    parallel_config=self.parallel_config,
                     dim=self.inner_dim,
                     num_attention_heads=num_attention_heads,
                     attention_head_dim=attention_head_dim,
@@ -722,6 +894,9 @@ class Flux2Transformer2DModel(nn.Module):
         joint_attention_kwargs = joint_attention_kwargs or {}
 
         num_txt_tokens = encoder_hidden_states.shape[1]
+        sp_size = self.parallel_config.sequence_parallel_size
+        if sp_size and sp_size > 1 and is_forward_context_available():
+            get_forward_context().split_text_embed_in_sp = False
 
         timestep = timestep.to(hidden_states.dtype) * 1000
         if guidance is not None:
@@ -742,18 +917,43 @@ class Flux2Transformer2DModel(nn.Module):
             txt_ids = txt_ids[0]
 
         if is_torch_npu_available():
-            freqs_cos_image, freqs_sin_image = self.pos_embed(img_ids.cpu())
-            image_rotary_emb = (freqs_cos_image.npu(), freqs_sin_image.npu())
-            freqs_cos_text, freqs_sin_text = self.pos_embed(txt_ids.cpu())
-            text_rotary_emb = (freqs_cos_text.npu(), freqs_sin_text.npu())
+            txt_freqs_cos, txt_freqs_sin, img_freqs_cos, img_freqs_sin = self.rope_prepare(img_ids.cpu(), txt_ids.cpu())
+            txt_freqs_cos = txt_freqs_cos.npu()
+            txt_freqs_sin = txt_freqs_sin.npu()
+            img_freqs_cos = img_freqs_cos.npu()
+            img_freqs_sin = img_freqs_sin.npu()
         else:
-            image_rotary_emb = self.pos_embed(img_ids)
-            text_rotary_emb = self.pos_embed(txt_ids)
+            txt_freqs_cos, txt_freqs_sin, img_freqs_cos, img_freqs_sin = self.rope_prepare(img_ids, txt_ids)
         concat_rotary_emb = (
-            torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0),
-            torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0),
+            torch.cat([txt_freqs_cos, img_freqs_cos], dim=0),
+            torch.cat([txt_freqs_sin, img_freqs_sin], dim=0),
         )
 
+        # Create separate masks for image and text portions for Ulysses SP joint attention
+        hidden_states_mask = None
+        encoder_hidden_states_mask = None
+        if is_forward_context_available():
+            ctx = get_forward_context()
+        else:
+            ctx = None
+        if ctx is not None and ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0:
+            batch_size = hidden_states.shape[0]
+            img_padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size
+            hidden_states_mask = torch.ones(
+                batch_size,
+                img_padded_seq_len,
+                dtype=torch.bool,
+                device=hidden_states.device,
+            )
+            hidden_states_mask[:, ctx.sp_original_seq_len :] = False
+            if hidden_states_mask.all():
+                hidden_states_mask = None
+
+        if hidden_states_mask is not None:
+            joint_attention_kwargs["hidden_states_mask"] = hidden_states_mask
+        if encoder_hidden_states_mask is not None:
+            joint_attention_kwargs["encoder_hidden_states_mask"] = encoder_hidden_states_mask
+
         for index_block, block in enumerate(self.transformer_blocks):
             encoder_hidden_states, hidden_states = block(
                 hidden_states=hidden_states,
@@ -773,6 +973,7 @@ class Flux2Transformer2DModel(nn.Module):
                 temb_mod_params=single_stream_mod,
                 image_rotary_emb=concat_rotary_emb,
                 joint_attention_kwargs=joint_attention_kwargs,
+                text_seq_len=num_txt_tokens,
             )
 
         hidden_states = hidden_states[:, num_txt_tokens:, ...]
diff --git a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py
index 8700309e5..948f922d4 100644
--- a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py
+++ b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py
@@ -398,7 +398,9 @@ class Flux2Pipeline(nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarM
             self._execution_device
         )
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, Flux2Transformer2DModel)
-        self.transformer = Flux2Transformer2DModel(quant_config=od_config.quantization_config, **transformer_kwargs)
+        self.transformer = Flux2Transformer2DModel(
+            quant_config=od_config.quantization_config, od_config=od_config, **transformer_kwargs
+        )
 
         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
         self.image_processor = Flux2ImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
diff --git a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py
index ddb32aa20..56a7066e3 100644
--- a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py
+++ b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py
@@ -15,6 +15,7 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     QKVParallelLinear,
+    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -33,6 +34,9 @@ from vllm_omni.diffusion.distributed.sp_plan import (
 )
 from vllm_omni.diffusion.forward_context import get_forward_context
 
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
 logger = init_logger(__name__)
 
 
@@ -230,11 +234,24 @@ class GlmImagePrepare(nn.Module):
 class GlmImageAdaLayerNormZero(nn.Module):
     """Adaptive LayerNorm with zero initialization for both image and text streams."""
 
-    def __init__(self, embedding_dim: int, dim: int) -> None:
+    def __init__(
+        self,
+        embedding_dim: int,
+        dim: int,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+    ) -> None:
         super().__init__()
         self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
         self.norm_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
-        self.linear = nn.Linear(embedding_dim, 12 * dim, bias=True)
+        self.linear = ReplicatedLinear(
+            embedding_dim,
+            12 * dim,
+            bias=True,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.linear",
+        )
 
     def forward(
         self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor
@@ -244,6 +261,8 @@ class GlmImageAdaLayerNormZero(nn.Module):
         norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to(dtype=dtype)
 
         emb = self.linear(temb)
+        if isinstance(emb, tuple):
+            emb = emb[0]
         (
             shift_msa,
             c_shift_msa,
@@ -286,14 +305,25 @@ class GlmImageAdaLayerNormContinuous(nn.Module):
         elementwise_affine: bool = True,
         eps: float = 1e-5,
         bias: bool = True,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
-        self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias)
+        self.linear = ReplicatedLinear(
+            conditioning_embedding_dim,
+            embedding_dim * 2,
+            bias=bias,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.linear",
+        )
         self.norm = nn.LayerNorm(embedding_dim, eps=eps, elementwise_affine=elementwise_affine)
 
     def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor:
         # NO SiLU here
         emb = self.linear(conditioning_embedding.to(x.dtype))
+        if isinstance(emb, tuple):
+            emb = emb[0]
         scale, shift = torch.chunk(emb, 2, dim=1)
         x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
         return x
@@ -469,6 +499,7 @@ class GlmImageAttention(nn.Module):
         out_bias: bool = True,
         eps: float = 1e-5,
         quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.dim = dim
@@ -485,6 +516,7 @@ class GlmImageAttention(nn.Module):
             bias=True,
             return_bias=False,
             quant_config=quant_config,
+            prefix=f"{prefix}.to_qkv",
         )
 
         # QK normalization (LayerNorm, not RMSNorm for GLM-Image)
@@ -501,6 +533,7 @@ class GlmImageAttention(nn.Module):
                     input_is_parallel=True,
                     return_bias=False,
                     quant_config=quant_config,
+                    prefix=f"{prefix}.to_out.0",
                 ),
                 nn.Dropout(0.0),
             ]
@@ -554,7 +587,7 @@ class GlmImageAttention(nn.Module):
         hidden_states_combined = torch.cat([encoder_hidden_states, hidden_states], dim=1)
 
         # QKV projection
-        qkv = self.to_qkv(hidden_states_combined)
+        qkv = self.to_qkv(hidden_states_combined.contiguous())
         q_size = self.to_qkv.num_heads * self.head_dim
         kv_size = self.to_qkv.num_kv_heads * self.head_dim
         query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1)
@@ -604,8 +637,12 @@ class GlmImageAttention(nn.Module):
             # Project combined [text, image] outputs, then split.
             # This keeps SP numerically aligned with the non-SP path.
             joint_hidden_states_out = joint_hidden_states_out.flatten(2, 3).to(dtype)
+            # Contiguous for FP8/W4A16 quantized RowParallelLinear
             for module in self.to_out:
-                joint_hidden_states_out = module(joint_hidden_states_out)
+                if isinstance(module, RowParallelLinear):
+                    joint_hidden_states_out = module(joint_hidden_states_out.contiguous())
+                else:
+                    joint_hidden_states_out = module(joint_hidden_states_out)
 
             encoder_hidden_states_out = joint_hidden_states_out[:, :text_seq_length, :]
             hidden_states_out = joint_hidden_states_out[:, text_seq_length:, :]
@@ -645,7 +682,10 @@ class GlmImageAttention(nn.Module):
 
             # Output projection
             for module in self.to_out:
-                hidden_states_out = module(hidden_states_out)
+                if isinstance(module, RowParallelLinear):
+                    hidden_states_out = module(hidden_states_out.contiguous())
+                else:
+                    hidden_states_out = module(hidden_states_out)
 
             # Split back to text and image
             encoder_hidden_states_out = hidden_states_out[:, :text_seq_length, :]
@@ -663,6 +703,7 @@ class ColumnParallelGELU(nn.Module):
         approximate: str = "none",
         bias: bool = True,
         quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.proj = ColumnParallelLinear(
@@ -672,6 +713,7 @@ class ColumnParallelGELU(nn.Module):
             gather_output=False,
             return_bias=False,
             quant_config=quant_config,
+            prefix=f"{prefix}.proj",
         )
         self.approximate = approximate
 
@@ -688,6 +730,7 @@ class ColumnParallelSiLU(nn.Module):
         *,
         bias: bool = True,
         quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.proj = ColumnParallelLinear(
@@ -697,6 +740,7 @@ class ColumnParallelSiLU(nn.Module):
             gather_output=False,
             return_bias=False,
             quant_config=quant_config,
+            prefix=f"{prefix}.proj",
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -714,6 +758,7 @@ class GlmImageFeedForward(nn.Module):
         bias: bool = True,
         activation_fn: str = "gelu",
         quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         inner_dim = inner_dim or int(dim * mult)
@@ -721,7 +766,7 @@ class GlmImageFeedForward(nn.Module):
 
         if activation_fn == "linear-silu":
             layers: list[nn.Module] = [
-                ColumnParallelSiLU(dim, inner_dim, bias=bias, quant_config=quant_config),
+                ColumnParallelSiLU(dim, inner_dim, bias=bias, quant_config=quant_config, prefix=f"{prefix}.net.0"),
                 nn.Identity(),
                 RowParallelLinear(
                     inner_dim,
@@ -730,12 +775,20 @@ class GlmImageFeedForward(nn.Module):
                     input_is_parallel=True,
                     return_bias=False,
                     quant_config=quant_config,
+                    prefix=f"{prefix}.net.2",
                 ),
             ]
         else:
             approximate = "tanh" if activation_fn == "gelu-approximate" else "none"
             layers = [
-                ColumnParallelGELU(dim, inner_dim, approximate=approximate, bias=bias, quant_config=quant_config),
+                ColumnParallelGELU(
+                    dim,
+                    inner_dim,
+                    approximate=approximate,
+                    bias=bias,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.net.0",
+                ),
                 nn.Identity(),
                 RowParallelLinear(
                     inner_dim,
@@ -744,6 +797,7 @@ class GlmImageFeedForward(nn.Module):
                     input_is_parallel=True,
                     return_bias=False,
                     quant_config=quant_config,
+                    prefix=f"{prefix}.net.2",
                 ),
             ]
 
@@ -751,7 +805,12 @@ class GlmImageFeedForward(nn.Module):
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         for module in self.net:
-            hidden_states = module(hidden_states)
+            if isinstance(module, ColumnParallelLinear):
+                hidden_states, _ = module(hidden_states)
+            elif isinstance(module, RowParallelLinear):
+                hidden_states = module(hidden_states.contiguous())
+            else:
+                hidden_states = module(hidden_states)
         return hidden_states
 
 
@@ -765,19 +824,21 @@ class GlmImageTransformerBlock(nn.Module):
         attention_head_dim: int = 40,
         time_embed_dim: int = 512,
         ffn_hidden_dim: int | None = None,
-        quant_config: "QuantizationConfig | None" = None,
         parallel_config: DiffusionParallelConfig | None = None,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
 
         # 1. Attention with AdaLN
-        self.norm1 = GlmImageAdaLayerNormZero(time_embed_dim, dim)
+        self.norm1 = GlmImageAdaLayerNormZero(time_embed_dim, dim, quant_config=quant_config, prefix=f"{prefix}.norm1")
         self.attn1 = GlmImageAttention(
             dim=dim,
             num_heads=num_attention_heads,
             head_dim=attention_head_dim,
-            quant_config=quant_config,
             parallel_config=parallel_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn1",
         )
 
         # 2. Feedforward
@@ -789,6 +850,7 @@ class GlmImageTransformerBlock(nn.Module):
             inner_dim=ffn_hidden_dim,
             activation_fn="gelu-approximate",
             quant_config=quant_config,
+            prefix=f"{prefix}.ff",
         )
 
     def forward(
@@ -959,6 +1021,7 @@ class GlmImageTransformer2DModel(CachedTransformer):
             inner_dim=inner_dim,
             activation_fn="gelu",
             quant_config=quant_config,
+            prefix="glyph_projector",
         )
         self.prior_token_embedding = nn.Embedding(prior_vq_quantizer_codebook_size, inner_dim)
         self.prior_projector = GlmImageFeedForward(
@@ -967,6 +1030,7 @@ class GlmImageTransformer2DModel(CachedTransformer):
             inner_dim=inner_dim,
             activation_fn="linear-silu",
             quant_config=quant_config,
+            prefix="prior_projector",
         )
 
         # Prepare module for SP (encapsulates patch embedding and RoPE for _sp_plan)
@@ -988,15 +1052,20 @@ class GlmImageTransformer2DModel(CachedTransformer):
                     attention_head_dim,
                     time_embed_dim,
                     ffn_hidden_dim=ffn_hidden_dim,
-                    quant_config=quant_config,
                     parallel_config=self.parallel_config,
+                    quant_config=quant_config,
+                    prefix=f"transformer_blocks.{i}",
                 )
-                for _ in range(num_layers)
+                for i in range(num_layers)
             ]
         )
 
         # 4. Output projection
-        self.norm_out = GlmImageAdaLayerNormContinuous(inner_dim, time_embed_dim, elementwise_affine=False)
+        # Final modulation feeds proj_out; quant_config is NOT applied here
+        # to avoid precision degradation in the final projection layer.
+        self.norm_out = GlmImageAdaLayerNormContinuous(
+            inner_dim, time_embed_dim, elementwise_affine=False, quant_config=None, prefix="norm_out"
+        )
         self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels, bias=True)
 
     def forward(
diff --git a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py b/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py
index 97cba18c2..102456d9f 100644
--- a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py
+++ b/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py
@@ -301,10 +301,7 @@ class GlmImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
 
         # Load transformer (DiT)
         logger.info("Loading GlmImageTransformer2DModel (DiT)...")
-        self.transformer = GlmImageTransformer2DModel(
-            od_config=od_config,
-            quant_config=od_config.quantization_config,
-        )
+        self.transformer = GlmImageTransformer2DModel(od_config=od_config, quant_config=od_config.quantization_config)
 
         # Weight sources for DiT loading
         self.weights_sources = [
@@ -689,15 +686,11 @@ class GlmImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
         prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
 
         # NOTE: DiffusionEngine does an internal warmup "dummy run" during
-        # initialization. That request has no request_id and does not carry
+        # initialization. That request carries the fixed dummy request_id and does not carry
         # Stage-0 (AR) outputs via req.extra. For GLM-Image, we allow that
         # specific warmup request to proceed by synthesizing minimal prior
         # tokens, while still raising a clear error for real requests.
-        is_dummy_warmup = (
-            not getattr(req, "request_id", None)
-            and prompt == "dummy run"
-            and (req.sampling_params.num_inference_steps == 1)
-        )
+        is_dummy_warmup = req.is_dummy_run()
 
         # Get pre-computed prompt embeddings if provided
         if isinstance(first_prompt, str):
diff --git a/vllm_omni/diffusion/models/helios/helios_transformer.py b/vllm_omni/diffusion/models/helios/helios_transformer.py
index 5e7934c3b..ed0a50fee 100644
--- a/vllm_omni/diffusion/models/helios/helios_transformer.py
+++ b/vllm_omni/diffusion/models/helios/helios_transformer.py
@@ -3,6 +3,7 @@
 # Adapted from Helios (https://github.com/BestWishYsh/Helios)
 
 import math
+from collections import OrderedDict
 from collections.abc import Iterable
 from functools import lru_cache
 from typing import TYPE_CHECKING, Any
@@ -447,21 +448,32 @@ class HeliosCrossAttention(nn.Module):
             causal=False,
         )
 
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-    ) -> torch.Tensor:
-        query = self.to_q(hidden_states)
-        query = self.norm_q(query)
-
+    def project_kv(self, encoder_hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         key = self.to_k(encoder_hidden_states)
         value = self.to_v(encoder_hidden_states)
         key = self.norm_k(key)
 
-        query = query.unflatten(2, (self.num_heads, self.head_dim))
         key = key.unflatten(2, (self.num_heads, self.head_dim))
         value = value.unflatten(2, (self.num_heads, self.head_dim))
+        return key, value
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None = None,
+        encoder_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        query = self.to_q(hidden_states)
+        query = self.norm_q(query)
+
+        if encoder_key_value is None:
+            if encoder_hidden_states is None:
+                raise ValueError("encoder_hidden_states is required when encoder_key_value is not provided.")
+            key, value = self.project_kv(encoder_hidden_states)
+        else:
+            key, value = encoder_key_value
+
+        query = query.unflatten(2, (self.num_heads, self.head_dim))
 
         hidden_states = self.attn(query, key, value)
         hidden_states = hidden_states.flatten(2, 3)
@@ -526,6 +538,7 @@ class HeliosTransformerBlock(nn.Module):
         temb: torch.Tensor,
         rotary_emb: torch.Tensor,
         original_context_length: int | None = None,
+        cross_attn_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
     ) -> torch.Tensor:
         if temb.ndim == 4:
             shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
@@ -556,12 +569,20 @@ class HeliosTransformerBlock(nn.Module):
                 hidden_states[:, history_seq_len:],
             )
             norm_hidden_states = self.norm2(current_hidden_states.float()).type_as(current_hidden_states)
-            attn_output = self.attn2(norm_hidden_states, encoder_hidden_states)
+            attn_output = self.attn2(
+                norm_hidden_states,
+                encoder_hidden_states,
+                encoder_key_value=cross_attn_key_value,
+            )
             current_hidden_states = current_hidden_states + attn_output
             hidden_states = torch.cat([history_hidden_states, current_hidden_states], dim=1)
         else:
             norm_hidden_states = self.norm2(hidden_states.float()).type_as(hidden_states)
-            attn_output = self.attn2(norm_hidden_states, encoder_hidden_states)
+            attn_output = self.attn2(
+                norm_hidden_states,
+                encoder_hidden_states,
+                encoder_key_value=cross_attn_key_value,
+            )
             hidden_states = hidden_states + attn_output
 
         # 3. Feed-forward
@@ -723,11 +744,80 @@ class HeliosTransformer3DModel(nn.Module):
         # 5. Output norm & projection
         self.norm_out = HeliosOutputNorm(inner_dim, eps)
         self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size))
+        self._projected_encoder_cache: OrderedDict[tuple, torch.Tensor] = OrderedDict()
+        self._cross_attn_kv_cache: OrderedDict[tuple, list[tuple[torch.Tensor, torch.Tensor]]] = OrderedDict()
+        self._cross_attn_cache_size = 2
 
     @property
     def dtype(self) -> torch.dtype:
         return next(self.parameters()).dtype
 
+    @staticmethod
+    def _tensor_cache_key(tensor: torch.Tensor) -> tuple:
+        try:
+            version = tensor._version
+        except RuntimeError:
+            version = None
+
+        return (
+            tensor.data_ptr(),
+            tuple(tensor.shape),
+            tuple(tensor.stride()),
+            tensor.dtype,
+            tensor.device.type,
+            tensor.device.index,
+            version,
+        )
+
+    @staticmethod
+    def _get_from_lru(cache: OrderedDict, key: tuple):
+        value = cache.get(key)
+        if value is not None:
+            cache.move_to_end(key)
+        return value
+
+    def _put_lru(self, cache: OrderedDict, key: tuple, value) -> None:
+        cache[key] = value
+        cache.move_to_end(key)
+        while len(cache) > self._cross_attn_cache_size:
+            cache.popitem(last=False)
+
+    def clear_cross_attention_cache(self) -> None:
+        self._projected_encoder_cache.clear()
+        self._cross_attn_kv_cache.clear()
+
+    def _cache_enabled(self) -> bool:
+        return not self.training and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+
+    def _project_encoder_hidden_states(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
+        if not self._cache_enabled():
+            return self.condition_embedder.text_embedder(encoder_hidden_states)
+
+        cache_key = self._tensor_cache_key(encoder_hidden_states)
+        cached = self._get_from_lru(self._projected_encoder_cache, cache_key)
+        if cached is not None:
+            return cached
+
+        projected = self.condition_embedder.text_embedder(encoder_hidden_states)
+        self._put_lru(self._projected_encoder_cache, cache_key, projected)
+        return projected
+
+    def _get_cross_attn_key_values(
+        self,
+        encoder_hidden_states: torch.Tensor,
+    ) -> list[tuple[torch.Tensor, torch.Tensor]] | None:
+        if not self._cache_enabled():
+            return None
+
+        cache_key = self._tensor_cache_key(encoder_hidden_states)
+        cached = self._get_from_lru(self._cross_attn_kv_cache, cache_key)
+        if cached is not None:
+            return cached
+
+        key_values = [block.attn2.project_kv(encoder_hidden_states) for block in self.blocks]
+        self._put_lru(self._cross_attn_kv_cache, cache_key, key_values)
+        return key_values
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -836,7 +926,12 @@ class HeliosTransformer3DModel(nn.Module):
                 .expand(batch_size, -1, history_context_length, -1)
             )
 
-        temb, timestep_proj, encoder_hidden_states = self.condition_embedder(timestep, encoder_hidden_states)
+        temb, timestep_proj, _ = self.condition_embedder(
+            timestep,
+            encoder_hidden_states,
+            is_return_encoder_hidden_states=False,
+        )
+        encoder_hidden_states = self._project_encoder_hidden_states(encoder_hidden_states)
         timestep_proj = timestep_proj.unflatten(-1, (6, -1))
 
         if indices_hidden_states is not None and not self.zero_history_timestep:
@@ -857,14 +952,17 @@ class HeliosTransformer3DModel(nn.Module):
         hidden_states = hidden_states.contiguous()
         encoder_hidden_states = encoder_hidden_states.contiguous()
         rotary_emb = rotary_emb.contiguous()
+        cross_attn_key_values = self._get_cross_attn_key_values(encoder_hidden_states)
 
-        for block in self.blocks:
+        for block_idx, block in enumerate(self.blocks):
+            cross_attn_key_value = None if cross_attn_key_values is None else cross_attn_key_values[block_idx]
             hidden_states = block(
                 hidden_states,
                 encoder_hidden_states,
                 timestep_proj,
                 rotary_emb,
                 original_context_length,
+                cross_attn_key_value,
             )
 
         # 7. Output normalization
diff --git a/vllm_omni/diffusion/models/hidream_image/__init__.py b/vllm_omni/diffusion/models/hidream_image/__init__.py
new file mode 100644
index 000000000..d8c52afdd
--- /dev/null
+++ b/vllm_omni/diffusion/models/hidream_image/__init__.py
@@ -0,0 +1,17 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""HiDream Image diffusion model components."""
+
+from vllm_omni.diffusion.models.hidream_image.hidream_image_transformer import (
+    HiDreamImageTransformer2DModel,
+)
+from vllm_omni.diffusion.models.hidream_image.pipeline_hidream_image import (
+    HiDreamImagePipeline,
+    get_hidream_image_post_process_func,
+)
+
+__all__ = [
+    "HiDreamImageTransformer2DModel",
+    "HiDreamImagePipeline",
+    "get_hidream_image_post_process_func",
+]
diff --git a/vllm_omni/diffusion/models/hidream_image/hidream_image_transformer.py b/vllm_omni/diffusion/models/hidream_image/hidream_image_transformer.py
new file mode 100644
index 000000000..5c12df14e
--- /dev/null
+++ b/vllm_omni/diffusion/models/hidream_image/hidream_image_transformer.py
@@ -0,0 +1,1042 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+from collections.abc import Iterable
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import (
+        QuantizationConfig,
+    )
+
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.utils import deprecate, logging
+from vllm.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_reduce,
+)
+from vllm.model_executor.layers.linear import (
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding
+
+logger = logging.get_logger(__name__)
+
+
+class HiDreamImageFeedForwardSwiGLU(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: int,
+        multiple_of: int = 256,
+        ffn_dim_multiplier: float | None = None,
+    ):
+        super().__init__()
+        hidden_dim = int(2 * hidden_dim / 3)
+        # custom dim factor multiplier
+        if ffn_dim_multiplier is not None:
+            hidden_dim = int(ffn_dim_multiplier * hidden_dim)
+        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w2(torch.nn.functional.silu(self.w1(x)) * self.w3(x))
+
+
+class HiDreamImagePooledEmbed(nn.Module):
+    def __init__(self, text_emb_dim, hidden_size):
+        super().__init__()
+        self.pooled_embedder = TimestepEmbedding(in_channels=text_emb_dim, time_embed_dim=hidden_size)
+
+    def forward(self, pooled_embed: torch.Tensor) -> torch.Tensor:
+        #
+        return self.pooled_embedder(pooled_embed.to(dtype=torch.bfloat16))
+
+
+class HiDreamImageTimestepEmbed(nn.Module):
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.time_proj = Timesteps(num_channels=frequency_embedding_size, flip_sin_to_cos=True, downscale_freq_shift=0)
+        self.timestep_embedder = TimestepEmbedding(in_channels=frequency_embedding_size, time_embed_dim=hidden_size)
+
+    def forward(self, timesteps: torch.Tensor, wdtype: torch.dtype | None = None) -> torch.Tensor:
+        t_emb = self.time_proj(timesteps).to(dtype=torch.bfloat16)
+        t_emb = self.timestep_embedder(t_emb)
+        return t_emb
+
+
+class HiDreamImageOutEmbed(nn.Module):
+    def __init__(self, hidden_size, patch_size, out_channels):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
+
+    def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor) -> torch.Tensor:
+        shift, scale = self.adaLN_modulation(temb).chunk(2, dim=1)
+        hidden_states = self.norm_final(hidden_states) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+        hidden_states = self.linear(hidden_states)
+        return hidden_states
+
+
+class HiDreamImagePatchEmbed(nn.Module):
+    def __init__(
+        self,
+        patch_size=2,
+        in_channels=4,
+        out_channels=1024,
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        self.out_channels = out_channels
+        self.proj = nn.Linear(in_channels * patch_size * patch_size, out_channels, bias=True)
+
+    def forward(self, latent) -> torch.Tensor:
+        latent = self.proj(latent.to(dtype=torch.bfloat16))
+        return latent
+
+
+def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
+    assert dim % 2 == 0, "The dimension must be even."
+
+    is_mps = pos.device.type == "mps"
+    is_npu = pos.device.type == "npu"
+
+    dtype = torch.float32 if (is_mps or is_npu) else torch.float64
+
+    scale = torch.arange(0, dim, 2, dtype=dtype, device=pos.device) / dim
+    omega = 1.0 / (theta**scale)
+
+    batch_size, seq_length = pos.shape
+    out = torch.einsum("...n,d->...nd", pos, omega)
+    cos_out = torch.cos(out)
+    sin_out = torch.sin(out)
+
+    stacked_out = torch.stack([cos_out, -sin_out, sin_out, cos_out], dim=-1)
+    out = stacked_out.view(batch_size, -1, dim // 2, 2, 2)
+    return out.float()
+
+
+class HiDreamImageEmbedND(nn.Module):
+    def __init__(self, theta: int, axes_dim: list[int]):
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: torch.Tensor) -> torch.Tensor:
+        n_axes = ids.shape[-1]
+        emb = torch.cat(
+            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
+            dim=-3,
+        )
+        return emb.unsqueeze(2)
+
+
+class DistributedRMSNorm(nn.Module):
+    """RMSNorm that computes global RMS across tensor parallel ranks."""
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        tp_size = get_tensor_model_parallel_world_size()
+
+        input_dtype = x.dtype
+        x_float = x.float()
+        local_sum_sq = (x_float**2).sum(dim=-1, keepdim=True)
+        local_count = x.shape[-1]
+
+        if tp_size > 1:
+            global_sum_sq = local_sum_sq.clone()
+            tensor_model_parallel_all_reduce(global_sum_sq)
+            global_count = local_count * tp_size
+        else:
+            global_sum_sq = local_sum_sq
+            global_count = local_count
+
+        rms = torch.sqrt(global_sum_sq / global_count + self.eps)
+
+        output = (x_float / rms) * self.weight.float()
+        return output.to(input_dtype)
+
+
+class HiDreamAttention(nn.Module):
+    def __init__(
+        self,
+        query_dim: int,
+        heads: int = 8,
+        dim_head: int = 64,
+        upcast_attention: bool = False,
+        upcast_softmax: bool = False,
+        scale_qk: bool = True,
+        eps: float = 1e-5,
+        processor=None,
+        out_dim: int = None,
+        single: bool = False,
+        quant_config: "QuantizationConfig | None" = None,
+    ):
+        super().__init__()
+        self.head_dim = dim_head
+
+        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+        self.query_dim = query_dim
+        self.upcast_attention = upcast_attention
+        self.upcast_softmax = upcast_softmax
+        self.out_dim = out_dim if out_dim is not None else query_dim
+
+        self.scale_qk = scale_qk
+        self.scale = dim_head**-0.5 if self.scale_qk else 1.0
+
+        self.heads = out_dim // dim_head if out_dim is not None else heads
+        self.sliceable_head_dim = heads
+        self.single = single
+
+        # Fused QKV projection using vLLM's optimized layer
+        self.to_qkv = QKVParallelLinear(
+            hidden_size=query_dim,
+            head_size=self.head_dim,
+            total_num_heads=self.heads,
+            quant_config=quant_config,
+            prefix="to_qkv",
+        )
+        self.query_num_heads = self.to_qkv.num_heads
+        self.kv_num_heads = self.to_qkv.num_kv_heads
+
+        self.to_out = RowParallelLinear(
+            self.inner_dim,
+            self.out_dim,
+            bias=True,
+            input_is_parallel=True,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix="to_out",
+        )
+
+        self.tp_inner_dim = self.query_num_heads * dim_head
+        if get_tensor_model_parallel_world_size() > 1:
+            self.q_rms_norm = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+            self.k_rms_norm = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+        else:
+            self.q_rms_norm = nn.RMSNorm(self.tp_inner_dim, eps)
+            self.k_rms_norm = nn.RMSNorm(self.tp_inner_dim, eps)
+
+        if not single:
+            self.to_qkv_t = QKVParallelLinear(
+                hidden_size=query_dim,
+                head_size=self.head_dim,
+                total_num_heads=self.heads,
+                quant_config=quant_config,
+                prefix="to_qkv_t",
+            )
+            self.to_out_t = RowParallelLinear(
+                self.inner_dim,
+                self.out_dim,
+                bias=True,
+                input_is_parallel=True,
+                return_bias=False,
+                quant_config=quant_config,
+                prefix="to_out_t",
+            )
+            if get_tensor_model_parallel_world_size() > 1:
+                self.q_rms_norm_t = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+                self.k_rms_norm_t = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+            else:
+                self.q_rms_norm_t = nn.RMSNorm(self.tp_inner_dim, eps)
+                self.k_rms_norm_t = nn.RMSNorm(self.tp_inner_dim, eps)
+        self.rope = RotaryEmbedding(is_neox_style=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        hidden_states_masks: torch.Tensor | None = None,
+        encoder_hidden_states: torch.Tensor | None = None,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        dtype = hidden_states.dtype
+        batch_size = hidden_states.shape[0]
+
+        qkv, _ = self.to_qkv(hidden_states)
+        q_size = self.to_qkv.num_heads * self.head_dim
+        kv_size = self.to_qkv.num_kv_heads * self.head_dim
+        query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1)
+
+        query_i = self.q_rms_norm(query).to(dtype=dtype)
+        key_i = self.k_rms_norm(key).to(dtype=dtype)
+        value_i = value
+
+        inner_dim = key_i.shape[-1]
+        head_dim = inner_dim // self.query_num_heads
+
+        query_i = query_i.view(batch_size, -1, self.query_num_heads, head_dim)
+        key_i = key_i.view(batch_size, -1, self.kv_num_heads, head_dim)
+        value_i = value_i.view(batch_size, -1, self.kv_num_heads, head_dim)
+        if hidden_states_masks is not None:
+            key_i = key_i * hidden_states_masks.view(batch_size, -1, 1, 1)
+
+        if not self.single:
+            qkv_t, _ = self.to_qkv_t(encoder_hidden_states)
+            q_size_t = self.to_qkv_t.num_heads * self.head_dim
+            kv_size_t = self.to_qkv_t.num_kv_heads * self.head_dim
+            query_t, key_t, value_t = qkv_t.split([q_size_t, kv_size_t, kv_size_t], dim=-1)
+            query_t = self.q_rms_norm_t(query_t).to(dtype=dtype)
+            key_t = self.k_rms_norm_t(key_t).to(dtype=dtype)
+            value_t = value_t
+
+            query_t = query_t.view(batch_size, -1, self.to_qkv_t.num_heads, head_dim)
+            key_t = key_t.view(batch_size, -1, self.to_qkv_t.num_kv_heads, head_dim)
+            value_t = value_t.view(batch_size, -1, self.to_qkv_t.num_kv_heads, head_dim)
+
+            num_image_tokens = query_i.shape[1]
+            num_text_tokens = query_t.shape[1]
+            query = torch.cat([query_i, query_t], dim=1)
+            key = torch.cat([key_i, key_t], dim=1)
+            value = torch.cat([value_i, value_t], dim=1)
+        else:
+            query = query_i
+            key = key_i
+            value = value_i
+
+        cos, sin = image_rotary_emb
+        cos = cos.squeeze(2).to(query.dtype)
+        sin = sin.squeeze(2).to(query.dtype)
+        if query.shape[-1] == image_rotary_emb[0].shape[-1] * 2:
+            query = self.rope(query, cos, sin)
+            key = self.rope(key, cos, sin)
+        else:
+            query_1, query_2 = query.chunk(2, dim=-1)
+            key_1, key_2 = key.chunk(2, dim=-1)
+            key_1 = self.rope(query_1, cos, sin)
+            key_1 = self.rope(key_1, cos, sin)
+            query = torch.cat([query_1, query_2], dim=-1)
+            key = torch.cat([key_1, key_2], dim=-1)
+
+        hidden_states = F.scaled_dot_product_attention(
+            query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.query_num_heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        if not self.single:
+            hidden_states_i, hidden_states_t = torch.split(hidden_states, [num_image_tokens, num_text_tokens], dim=1)
+            hidden_states_i = self.to_out(hidden_states_i)
+            hidden_states_t = self.to_out_t(hidden_states_t)
+            return hidden_states_i, hidden_states_t
+        else:
+            hidden_states = self.to_out(hidden_states)
+            return hidden_states
+
+
+# Modified from https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/model.py
+class MoEGate(nn.Module):
+    def __init__(
+        self,
+        embed_dim,
+        num_routed_experts=4,
+        num_activated_experts=2,
+        aux_loss_alpha=0.01,
+        _force_inference_output=False,
+    ):
+        super().__init__()
+        self.top_k = num_activated_experts
+        self.n_routed_experts = num_routed_experts
+
+        self.scoring_func = "softmax"
+        self.alpha = aux_loss_alpha
+        self.seq_aux = False
+
+        # topk selection algorithm
+        self.norm_topk_prob = False
+        self.gating_dim = embed_dim
+        self.weight = nn.Parameter(torch.randn(self.n_routed_experts, self.gating_dim) / embed_dim**0.5)
+
+        self._force_inference_output = _force_inference_output
+
+    def forward(self, hidden_states):
+        bsz, seq_len, h = hidden_states.shape
+        ### compute gating score
+        hidden_states = hidden_states.view(-1, h)
+        logits = F.linear(hidden_states, self.weight, None)
+        if self.scoring_func == "softmax":
+            scores = logits.softmax(dim=-1)
+        else:
+            raise NotImplementedError(f"insupportable scoring function for MoE gating: {self.scoring_func}")
+
+        ### select top-k experts
+        topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
+
+        ### norm gate to sum 1
+        if self.top_k > 1 and self.norm_topk_prob:
+            denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
+            topk_weight = topk_weight / denominator
+
+        ### expert-level computation auxiliary loss
+        if self.training and self.alpha > 0.0 and not self._force_inference_output:
+            scores_for_aux = scores
+            aux_topk = self.top_k
+            # always compute aux loss based on the naive greedy topk method
+            topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
+            if self.seq_aux:
+                scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)
+                ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device)
+                ce.scatter_add_(
+                    1, topk_idx_for_aux_loss, torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device)
+                ).div_(seq_len * aux_topk / self.n_routed_experts)
+                aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha
+            else:
+                mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)
+                ce = mask_ce.float().mean(0)
+
+                Pi = scores_for_aux.mean(0)
+                fi = ce * self.n_routed_experts
+                aux_loss = (Pi * fi).sum() * self.alpha
+        else:
+            aux_loss = None
+        return topk_idx, topk_weight, aux_loss
+
+
+# Modified from https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/model.py
+class MOEFeedForwardSwiGLU(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: int,
+        num_routed_experts: int,
+        num_activated_experts: int,
+        _force_inference_output: bool = False,
+    ):
+        super().__init__()
+        self.shared_experts = HiDreamImageFeedForwardSwiGLU(dim, hidden_dim // 2)
+        self.experts = nn.ModuleList(
+            [HiDreamImageFeedForwardSwiGLU(dim, hidden_dim) for i in range(num_routed_experts)]
+        )
+        self._force_inference_output = _force_inference_output
+        self.gate = MoEGate(
+            embed_dim=dim,
+            num_routed_experts=num_routed_experts,
+            num_activated_experts=num_activated_experts,
+            _force_inference_output=_force_inference_output,
+        )
+        self.num_activated_experts = num_activated_experts
+
+    def forward(self, x):
+        wtype = x.dtype
+        identity = x
+        orig_shape = x.shape
+        topk_idx, topk_weight, aux_loss = self.gate(x)
+        x = x.view(-1, x.shape[-1])
+        flat_topk_idx = topk_idx.view(-1)
+        if self.training and not self._force_inference_output:
+            x = x.repeat_interleave(self.num_activated_experts, dim=0)
+            y = torch.empty_like(x, dtype=wtype)
+            for i, expert in enumerate(self.experts):
+                y[flat_topk_idx == i] = expert(x[flat_topk_idx == i]).to(dtype=wtype)
+            y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
+            y = y.view(*orig_shape).to(dtype=wtype)
+        else:
+            y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)
+        y = y + self.shared_experts(identity)
+        return y
+
+    @torch.no_grad()
+    def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
+        expert_cache = torch.zeros_like(x)
+        idxs = flat_expert_indices.argsort()
+        tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0)
+        token_idxs = idxs // self.num_activated_experts
+
+        # TODO: This MoE expert routing is sequential and can be a bottleneck for models
+        # with MoE in every transformer block. Consider optimizing this in the future
+        # using techniques like grouped GEMM or batched expert execution.
+        for i, end_idx in enumerate(tokens_per_expert):
+            start_idx = 0 if i == 0 else tokens_per_expert[i - 1]
+            if start_idx == end_idx:
+                continue
+            expert = self.experts[i]
+            exp_token_idx = token_idxs[start_idx:end_idx]
+            expert_tokens = x[exp_token_idx]
+            expert_out = expert(expert_tokens)
+            expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
+
+            # for fp16 and other dtype
+            expert_cache = expert_cache.to(expert_out.dtype)
+            expert_cache.scatter_reduce_(0, exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]), expert_out, reduce="sum")
+        return expert_cache
+
+
+class TextProjection(nn.Module):
+    def __init__(self, in_features, hidden_size):
+        super().__init__()
+        self.linear = nn.Linear(in_features=in_features, out_features=hidden_size, bias=False)
+
+    def forward(self, caption):
+        hidden_states = self.linear(caption)
+        return hidden_states
+
+
+class HiDreamImageSingleTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        num_routed_experts: int = 4,
+        num_activated_experts: int = 2,
+        _force_inference_output: bool = False,
+    ):
+        super().__init__()
+        self.num_attention_heads = num_attention_heads
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim, bias=True))
+
+        # 1. Attention
+        self.norm1_i = nn.LayerNorm(dim, eps=1e-06, elementwise_affine=False)
+        self.attn1 = HiDreamAttention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            single=True,
+        )
+
+        # 3. Feed-forward
+        self.norm3_i = nn.LayerNorm(dim, eps=1e-06, elementwise_affine=False)
+        if num_routed_experts > 0:
+            self.ff_i = MOEFeedForwardSwiGLU(
+                dim=dim,
+                hidden_dim=4 * dim,
+                num_routed_experts=num_routed_experts,
+                num_activated_experts=num_activated_experts,
+                _force_inference_output=_force_inference_output,
+            )
+        else:
+            self.ff_i = HiDreamImageFeedForwardSwiGLU(dim=dim, hidden_dim=4 * dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        hidden_states_masks: torch.Tensor | None = None,
+        encoder_hidden_states: torch.Tensor | None = None,
+        temb: torch.Tensor | None = None,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        wtype = hidden_states.dtype
+        shift_msa_i, scale_msa_i, gate_msa_i, shift_mlp_i, scale_mlp_i, gate_mlp_i = self.adaLN_modulation(temb)[
+            :, None
+        ].chunk(6, dim=-1)
+
+        # 1. MM-Attention
+        norm_hidden_states = self.norm1_i(hidden_states).to(dtype=wtype)
+        norm_hidden_states = norm_hidden_states * (1 + scale_msa_i) + shift_msa_i
+        attn_output_i = self.attn1(
+            norm_hidden_states,
+            hidden_states_masks,
+            image_rotary_emb=image_rotary_emb,
+        )
+        hidden_states = gate_msa_i * attn_output_i + hidden_states
+
+        # 2. Feed-forward
+        norm_hidden_states = self.norm3_i(hidden_states).to(dtype=wtype)
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp_i) + shift_mlp_i
+        ff_output_i = gate_mlp_i * self.ff_i(norm_hidden_states.to(dtype=wtype))
+        hidden_states = ff_output_i + hidden_states
+        return hidden_states
+
+
+class HiDreamImageTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        num_routed_experts: int = 4,
+        num_activated_experts: int = 2,
+        _force_inference_output: bool = False,
+    ):
+        super().__init__()
+        self.num_attention_heads = num_attention_heads
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, 12 * dim, bias=True))
+
+        # 1. Attention
+        self.norm1_i = nn.LayerNorm(dim, eps=1e-06, elementwise_affine=False)
+        self.norm1_t = nn.LayerNorm(dim, eps=1e-06, elementwise_affine=False)
+        self.attn1 = HiDreamAttention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            single=False,
+        )
+
+        # 3. Feed-forward
+        self.norm3_i = nn.LayerNorm(dim, eps=1e-06, elementwise_affine=False)
+        if num_routed_experts > 0:
+            self.ff_i = MOEFeedForwardSwiGLU(
+                dim=dim,
+                hidden_dim=4 * dim,
+                num_routed_experts=num_routed_experts,
+                num_activated_experts=num_activated_experts,
+                _force_inference_output=_force_inference_output,
+            )
+        else:
+            self.ff_i = HiDreamImageFeedForwardSwiGLU(dim=dim, hidden_dim=4 * dim)
+        self.norm3_t = nn.LayerNorm(dim, eps=1e-06, elementwise_affine=False)
+        self.ff_t = HiDreamImageFeedForwardSwiGLU(dim=dim, hidden_dim=4 * dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        hidden_states_masks: torch.Tensor | None = None,
+        encoder_hidden_states: torch.Tensor | None = None,
+        temb: torch.Tensor | None = None,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        wtype = hidden_states.dtype
+        (
+            shift_msa_i,
+            scale_msa_i,
+            gate_msa_i,
+            shift_mlp_i,
+            scale_mlp_i,
+            gate_mlp_i,
+            shift_msa_t,
+            scale_msa_t,
+            gate_msa_t,
+            shift_mlp_t,
+            scale_mlp_t,
+            gate_mlp_t,
+        ) = self.adaLN_modulation(temb)[:, None].chunk(12, dim=-1)
+
+        # 1. MM-Attention
+        norm_hidden_states = self.norm1_i(hidden_states).to(dtype=wtype)
+        norm_hidden_states = norm_hidden_states * (1 + scale_msa_i) + shift_msa_i
+        norm_encoder_hidden_states = self.norm1_t(encoder_hidden_states).to(dtype=wtype)
+        norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + scale_msa_t) + shift_msa_t
+
+        attn_output_i, attn_output_t = self.attn1(
+            norm_hidden_states,
+            hidden_states_masks,
+            norm_encoder_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+        )
+
+        hidden_states = gate_msa_i * attn_output_i + hidden_states
+        encoder_hidden_states = gate_msa_t * attn_output_t + encoder_hidden_states
+
+        # 2. Feed-forward
+        norm_hidden_states = self.norm3_i(hidden_states).to(dtype=wtype)
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp_i) + shift_mlp_i
+        norm_encoder_hidden_states = self.norm3_t(encoder_hidden_states).to(dtype=wtype)
+        norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + scale_mlp_t) + shift_mlp_t
+
+        ff_output_i = gate_mlp_i * self.ff_i(norm_hidden_states)
+        ff_output_t = gate_mlp_t * self.ff_t(norm_encoder_hidden_states)
+        hidden_states = ff_output_i + hidden_states
+        encoder_hidden_states = ff_output_t + encoder_hidden_states
+        return hidden_states, encoder_hidden_states
+
+
+class HiDreamBlock(nn.Module):
+    def __init__(self, block: HiDreamImageTransformerBlock | HiDreamImageSingleTransformerBlock):
+        super().__init__()
+        self.block = block
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        hidden_states_masks: torch.Tensor | None = None,
+        encoder_hidden_states: torch.Tensor | None = None,
+        temb: torch.Tensor | None = None,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        return self.block(
+            hidden_states=hidden_states,
+            hidden_states_masks=hidden_states_masks,
+            encoder_hidden_states=encoder_hidden_states,
+            temb=temb,
+            image_rotary_emb=image_rotary_emb,
+        )
+
+
+class HiDreamImageTransformer2DModel(nn.Module):
+    def __init__(
+        self,
+        od_config: OmniDiffusionConfig = None,
+        patch_size: int | None = None,
+        in_channels: int = 64,
+        out_channels: int | None = None,
+        num_layers: int = 16,
+        num_single_layers: int = 32,
+        attention_head_dim: int = 128,
+        num_attention_heads: int = 20,
+        caption_channels: list[int] = None,
+        text_emb_dim: int = 2048,
+        num_routed_experts: int = 4,
+        num_activated_experts: int = 2,
+        axes_dims_rope: tuple[int, int] = (32, 32),
+        max_resolution: tuple[int, int] = (128, 128),
+        llama_layers: list[int] = None,
+        force_inference_output: bool = False,
+        quant_config: "QuantizationConfig | None" = None,
+    ):
+        super().__init__()
+        self.parallel_config = od_config.parallel_config
+
+        self.in_channels = in_channels
+        self.patch_size = patch_size
+        self.force_inference_output = force_inference_output
+        self.llama_layers = llama_layers
+
+        self.out_channels = out_channels or in_channels
+        self.inner_dim = num_attention_heads * attention_head_dim
+
+        self.t_embedder = HiDreamImageTimestepEmbed(self.inner_dim)
+        self.p_embedder = HiDreamImagePooledEmbed(text_emb_dim, self.inner_dim)
+        self.x_embedder = HiDreamImagePatchEmbed(
+            patch_size=patch_size,
+            in_channels=in_channels,
+            out_channels=self.inner_dim,
+        )
+        self.pe_embedder = HiDreamImageEmbedND(theta=10000, axes_dim=axes_dims_rope)
+
+        self.double_stream_blocks = nn.ModuleList(
+            [
+                HiDreamBlock(
+                    HiDreamImageTransformerBlock(
+                        dim=self.inner_dim,
+                        num_attention_heads=num_attention_heads,
+                        attention_head_dim=attention_head_dim,
+                        num_routed_experts=num_routed_experts,
+                        num_activated_experts=num_activated_experts,
+                        _force_inference_output=force_inference_output,
+                    )
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        self.single_stream_blocks = nn.ModuleList(
+            [
+                HiDreamBlock(
+                    HiDreamImageSingleTransformerBlock(
+                        dim=self.inner_dim,
+                        num_attention_heads=num_attention_heads,
+                        attention_head_dim=attention_head_dim,
+                        num_routed_experts=num_routed_experts,
+                        num_activated_experts=num_activated_experts,
+                        _force_inference_output=force_inference_output,
+                    )
+                )
+                for _ in range(num_single_layers)
+            ]
+        )
+
+        self.final_layer = HiDreamImageOutEmbed(self.inner_dim, patch_size, self.out_channels)
+
+        num_llama_projections = len(self.llama_layers) if self.llama_layers is not None else 0
+        caption_channels = [caption_channels[1]] * num_llama_projections + [caption_channels[0]]
+        caption_projection = []
+        for caption_channel in caption_channels:
+            caption_projection.append(TextProjection(in_features=caption_channel, hidden_size=self.inner_dim))
+        self.caption_projection = nn.ModuleList(caption_projection)
+        self.max_seq = max_resolution[0] * max_resolution[1] // (patch_size * patch_size)
+
+    def unpatchify(self, x: torch.Tensor, img_sizes: list[tuple[int, int]], is_training: bool) -> list[torch.Tensor]:
+        if is_training and not self.force_inference_output:
+            B, S, F = x.shape
+            C = F // (self.patch_size * self.patch_size)
+            x = (
+                x.reshape(B, S, self.patch_size, self.patch_size, C)
+                .permute(0, 4, 1, 2, 3)
+                .reshape(B, C, S, self.patch_size * self.patch_size)
+            )
+        else:
+            x_arr = []
+            p1 = self.patch_size
+            p2 = self.patch_size
+            for i, img_size in enumerate(img_sizes):
+                pH, pW = img_size
+                t = x[i, : pH * pW].reshape(1, pH, pW, -1)
+                F_token = t.shape[-1]
+                C = F_token // (p1 * p2)
+                t = t.reshape(1, pH, pW, p1, p2, C)
+                t = t.permute(0, 5, 1, 3, 2, 4)
+                t = t.reshape(1, C, pH * p1, pW * p2)
+                x_arr.append(t)
+            x = torch.cat(x_arr, dim=0)
+        return x
+
+    def patchify(self, hidden_states):
+        batch_size, channels, height, width = hidden_states.shape
+        patch_size = self.patch_size
+        patch_height, patch_width = height // patch_size, width // patch_size
+        device = hidden_states.device
+        dtype = hidden_states.dtype
+
+        # create img_sizes
+        img_sizes = torch.tensor([patch_height, patch_width], dtype=torch.int64, device=device).reshape(-1)
+        img_sizes = img_sizes.unsqueeze(0).repeat(batch_size, 1)
+
+        # create hidden_states_masks
+        if hidden_states.shape[-2] != hidden_states.shape[-1]:
+            hidden_states_masks = torch.zeros((batch_size, self.max_seq), dtype=dtype, device=device)
+            hidden_states_masks[:, : patch_height * patch_width] = 1.0
+        else:
+            hidden_states_masks = None
+
+        # create img_ids
+        img_ids = torch.zeros(patch_height, patch_width, 3, device=device)
+        row_indices = torch.arange(patch_height, device=device)[:, None]
+        col_indices = torch.arange(patch_width, device=device)[None, :]
+        img_ids[..., 1] = img_ids[..., 1] + row_indices
+        img_ids[..., 2] = img_ids[..., 2] + col_indices
+        img_ids = img_ids.reshape(patch_height * patch_width, -1)
+
+        if hidden_states.shape[-2] != hidden_states.shape[-1]:
+            # Handle non-square latents
+            img_ids_pad = torch.zeros(self.max_seq, 3, device=device)
+            img_ids_pad[: patch_height * patch_width, :] = img_ids
+            img_ids = img_ids_pad.unsqueeze(0).repeat(batch_size, 1, 1)
+        else:
+            img_ids = img_ids.unsqueeze(0).repeat(batch_size, 1, 1)
+
+        # patchify hidden_states
+        if hidden_states.shape[-2] != hidden_states.shape[-1]:
+            # Handle non-square latents
+            out = torch.zeros(
+                (batch_size, channels, self.max_seq, patch_size * patch_size),
+                dtype=dtype,
+                device=device,
+            )
+            hidden_states = hidden_states.reshape(
+                batch_size, channels, patch_height, patch_size, patch_width, patch_size
+            )
+            hidden_states = hidden_states.permute(0, 1, 2, 4, 3, 5)
+            hidden_states = hidden_states.reshape(
+                batch_size, channels, patch_height * patch_width, patch_size * patch_size
+            )
+            out[:, :, 0 : patch_height * patch_width] = hidden_states
+            hidden_states = out
+            hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
+                batch_size, self.max_seq, patch_size * patch_size * channels
+            )
+
+        else:
+            # Handle square latents
+            hidden_states = hidden_states.reshape(
+                batch_size, channels, patch_height, patch_size, patch_width, patch_size
+            )
+            hidden_states = hidden_states.permute(0, 2, 4, 3, 5, 1)
+            hidden_states = hidden_states.reshape(
+                batch_size, patch_height * patch_width, patch_size * patch_size * channels
+            )
+
+        return hidden_states, hidden_states_masks, img_sizes, img_ids
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timesteps: torch.LongTensor = None,
+        encoder_hidden_states_t5: torch.Tensor = None,
+        encoder_hidden_states_llama3: torch.Tensor = None,
+        pooled_embeds: torch.Tensor = None,
+        img_ids: torch.Tensor | None = None,
+        img_sizes: list[tuple[int, int]] | None = None,
+        hidden_states_masks: torch.Tensor | None = None,
+        return_dict: bool = True,
+        **kwargs,
+    ) -> tuple[torch.Tensor] | Transformer2DModelOutput:
+        encoder_hidden_states = kwargs.get("encoder_hidden_states", None)
+
+        if encoder_hidden_states is not None:
+            deprecation_message = (
+                "The `encoder_hidden_states` argument is deprecated."
+                "Please use `encoder_hidden_states_t5` and `encoder_hidden_states_llama3` instead."
+            )
+            deprecate("encoder_hidden_states", "0.35.0", deprecation_message)
+            encoder_hidden_states_t5 = encoder_hidden_states[0]
+            encoder_hidden_states_llama3 = encoder_hidden_states[1]
+
+        if img_ids is not None and img_sizes is not None and hidden_states_masks is None:
+            deprecation_message = (
+                "Passing `img_ids` and `img_sizes` with unpachified `hidden_states` is deprecated and will be ignored."
+            )
+            deprecate("img_ids", "0.35.0", deprecation_message)
+
+        if hidden_states_masks is not None and (img_ids is None or img_sizes is None):
+            raise ValueError("if `hidden_states_masks` is passed, `img_ids` and `img_sizes` must also be passed.")
+        elif hidden_states_masks is not None and hidden_states.ndim != 3:
+            raise ValueError(
+                "if `hidden_states_masks` is passed, `hidden_states` must be a 3D tensors with shape"
+                "(batch_size, patch_height * patch_width, patch_size * patch_size * channels)"
+            )
+
+        # spatial forward
+        batch_size = hidden_states.shape[0]
+        hidden_states_type = hidden_states.dtype
+
+        # Patchify the input
+        if hidden_states_masks is None:
+            hidden_states, hidden_states_masks, img_sizes, img_ids = self.patchify(hidden_states)
+
+        # Embed the hidden states
+        hidden_states = self.x_embedder(hidden_states)
+
+        # 0. time
+        timesteps = self.t_embedder(timesteps, hidden_states_type)
+        p_embedder = self.p_embedder(pooled_embeds)
+        temb = timesteps + p_embedder
+
+        encoder_hidden_states = [encoder_hidden_states_llama3[k] for k in self.llama_layers]
+
+        if self.caption_projection is not None:
+            new_encoder_hidden_states = []
+            for i, enc_hidden_state in enumerate(encoder_hidden_states):
+                enc_hidden_state = self.caption_projection[i](enc_hidden_state)
+                enc_hidden_state = enc_hidden_state.view(batch_size, -1, hidden_states.shape[-1])
+                new_encoder_hidden_states.append(enc_hidden_state)
+            encoder_hidden_states = new_encoder_hidden_states
+            encoder_hidden_states_t5 = self.caption_projection[-1](encoder_hidden_states_t5)
+            encoder_hidden_states_t5 = encoder_hidden_states_t5.view(batch_size, -1, hidden_states.shape[-1])
+            encoder_hidden_states.append(encoder_hidden_states_t5)
+
+        txt_ids = torch.zeros(
+            batch_size,
+            encoder_hidden_states[-1].shape[1] + encoder_hidden_states[-2].shape[1] + encoder_hidden_states[0].shape[1],
+            3,
+            device=img_ids.device,
+            dtype=img_ids.dtype,
+        )
+        image_rotary_emb = self.pe_embedder(img_ids)
+        text_rotary_emb = self.pe_embedder(txt_ids)
+        concat_rotary_emb = (
+            torch.cat([image_rotary_emb[..., 0, 0], text_rotary_emb[..., 0, 0]], dim=1),
+            torch.cat([image_rotary_emb[..., 1, 0], text_rotary_emb[..., 1, 0]], dim=1),
+        )
+
+        # 2. Blocks
+        block_id = 0
+        initial_encoder_hidden_states = torch.cat([encoder_hidden_states[-1], encoder_hidden_states[-2]], dim=1)
+        initial_encoder_hidden_states_seq_len = initial_encoder_hidden_states.shape[1]
+        for bid, block in enumerate(self.double_stream_blocks):
+            cur_llama31_encoder_hidden_states = encoder_hidden_states[block_id]
+            cur_encoder_hidden_states = torch.cat(
+                [initial_encoder_hidden_states, cur_llama31_encoder_hidden_states], dim=1
+            )
+            hidden_states, initial_encoder_hidden_states = block(
+                hidden_states=hidden_states,
+                hidden_states_masks=hidden_states_masks,
+                encoder_hidden_states=cur_encoder_hidden_states,
+                temb=temb,
+                image_rotary_emb=concat_rotary_emb,
+            )
+            initial_encoder_hidden_states = initial_encoder_hidden_states[:, :initial_encoder_hidden_states_seq_len]
+            block_id += 1
+
+        image_tokens_seq_len = hidden_states.shape[1]
+        hidden_states = torch.cat([hidden_states, initial_encoder_hidden_states], dim=1)
+        hidden_states_seq_len = hidden_states.shape[1]
+        if hidden_states_masks is not None:
+            encoder_attention_mask_ones = torch.ones(
+                (batch_size, initial_encoder_hidden_states.shape[1] + cur_llama31_encoder_hidden_states.shape[1]),
+                device=hidden_states_masks.device,
+                dtype=hidden_states_masks.dtype,
+            )
+            hidden_states_masks = torch.cat([hidden_states_masks, encoder_attention_mask_ones], dim=1)
+
+        for bid, block in enumerate(self.single_stream_blocks):
+            cur_llama31_encoder_hidden_states = encoder_hidden_states[block_id]
+            hidden_states = torch.cat([hidden_states, cur_llama31_encoder_hidden_states], dim=1)
+
+            hidden_states = block(
+                hidden_states=hidden_states,
+                hidden_states_masks=hidden_states_masks,
+                encoder_hidden_states=None,
+                temb=temb,
+                image_rotary_emb=concat_rotary_emb,
+            )
+            hidden_states = hidden_states[:, :hidden_states_seq_len]
+            block_id += 1
+
+        hidden_states = hidden_states[:, :image_tokens_seq_len, ...]
+        output = self.final_layer(hidden_states, temb)
+        output = self.unpatchify(output, img_sizes, self.training)
+        if hidden_states_masks is not None:
+            hidden_states_masks = hidden_states_masks[:, :image_tokens_seq_len]
+
+        if not return_dict:
+            return (output,)
+        return Transformer2DModelOutput(sample=output)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        tp_rank = get_tensor_model_parallel_rank()
+        tp_size = get_tensor_model_parallel_world_size()
+
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            # self-attn
+            (".to_qkv", ".to_q", "q"),
+            (".to_qkv", ".to_k", "k"),
+            (".to_qkv", ".to_v", "v"),
+            # self-attn no single
+            (".to_qkv_t", ".to_q_t", "q"),
+            (".to_qkv_t", ".to_k_t", "k"),
+            (".to_qkv_t", ".to_v_t", "v"),
+            # cross-attn
+            (".add_kv_proj", ".add_q_proj", "q"),
+            (".add_kv_proj", ".add_k_proj", "k"),
+            (".add_kv_proj", ".add_v_proj", "v"),
+        ]
+
+        params_dict = dict(self.named_parameters())
+
+        # we need to load the buffers for beta and eps (XIELU)
+        for name, buffer in self.named_buffers():
+            if name.endswith(".beta") or name.endswith(".eps"):
+                params_dict[name] = buffer
+
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            original_name = name
+            lookup_name = name
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in original_name:
+                    continue
+                lookup_name = original_name.replace(weight_name, param_name)
+                param = params_dict[lookup_name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                if lookup_name not in params_dict and ".to_out.0." in lookup_name:
+                    lookup_name = lookup_name.replace(".to_out.0.", ".to_out.")
+                if lookup_name not in params_dict:
+                    continue
+                param = params_dict[lookup_name]
+                if tp_size > 1 and any(
+                    norm_name in lookup_name
+                    for norm_name in [
+                        ".attn1.q_rms_norm",
+                        ".attn1.k_rms_norm",
+                    ]
+                ):
+                    shard_size = loaded_weight.shape[0] // tp_size
+                    loaded_weight = loaded_weight[tp_rank * shard_size : (tp_rank + 1) * shard_size]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(original_name)
+            loaded_params.add(lookup_name)
+        return loaded_params
diff --git a/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py b/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py
new file mode 100644
index 000000000..39869a00d
--- /dev/null
+++ b/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py
@@ -0,0 +1,1027 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import inspect
+import json
+import math
+import os
+from collections.abc import Callable, Iterable
+from typing import Any
+
+import numpy as np
+import torch
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models import AutoencoderKL
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import deprecate, logging
+from diffusers.utils.torch_utils import randn_tensor
+from torch import nn
+from transformers import (
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+    LlamaForCausalLM,
+    PreTrainedTokenizerFast,
+    T5EncoderModel,
+    T5Tokenizer,
+)
+from vllm.model_executor.models.utils import AutoWeightsLoader
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.hidream_image import HiDreamImageTransformer2DModel
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
+from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+
+logger = logging.get_logger(__name__)
+
+
+def get_hidream_image_post_process_func(
+    od_config: OmniDiffusionConfig,
+):
+    model_name = od_config.model
+    if os.path.exists(model_name):
+        model_path = model_name
+    else:
+        model_path = download_weights_from_hf_specific(model_name, None, ["*"])
+    vae_config_path = os.path.join(model_path, "vae/config.json")
+    if not os.path.exists(vae_config_path):
+        raise FileNotFoundError(
+            f"VAE config not found at {vae_config_path}. "
+            "Please ensure the model path contains a valid VAE configuration."
+        )
+    with open(vae_config_path) as f:
+        vae_config = json.load(f)
+        vae_scale_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) if "block_out_channels" in vae_config else 8
+
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2)
+
+    def post_process_func(
+        images: torch.Tensor,
+    ):
+        return image_processor.postprocess(images)
+
+    return post_process_func
+
+
+# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: int | None = None,
+    device: str | torch.device | None = None,
+    timesteps: list[int] | None = None,
+    sigmas: list[float] | None = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`list[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`list[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        timesteps (`torch.Tensor`): The timestep schedule from the scheduler.
+        num_inference_steps (`int`): The number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class HiDreamI1ImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin, ProgressBarMixin):
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.od_config = od_config
+        dtype = getattr(od_config, "dtype", torch.bfloat16)
+        llama_path = self.od_config.extras["auxiliary_text_encoder"]
+        if llama_path is None:
+            logger.warning(
+                f"auxiliary_text_encoder is not provided. "
+                f"Attempting to load default LLaMA model: {self.od_config.extras['default_llama_model_id']}"
+            )
+            llama_path = self.od_config.extras["default_llama_model_id"]
+        logger.info(f"Loading LLAMA model from: {llama_path}")
+
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder="transformer",
+                revision=None,
+                prefix="transformer.",
+                fall_back_to_pt=True,
+            ),
+        ]
+
+        self.device = get_local_device()
+        model = od_config.model
+        # Check if model is a local path
+        local_files_only = os.path.exists(model)
+
+        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            model, subfolder="scheduler", local_files_only=local_files_only
+        )
+        self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
+            self.device
+        )
+        self.text_encoder = CLIPTextModelWithProjection.from_pretrained(
+            model, subfolder="text_encoder", local_files_only=local_files_only
+        )
+        self.tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
+        self.text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(
+            model, subfolder="text_encoder_2", local_files_only=local_files_only
+        )
+        self.tokenizer_2 = CLIPTokenizer.from_pretrained(
+            model, subfolder="tokenizer_2", local_files_only=local_files_only
+        )
+        self.text_encoder_3 = T5EncoderModel.from_pretrained(
+            model, subfolder="text_encoder_3", local_files_only=local_files_only
+        )
+        self.tokenizer_3 = T5Tokenizer.from_pretrained(
+            model, subfolder="tokenizer_3", local_files_only=local_files_only
+        )
+        self.text_encoder_4 = LlamaForCausalLM.from_pretrained(llama_path, output_hidden_states=True, dtype=dtype).to(
+            self.device
+        )
+        self.tokenizer_4 = PreTrainedTokenizerFast.from_pretrained(llama_path, use_fast=False)
+        transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, HiDreamImageTransformer2DModel)
+        self.transformer = HiDreamImageTransformer2DModel(
+            od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs
+        )
+
+        self.stage = None
+
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
+        # HiDreamImage latents are turned into 2x2 patches and packed.
+        # This means the latent width and height has to be divisible
+        # by the patch size. So the vae scale factor is multiplied by the patch size to account for this
+        self.default_sample_size = 128
+        if getattr(self, "tokenizer_4", None) is not None:
+            self.tokenizer_4.pad_token = self.tokenizer_4.eos_token
+
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
+        )
+
+    def check_inputs(
+        self,
+        prompt,
+        prompt_2,
+        prompt_3,
+        prompt_4,
+        negative_prompt=None,
+        negative_prompt_2=None,
+        negative_prompt_3=None,
+        negative_prompt_4=None,
+        prompt_embeds_t5=None,
+        prompt_embeds_llama3=None,
+        negative_prompt_embeds_t5=None,
+        negative_prompt_embeds_llama3=None,
+        pooled_prompt_embeds=None,
+        negative_pooled_prompt_embeds=None,
+        callback_on_step_end_tensor_inputs=None,
+    ):
+        if prompt is not None and pooled_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `pooled_prompt_embeds`: {pooled_prompt_embeds}. "
+                "Please make sure to only forward one of the two."
+            )
+        elif prompt_2 is not None and pooled_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_2`: {prompt_2} and `pooled_prompt_embeds`: {pooled_prompt_embeds}."
+                "Please make sure to only forward one of the two."
+            )
+        elif prompt_3 is not None and prompt_embeds_t5 is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_3`: {prompt_3} and `prompt_embeds_t5`: {prompt_embeds_t5}."
+                "Please make sure to only forward one of the two."
+            )
+        elif prompt_4 is not None and prompt_embeds_llama3 is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_4`: {prompt_4} and `prompt_embeds_llama3`: {prompt_embeds_llama3}."
+                "Please make sure to only forward one of the two."
+            )
+        elif prompt is None and pooled_prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `pooled_prompt_embeds`."
+                "Cannot leave both `prompt` and `pooled_prompt_embeds` undefined."
+            )
+        elif prompt is None and prompt_embeds_t5 is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds_t5`."
+                "Cannot leave both `prompt` and `prompt_embeds_t5` undefined."
+            )
+        elif prompt is None and prompt_embeds_llama3 is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds_llama3`."
+                "Cannot leave both `prompt` and `prompt_embeds_llama3` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+        elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
+            raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
+        elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)):
+            raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}")
+        elif prompt_4 is not None and (not isinstance(prompt_4, str) and not isinstance(prompt_4, list)):
+            raise ValueError(f"`prompt_4` has to be of type `str` or `list` but is {type(prompt_4)}")
+
+        if negative_prompt is not None and negative_pooled_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_pooled_prompt_embeds`:"
+                f" {negative_pooled_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_2 is not None and negative_pooled_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_pooled_prompt_embeds`:"
+                f" {negative_pooled_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_3 is not None and negative_prompt_embeds_t5 is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} and `negative_prompt_embeds_t5`:"
+                f" {negative_prompt_embeds_t5}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_4 is not None and negative_prompt_embeds_llama3 is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_4`: {negative_prompt_4} and `negative_prompt_embeds_llama3`:"
+                f" {negative_prompt_embeds_llama3}. Please make sure to only forward one of the two."
+            )
+
+        if pooled_prompt_embeds is not None and negative_pooled_prompt_embeds is not None:
+            if pooled_prompt_embeds.shape != negative_pooled_prompt_embeds.shape:
+                raise ValueError(
+                    "`pooled_prompt_embeds` and `negative_pooled_prompt_embeds`"
+                    " must have the same shape when passed directly, but"
+                    f" got: `pooled_prompt_embeds` {pooled_prompt_embeds.shape} != `negative_pooled_prompt_embeds`"
+                    f" {negative_pooled_prompt_embeds.shape}."
+                )
+        if prompt_embeds_t5 is not None and negative_prompt_embeds_t5 is not None:
+            if prompt_embeds_t5.shape != negative_prompt_embeds_t5.shape:
+                raise ValueError(
+                    "`prompt_embeds_t5` and `negative_prompt_embeds_t5`"
+                    " must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds_t5` {prompt_embeds_t5.shape} != `negative_prompt_embeds_t5`"
+                    f" {negative_prompt_embeds_t5.shape}."
+                )
+        if prompt_embeds_llama3 is not None and negative_prompt_embeds_llama3 is not None:
+            if prompt_embeds_llama3.shape != negative_prompt_embeds_llama3.shape:
+                raise ValueError(
+                    "`prompt_embeds_llama3` and `negative_prompt_embeds_llama3`"
+                    " must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds_llama3` {prompt_embeds_llama3.shape} != `negative_prompt_embeds_llama3`"
+                    f" {negative_prompt_embeds_llama3.shape}."
+                )
+
+    def _get_t5_prompt_embeds(
+        self,
+        prompt: str | list[str] = None,
+        max_sequence_length: int = 128,
+        dtype: torch.dtype | None = None,
+    ):
+        dtype = dtype or self.text_encoder_3.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        text_inputs = self.tokenizer_3(
+            prompt,
+            padding="max_length",
+            max_length=min(max_sequence_length, self.tokenizer_3.model_max_length),
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        attention_mask = text_inputs.attention_mask
+        untruncated_ids = self.tokenizer_3(prompt, padding="longest", return_tensors="pt").input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer_3.batch_decode(
+                untruncated_ids[:, min(max_sequence_length, self.tokenizer_3.model_max_length) - 1 : -1]
+            )
+            logger.warning(
+                "The following part of your input was truncated because `max_sequence_length` is set to "
+                f" {min(max_sequence_length, self.tokenizer_3.model_max_length)} tokens: {removed_text}"
+            )
+
+        prompt_embeds = self.text_encoder_3(
+            text_input_ids.to(self.device), attention_mask=attention_mask.to(self.device)
+        )[0]
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device)
+        return prompt_embeds
+
+    def _get_clip_prompt_embeds(
+        self,
+        tokenizer,
+        text_encoder,
+        prompt: str | list[str],
+        max_sequence_length: int = 128,
+        dtype: torch.dtype | None = None,
+    ):
+        dtype = dtype or text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        model_max_length = self.tokenizer.model_max_length
+        text_inputs = tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=min(max_sequence_length, model_max_length),
+            truncation=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = tokenizer.batch_decode(untruncated_ids[:, model_max_length - 1 : -1])
+            logger.warning(
+                "The following part of your input was truncated because CLIP can only handle sequences up to"
+                f" {model_max_length} tokens: {removed_text}"
+            )
+        prompt_embeds = text_encoder(text_input_ids.to(self.device), output_hidden_states=True)
+
+        # Use pooled output of CLIPTextModel
+        prompt_embeds = prompt_embeds[0]
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device)
+        return prompt_embeds
+
+    def _get_llama3_prompt_embeds(
+        self,
+        prompt: str | list[str] = None,
+        max_sequence_length: int = 128,
+        dtype: torch.dtype | None = None,
+    ):
+        dtype = dtype or self.text_encoder_4.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        text_inputs = self.tokenizer_4(
+            prompt,
+            padding="max_length",
+            max_length=min(max_sequence_length, self.tokenizer_4.model_max_length),
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        attention_mask = text_inputs.attention_mask
+        untruncated_ids = self.tokenizer_4(prompt, padding="longest", return_tensors="pt").input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer_4.batch_decode(
+                untruncated_ids[:, min(max_sequence_length, self.tokenizer_4.model_max_length) - 1 : -1]
+            )
+            logger.warning(
+                "The following part of your input was truncated because `max_sequence_length` is set to "
+                f" {min(max_sequence_length, self.tokenizer_4.model_max_length)} tokens: {removed_text}"
+            )
+
+        outputs = self.text_encoder_4(
+            text_input_ids.to(self.device),
+            attention_mask=attention_mask.to(self.device),
+            output_hidden_states=True,
+        )
+
+        prompt_embeds = outputs.hidden_states[1:]
+        prompt_embeds = torch.stack(prompt_embeds, dim=0)
+        return prompt_embeds
+
+    def encode_prompt(
+        self,
+        prompt: str | list[str] | None = None,
+        prompt_2: str | list[str] | None = None,
+        prompt_3: str | list[str] | None = None,
+        prompt_4: str | list[str] | None = None,
+        dtype: torch.dtype | None = None,
+        num_images_per_prompt: int = 1,
+        do_classifier_free_guidance: bool = True,
+        negative_prompt: str | list[str] | None = None,
+        negative_prompt_2: str | list[str] | None = None,
+        negative_prompt_3: str | list[str] | None = None,
+        negative_prompt_4: str | list[str] | None = None,
+        prompt_embeds_t5: list[torch.FloatTensor] | None = None,
+        prompt_embeds_llama3: list[torch.FloatTensor] | None = None,
+        negative_prompt_embeds_t5: list[torch.FloatTensor] | None = None,
+        negative_prompt_embeds_llama3: list[torch.FloatTensor] | None = None,
+        pooled_prompt_embeds: torch.FloatTensor | None = None,
+        negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+        max_sequence_length: int = 128,
+        lora_scale: float | None = None,
+    ):
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        if prompt is not None:
+            batch_size = len(prompt)
+        else:
+            batch_size = pooled_prompt_embeds.shape[0]
+
+        if pooled_prompt_embeds is None:
+            pooled_prompt_embeds_1 = self._get_clip_prompt_embeds(
+                self.tokenizer, self.text_encoder, prompt, max_sequence_length, dtype
+            )
+
+        if do_classifier_free_guidance and negative_pooled_prompt_embeds is None:
+            negative_prompt = negative_prompt or ""
+            negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+
+            if len(negative_prompt) > 1 and len(negative_prompt) != batch_size:
+                raise ValueError(f"negative_prompt must be of length 1 or {batch_size}")
+
+            negative_pooled_prompt_embeds_1 = self._get_clip_prompt_embeds(
+                self.tokenizer, self.text_encoder, negative_prompt, max_sequence_length, dtype
+            )
+
+            if negative_pooled_prompt_embeds_1.shape[0] == 1 and batch_size > 1:
+                negative_pooled_prompt_embeds_1 = negative_pooled_prompt_embeds_1.repeat(batch_size, 1)
+
+        if pooled_prompt_embeds is None:
+            prompt_2 = prompt_2 or prompt
+            prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+            if len(prompt_2) > 1 and len(prompt_2) != batch_size:
+                raise ValueError(f"prompt_2 must be of length 1 or {batch_size}")
+
+            pooled_prompt_embeds_2 = self._get_clip_prompt_embeds(
+                self.tokenizer_2, self.text_encoder_2, prompt_2, max_sequence_length, dtype
+            )
+
+            if pooled_prompt_embeds_2.shape[0] == 1 and batch_size > 1:
+                pooled_prompt_embeds_2 = pooled_prompt_embeds_2.repeat(batch_size, 1)
+
+        if do_classifier_free_guidance and negative_pooled_prompt_embeds is None:
+            negative_prompt_2 = negative_prompt_2 or negative_prompt
+            negative_prompt_2 = [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
+
+            if len(negative_prompt_2) > 1 and len(negative_prompt_2) != batch_size:
+                raise ValueError(f"negative_prompt_2 must be of length 1 or {batch_size}")
+
+            negative_pooled_prompt_embeds_2 = self._get_clip_prompt_embeds(
+                self.tokenizer_2, self.text_encoder_2, negative_prompt_2, max_sequence_length, dtype
+            )
+
+            if negative_pooled_prompt_embeds_2.shape[0] == 1 and batch_size > 1:
+                negative_pooled_prompt_embeds_2 = negative_pooled_prompt_embeds_2.repeat(batch_size, 1)
+
+        if pooled_prompt_embeds is None:
+            pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_1, pooled_prompt_embeds_2], dim=-1)
+
+        if do_classifier_free_guidance and negative_pooled_prompt_embeds is None:
+            negative_pooled_prompt_embeds = torch.cat(
+                [negative_pooled_prompt_embeds_1, negative_pooled_prompt_embeds_2], dim=-1
+            )
+
+        if prompt_embeds_t5 is None:
+            prompt_3 = prompt_3 or prompt
+            prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3
+
+            if len(prompt_3) > 1 and len(prompt_3) != batch_size:
+                raise ValueError(f"prompt_3 must be of length 1 or {batch_size}")
+
+            prompt_embeds_t5 = self._get_t5_prompt_embeds(prompt_3, max_sequence_length, dtype)
+
+            if prompt_embeds_t5.shape[0] == 1 and batch_size > 1:
+                prompt_embeds_t5 = prompt_embeds_t5.repeat(batch_size, 1, 1)
+
+        if do_classifier_free_guidance and negative_prompt_embeds_t5 is None:
+            negative_prompt_3 = negative_prompt_3 or negative_prompt
+            negative_prompt_3 = [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3
+
+            if len(negative_prompt_3) > 1 and len(negative_prompt_3) != batch_size:
+                raise ValueError(f"negative_prompt_3 must be of length 1 or {batch_size}")
+
+            negative_prompt_embeds_t5 = self._get_t5_prompt_embeds(negative_prompt_3, max_sequence_length, dtype)
+
+            if negative_prompt_embeds_t5.shape[0] == 1 and batch_size > 1:
+                negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(batch_size, 1, 1)
+
+        if prompt_embeds_llama3 is None:
+            prompt_4 = prompt_4 or prompt
+            prompt_4 = [prompt_4] if isinstance(prompt_4, str) else prompt_4
+
+            if len(prompt_4) > 1 and len(prompt_4) != batch_size:
+                raise ValueError(f"prompt_4 must be of length 1 or {batch_size}")
+
+            prompt_embeds_llama3 = self._get_llama3_prompt_embeds(prompt_4, max_sequence_length, dtype)
+
+            if prompt_embeds_llama3.shape[0] == 1 and batch_size > 1:
+                prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, batch_size, 1, 1)
+
+        if do_classifier_free_guidance and negative_prompt_embeds_llama3 is None:
+            negative_prompt_4 = negative_prompt_4 or negative_prompt
+            negative_prompt_4 = [negative_prompt_4] if isinstance(negative_prompt_4, str) else negative_prompt_4
+
+            if len(negative_prompt_4) > 1 and len(negative_prompt_4) != batch_size:
+                raise ValueError(f"negative_prompt_4 must be of length 1 or {batch_size}")
+
+            negative_prompt_embeds_llama3 = self._get_llama3_prompt_embeds(
+                negative_prompt_4, max_sequence_length, dtype
+            )
+
+            if negative_prompt_embeds_llama3.shape[0] == 1 and batch_size > 1:
+                negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, batch_size, 1, 1)
+
+        # duplicate pooled_prompt_embeds for each generation per prompt
+        pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt)
+        pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+        # duplicate t5_prompt_embeds for batch_size and num_images_per_prompt
+        bs_embed, seq_len, _ = prompt_embeds_t5.shape
+        if bs_embed == 1 and batch_size > 1:
+            prompt_embeds_t5 = prompt_embeds_t5.repeat(batch_size, 1, 1)
+        elif bs_embed > 1 and bs_embed != batch_size:
+            raise ValueError(f"cannot duplicate prompt_embeds_t5 of batch size {bs_embed}")
+        prompt_embeds_t5 = prompt_embeds_t5.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds_t5 = prompt_embeds_t5.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        # duplicate llama3_prompt_embeds for batch_size and num_images_per_prompt
+        _, bs_embed, seq_len, dim = prompt_embeds_llama3.shape
+        if bs_embed == 1 and batch_size > 1:
+            prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, batch_size, 1, 1)
+        elif bs_embed > 1 and bs_embed != batch_size:
+            raise ValueError(f"cannot duplicate prompt_embeds_llama3 of batch size {bs_embed}")
+        prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, 1, num_images_per_prompt, 1)
+        prompt_embeds_llama3 = prompt_embeds_llama3.view(-1, batch_size * num_images_per_prompt, seq_len, dim)
+
+        if do_classifier_free_guidance:
+            # duplicate negative_pooled_prompt_embeds for batch_size and num_images_per_prompt
+            bs_embed, seq_len = negative_pooled_prompt_embeds.shape
+            if bs_embed == 1 and batch_size > 1:
+                negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(batch_size, 1)
+            elif bs_embed > 1 and bs_embed != batch_size:
+                raise ValueError(f"cannot duplicate negative_pooled_prompt_embeds of batch size {bs_embed}")
+            negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt)
+            negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+            # duplicate negative_t5_prompt_embeds for batch_size and num_images_per_prompt
+            bs_embed, seq_len, _ = negative_prompt_embeds_t5.shape
+            if bs_embed == 1 and batch_size > 1:
+                negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(batch_size, 1, 1)
+            elif bs_embed > 1 and bs_embed != batch_size:
+                raise ValueError(f"cannot duplicate negative_prompt_embeds_t5 of batch size {bs_embed}")
+            negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(1, num_images_per_prompt, 1)
+            negative_prompt_embeds_t5 = negative_prompt_embeds_t5.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+            # duplicate negative_prompt_embeds_llama3 for batch_size and num_images_per_prompt
+            _, bs_embed, seq_len, dim = negative_prompt_embeds_llama3.shape
+            if bs_embed == 1 and batch_size > 1:
+                negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, batch_size, 1, 1)
+            elif bs_embed > 1 and bs_embed != batch_size:
+                raise ValueError(f"cannot duplicate negative_prompt_embeds_llama3 of batch size {bs_embed}")
+            negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, 1, num_images_per_prompt, 1)
+            negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.view(
+                -1, batch_size * num_images_per_prompt, seq_len, dim
+            )
+
+        return (
+            prompt_embeds_t5,
+            negative_prompt_embeds_t5,
+            prompt_embeds_llama3,
+            negative_prompt_embeds_llama3,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        )
+
+    def enable_vae_slicing(self):
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        depr_message = (
+            f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this"
+            " method will be removed in a future version. Please use `pipe.vae.enable_slicing()`."
+        )
+        deprecate(
+            "enable_vae_slicing",
+            "0.40.0",
+            depr_message,
+        )
+        self.vae.enable_slicing()
+
+    def disable_vae_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        depr_message = (
+            f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this"
+            " method will be removed in a future version. Please use `pipe.vae.disable_slicing()`."
+        )
+        deprecate(
+            "disable_vae_slicing",
+            "0.40.0",
+            depr_message,
+        )
+        self.vae.disable_slicing()
+
+    def enable_vae_tiling(self):
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+        """
+        depr_message = (
+            f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}`"
+            " is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`."
+        )
+        deprecate(
+            "enable_vae_tiling",
+            "0.40.0",
+            depr_message,
+        )
+        self.vae.enable_tiling()
+
+    def disable_vae_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        depr_message = (
+            f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is"
+            " deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`."
+        )
+        deprecate(
+            "disable_vae_tiling",
+            "0.40.0",
+            depr_message,
+        )
+        self.vae.disable_tiling()
+
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        generator,
+        latents=None,
+    ):
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (self.vae_scale_factor * 2))
+        width = 2 * (int(width) // (self.vae_scale_factor * 2))
+
+        shape = (batch_size, num_channels_latents, height, width)
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=self.device, dtype=dtype)
+        else:
+            if latents.shape != shape:
+                raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
+            latents = latents.to(self.device)
+        return latents
+
+    def prepare_timesteps(self, num_inference_steps, sigmas, image_seq_len):
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.get("base_image_seq_len", 256),
+            self.scheduler.config.get("max_image_seq_len", 4096),
+            self.scheduler.config.get("base_shift", 0.5),
+            self.scheduler.config.get("max_shift", 1.15),
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            self.device,
+            sigmas=sigmas,
+            mu=mu,
+        )
+        return timesteps, num_inference_steps
+
+    def _extract_prompts(self, prompts):
+        """Extract prompt and negative_prompt from OmniPromptType list."""
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts] or None
+        if all(isinstance(p, str) or p.get("negative_prompt") is None for p in prompts):
+            negative_prompt = None
+        elif prompts:
+            negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in prompts]
+        else:
+            negative_prompt = None
+        return prompt, negative_prompt
+
+    def check_cfg_parallel_validity(self, true_cfg_scale: float, has_neg_prompt: bool):
+        if get_classifier_free_guidance_world_size() == 1:
+            return True
+
+        if true_cfg_scale <= 1:
+            logger.warning("CFG parallel is NOT working correctly when true_cfg_scale <= 1.")
+            return False
+
+        if not has_neg_prompt:
+            logger.warning(
+                "CFG parallel is NOT working correctly when there is no negative prompt or negative prompt embeddings."
+            )
+            return False
+        return True
+
+    def diffuse(
+        self,
+        prompt_embeds_t5: torch.Tensor,
+        prompt_embeds_llama3: torch.Tensor,
+        pooled_prompt_embeds: torch.Tensor,
+        latents: torch.Tensor,
+        timesteps: torch.Tensor,
+        do_true_cfg: bool,
+    ) -> torch.Tensor:
+        with self.progress_bar(total=len(timesteps)) as pbar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latent_model_input.shape[0])
+                noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    timesteps=timestep,
+                    encoder_hidden_states_t5=prompt_embeds_t5,
+                    encoder_hidden_states_llama3=prompt_embeds_llama3,
+                    pooled_embeds=pooled_prompt_embeds,
+                    return_dict=False,
+                )[0]
+                noise_pred = -noise_pred
+
+                # TODO: Modify CFG guidance
+                # perform guidance
+                if self.do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
+
+                pbar.update()
+
+        return latents
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale > 1
+
+    @property
+    def attention_kwargs(self):
+        return self._attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    def forward(
+        self,
+        req: OmniDiffusionRequest,
+        prompt: str | list[str] = None,
+        prompt_2: str | list[str] | None = None,
+        prompt_3: str | list[str] | None = None,
+        prompt_4: str | list[str] | None = None,
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int = 50,
+        sigmas: list[float] | None = None,
+        guidance_scale: float = 5.0,
+        negative_prompt: str | list[str] | None = None,
+        negative_prompt_2: str | list[str] | None = None,
+        negative_prompt_3: str | list[str] | None = None,
+        negative_prompt_4: str | list[str] | None = None,
+        num_images_per_prompt: int | None = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.FloatTensor | None = None,
+        prompt_embeds_t5: torch.FloatTensor | None = None,
+        prompt_embeds_llama3: torch.FloatTensor | None = None,
+        negative_prompt_embeds_t5: torch.FloatTensor | None = None,
+        negative_prompt_embeds_llama3: torch.FloatTensor | None = None,
+        pooled_prompt_embeds: torch.FloatTensor | None = None,
+        negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+        output_type: str | None = "pil",
+        return_dict: bool = True,
+        attention_kwargs: dict[str, Any] | None = None,
+        callback_on_step_end: Callable[[int, int], None] | None = None,
+        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+        max_sequence_length: int = 128,
+        **kwargs,
+    ):
+        extracted_prompt, negative_prompt = self._extract_prompts(req.prompts)
+        prompt = extracted_prompt or prompt
+
+        height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
+        width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
+
+        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        sigmas = req.sampling_params.sigmas or sigmas
+        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
+        generator = req.sampling_params.generator or generator
+        true_cfg_scale = req.sampling_params.true_cfg_scale or guidance_scale
+        if req.sampling_params.guidance_scale_provided:
+            guidance_scale = req.sampling_params.guidance_scale
+        num_images_per_prompt = (
+            req.sampling_params.num_outputs_per_prompt
+            if req.sampling_params.num_outputs_per_prompt > 0
+            else num_images_per_prompt
+        )
+
+        prompt_embeds = kwargs.get("prompt_embeds", None)
+        negative_prompt_embeds = kwargs.get("negative_prompt_embeds", None)
+
+        if prompt_embeds is not None:
+            deprecation_message = (
+                "The `prompt_embeds` argument is deprecated."
+                " Please use `prompt_embeds_t5` and `prompt_embeds_llama3` instead."
+            )
+            deprecate("prompt_embeds", "0.35.0", deprecation_message)
+            prompt_embeds_t5 = prompt_embeds[0]
+            prompt_embeds_llama3 = prompt_embeds[1]
+
+        if negative_prompt_embeds is not None:
+            deprecation_message = (
+                "The `negative_prompt_embeds` argument is deprecated."
+                "Please use `negative_prompt_embeds_t5` and `negative_prompt_embeds_llama3` instead."
+            )
+            deprecate("negative_prompt_embeds", "0.35.0", deprecation_message)
+            negative_prompt_embeds_t5 = negative_prompt_embeds[0]
+            negative_prompt_embeds_llama3 = negative_prompt_embeds[1]
+
+        division = self.vae_scale_factor * 2
+        S_max = (self.default_sample_size * self.vae_scale_factor) ** 2
+        scale = S_max / (width * height)
+        scale = math.sqrt(scale)
+        width, height = int(width * scale // division * division), int(height * scale // division * division)
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            prompt_3,
+            prompt_4,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            negative_prompt_3=negative_prompt_3,
+            negative_prompt_4=negative_prompt_4,
+            prompt_embeds_t5=prompt_embeds_t5,
+            prompt_embeds_llama3=prompt_embeds_llama3,
+            negative_prompt_embeds_t5=negative_prompt_embeds_t5,
+            negative_prompt_embeds_llama3=negative_prompt_embeds_llama3,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._attention_kwargs = attention_kwargs
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        elif pooled_prompt_embeds is not None:
+            batch_size = pooled_prompt_embeds.shape[0]
+
+        # TODO: CFG guidance configuration
+        has_neg_prompt = negative_prompt is not None or (
+            negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None
+        )
+        do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
+        self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt)
+
+        # 3. Encode prompt
+        lora_scale = self.attention_kwargs.get("scale", None) if self.attention_kwargs is not None else None
+        (
+            prompt_embeds_t5,
+            negative_prompt_embeds_t5,
+            prompt_embeds_llama3,
+            negative_prompt_embeds_llama3,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_3=prompt_3,
+            prompt_4=prompt_4,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            negative_prompt_3=negative_prompt_3,
+            negative_prompt_4=negative_prompt_4,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            prompt_embeds_t5=prompt_embeds_t5,
+            prompt_embeds_llama3=prompt_embeds_llama3,
+            negative_prompt_embeds_t5=negative_prompt_embeds_t5,
+            negative_prompt_embeds_llama3=negative_prompt_embeds_llama3,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            lora_scale=lora_scale,
+        )
+
+        if self.do_classifier_free_guidance:
+            prompt_embeds_t5 = torch.cat([negative_prompt_embeds_t5, prompt_embeds_t5], dim=0)
+            prompt_embeds_llama3 = torch.cat([negative_prompt_embeds_llama3, prompt_embeds_llama3], dim=1)
+            pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0)
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            pooled_prompt_embeds.dtype,
+            generator,
+            latents,
+        )
+
+        # 5. Prepare timesteps
+        timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1])
+        self._num_timesteps = len(timesteps)
+
+        # 6. Denoising loop
+        latents = self.diffuse(
+            prompt_embeds_t5,
+            prompt_embeds_llama3,
+            pooled_prompt_embeds,
+            latents,
+            timesteps,
+            do_true_cfg,
+        )
+
+        if output_type == "latent":
+            image = latents
+
+        else:
+            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+
+            image = self.vae.decode(latents, return_dict=False)[0]
+
+        return DiffusionOutput(
+            output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)
diff --git a/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_image.py b/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_image.py
new file mode 100644
index 000000000..0c5486ae5
--- /dev/null
+++ b/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_image.py
@@ -0,0 +1,1027 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import inspect
+import json
+import math
+import os
+from collections.abc import Callable, Iterable
+from typing import Any
+
+import numpy as np
+import torch
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models import AutoencoderKL
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import deprecate, logging
+from diffusers.utils.torch_utils import randn_tensor
+from torch import nn
+from transformers import (
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+    LlamaForCausalLM,
+    PreTrainedTokenizerFast,
+    T5EncoderModel,
+    T5Tokenizer,
+)
+from vllm.model_executor.models.utils import AutoWeightsLoader
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.hidream_image import HiDreamImageTransformer2DModel
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
+from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+
+logger = logging.get_logger(__name__)
+
+
+def get_hidream_image_post_process_func(
+    od_config: OmniDiffusionConfig,
+):
+    model_name = od_config.model
+    if os.path.exists(model_name):
+        model_path = model_name
+    else:
+        model_path = download_weights_from_hf_specific(model_name, None, ["*"])
+    vae_config_path = os.path.join(model_path, "vae/config.json")
+    if not os.path.exists(vae_config_path):
+        raise FileNotFoundError(
+            f"VAE config not found at {vae_config_path}. "
+            "Please ensure the model path contains a valid VAE configuration."
+        )
+    with open(vae_config_path) as f:
+        vae_config = json.load(f)
+        vae_scale_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) if "block_out_channels" in vae_config else 8
+
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2)
+
+    def post_process_func(
+        images: torch.Tensor,
+    ):
+        return image_processor.postprocess(images)
+
+    return post_process_func
+
+
+# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: int | None = None,
+    device: str | torch.device | None = None,
+    timesteps: list[int] | None = None,
+    sigmas: list[float] | None = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`list[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`list[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        timesteps (`torch.Tensor`): The timestep schedule from the scheduler.
+        num_inference_steps (`int`): The number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class HiDreamImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin, ProgressBarMixin):
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.od_config = od_config
+        dtype = getattr(od_config, "dtype", torch.bfloat16)
+        llama_path = self.od_config.extras["auxiliary_text_encoder"]
+        if llama_path is None:
+            logger.warning(
+                f"auxiliary_text_encoder is not provided. "
+                f"Attempting to load default LLaMA model: {self.od_config.extras['default_llama_model_id']}"
+            )
+            llama_path = self.od_config.extras["default_llama_model_id"]
+        logger.info(f"Loading LLAMA model from: {llama_path}")
+
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder="transformer",
+                revision=None,
+                prefix="transformer.",
+                fall_back_to_pt=True,
+            ),
+        ]
+
+        self.device = get_local_device()
+        model = od_config.model
+        # Check if model is a local path
+        local_files_only = os.path.exists(model)
+
+        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            model, subfolder="scheduler", local_files_only=local_files_only
+        )
+        self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
+            self.device
+        )
+        self.text_encoder = CLIPTextModelWithProjection.from_pretrained(
+            model, subfolder="text_encoder", local_files_only=local_files_only
+        )
+        self.tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
+        self.text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(
+            model, subfolder="text_encoder_2", local_files_only=local_files_only
+        )
+        self.tokenizer_2 = CLIPTokenizer.from_pretrained(
+            model, subfolder="tokenizer_2", local_files_only=local_files_only
+        )
+        self.text_encoder_3 = T5EncoderModel.from_pretrained(
+            model, subfolder="text_encoder_3", local_files_only=local_files_only
+        )
+        self.tokenizer_3 = T5Tokenizer.from_pretrained(
+            model, subfolder="tokenizer_3", local_files_only=local_files_only
+        )
+        self.text_encoder_4 = LlamaForCausalLM.from_pretrained(llama_path, output_hidden_states=True, dtype=dtype).to(
+            self.device
+        )
+        self.tokenizer_4 = PreTrainedTokenizerFast.from_pretrained(llama_path, use_fast=False)
+        transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, HiDreamImageTransformer2DModel)
+        self.transformer = HiDreamImageTransformer2DModel(
+            od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs
+        )
+
+        self.stage = None
+
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
+        # HiDreamImage latents are turned into 2x2 patches and packed.
+        # This means the latent width and height has to be divisible
+        # by the patch size. So the vae scale factor is multiplied by the patch size to account for this
+        self.default_sample_size = 128
+        if getattr(self, "tokenizer_4", None) is not None:
+            self.tokenizer_4.pad_token = self.tokenizer_4.eos_token
+
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
+        )
+
+    def check_inputs(
+        self,
+        prompt,
+        prompt_2,
+        prompt_3,
+        prompt_4,
+        negative_prompt=None,
+        negative_prompt_2=None,
+        negative_prompt_3=None,
+        negative_prompt_4=None,
+        prompt_embeds_t5=None,
+        prompt_embeds_llama3=None,
+        negative_prompt_embeds_t5=None,
+        negative_prompt_embeds_llama3=None,
+        pooled_prompt_embeds=None,
+        negative_pooled_prompt_embeds=None,
+        callback_on_step_end_tensor_inputs=None,
+    ):
+        if prompt is not None and pooled_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `pooled_prompt_embeds`: {pooled_prompt_embeds}. "
+                "Please make sure to only forward one of the two."
+            )
+        elif prompt_2 is not None and pooled_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_2`: {prompt_2} and `pooled_prompt_embeds`: {pooled_prompt_embeds}."
+                "Please make sure to only forward one of the two."
+            )
+        elif prompt_3 is not None and prompt_embeds_t5 is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_3`: {prompt_3} and `prompt_embeds_t5`: {prompt_embeds_t5}."
+                "Please make sure to only forward one of the two."
+            )
+        elif prompt_4 is not None and prompt_embeds_llama3 is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_4`: {prompt_4} and `prompt_embeds_llama3`: {prompt_embeds_llama3}."
+                "Please make sure to only forward one of the two."
+            )
+        elif prompt is None and pooled_prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `pooled_prompt_embeds`."
+                "Cannot leave both `prompt` and `pooled_prompt_embeds` undefined."
+            )
+        elif prompt is None and prompt_embeds_t5 is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds_t5`."
+                "Cannot leave both `prompt` and `prompt_embeds_t5` undefined."
+            )
+        elif prompt is None and prompt_embeds_llama3 is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds_llama3`."
+                "Cannot leave both `prompt` and `prompt_embeds_llama3` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+        elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
+            raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
+        elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)):
+            raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}")
+        elif prompt_4 is not None and (not isinstance(prompt_4, str) and not isinstance(prompt_4, list)):
+            raise ValueError(f"`prompt_4` has to be of type `str` or `list` but is {type(prompt_4)}")
+
+        if negative_prompt is not None and negative_pooled_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_pooled_prompt_embeds`:"
+                f" {negative_pooled_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_2 is not None and negative_pooled_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_pooled_prompt_embeds`:"
+                f" {negative_pooled_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_3 is not None and negative_prompt_embeds_t5 is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} and `negative_prompt_embeds_t5`:"
+                f" {negative_prompt_embeds_t5}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_4 is not None and negative_prompt_embeds_llama3 is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_4`: {negative_prompt_4} and `negative_prompt_embeds_llama3`:"
+                f" {negative_prompt_embeds_llama3}. Please make sure to only forward one of the two."
+            )
+
+        if pooled_prompt_embeds is not None and negative_pooled_prompt_embeds is not None:
+            if pooled_prompt_embeds.shape != negative_pooled_prompt_embeds.shape:
+                raise ValueError(
+                    "`pooled_prompt_embeds` and `negative_pooled_prompt_embeds`"
+                    " must have the same shape when passed directly, but"
+                    f" got: `pooled_prompt_embeds` {pooled_prompt_embeds.shape} != `negative_pooled_prompt_embeds`"
+                    f" {negative_pooled_prompt_embeds.shape}."
+                )
+        if prompt_embeds_t5 is not None and negative_prompt_embeds_t5 is not None:
+            if prompt_embeds_t5.shape != negative_prompt_embeds_t5.shape:
+                raise ValueError(
+                    "`prompt_embeds_t5` and `negative_prompt_embeds_t5`"
+                    " must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds_t5` {prompt_embeds_t5.shape} != `negative_prompt_embeds_t5`"
+                    f" {negative_prompt_embeds_t5.shape}."
+                )
+        if prompt_embeds_llama3 is not None and negative_prompt_embeds_llama3 is not None:
+            if prompt_embeds_llama3.shape != negative_prompt_embeds_llama3.shape:
+                raise ValueError(
+                    "`prompt_embeds_llama3` and `negative_prompt_embeds_llama3`"
+                    " must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds_llama3` {prompt_embeds_llama3.shape} != `negative_prompt_embeds_llama3`"
+                    f" {negative_prompt_embeds_llama3.shape}."
+                )
+
+    def _get_t5_prompt_embeds(
+        self,
+        prompt: str | list[str] = None,
+        max_sequence_length: int = 128,
+        dtype: torch.dtype | None = None,
+    ):
+        dtype = dtype or self.text_encoder_3.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        text_inputs = self.tokenizer_3(
+            prompt,
+            padding="max_length",
+            max_length=min(max_sequence_length, self.tokenizer_3.model_max_length),
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        attention_mask = text_inputs.attention_mask
+        untruncated_ids = self.tokenizer_3(prompt, padding="longest", return_tensors="pt").input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer_3.batch_decode(
+                untruncated_ids[:, min(max_sequence_length, self.tokenizer_3.model_max_length) - 1 : -1]
+            )
+            logger.warning(
+                "The following part of your input was truncated because `max_sequence_length` is set to "
+                f" {min(max_sequence_length, self.tokenizer_3.model_max_length)} tokens: {removed_text}"
+            )
+
+        prompt_embeds = self.text_encoder_3(
+            text_input_ids.to(self.device), attention_mask=attention_mask.to(self.device)
+        )[0]
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device)
+        return prompt_embeds
+
+    def _get_clip_prompt_embeds(
+        self,
+        tokenizer,
+        text_encoder,
+        prompt: str | list[str],
+        max_sequence_length: int = 128,
+        dtype: torch.dtype | None = None,
+    ):
+        dtype = dtype or text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        model_max_length = self.tokenizer.model_max_length
+        text_inputs = tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=min(max_sequence_length, model_max_length),
+            truncation=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = tokenizer.batch_decode(untruncated_ids[:, model_max_length - 1 : -1])
+            logger.warning(
+                "The following part of your input was truncated because CLIP can only handle sequences up to"
+                f" {model_max_length} tokens: {removed_text}"
+            )
+        prompt_embeds = text_encoder(text_input_ids.to(self.device), output_hidden_states=True)
+
+        # Use pooled output of CLIPTextModel
+        prompt_embeds = prompt_embeds[0]
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device)
+        return prompt_embeds
+
+    def _get_llama3_prompt_embeds(
+        self,
+        prompt: str | list[str] = None,
+        max_sequence_length: int = 128,
+        dtype: torch.dtype | None = None,
+    ):
+        dtype = dtype or self.text_encoder_4.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        text_inputs = self.tokenizer_4(
+            prompt,
+            padding="max_length",
+            max_length=min(max_sequence_length, self.tokenizer_4.model_max_length),
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        attention_mask = text_inputs.attention_mask
+        untruncated_ids = self.tokenizer_4(prompt, padding="longest", return_tensors="pt").input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer_4.batch_decode(
+                untruncated_ids[:, min(max_sequence_length, self.tokenizer_4.model_max_length) - 1 : -1]
+            )
+            logger.warning(
+                "The following part of your input was truncated because `max_sequence_length` is set to "
+                f" {min(max_sequence_length, self.tokenizer_4.model_max_length)} tokens: {removed_text}"
+            )
+
+        outputs = self.text_encoder_4(
+            text_input_ids.to(self.device),
+            attention_mask=attention_mask.to(self.device),
+            output_hidden_states=True,
+        )
+
+        prompt_embeds = outputs.hidden_states[1:]
+        prompt_embeds = torch.stack(prompt_embeds, dim=0)
+        return prompt_embeds
+
+    def encode_prompt(
+        self,
+        prompt: str | list[str] | None = None,
+        prompt_2: str | list[str] | None = None,
+        prompt_3: str | list[str] | None = None,
+        prompt_4: str | list[str] | None = None,
+        dtype: torch.dtype | None = None,
+        num_images_per_prompt: int = 1,
+        do_classifier_free_guidance: bool = True,
+        negative_prompt: str | list[str] | None = None,
+        negative_prompt_2: str | list[str] | None = None,
+        negative_prompt_3: str | list[str] | None = None,
+        negative_prompt_4: str | list[str] | None = None,
+        prompt_embeds_t5: list[torch.FloatTensor] | None = None,
+        prompt_embeds_llama3: list[torch.FloatTensor] | None = None,
+        negative_prompt_embeds_t5: list[torch.FloatTensor] | None = None,
+        negative_prompt_embeds_llama3: list[torch.FloatTensor] | None = None,
+        pooled_prompt_embeds: torch.FloatTensor | None = None,
+        negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+        max_sequence_length: int = 128,
+        lora_scale: float | None = None,
+    ):
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        if prompt is not None:
+            batch_size = len(prompt)
+        else:
+            batch_size = pooled_prompt_embeds.shape[0]
+
+        if pooled_prompt_embeds is None:
+            pooled_prompt_embeds_1 = self._get_clip_prompt_embeds(
+                self.tokenizer, self.text_encoder, prompt, max_sequence_length, dtype
+            )
+
+        if do_classifier_free_guidance and negative_pooled_prompt_embeds is None:
+            negative_prompt = negative_prompt or ""
+            negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+
+            if len(negative_prompt) > 1 and len(negative_prompt) != batch_size:
+                raise ValueError(f"negative_prompt must be of length 1 or {batch_size}")
+
+            negative_pooled_prompt_embeds_1 = self._get_clip_prompt_embeds(
+                self.tokenizer, self.text_encoder, negative_prompt, max_sequence_length, dtype
+            )
+
+            if negative_pooled_prompt_embeds_1.shape[0] == 1 and batch_size > 1:
+                negative_pooled_prompt_embeds_1 = negative_pooled_prompt_embeds_1.repeat(batch_size, 1)
+
+        if pooled_prompt_embeds is None:
+            prompt_2 = prompt_2 or prompt
+            prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+            if len(prompt_2) > 1 and len(prompt_2) != batch_size:
+                raise ValueError(f"prompt_2 must be of length 1 or {batch_size}")
+
+            pooled_prompt_embeds_2 = self._get_clip_prompt_embeds(
+                self.tokenizer_2, self.text_encoder_2, prompt_2, max_sequence_length, dtype
+            )
+
+            if pooled_prompt_embeds_2.shape[0] == 1 and batch_size > 1:
+                pooled_prompt_embeds_2 = pooled_prompt_embeds_2.repeat(batch_size, 1)
+
+        if do_classifier_free_guidance and negative_pooled_prompt_embeds is None:
+            negative_prompt_2 = negative_prompt_2 or negative_prompt
+            negative_prompt_2 = [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
+
+            if len(negative_prompt_2) > 1 and len(negative_prompt_2) != batch_size:
+                raise ValueError(f"negative_prompt_2 must be of length 1 or {batch_size}")
+
+            negative_pooled_prompt_embeds_2 = self._get_clip_prompt_embeds(
+                self.tokenizer_2, self.text_encoder_2, negative_prompt_2, max_sequence_length, dtype
+            )
+
+            if negative_pooled_prompt_embeds_2.shape[0] == 1 and batch_size > 1:
+                negative_pooled_prompt_embeds_2 = negative_pooled_prompt_embeds_2.repeat(batch_size, 1)
+
+        if pooled_prompt_embeds is None:
+            pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_1, pooled_prompt_embeds_2], dim=-1)
+
+        if do_classifier_free_guidance and negative_pooled_prompt_embeds is None:
+            negative_pooled_prompt_embeds = torch.cat(
+                [negative_pooled_prompt_embeds_1, negative_pooled_prompt_embeds_2], dim=-1
+            )
+
+        if prompt_embeds_t5 is None:
+            prompt_3 = prompt_3 or prompt
+            prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3
+
+            if len(prompt_3) > 1 and len(prompt_3) != batch_size:
+                raise ValueError(f"prompt_3 must be of length 1 or {batch_size}")
+
+            prompt_embeds_t5 = self._get_t5_prompt_embeds(prompt_3, max_sequence_length, dtype)
+
+            if prompt_embeds_t5.shape[0] == 1 and batch_size > 1:
+                prompt_embeds_t5 = prompt_embeds_t5.repeat(batch_size, 1, 1)
+
+        if do_classifier_free_guidance and negative_prompt_embeds_t5 is None:
+            negative_prompt_3 = negative_prompt_3 or negative_prompt
+            negative_prompt_3 = [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3
+
+            if len(negative_prompt_3) > 1 and len(negative_prompt_3) != batch_size:
+                raise ValueError(f"negative_prompt_3 must be of length 1 or {batch_size}")
+
+            negative_prompt_embeds_t5 = self._get_t5_prompt_embeds(negative_prompt_3, max_sequence_length, dtype)
+
+            if negative_prompt_embeds_t5.shape[0] == 1 and batch_size > 1:
+                negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(batch_size, 1, 1)
+
+        if prompt_embeds_llama3 is None:
+            prompt_4 = prompt_4 or prompt
+            prompt_4 = [prompt_4] if isinstance(prompt_4, str) else prompt_4
+
+            if len(prompt_4) > 1 and len(prompt_4) != batch_size:
+                raise ValueError(f"prompt_4 must be of length 1 or {batch_size}")
+
+            prompt_embeds_llama3 = self._get_llama3_prompt_embeds(prompt_4, max_sequence_length, dtype)
+
+            if prompt_embeds_llama3.shape[0] == 1 and batch_size > 1:
+                prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, batch_size, 1, 1)
+
+        if do_classifier_free_guidance and negative_prompt_embeds_llama3 is None:
+            negative_prompt_4 = negative_prompt_4 or negative_prompt
+            negative_prompt_4 = [negative_prompt_4] if isinstance(negative_prompt_4, str) else negative_prompt_4
+
+            if len(negative_prompt_4) > 1 and len(negative_prompt_4) != batch_size:
+                raise ValueError(f"negative_prompt_4 must be of length 1 or {batch_size}")
+
+            negative_prompt_embeds_llama3 = self._get_llama3_prompt_embeds(
+                negative_prompt_4, max_sequence_length, dtype
+            )
+
+            if negative_prompt_embeds_llama3.shape[0] == 1 and batch_size > 1:
+                negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, batch_size, 1, 1)
+
+        # duplicate pooled_prompt_embeds for each generation per prompt
+        pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt)
+        pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+        # duplicate t5_prompt_embeds for batch_size and num_images_per_prompt
+        bs_embed, seq_len, _ = prompt_embeds_t5.shape
+        if bs_embed == 1 and batch_size > 1:
+            prompt_embeds_t5 = prompt_embeds_t5.repeat(batch_size, 1, 1)
+        elif bs_embed > 1 and bs_embed != batch_size:
+            raise ValueError(f"cannot duplicate prompt_embeds_t5 of batch size {bs_embed}")
+        prompt_embeds_t5 = prompt_embeds_t5.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds_t5 = prompt_embeds_t5.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        # duplicate llama3_prompt_embeds for batch_size and num_images_per_prompt
+        _, bs_embed, seq_len, dim = prompt_embeds_llama3.shape
+        if bs_embed == 1 and batch_size > 1:
+            prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, batch_size, 1, 1)
+        elif bs_embed > 1 and bs_embed != batch_size:
+            raise ValueError(f"cannot duplicate prompt_embeds_llama3 of batch size {bs_embed}")
+        prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, 1, num_images_per_prompt, 1)
+        prompt_embeds_llama3 = prompt_embeds_llama3.view(-1, batch_size * num_images_per_prompt, seq_len, dim)
+
+        if do_classifier_free_guidance:
+            # duplicate negative_pooled_prompt_embeds for batch_size and num_images_per_prompt
+            bs_embed, seq_len = negative_pooled_prompt_embeds.shape
+            if bs_embed == 1 and batch_size > 1:
+                negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(batch_size, 1)
+            elif bs_embed > 1 and bs_embed != batch_size:
+                raise ValueError(f"cannot duplicate negative_pooled_prompt_embeds of batch size {bs_embed}")
+            negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt)
+            negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+            # duplicate negative_t5_prompt_embeds for batch_size and num_images_per_prompt
+            bs_embed, seq_len, _ = negative_prompt_embeds_t5.shape
+            if bs_embed == 1 and batch_size > 1:
+                negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(batch_size, 1, 1)
+            elif bs_embed > 1 and bs_embed != batch_size:
+                raise ValueError(f"cannot duplicate negative_prompt_embeds_t5 of batch size {bs_embed}")
+            negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(1, num_images_per_prompt, 1)
+            negative_prompt_embeds_t5 = negative_prompt_embeds_t5.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+            # duplicate negative_prompt_embeds_llama3 for batch_size and num_images_per_prompt
+            _, bs_embed, seq_len, dim = negative_prompt_embeds_llama3.shape
+            if bs_embed == 1 and batch_size > 1:
+                negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, batch_size, 1, 1)
+            elif bs_embed > 1 and bs_embed != batch_size:
+                raise ValueError(f"cannot duplicate negative_prompt_embeds_llama3 of batch size {bs_embed}")
+            negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, 1, num_images_per_prompt, 1)
+            negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.view(
+                -1, batch_size * num_images_per_prompt, seq_len, dim
+            )
+
+        return (
+            prompt_embeds_t5,
+            negative_prompt_embeds_t5,
+            prompt_embeds_llama3,
+            negative_prompt_embeds_llama3,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        )
+
+    def enable_vae_slicing(self):
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        depr_message = (
+            f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this"
+            " method will be removed in a future version. Please use `pipe.vae.enable_slicing()`."
+        )
+        deprecate(
+            "enable_vae_slicing",
+            "0.40.0",
+            depr_message,
+        )
+        self.vae.enable_slicing()
+
+    def disable_vae_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        depr_message = (
+            f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this"
+            " method will be removed in a future version. Please use `pipe.vae.disable_slicing()`."
+        )
+        deprecate(
+            "disable_vae_slicing",
+            "0.40.0",
+            depr_message,
+        )
+        self.vae.disable_slicing()
+
+    def enable_vae_tiling(self):
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+        """
+        depr_message = (
+            f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}`"
+            " is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`."
+        )
+        deprecate(
+            "enable_vae_tiling",
+            "0.40.0",
+            depr_message,
+        )
+        self.vae.enable_tiling()
+
+    def disable_vae_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        depr_message = (
+            f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is"
+            " deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`."
+        )
+        deprecate(
+            "disable_vae_tiling",
+            "0.40.0",
+            depr_message,
+        )
+        self.vae.disable_tiling()
+
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        generator,
+        latents=None,
+    ):
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (self.vae_scale_factor * 2))
+        width = 2 * (int(width) // (self.vae_scale_factor * 2))
+
+        shape = (batch_size, num_channels_latents, height, width)
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=self.device, dtype=dtype)
+        else:
+            if latents.shape != shape:
+                raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
+            latents = latents.to(self.device)
+        return latents
+
+    def prepare_timesteps(self, num_inference_steps, sigmas, image_seq_len):
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.get("base_image_seq_len", 256),
+            self.scheduler.config.get("max_image_seq_len", 4096),
+            self.scheduler.config.get("base_shift", 0.5),
+            self.scheduler.config.get("max_shift", 1.15),
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            self.device,
+            sigmas=sigmas,
+            mu=mu,
+        )
+        return timesteps, num_inference_steps
+
+    def _extract_prompts(self, prompts):
+        """Extract prompt and negative_prompt from OmniPromptType list."""
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts] or None
+        if all(isinstance(p, str) or p.get("negative_prompt") is None for p in prompts):
+            negative_prompt = None
+        elif prompts:
+            negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in prompts]
+        else:
+            negative_prompt = None
+        return prompt, negative_prompt
+
+    def check_cfg_parallel_validity(self, true_cfg_scale: float, has_neg_prompt: bool):
+        if get_classifier_free_guidance_world_size() == 1:
+            return True
+
+        if true_cfg_scale <= 1:
+            logger.warning("CFG parallel is NOT working correctly when true_cfg_scale <= 1.")
+            return False
+
+        if not has_neg_prompt:
+            logger.warning(
+                "CFG parallel is NOT working correctly when there is no negative prompt or negative prompt embeddings."
+            )
+            return False
+        return True
+
+    def diffuse(
+        self,
+        prompt_embeds_t5: torch.Tensor,
+        prompt_embeds_llama3: torch.Tensor,
+        pooled_prompt_embeds: torch.Tensor,
+        latents: torch.Tensor,
+        timesteps: torch.Tensor,
+        do_true_cfg: bool,
+    ) -> torch.Tensor:
+        with self.progress_bar(total=len(timesteps)) as pbar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latent_model_input.shape[0])
+                noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    timesteps=timestep,
+                    encoder_hidden_states_t5=prompt_embeds_t5,
+                    encoder_hidden_states_llama3=prompt_embeds_llama3,
+                    pooled_embeds=pooled_prompt_embeds,
+                    return_dict=False,
+                )[0]
+                noise_pred = -noise_pred
+
+                # TODO: Modify CFG guidance
+                # perform guidance
+                if self.do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
+
+                pbar.update()
+
+        return latents
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale > 1
+
+    @property
+    def attention_kwargs(self):
+        return self._attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    def forward(
+        self,
+        req: OmniDiffusionRequest,
+        prompt: str | list[str] = None,
+        prompt_2: str | list[str] | None = None,
+        prompt_3: str | list[str] | None = None,
+        prompt_4: str | list[str] | None = None,
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int = 50,
+        sigmas: list[float] | None = None,
+        guidance_scale: float = 5.0,
+        negative_prompt: str | list[str] | None = None,
+        negative_prompt_2: str | list[str] | None = None,
+        negative_prompt_3: str | list[str] | None = None,
+        negative_prompt_4: str | list[str] | None = None,
+        num_images_per_prompt: int | None = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.FloatTensor | None = None,
+        prompt_embeds_t5: torch.FloatTensor | None = None,
+        prompt_embeds_llama3: torch.FloatTensor | None = None,
+        negative_prompt_embeds_t5: torch.FloatTensor | None = None,
+        negative_prompt_embeds_llama3: torch.FloatTensor | None = None,
+        pooled_prompt_embeds: torch.FloatTensor | None = None,
+        negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+        output_type: str | None = "pil",
+        return_dict: bool = True,
+        attention_kwargs: dict[str, Any] | None = None,
+        callback_on_step_end: Callable[[int, int], None] | None = None,
+        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+        max_sequence_length: int = 128,
+        **kwargs,
+    ):
+        extracted_prompt, negative_prompt = self._extract_prompts(req.prompts)
+        prompt = extracted_prompt or prompt
+
+        height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
+        width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
+
+        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        sigmas = req.sampling_params.sigmas or sigmas
+        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
+        generator = req.sampling_params.generator or generator
+        true_cfg_scale = req.sampling_params.true_cfg_scale or guidance_scale
+        if req.sampling_params.guidance_scale_provided:
+            guidance_scale = req.sampling_params.guidance_scale
+        num_images_per_prompt = (
+            req.sampling_params.num_outputs_per_prompt
+            if req.sampling_params.num_outputs_per_prompt > 0
+            else num_images_per_prompt
+        )
+
+        prompt_embeds = kwargs.get("prompt_embeds", None)
+        negative_prompt_embeds = kwargs.get("negative_prompt_embeds", None)
+
+        if prompt_embeds is not None:
+            deprecation_message = (
+                "The `prompt_embeds` argument is deprecated."
+                " Please use `prompt_embeds_t5` and `prompt_embeds_llama3` instead."
+            )
+            deprecate("prompt_embeds", "0.35.0", deprecation_message)
+            prompt_embeds_t5 = prompt_embeds[0]
+            prompt_embeds_llama3 = prompt_embeds[1]
+
+        if negative_prompt_embeds is not None:
+            deprecation_message = (
+                "The `negative_prompt_embeds` argument is deprecated."
+                "Please use `negative_prompt_embeds_t5` and `negative_prompt_embeds_llama3` instead."
+            )
+            deprecate("negative_prompt_embeds", "0.35.0", deprecation_message)
+            negative_prompt_embeds_t5 = negative_prompt_embeds[0]
+            negative_prompt_embeds_llama3 = negative_prompt_embeds[1]
+
+        division = self.vae_scale_factor * 2
+        S_max = (self.default_sample_size * self.vae_scale_factor) ** 2
+        scale = S_max / (width * height)
+        scale = math.sqrt(scale)
+        width, height = int(width * scale // division * division), int(height * scale // division * division)
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            prompt_3,
+            prompt_4,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            negative_prompt_3=negative_prompt_3,
+            negative_prompt_4=negative_prompt_4,
+            prompt_embeds_t5=prompt_embeds_t5,
+            prompt_embeds_llama3=prompt_embeds_llama3,
+            negative_prompt_embeds_t5=negative_prompt_embeds_t5,
+            negative_prompt_embeds_llama3=negative_prompt_embeds_llama3,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._attention_kwargs = attention_kwargs
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        elif pooled_prompt_embeds is not None:
+            batch_size = pooled_prompt_embeds.shape[0]
+
+        # TODO: CFG guidance configuration
+        has_neg_prompt = negative_prompt is not None or (
+            negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None
+        )
+        do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
+        self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt)
+
+        # 3. Encode prompt
+        lora_scale = self.attention_kwargs.get("scale", None) if self.attention_kwargs is not None else None
+        (
+            prompt_embeds_t5,
+            negative_prompt_embeds_t5,
+            prompt_embeds_llama3,
+            negative_prompt_embeds_llama3,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_3=prompt_3,
+            prompt_4=prompt_4,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            negative_prompt_3=negative_prompt_3,
+            negative_prompt_4=negative_prompt_4,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            prompt_embeds_t5=prompt_embeds_t5,
+            prompt_embeds_llama3=prompt_embeds_llama3,
+            negative_prompt_embeds_t5=negative_prompt_embeds_t5,
+            negative_prompt_embeds_llama3=negative_prompt_embeds_llama3,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            lora_scale=lora_scale,
+        )
+
+        if self.do_classifier_free_guidance:
+            prompt_embeds_t5 = torch.cat([negative_prompt_embeds_t5, prompt_embeds_t5], dim=0)
+            prompt_embeds_llama3 = torch.cat([negative_prompt_embeds_llama3, prompt_embeds_llama3], dim=1)
+            pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0)
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            pooled_prompt_embeds.dtype,
+            generator,
+            latents,
+        )
+
+        # 5. Prepare timesteps
+        timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1])
+        self._num_timesteps = len(timesteps)
+
+        # 6. Denoising loop
+        latents = self.diffuse(
+            prompt_embeds_t5,
+            prompt_embeds_llama3,
+            pooled_prompt_embeds,
+            latents,
+            timesteps,
+            do_true_cfg,
+        )
+
+        if output_type == "latent":
+            image = latents
+
+        else:
+            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+
+            image = self.vae.decode(latents, return_dict=False)[0]
+
+        return DiffusionOutput(
+            output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)
diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py
index c39965f67..80bb1d7b6 100644
--- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py
+++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py
@@ -891,6 +891,7 @@ class ImageKVCacheManager:
         # cache related
         self.image_token_len: int = image_token_len
         self.image_kv_cache_map: tuple[torch.Tensor, torch.Tensor] | None = None
+        self.image_kv_cache_lens: torch.Tensor | None = None
         self._injected_ar_kv: list[tuple[torch.Tensor, torch.Tensor]] | None = None
 
         self.sp_size = get_sequence_parallel_world_size()
@@ -904,12 +905,23 @@ class ImageKVCacheManager:
             prefix=f"{prefix}.attn" if prefix else "",
         )
 
+    @staticmethod
+    def _get_current_starts(
+        gen_timestep_scatter_index: torch.Tensor | None,
+        ar_kv_len: int = 0,
+    ) -> torch.Tensor:
+        assert gen_timestep_scatter_index is not None, (
+            "`gen_timestep_scatter_index` is required to locate the generated image timestep token for image KV reuse."
+        )
+        return gen_timestep_scatter_index[:, -1].to(dtype=torch.long) + ar_kv_len
+
     def _cache_prompt_kv(
         self,
         key: torch.Tensor,
         value: torch.Tensor,
         seq_len: int,
         shard_image_size: int | None = None,
+        gen_timestep_scatter_index: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Cache prompt KV on first_step. Consumes _injected_ar_kv.
 
@@ -939,12 +951,26 @@ class ImageKVCacheManager:
             key = torch.cat(new_keys, dim=0)
             value = torch.cat(new_values, dim=0)
 
-        image_size = shard_image_size if self.sp_size > 1 else (self.image_token_len + 1)
-        cached_prompt_len = seq_len - image_size
+        cached_prompt_lens = self._get_current_starts(gen_timestep_scatter_index, ar_kv_len)
+        assert torch.all(cached_prompt_lens <= key.shape[1]), (
+            f"cached_prompt_lens({cached_prompt_lens.tolist()}) must be <= key length({key.shape[1]})"
+        )
+        if shard_image_size is not None:
+            expected_prompt_len = seq_len - shard_image_size
+            assert torch.all(cached_prompt_lens == expected_prompt_len), (
+                f"cached_prompt_lens({cached_prompt_lens.tolist()}) != "
+                f"seq_len({seq_len}) - shard_image_size({shard_image_size})"
+            )
 
-        cached_key = key[:, :cached_prompt_len].reshape(-1, num_kv_heads, head_dim)
-        cached_value = value[:, :cached_prompt_len].reshape(-1, num_kv_heads, head_dim)
+        max_cached_prompt_len = int(cached_prompt_lens.max().item())
+        cached_key = key.new_zeros(bs, max_cached_prompt_len, num_kv_heads, head_dim)
+        cached_value = value.new_zeros(bs, max_cached_prompt_len, num_kv_heads, head_dim)
+        for b in range(bs):
+            cached_prompt_len = int(cached_prompt_lens[b].item())
+            cached_key[b, :cached_prompt_len] = key[b, :cached_prompt_len]
+            cached_value[b, :cached_prompt_len] = value[b, :cached_prompt_len]
         self.image_kv_cache_map = (cached_key, cached_value)
+        self.image_kv_cache_lens = cached_prompt_lens
         return key, value
 
     def _reuse_prompt_kv(
@@ -954,42 +980,36 @@ class ImageKVCacheManager:
         seq_len: int,
         bs: int,
         shard_image_size: int | None = None,
+        position_ids: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Reuse cached prompt KV in subsequent denoising steps.
 
-        Non-SP: returns [cached_prompt, new_image, zero_eoi] shaped [bs, seq_len, ...].
-        SP: returns cached prompt only shaped [bs, cached_prompt_len, ...] (used as joint_text).
+        Non-SP: concatenates cached prompt KV with current image KV.
+        SP: returns cached prompt KV only (used as joint_text).
         """
         cached_key, cached_value = self.image_kv_cache_map
-        _, q_len, num_kv_heads, head_dim = key.shape
-        cached_prompt_len = cached_key.shape[0] // bs
+        _, q_len, _, _ = key.shape
+        assert cached_key.dim() == 4 and cached_key.shape[0] == bs
+        max_cached_prompt_len = cached_key.shape[1]
 
         if shard_image_size is not None:
-            assert cached_prompt_len == seq_len - q_len
-            result_k = cached_key.reshape(bs, cached_prompt_len, num_kv_heads, head_dim)
-            result_v = cached_value.reshape(bs, cached_prompt_len, num_kv_heads, head_dim)
-            return result_k.contiguous(), result_v.contiguous()
+            assert max_cached_prompt_len + q_len == seq_len, (
+                f"max_cached_prompt_len({max_cached_prompt_len}) + q_len({q_len}) != seq_len({seq_len})"
+            )
+            return cached_key.contiguous(), cached_value.contiguous()
 
-        assert cached_prompt_len + 1 == seq_len - q_len, f"{cached_prompt_len + 1} != {seq_len - q_len}"
-
-        key = key.reshape(-1, num_kv_heads, head_dim)
-        value = value.reshape(-1, num_kv_heads, head_dim)
-        zero_eoi = torch.zeros(1, num_kv_heads, head_dim, device=key.device, dtype=key.dtype)
-
-        new_key, new_value = [], []
-        for b in range(bs):
-            cache_base = b * cached_prompt_len
-            img_base = b * q_len
-            new_key.append(cached_key[cache_base : cache_base + cached_prompt_len])
-            new_key.append(key[img_base : img_base + q_len])
-            new_key.append(zero_eoi)
-            new_value.append(cached_value[cache_base : cache_base + cached_prompt_len])
-            new_value.append(value[img_base : img_base + q_len])
-            new_value.append(zero_eoi)
-
-        new_key = torch.cat(new_key).reshape(bs, seq_len, num_kv_heads, head_dim)
-        new_value = torch.cat(new_value).reshape(bs, seq_len, num_kv_heads, head_dim)
+        assert max_cached_prompt_len + q_len == seq_len, (
+            f"max_cached_prompt_len({max_cached_prompt_len}) + q_len({q_len}) != seq_len({seq_len})"
+        )
+        assert self.image_kv_cache_lens is not None
+        if position_ids is not None:
+            assert position_ids.shape == (bs, q_len)
+            assert torch.all(position_ids[:, 0] == self.image_kv_cache_lens), (
+                "The first current position must immediately follow each sample's cached prompt KV."
+            )
+        new_key = torch.cat([cached_key, key], dim=1)
+        new_value = torch.cat([cached_value, value], dim=1)
         return new_key.contiguous(), new_value.contiguous()
 
     def _build_neg_ar_kv(
@@ -1032,7 +1052,7 @@ class ImageKVCacheManager:
 
         query_lens = kwargs.get("query_lens")
         seq_lens = kwargs.get("seq_lens")
-        shard_image_size = kwargs.get("shard_image_size")
+        shard_image_size = kwargs.get("shard_image_size") if self.sp_size > 1 else None
         bs = len(query_lens)
         q_len = query_lens[0]
         seq_len = seq_lens[0]
@@ -1058,7 +1078,14 @@ class ImageKVCacheManager:
                 value = value[:, :0, :, :]
         elif first_step:
             self.image_kv_cache_map = None  # reset first
-            key, value = self._cache_prompt_kv(key, value, seq_len, shard_image_size)
+            self.image_kv_cache_lens = None
+            key, value = self._cache_prompt_kv(
+                key,
+                value,
+                seq_len,
+                shard_image_size,
+                kwargs.get("gen_timestep_scatter_index"),
+            )
             if self.sp_size > 1:
                 local_prompt_len = seq_len - shard_image_size
                 join_query_len = query.shape[1] - shard_image_size
@@ -1070,7 +1097,7 @@ class ImageKVCacheManager:
                 value = value[:, local_prompt_len:, :, :]
         else:
             if self.sp_size <= 1:
-                key, value = self._reuse_prompt_kv(key, value, seq_len, bs)
+                key, value = self._reuse_prompt_kv(key, value, seq_len, bs, position_ids=kwargs.get("position_ids"))
             else:
                 joint_text_query = query[:, :0, :, :]
                 joint_text_key, joint_text_value = self._reuse_prompt_kv(key, value, seq_len, bs, shard_image_size)
@@ -1814,6 +1841,7 @@ class HunyuanImage3DecoderLayer(nn.Module):
         hidden_states, self_attn_weights, present_key_value = self.self_attn(
             hidden_states=hidden_states,
             attention_mask=attention_mask,
+            position_ids=position_ids,
             custom_pos_emb=custom_pos_emb,
             past_key_value=past_key_value,
             output_attentions=output_attentions,
@@ -1878,12 +1906,21 @@ class HunyuanImagePreprocessor(nn.Module):
         image_token_len: int,
         is_first_step: bool = False,
         uncond_cfg_prefill: bool = False,
+        gen_timestep_scatter_index: torch.Tensor | None = None,
     ):
         # ---------- compute prompt length ----------
         if uncond_cfg_prefill:
             prompt_len = hidden_states.shape[1]
         elif is_first_step:
-            prompt_len = hidden_states.shape[1] - image_token_len - 1
+            assert gen_timestep_scatter_index is not None, (
+                "`gen_timestep_scatter_index` is required to split prompt and "
+                "generated image tokens for sequence parallel image generation."
+            )
+            prompt_lens = gen_timestep_scatter_index[:, -1]
+            assert torch.all(prompt_lens == prompt_lens[0]), (
+                "Sequence parallel image generation requires the generated timestep position to match across the batch."
+            )
+            prompt_len = int(prompt_lens[0].item())
         else:
             prompt_len = 0
 
@@ -2073,7 +2110,7 @@ class HunyuanImage3Model(nn.Module):
             return False
 
         def is_scalar_quant_scale(name: str, tensor: torch.Tensor) -> bool:
-            return tensor.numel() == 1 and name.endswith((".input_scale", ".weight_scale"))
+            return tensor.numel() == 1 and name.endswith((".input_scale", ".weight_scale", ".weight_scale_2"))
 
         def load_split_param(
             name: str,
@@ -2327,6 +2364,7 @@ class HunyuanImage3Model(nn.Module):
                 num_image_tokens,
                 first_step,
                 uncond_cfg_prefill=uncond_cfg_prefill,
+                gen_timestep_scatter_index=gen_timestep_scatter_index,
             )
             assert len(set(query_lens)) == 1 and len(set(seq_lens)) == 1, (
                 "query_lens and seq_lens must be the same for sequence parallel"
@@ -2334,15 +2372,15 @@ class HunyuanImage3Model(nn.Module):
             prompt_size = text_hidden_states.shape[1]
             shard_image_size = image_hidden_states.shape[1]
 
-            if first_step:  # [image tokens, last token]
-                shard_padding_size = shard_image_size * sp_world_size - (num_image_tokens + 1)
-            else:  # [image tokens, last token]
+            if first_step:
+                shard_padding_size = shard_image_size * sp_world_size - (origin_query_len - prompt_size)
+            else:
                 shard_padding_size = shard_image_size * sp_world_size - num_image_tokens
             if first_step:
                 seq_lens = [prompt_size + shard_image_size + ar_kv_reuse_len for _ in seq_lens]
             else:
                 seq_lens = [x - y for x, y in zip(seq_lens, query_lens)]
-                seq_lens = [seq_len + shard_image_size - 1 for seq_len in seq_lens]
+                seq_lens = [seq_len + shard_image_size for seq_len in seq_lens]
             query_lens = [prompt_size + shard_image_size for _ in query_lens]
 
             if prompt_size > 0:
@@ -2357,10 +2395,6 @@ class HunyuanImage3Model(nn.Module):
                 position_ids = image_position_ids
                 custom_pos_emb = (image_custom_pos_emb_sin, image_custom_pos_emb_cos)
 
-            if not first_step:
-                assert torch.all(~attention_mask[..., -1]), "The last token should not be attended to"
-                attention_mask = attention_mask[..., :-1]  # we attention without last token
-
             if shard_padding_size > 0:
                 B, H, Q, K = attention_mask.shape
                 pad = shard_padding_size
@@ -3091,15 +3125,11 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
                 # Scheduler step (all ranks compute locally in CFG parallel)
                 latents = self.scheduler.step(pred, t, latents, **_scheduler_step_extra_kwargs, return_dict=False)[0]
                 if i != len(timesteps) - 1 and should_compute:
-                    offset = model_kwargs.get("ar_kv_reuse_offset", 0)
                     model_kwargs = self.model._update_model_kwargs_for_generation(  # noqa
                         model_output,
                         model_kwargs,
                     )
-                    if input_ids.shape[1] != model_kwargs["position_ids"].shape[1]:
-                        # Select using relative positions
-                        index = model_kwargs["position_ids"] - offset
-                        input_ids = torch.gather(input_ids, 1, index=index)
+                    input_ids = None
                     attention_mask = model_kwargs.get("attention_mask")
                     b, _, q_len1, seq_len = attention_mask.shape
                     query_lens = [q_len1] * b
diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
index 677d67796..b4a588df9 100644
--- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
+++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
@@ -1001,7 +1001,7 @@ class HunyuanImage3Pipeline(
         if inputs_embeds is not None and past_key_values is None:
             model_inputs = {"inputs_embeds": inputs_embeds}
         else:
-            if input_ids.shape[1] != kwargs["position_ids"].shape[1]:  # in decode steps
+            if input_ids is not None and input_ids.shape[1] != kwargs["position_ids"].shape[1]:  # in decode steps
                 input_ids = torch.gather(input_ids, dim=1, index=kwargs["position_ids"])
             model_inputs = {"input_ids": input_ids}
 
@@ -1082,12 +1082,27 @@ class HunyuanImage3Pipeline(
 
                 # attention mask
                 mask_list = []
-                for attention_mask_i, position_ids_i in zip(
-                    model_kwargs["attention_mask"], updated_model_kwargs["position_ids"]
+                current_starts = timestep_position_ids.reshape(-1)
+                max_current_start = int(current_starts.max().item())
+                for attention_mask_i, position_ids_i, current_start_i in zip(
+                    model_kwargs["attention_mask"], updated_model_kwargs["position_ids"], current_starts
                 ):
-                    mask_list.append(
-                        torch.index_select(attention_mask_i, dim=1, index=position_ids_i.reshape(-1) - offset)
+                    query_mask = torch.index_select(attention_mask_i, dim=1, index=position_ids_i.reshape(-1) - offset)
+                    current_start = int(current_start_i.item())
+                    prefix_mask = torch.index_select(
+                        query_mask,
+                        dim=2,
+                        index=torch.arange(current_start, device=image_mask.device),
                     )
+                    if current_start < max_current_start:
+                        prefix_pad = query_mask.new_zeros(
+                            query_mask.shape[0],
+                            query_mask.shape[1],
+                            max_current_start - current_start,
+                        )
+                        prefix_mask = torch.cat([prefix_mask, prefix_pad], dim=2)
+                    current_mask = torch.index_select(query_mask, dim=2, index=position_ids_i.reshape(-1))
+                    mask_list.append(torch.cat([prefix_mask, current_mask], dim=2))
                 attention_mask = torch.stack(mask_list, dim=0)
                 updated_model_kwargs["attention_mask"] = attention_mask
                 updated_model_kwargs["gen_timestep_scatter_index"] = model_kwargs["gen_timestep_scatter_index"]
@@ -1219,9 +1234,14 @@ class HunyuanImage3Pipeline(
         )
         custom_pos_emb = self.get_pos_emb(custom_pos_emb, position_ids)
 
-        inputs_embeds = self.model.embed_tokens(input_ids)
-
-        bsz, seq_len, n_embd = inputs_embeds.shape
+        if input_ids is not None:
+            inputs_embeds = self.model.embed_tokens(input_ids)
+            bsz, seq_len, n_embd = inputs_embeds.shape
+        else:
+            inputs_embeds = None
+            bsz = images.shape[0] if isinstance(images, torch.Tensor) else len(images)
+            seq_len = 0
+            n_embd = self.config.hidden_size
 
         # Instantiate placeholder tokens: <timestep>, <img> for the gen image
         if mode == "gen_text":
@@ -1232,6 +1252,7 @@ class HunyuanImage3Pipeline(
             token_h, token_w = None, None
         else:
             if first_step:
+                assert inputs_embeds is not None
                 inputs_embeds, token_h, token_w = self.instantiate_vae_image_tokens(
                     inputs_embeds, images, timestep, image_mask
                 )
@@ -1253,6 +1274,8 @@ class HunyuanImage3Pipeline(
             inputs_embeds = self.instantiate_vit_image_tokens(
                 inputs_embeds, cond_vit_images, cond_vit_image_mask, vit_kwargs
             )
+        assert inputs_embeds is not None
+        bsz, seq_len, n_embd = inputs_embeds.shape
 
         # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
         from vllm.forward_context import set_forward_context
@@ -1291,7 +1314,7 @@ class HunyuanImage3Pipeline(
             diffusion_prediction = None
         else:
             logits = None
-            hidden_states = hidden_states.to(input_ids.device)
+            hidden_states = hidden_states.to(inputs_embeds.device)
             assert hidden_states.numel() == bsz * seq_len * n_embd, (
                 f"Shape mismatch: {hidden_states.shape} cannot reshape to ({bsz}, {seq_len}, {n_embd})"
             )
diff --git a/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py b/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py
index b178b021f..b64fe2e97 100644
--- a/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py
+++ b/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py
@@ -146,6 +146,7 @@ def resolve_stop_token_ids(
     task: str = "it2i",
     bot_task: str | None | _DefaultBotTask = _DEFAULT_BOT_TASK,
     tokenizer: Any | None = None,
+    image_size: str | None = None,
 ) -> list[int]:
     """AR stop-token ids for a given (task, bot_task) generation request.
 
@@ -171,6 +172,14 @@ def resolve_stop_token_ids(
     if bot_task not in _BOT_TASK_PRESETS:
         raise ValueError(f"Unknown bot_task {bot_task!r}. Choose from: {available_bot_tasks()}")
     if task in ("it2i", "t2i"):
+        need_ratio = image_size is None or image_size == "auto"
+        if not need_ratio:
+            end_of_think_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["</think>"]
+            end_of_recaption_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["</recaption>"]
+            if bot_task in ("think_recaption", "recaption"):
+                return [end_of_recaption_id]
+            if bot_task == "think":
+                return [end_of_think_id, end_of_recaption_id]
         # Main ratio range: <img_ratio_0> .. <img_ratio_32>.
         start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_0>"]
         end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_32>"]
@@ -219,7 +228,7 @@ def build_prompt(
     effective_sys_type = sys_type or preset_sys_type
 
     system_prompt = get_system_prompt(effective_sys_type, bot_task, custom_system_prompt)
-    sys_text = system_prompt or ""
+    sys_text = system_prompt.strip() if system_prompt is not None else ""
 
     has_image_input = task in ("i2t", "it2i")
     if has_image_input:
@@ -288,7 +297,7 @@ def build_prompt_tokens(
         )
 
     system_prompt = get_system_prompt(effective_sys_type, bot_task, custom_system_prompt)
-    sys_text = system_prompt or ""
+    sys_text = system_prompt.strip() if system_prompt is not None else ""
 
     ids: list[int] = [bos_id]
     if sys_text:
diff --git a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py b/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py
index 30ba62184..1f06799c3 100644
--- a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py
+++ b/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py
@@ -762,7 +762,13 @@ class HunyuanVideo15Transformer3DModel(nn.Module):
         encoder_hidden_states = torch.stack(new_encoder_hidden_states)
         encoder_attention_mask = torch.stack(new_encoder_attention_mask)
 
-        # Create explicit attn_mask for image tokens when SP auto_pad is active.
+        max_valid_encoder_tokens = int(encoder_attention_mask.sum(dim=1).max().item())
+        if max_valid_encoder_tokens < encoder_attention_mask.shape[1]:
+            encoder_hidden_states = encoder_hidden_states[:, :max_valid_encoder_tokens]
+            encoder_attention_mask = encoder_attention_mask[:, :max_valid_encoder_tokens]
+        if encoder_attention_mask.all():
+            encoder_attention_mask = None
+
         ctx = get_forward_context()
         hidden_states_mask = None
         if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0:
@@ -830,9 +836,10 @@ class HunyuanVideo15Transformer3DModel(nn.Module):
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in original_name:
                     continue
-                lookup_name = original_name.replace(weight_name, param_name)
-                if lookup_name not in params_dict:
-                    break
+                maybe_lookup_name = original_name.replace(weight_name, param_name)
+                if maybe_lookup_name not in params_dict:
+                    continue
+                lookup_name = maybe_lookup_name
                 param = params_dict[lookup_name]
                 weight_loader = param.weight_loader
                 weight_loader(param, loaded_weight, shard_id)
diff --git a/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py b/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py
index 6445bfee2..3688aa7c2 100644
--- a/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py
+++ b/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py
@@ -11,7 +11,6 @@ from typing import Any
 
 import numpy as np
 import torch
-from diffusers import AutoencoderKLHunyuanVideo15
 from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
 from diffusers.utils.torch_utils import randn_tensor
 from diffusers.video_processor import VideoProcessor
@@ -20,6 +19,9 @@ from transformers import AutoConfig, ByT5Tokenizer, Qwen2_5_VLTextModel, Qwen2To
 from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan_video_15 import (
+    DistributedAutoencoderKLHunyuanVideo15,
+)
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -110,7 +112,7 @@ class HunyuanVideo15Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, Diff
         t5_config = AutoConfig.from_pretrained(model, subfolder="text_encoder_2", local_files_only=local_files_only)
         self.text_encoder_2 = T5EncoderModel(t5_config, prefix="text_encoder_2").to(dtype=dtype, device=self.device)
 
-        self.vae = AutoencoderKLHunyuanVideo15.from_pretrained(
+        self.vae = DistributedAutoencoderKLHunyuanVideo15.from_pretrained(
             model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
         ).to(self.device)
 
diff --git a/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py b/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py
index c1acd1a89..ad6339c8c 100644
--- a/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py
+++ b/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py
@@ -11,7 +11,6 @@ from typing import Any, cast
 import numpy as np
 import PIL.Image
 import torch
-from diffusers import AutoencoderKLHunyuanVideo15
 from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
 from diffusers.utils.torch_utils import randn_tensor
 from diffusers.video_processor import VideoProcessor
@@ -27,6 +26,9 @@ from transformers import (
 from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan_video_15 import (
+    DistributedAutoencoderKLHunyuanVideo15,
+)
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -139,7 +141,7 @@ class HunyuanVideo15I2VPipeline(
             model, subfolder="feature_extractor", local_files_only=local_files_only
         )
 
-        self.vae = AutoencoderKLHunyuanVideo15.from_pretrained(
+        self.vae = DistributedAutoencoderKLHunyuanVideo15.from_pretrained(
             model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
         ).to(self.device)
 
diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py
index 2a65ad6b0..a1d72cf66 100644
--- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py
+++ b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py
@@ -67,13 +67,20 @@ def _make_rms_norm(hidden_size: int, *, eps: float, elementwise_affine: bool) ->
 
 def apply_interleaved_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
     cos, sin = freqs
-    x_real, x_imag = x.unflatten(2, (-1, 2)).unbind(-1)  # [B, S, C // 2]
+    # Concrete pair count instead of -1 keeps SDPA shape static under torch.compile.
+    x_real, x_imag = x.unflatten(2, (x.shape[2] // 2, 2)).unbind(-1)  # [B, S, C // 2]
     x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2)
     out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
     return out
 
 
-def apply_split_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+def apply_split_rotary_emb(
+    x: torch.Tensor,
+    freqs: tuple[torch.Tensor, torch.Tensor],
+    *,
+    head_dim: int,
+) -> torch.Tensor:
+    # `head_dim` is plumbed in (not inferred via `-1`) so SDPA shape stays static under torch.compile.
     cos, sin = freqs
 
     x_dtype = x.dtype
@@ -83,7 +90,7 @@ def apply_split_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Ten
         # The cos/sin batch dim may only be broadcastable, so take batch size from x
         b = x.shape[0]
         _, h, t, _ = cos.shape
-        x = x.reshape(b, t, h, -1).swapaxes(1, 2)
+        x = x.reshape(b, t, h, head_dim).swapaxes(1, 2)
         needs_reshape = True
 
     # Split last dim (2*r) into (d=2, r)
@@ -110,7 +117,7 @@ def apply_split_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Ten
     out = out.reshape(*out.shape[:-2], last)
 
     if needs_reshape:
-        out = out.swapaxes(1, 2).reshape(b, t, -1)
+        out = out.swapaxes(1, 2).reshape(b, t, h * head_dim)
 
     out = out.to(dtype=x_dtype)
     return out
@@ -488,12 +495,14 @@ class LTX2AudioVideoAttnProcessor:
                     key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb
                 )
             elif attn.rope_type == "split":
-                query = apply_split_rotary_emb(query, query_rotary_emb)
-                key = apply_split_rotary_emb(key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb)
+                query = apply_split_rotary_emb(query, query_rotary_emb, head_dim=attn.head_dim)
+                key = apply_split_rotary_emb(
+                    key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb, head_dim=attn.head_dim
+                )
 
-        query = query.unflatten(2, (attn.heads, -1))
-        key = key.unflatten(2, (attn.heads, -1))
-        value = value.unflatten(2, (attn.heads, -1))
+        query = query.unflatten(2, (attn.heads, attn.head_dim))
+        key = key.unflatten(2, (attn.heads, attn.head_dim))
+        value = value.unflatten(2, (attn.heads, attn.head_dim))
 
         attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
         hidden_states = attn.attn(query, key, value, attn_metadata)
@@ -502,7 +511,7 @@ class LTX2AudioVideoAttnProcessor:
 
         # LTX-2.3: per-head gated attention
         if attn.to_gate_logits is not None:
-            hidden_states = hidden_states.unflatten(2, (attn.heads, -1))  # [B, T, H, D]
+            hidden_states = hidden_states.unflatten(2, (attn.heads, attn.head_dim))  # [B, T, H, D]
             # 2.0 * sigmoid so zero-init gates produce 1.0 (identity)
             gates = 2.0 * torch.sigmoid(gate_logits)  # [B, T, H]
             hidden_states = hidden_states * gates.unsqueeze(-1)
@@ -1391,9 +1400,11 @@ class LTX2AudioVideoRotaryPosEmbed(nn.Module):
             # Reshape freqs to be compatible with multi-head attention
             b = cos_freq.shape[0]
             t = cos_freq.shape[1]
+            # Concrete per-head dim instead of -1 — see apply_split_rotary_emb
+            r = self.dim // self.num_attention_heads // 2
 
-            cos_freq = cos_freq.reshape(b, t, self.num_attention_heads, -1)
-            sin_freq = sin_freq.reshape(b, t, self.num_attention_heads, -1)
+            cos_freq = cos_freq.reshape(b, t, self.num_attention_heads, r)
+            sin_freq = sin_freq.reshape(b, t, self.num_attention_heads, r)
 
             cos_freqs = torch.swapaxes(cos_freq, 1, 2)  # (B,H,T,D//2)
             sin_freqs = torch.swapaxes(sin_freq, 1, 2)  # (B,H,T,D//2)
@@ -1939,10 +1950,6 @@ class LTX2VideoTransformer3DModel(nn.Module):
         # 5. Run transformer blocks
         for block in self.transformer_blocks:
             block_kwargs = {
-                "hidden_states": hidden_states,
-                "audio_hidden_states": audio_hidden_states,
-                "encoder_hidden_states": encoder_hidden_states,
-                "audio_encoder_hidden_states": audio_encoder_hidden_states,
                 "temb": temb,
                 "temb_audio": temb_audio,
                 "temb_ca_scale_shift": video_cross_attn_scale_shift,
@@ -1958,11 +1965,27 @@ class LTX2VideoTransformer3DModel(nn.Module):
                 "encoder_attention_mask": encoder_attention_mask,
                 "audio_encoder_attention_mask": audio_encoder_attention_mask,
             }
+            # NOTE: We intentionally pass the first 4 args here positionally,
+            # because CacheDiT expects the first 2 args to be hidden_states
+            # and encoder_hidden_states, so passing them as kwargs will cause
+            # positional / keywords arg collisions.
             if torch.is_grad_enabled() and self.gradient_checkpointing:
-                hidden_states, audio_hidden_states = self._gradient_checkpointing_func(block, **block_kwargs)
+                hidden_states, audio_hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    audio_hidden_states,
+                    encoder_hidden_states,
+                    audio_encoder_hidden_states,
+                    **block_kwargs,
+                )
             else:
-                hidden_states, audio_hidden_states = block(**block_kwargs)
-
+                hidden_states, audio_hidden_states = block(
+                    hidden_states,
+                    audio_hidden_states,
+                    encoder_hidden_states,
+                    audio_encoder_hidden_states,
+                    **block_kwargs,
+                )
         # 6. Output layers (including unpatchification)
         scale_shift_values = self.scale_shift_table[None, None] + embedded_timestep[:, :, None]
         shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py
index c92fd687f..82f59e7ff 100644
--- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py
+++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py
@@ -157,6 +157,9 @@ class _VideoAudioScheduler:
 
 
 class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
+    # Audio is diffused jointly with video; warmup must size audio tokens.
+    dummy_run_num_frames = 2
+
     def __init__(
         self,
         *,
@@ -1118,6 +1121,8 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
 class LTX2TwoStagesPipeline(nn.Module, SupportsComponentDiscovery):
     """LTX2TwoStagesPipeline is for two stages image to video generation"""
 
+    dummy_run_num_frames = 2
+
     _dit_modules: ClassVar[list[str]] = ["pipe.transformer"]
     _encoder_modules: ClassVar[list[str]] = ["pipe.text_encoder"]
     _vae_modules: ClassVar[list[str]] = ["pipe.vae", "pipe.audio_vae"]
diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py
index ec3b495f5..e6985e186 100644
--- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py
+++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py
@@ -37,6 +37,12 @@ from vllm.logger import init_logger
 from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_cfg_group,
+    get_classifier_free_guidance_rank,
+    get_classifier_free_guidance_world_size,
+)
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
@@ -44,6 +50,7 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest
 
 from .pipeline_ltx2 import (
     _get_prompt_field,
+    _VideoAudioScheduler,
     calculate_shift,
     create_transformer_from_config,
     load_transformer_config,
@@ -104,7 +111,7 @@ def get_ltx2_post_process_func(od_config: OmniDiffusionConfig):
     return post_process_func
 
 
-class LTX23Pipeline(nn.Module, ProgressBarMixin):
+class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
     """Fully independent LTX-2.3 pipeline.
 
     Key differences from LTX2Pipeline:
@@ -115,6 +122,9 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin):
     - CPU offloading: text encoder, connectors, VAE, vocoder stay on CPU
     """
 
+    # Audio is diffused jointly with video; warmup must size audio tokens.
+    dummy_run_num_frames = 2
+
     def __init__(
         self,
         *,
@@ -445,6 +455,26 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin):
     def _unpad_audio_latents(latents: torch.Tensor, num_frames: int) -> torch.Tensor:
         return latents[:, :num_frames]
 
+    @staticmethod
+    def _get_sp_padded_audio_latent_length(audio_latent_length: int, sp_size: int) -> int:
+        if sp_size > 1:
+            audio_latent_length += (sp_size - (audio_latent_length % sp_size)) % sp_size
+        return audio_latent_length
+
+    def _resolve_audio_latent_length(self, audio_latent_length: int, audio_latents: torch.Tensor | None) -> int:
+        if audio_latents is None or audio_latents.ndim != 4:
+            return audio_latent_length
+
+        provided_latent_length = audio_latents.shape[2]
+        sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1) or 1
+        padded_latent_length = self._get_sp_padded_audio_latent_length(audio_latent_length, int(sp_size))
+
+        # Keep requested duration semantics when callers pass 4D latents that
+        # are already padded for SP; other 4D lengths retain shape inference.
+        if provided_latent_length in {audio_latent_length, padded_latent_length}:
+            return audio_latent_length
+        return provided_latent_length
+
     # ------------------------------------------------------------------
     # Latent preparation
     # ------------------------------------------------------------------
@@ -497,9 +527,11 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin):
         latents: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, int, int]:
         original_latent_length = audio_latent_length
-        padded_latent_length = original_latent_length
         latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
 
+        sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1) or 1
+        padded_latent_length = self._get_sp_padded_audio_latent_length(original_latent_length, int(sp_size))
+
         if latents is not None:
             if latents.ndim == 4:
                 latents = self._pack_audio_latents(latents)
@@ -508,9 +540,32 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin):
             latents = self._normalize_audio_latents(latents, self.audio_vae.latents_mean, self.audio_vae.latents_std)
             noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype)
             latents = noise_scale * noise + (1 - noise_scale) * latents
+
+            if latents.shape[1] not in {original_latent_length, padded_latent_length}:
+                raise ValueError(
+                    "Provided `audio_latents` has incompatible audio frame count "
+                    f"{latents.shape[1]}; expected {original_latent_length} or {padded_latent_length}."
+                )
+
+            if latents.shape[1] == original_latent_length and padded_latent_length > original_latent_length:
+                padding = torch.zeros(
+                    latents.shape[0],
+                    padded_latent_length - original_latent_length,
+                    latents.shape[2],
+                    dtype=latents.dtype,
+                    device=latents.device,
+                )
+                latents = torch.cat([latents, padding], dim=1)
+
             return latents.to(device=device, dtype=dtype), original_latent_length, padded_latent_length
 
         shape = (batch_size, num_channels_latents, padded_latent_length, latent_mel_bins)
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
         latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
         latents = self._pack_audio_latents(latents)
         return latents, original_latent_length, padded_latent_length
@@ -593,6 +648,121 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin):
             return cache_context(context_name)
         return nullcontext()
 
+    # ------------------------------------------------------------------
+    # CFG helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _combine_x0_space_cfg(
+        sample: torch.Tensor,
+        positive_noise_pred: torch.Tensor,
+        negative_noise_pred: torch.Tensor,
+        sigma: torch.Tensor,
+        guidance_scale: float,
+    ) -> torch.Tensor:
+        x0_cond = sample - positive_noise_pred * sigma
+        x0_uncond = sample - negative_noise_pred * sigma
+        x0_guided = x0_cond + (guidance_scale - 1) * (x0_cond - x0_uncond)
+        return (sample - x0_guided) / sigma
+
+    def predict_noise(self, **kwargs):
+        with self._transformer_cache_context("cond_uncond"):
+            noise_pred_video, noise_pred_audio = self.transformer(**kwargs)
+        return noise_pred_video.float(), noise_pred_audio.float()
+
+    def combine_cfg_noise(
+        self,
+        positive_noise_pred,
+        negative_noise_pred,
+        true_cfg_scale,
+        cfg_normalize=False,
+        *,
+        video_latents: torch.Tensor | None = None,
+        audio_latents: torch.Tensor | None = None,
+        video_sigma: torch.Tensor | None = None,
+        audio_sigma: torch.Tensor | None = None,
+    ):
+        if video_latents is None or audio_latents is None or video_sigma is None or audio_sigma is None:
+            raise ValueError("LTX23Pipeline applies CFG in x0-space and requires video/audio latents and sigmas.")
+
+        video_pos, audio_pos = positive_noise_pred
+        video_neg, audio_neg = negative_noise_pred
+        video_combined = self._combine_x0_space_cfg(
+            video_latents,
+            video_pos,
+            video_neg,
+            video_sigma,
+            true_cfg_scale,
+        )
+        audio_combined = self._combine_x0_space_cfg(
+            audio_latents,
+            audio_pos,
+            audio_neg,
+            audio_sigma,
+            true_cfg_scale,
+        )
+        if cfg_normalize:
+            video_combined = self.cfg_normalize_function(video_pos, video_combined)
+            audio_combined = self.cfg_normalize_function(audio_pos, audio_combined)
+        return video_combined, audio_combined
+
+    def predict_noise_with_parallel_cfg(
+        self,
+        true_cfg_scale: float,
+        positive_kwargs: dict[str, Any],
+        negative_kwargs: dict[str, Any],
+        cfg_normalize: bool = True,
+        output_slice: int | None = None,
+        *,
+        video_latents: torch.Tensor | None = None,
+        audio_latents: torch.Tensor | None = None,
+        video_sigma: torch.Tensor | None = None,
+        audio_sigma: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        def maybe_slice(pred: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+            if output_slice is None:
+                return pred
+            return pred[0][:, :output_slice], pred[1][:, :output_slice]
+
+        cfg_world_size = get_classifier_free_guidance_world_size()
+        if cfg_world_size != 2:
+            raise ValueError(f"LTX23Pipeline parallel CFG requires cfg_parallel_size 2, but got {cfg_world_size}.")
+
+        cfg_group = get_cfg_group()
+        cfg_rank = get_classifier_free_guidance_rank()
+        branch_kwargs = positive_kwargs if cfg_rank == 0 else negative_kwargs
+        local_video_pred, local_audio_pred = maybe_slice(self.predict_noise(**branch_kwargs))
+
+        gathered_video = cfg_group.all_gather(local_video_pred, separate_tensors=True)
+        gathered_audio = cfg_group.all_gather(local_audio_pred, separate_tensors=True)
+        positive_noise_pred = (gathered_video[0], gathered_audio[0])
+        negative_noise_pred = (gathered_video[1], gathered_audio[1])
+
+        return self.combine_cfg_noise(
+            positive_noise_pred,
+            negative_noise_pred,
+            true_cfg_scale,
+            cfg_normalize,
+            video_latents=video_latents,
+            audio_latents=audio_latents,
+            video_sigma=video_sigma,
+            audio_sigma=audio_sigma,
+        )
+
+    def _synchronize_cfg_parallel_step_output(
+        self,
+        latents: tuple[torch.Tensor, torch.Tensor],
+        do_true_cfg: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if not (do_true_cfg and get_classifier_free_guidance_world_size() > 1):
+            return latents
+
+        latents = tuple(tensor.contiguous() for tensor in latents)
+        device = next((tensor.device for tensor in latents if tensor.is_cuda), None)
+        if device is not None:
+            torch.cuda.current_stream(device).synchronize()
+        return latents
+
     # ------------------------------------------------------------------
     # Forward pass
     # ------------------------------------------------------------------
@@ -711,6 +881,12 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin):
         self._attention_kwargs = attention_kwargs
         self._interrupt = False
         self._current_timestep = None
+        cfg_world_size = get_classifier_free_guidance_world_size()
+        if self.do_classifier_free_guidance and cfg_world_size not in (1, 2):
+            raise ValueError(
+                f"LTX23Pipeline supports CFG parallelism with cfg_parallel_size 1 or 2, but got {cfg_world_size}."
+            )
+        cfg_parallel_ready = self.do_classifier_free_guidance and cfg_world_size > 1
 
         if prompt is not None and isinstance(prompt, str):
             batch_size = 1
@@ -757,6 +933,21 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin):
         if torch.cuda.is_available():
             torch.accelerator.empty_cache()
 
+        positive_connector_prompt_embeds = connector_prompt_embeds
+        positive_connector_audio_prompt_embeds = connector_audio_prompt_embeds
+        positive_connector_attention_mask = connector_attention_mask
+        negative_connector_prompt_embeds = None
+        negative_connector_audio_prompt_embeds = None
+        negative_connector_attention_mask = None
+        if self.do_classifier_free_guidance:
+            split_batch = batch_size * num_videos_per_prompt
+            negative_connector_prompt_embeds = connector_prompt_embeds[:split_batch]
+            positive_connector_prompt_embeds = connector_prompt_embeds[split_batch:]
+            negative_connector_audio_prompt_embeds = connector_audio_prompt_embeds[:split_batch]
+            positive_connector_audio_prompt_embeds = connector_audio_prompt_embeds[split_batch:]
+            negative_connector_attention_mask = connector_attention_mask[:split_batch]
+            positive_connector_attention_mask = connector_attention_mask[split_batch:]
+
         # ---- Prepare latents ----
         latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
         latent_height = height // self.vae_spatial_compression_ratio
@@ -783,8 +974,7 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin):
             self.audio_sampling_rate / self.audio_hop_length / float(self.audio_vae_temporal_compression_ratio)
         )
         audio_num_frames = round(duration_s * audio_latents_per_second)
-        if audio_latents is not None and audio_latents.ndim == 4:
-            _, _, audio_num_frames, _ = audio_latents.shape
+        audio_num_frames = self._resolve_audio_latent_length(audio_num_frames, audio_latents)
 
         num_mel_bins = self.audio_vae.config.mel_bins if self.audio_vae is not None else 64
         latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
@@ -813,6 +1003,7 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin):
             self.scheduler.config.get("max_shift", 2.05),
         )
         audio_scheduler = copy.deepcopy(self.scheduler)
+        video_audio_scheduler = _VideoAudioScheduler(self.scheduler, audio_scheduler)
         _ = retrieve_timesteps(audio_scheduler, num_inference_steps, device, timesteps, sigmas=sigmas, mu=mu)
         timesteps, num_inference_steps = retrieve_timesteps(
             self.scheduler,
@@ -839,9 +1030,9 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin):
             audio_latents.device,
         )
 
-        # ---- CFG: duplicate coords for batch=2 ----
+        # ---- CFG: duplicate coords for single-rank batch=2 CFG ----
         # Connector outputs are already batch=2 (neg+pos concatenated before connector call)
-        if self.do_classifier_free_guidance:
+        if self.do_classifier_free_guidance and not cfg_parallel_ready:
             video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1))
             audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1))
 
@@ -856,60 +1047,112 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin):
 
                 self._current_timestep = t
 
-                # Duplicate latents for CFG (uncond + cond)
-                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
-                latent_model_input = latent_model_input.to(connector_prompt_embeds.dtype)
-                audio_latent_model_input = (
-                    torch.cat([audio_latents] * 2) if self.do_classifier_free_guidance else audio_latents
-                )
-                audio_latent_model_input = audio_latent_model_input.to(connector_prompt_embeds.dtype)
-                ts = t.expand(latent_model_input.shape[0])
-
-                with self._transformer_cache_context("cond_uncond"):
-                    noise_pred_video, noise_pred_audio = self.transformer(
-                        hidden_states=latent_model_input,
-                        audio_hidden_states=audio_latent_model_input,
-                        encoder_hidden_states=connector_prompt_embeds,
-                        audio_encoder_hidden_states=connector_audio_prompt_embeds,
-                        timestep=ts,
-                        sigma=ts,  # LTX-2.3: sigma for prompt_adaln
-                        encoder_attention_mask=connector_attention_mask,
-                        audio_encoder_attention_mask=connector_attention_mask,
-                        num_frames=latent_num_frames,
-                        height=latent_height,
-                        width=latent_width,
-                        fps=frame_rate,
-                        audio_num_frames=padded_audio_num_frames,
-                        video_coords=video_coords,
-                        audio_coords=audio_coords,
-                        attention_kwargs=attention_kwargs,
-                        return_dict=False,
+                if cfg_parallel_ready:
+                    latent_model_input = latents.to(positive_connector_prompt_embeds.dtype)
+                    audio_latent_model_input = audio_latents.to(positive_connector_prompt_embeds.dtype)
+                    ts = t.expand(latent_model_input.shape[0])
+                    positive_kwargs = {
+                        "hidden_states": latent_model_input,
+                        "audio_hidden_states": audio_latent_model_input,
+                        "encoder_hidden_states": positive_connector_prompt_embeds,
+                        "audio_encoder_hidden_states": positive_connector_audio_prompt_embeds,
+                        "timestep": ts,
+                        "sigma": ts,
+                        "encoder_attention_mask": positive_connector_attention_mask,
+                        "audio_encoder_attention_mask": positive_connector_attention_mask,
+                        "num_frames": latent_num_frames,
+                        "height": latent_height,
+                        "width": latent_width,
+                        "fps": frame_rate,
+                        "audio_num_frames": padded_audio_num_frames,
+                        "video_coords": video_coords,
+                        "audio_coords": audio_coords,
+                        "attention_kwargs": attention_kwargs,
+                        "return_dict": False,
+                    }
+                    negative_kwargs = {
+                        **positive_kwargs,
+                        "encoder_hidden_states": negative_connector_prompt_embeds,
+                        "audio_encoder_hidden_states": negative_connector_audio_prompt_embeds,
+                        "encoder_attention_mask": negative_connector_attention_mask,
+                        "audio_encoder_attention_mask": negative_connector_attention_mask,
+                    }
+                    noise_pred_video, noise_pred_audio = self.predict_noise_with_parallel_cfg(
+                        true_cfg_scale=guidance_scale,
+                        positive_kwargs=positive_kwargs,
+                        negative_kwargs=negative_kwargs,
+                        cfg_normalize=False,
+                        video_latents=latents,
+                        audio_latents=audio_latents,
+                        video_sigma=self.scheduler.sigmas[i],
+                        audio_sigma=audio_scheduler.sigmas[i],
                     )
 
-                noise_pred_video = noise_pred_video.float()
-                noise_pred_audio = noise_pred_audio.float()
+                    latents, audio_latents = self.scheduler_step_maybe_with_cfg(
+                        (noise_pred_video, noise_pred_audio),
+                        (t, t),
+                        (latents, audio_latents),
+                        do_true_cfg=self.do_classifier_free_guidance,
+                        per_request_scheduler=video_audio_scheduler,
+                    )
+                    latents, audio_latents = self._synchronize_cfg_parallel_step_output(
+                        (latents, audio_latents),
+                        do_true_cfg=self.do_classifier_free_guidance,
+                    )
+                else:
+                    latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
+                    latent_model_input = latent_model_input.to(connector_prompt_embeds.dtype)
+                    audio_latent_model_input = (
+                        torch.cat([audio_latents] * 2) if self.do_classifier_free_guidance else audio_latents
+                    )
+                    audio_latent_model_input = audio_latent_model_input.to(connector_prompt_embeds.dtype)
+                    ts = t.expand(latent_model_input.shape[0])
 
-                # CFG in x0-space (delta formulation matching diffusers)
-                if self.do_classifier_free_guidance:
-                    noise_pred_video_uncond, noise_pred_video_cond = noise_pred_video.chunk(2)
-                    # Convert velocity to x0: x0 = sample - velocity * sigma
-                    x0_video_cond = latents - noise_pred_video_cond * self.scheduler.sigmas[i]
-                    x0_video_uncond = latents - noise_pred_video_uncond * self.scheduler.sigmas[i]
-                    video_cfg_delta = (guidance_scale - 1) * (x0_video_cond - x0_video_uncond)
-                    x0_video_guided = x0_video_cond + video_cfg_delta
+                    with self._transformer_cache_context("cond_uncond"):
+                        noise_pred_video, noise_pred_audio = self.transformer(
+                            hidden_states=latent_model_input,
+                            audio_hidden_states=audio_latent_model_input,
+                            encoder_hidden_states=connector_prompt_embeds,
+                            audio_encoder_hidden_states=connector_audio_prompt_embeds,
+                            timestep=ts,
+                            sigma=ts,  # LTX-2.3: sigma for prompt_adaln
+                            encoder_attention_mask=connector_attention_mask,
+                            audio_encoder_attention_mask=connector_attention_mask,
+                            num_frames=latent_num_frames,
+                            height=latent_height,
+                            width=latent_width,
+                            fps=frame_rate,
+                            audio_num_frames=padded_audio_num_frames,
+                            video_coords=video_coords,
+                            audio_coords=audio_coords,
+                            attention_kwargs=attention_kwargs,
+                            return_dict=False,
+                        )
 
-                    noise_pred_audio_uncond, noise_pred_audio_cond = noise_pred_audio.chunk(2)
-                    x0_audio_cond = audio_latents - noise_pred_audio_cond * audio_scheduler.sigmas[i]
-                    x0_audio_uncond = audio_latents - noise_pred_audio_uncond * audio_scheduler.sigmas[i]
-                    audio_cfg_delta = (guidance_scale - 1) * (x0_audio_cond - x0_audio_uncond)
-                    x0_audio_guided = x0_audio_cond + audio_cfg_delta
+                    noise_pred_video = noise_pred_video.float()
+                    noise_pred_audio = noise_pred_audio.float()
 
-                    # Convert x0 back to velocity: v = (sample - x0) / sigma
-                    noise_pred_video = (latents - x0_video_guided) / self.scheduler.sigmas[i]
-                    noise_pred_audio = (audio_latents - x0_audio_guided) / audio_scheduler.sigmas[i]
+                    if self.do_classifier_free_guidance:
+                        noise_pred_video_uncond, noise_pred_video_cond = noise_pred_video.chunk(2)
+                        noise_pred_video = self._combine_x0_space_cfg(
+                            latents,
+                            noise_pred_video_cond,
+                            noise_pred_video_uncond,
+                            self.scheduler.sigmas[i],
+                            guidance_scale,
+                        )
 
-                latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0]
-                audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
+                        noise_pred_audio_uncond, noise_pred_audio_cond = noise_pred_audio.chunk(2)
+                        noise_pred_audio = self._combine_x0_space_cfg(
+                            audio_latents,
+                            noise_pred_audio_cond,
+                            noise_pred_audio_uncond,
+                            audio_scheduler.sigmas[i],
+                            guidance_scale,
+                        )
+
+                    latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0]
+                    audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
 
                 pbar.update()
 
diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py
index af4135d65..e93f4e744 100644
--- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py
+++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py
@@ -735,6 +735,7 @@ class LTX2ImageToVideoTwoStagesPipeline(nn.Module, SupportsComponentDiscovery):
     """LTXImageToVideoTwoStagesPipeline is for two stages image to video generation"""
 
     support_image_input = True
+    dummy_run_num_frames = 2
 
     _dit_modules: ClassVar[list[str]] = ["pipe.transformer"]
     _encoder_modules: ClassVar[list[str]] = ["pipe.text_encoder"]
diff --git a/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py b/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py
index 7f073ab87..65da889f1 100644
--- a/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py
+++ b/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py
@@ -2130,6 +2130,28 @@ class MagiHumanPipeline(nn.Module, ProgressBarMixin, DiffusionPipelineProfilerMi
 
         return resample(audio_np, target_len)
 
+    def encode_prompt(
+        self,
+        prompt: str,
+        target_length: int | None = None,
+    ) -> tuple[torch.Tensor, int]:
+        """Encode *prompt* with the T5-Gemma text encoder and pad to fixed length.
+
+        This is the single text-encoder entrypoint so the runner-level
+        prompt-embedding cache (see
+        ``vllm_omni/diffusion/cache/prompt_embed_cache.py``) can transparently
+        memoize results when the same prompt is submitted repeatedly.
+
+        Returns:
+            ``(context, original_context_len)`` matching
+            :func:`_get_padded_t5_gemma_embedding`.
+        """
+        return _get_padded_t5_gemma_embedding(
+            prompt,
+            self.text_encoder,
+            target_length if target_length is not None else self.t5_gemma_target_length,
+        )
+
     @torch.inference_mode()
     def forward(
         self,
@@ -2200,11 +2222,7 @@ class MagiHumanPipeline(nn.Module, ProgressBarMixin, DiffusionPipelineProfilerMi
             device=device,
         )
 
-        context, original_context_len = _get_padded_t5_gemma_embedding(
-            prompt,
-            self.text_encoder,
-            self.t5_gemma_target_length,
-        )
+        context, original_context_len = self.encode_prompt(prompt)
 
         if image_path is not None:
             br_image = self._encode_image(load_image(image_path), br_height, br_width)
diff --git a/vllm_omni/diffusion/models/ming_flash_omni/__init__.py b/vllm_omni/diffusion/models/ming_flash_omni/__init__.py
new file mode 100644
index 000000000..d10aed623
--- /dev/null
+++ b/vllm_omni/diffusion/models/ming_flash_omni/__init__.py
@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 The vLLM-Omni team.
+
+from .byte5_encoder import MingByT5Encoder
+from .condition_encoder import MingConditionEncoder
+from .ming_zimage_transformer import MingZImageTransformer2DModel
+from .pipeline_ming_imagegen import (
+    MingImagePipeline,
+    get_ming_image_post_process_func,
+)
+from .t5_block_mapper import T5EncoderBlockByT5Mapper
+
+__all__ = [
+    "MingByT5Encoder",
+    "MingConditionEncoder",
+    "MingImagePipeline",
+    "MingZImageTransformer2DModel",
+    "T5EncoderBlockByT5Mapper",
+    "get_ming_image_post_process_func",
+]
diff --git a/vllm_omni/diffusion/models/ming_flash_omni/byte5_encoder.py b/vllm_omni/diffusion/models/ming_flash_omni/byte5_encoder.py
new file mode 100644
index 000000000..e0748abcc
--- /dev/null
+++ b/vllm_omni/diffusion/models/ming_flash_omni/byte5_encoder.py
@@ -0,0 +1,223 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 The vLLM-Omni team.
+
+"""ByT5 glyph/text encoder for Ming-flash-omni-2.0 image generation.
+
+Ported from Ming's ``load_byte5_and_byte5_tokenizer`` +
+``T5EncoderBlockByT5Mapper`` + ``get_condition_embeds_for_image_gen`` byte5
+branch. The released checkpoint's byte5 weights were trained with font/color
+special tokens, so we replicate that vocabulary extension before loading —
+otherwise ``byte5_model.pt`` would shape-mismatch at the embedding layer.
+
+Typical forward: takes a list of user-supplied prompt strings (possibly with
+``<cn-font-N>`` / ``<color-N>`` markers), returns ``[B, byte5_max_length,
+diffusion_c_input_dim]`` features ready to be concatenated onto cap_feats
+along the sequence dimension.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from types import SimpleNamespace
+
+import torch
+import torch.nn as nn
+from transformers import AutoTokenizer, T5ForConditionalGeneration
+
+from vllm_omni.diffusion.models.ming_flash_omni.t5_block_mapper import (
+    T5EncoderBlockByT5Mapper,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _add_multilingual_special_tokens(
+    tokenizer,
+    text_encoder: nn.Module,
+    font_ann_path: Path,
+    color_ann_path: Path,
+    add_font: bool,
+    add_color: bool,
+    add_align: bool = False,
+) -> None:
+    """Extend the byte5 vocab with per-language font + color markers.
+
+    Mirrors ``add_special_token_multilingual`` in Ming's bizgen utils. The token
+    set must match the one the checkpoint was trained with, otherwise the
+    resized embedding table shape won't line up with the shipped weights.
+    """
+    idx_font_dict = json.loads(font_ann_path.read_text())
+    idx_color_dict = json.loads(color_ann_path.read_text())
+
+    font_tokens: list[str] = []
+    for font_code in idx_font_dict:
+        prefix = font_code[:3]
+        if prefix in ("cn-", "en-", "jp-", "kr-"):
+            font_tokens.append(f"<{prefix}font-{idx_font_dict[font_code]}>")
+        else:
+            font_tokens.append(f"<font-{idx_font_dict[font_code]}>")
+    color_tokens = [f"<color-{i}>" for i in range(len(idx_color_dict))]
+    align_tokens = [f"<align-{i}>" for i in range(3)]
+
+    extra: list[str] = []
+    if add_color:
+        extra += color_tokens
+    if add_font:
+        extra += font_tokens
+    if add_align:
+        extra += align_tokens
+    tokenizer.add_tokens(extra, special_tokens=True)
+    text_encoder.resize_token_embeddings(len(tokenizer))
+
+
+class MingByT5Encoder(nn.Module):
+    """Bundles byte5 tokenizer + T5 encoder + T5EncoderBlockByT5Mapper.
+
+    Build with ``MingByT5Encoder.from_checkpoint(<model>/byte5)`` when the
+    checkpoint ships byte5 weights; otherwise callers can skip this and the
+    pipeline falls back to no-byte5 conditioning.
+    """
+
+    def __init__(
+        self,
+        tokenizer,
+        text_encoder: nn.Module,
+        mapper: T5EncoderBlockByT5Mapper,
+        max_length: int,
+    ) -> None:
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.text_encoder = text_encoder
+        self.mapper = mapper
+        self.max_length = max_length
+
+    @classmethod
+    def from_checkpoint(
+        cls,
+        byte5_dir: Path,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> MingByT5Encoder:
+        # Wrap the whole load in fork_rng so any nn.init.normal_ inside
+        # T5ForConditionalGeneration.from_pretrained / vocab-resize cannot
+        # advance the default CPU/CUDA generator state. Without this, the
+        # diffusion pipeline's seeded noise becomes order-dependent across
+        # requests (same-seed replays diverge).
+        cuda_devs = list(range(torch.accelerator.device_count())) if torch.cuda.is_available() else []
+        with torch.random.fork_rng(devices=cuda_devs, enabled=True):
+            return cls._from_checkpoint_impl(byte5_dir, device=device, dtype=dtype)
+
+    @classmethod
+    def _from_checkpoint_impl(
+        cls,
+        byte5_dir: Path,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> MingByT5Encoder:
+        byte5_dir = Path(byte5_dir)
+        # Ming checkpoint uses ``byt5`` (no 'e') in filenames and JSON keys
+        # variable names below keep the ``byte5_`` spelling for readability.
+        cfg_path = byte5_dir / "byt5.json"
+        cfg_raw = json.loads(cfg_path.read_text())
+        cfg = SimpleNamespace(**cfg_raw)
+        byte5_config = cfg.byt5_config
+        mapper_config = cfg.byt5_mapper_config
+        max_length = int(cfg.byt5_max_length)
+
+        # ---- Tokenizer + T5 encoder (base).
+        ckpt_key = byte5_config.get("byt5_ckpt_path")
+        byte5_ckpt_path = byte5_dir / ckpt_key.lstrip("./")
+        tokenizer = AutoTokenizer.from_pretrained(byte5_ckpt_path, local_files_only=True)
+        text_encoder = T5ForConditionalGeneration.from_pretrained(byte5_ckpt_path, local_files_only=True).get_encoder()
+
+        # ---- Extend vocab with font/color markers so the shipped weights load.
+        if byte5_config.get("special_token"):
+            if not byte5_config.get("multilingual", True):
+                raise NotImplementedError(
+                    "Non-multilingual byte5 vocab extension is not ported; "
+                    "the released Ming checkpoint uses multilingual=True."
+                )
+            _add_multilingual_special_tokens(
+                tokenizer,
+                text_encoder,
+                font_ann_path=byte5_dir / byte5_config["font_ann_path"].lstrip("./"),
+                color_ann_path=byte5_dir / byte5_config["color_ann_path"].lstrip("./"),
+                add_font=bool(byte5_config.get("font_special_token")),
+                add_color=bool(byte5_config.get("color_special_token")),
+            )
+
+        # ---- Load byt5 text-encoder weights (base.pt wraps the backbone in a
+        # trainable-module container; byt5_model.pt has the top-level encoder
+        # state). Follow Ming's two-step load.
+        base_state = torch.load(byte5_dir / "byt5_model" / "base.pt", map_location="cpu", weights_only=False)
+        prefix = "module.text_tower.encoder."
+        base_filtered = {
+            name[len(prefix) :]: state for name, state in base_state["state_dict"].items() if name.startswith(prefix)
+        }
+        text_encoder.load_state_dict(base_filtered, strict=True)
+        del base_state, base_filtered
+
+        encoder_state = torch.load(byte5_dir / "byt5_model" / "byt5_model.pt", map_location="cpu", weights_only=False)
+        text_encoder.load_state_dict(encoder_state)
+        del encoder_state
+
+        text_encoder.to(device=device, dtype=dtype).eval()
+
+        # ---- Mapper.
+        mapper = T5EncoderBlockByT5Mapper(
+            byte5_config=text_encoder.config,
+            num_layers=int(mapper_config["num_layers"]),
+            sdxl_channels=int(mapper_config["sdxl_channels"]),
+        )
+        mapper_state = torch.load(byte5_dir / "byt5_mapper" / "byt5_mapper.pt", map_location="cpu", weights_only=False)
+        # mapper now uses vllm-omni TP-aware T5 layers (fused qkv_proj / wi);
+        mapper.load_weights(mapper_state.items())
+        del mapper_state
+        mapper.to(device=device, dtype=dtype).eval()
+
+        logger.info(
+            "[MingByT5Encoder] ready: d_model=%d mapper_layers=%d sdxl_channels=%d max_length=%d vocab=%d",
+            text_encoder.config.d_model,
+            mapper_config["num_layers"],
+            mapper_config["sdxl_channels"],
+            max_length,
+            len(tokenizer),
+        )
+        return cls(tokenizer, text_encoder, mapper, max_length)
+
+    @torch.inference_mode()
+    def forward(self, texts: list[str]) -> torch.Tensor:
+        """Tokenize → T5 encode → mapper; masks out padded positions.
+
+        Returns ``[B, max_length, sdxl_channels]``. Padded positions are zeroed
+        so the downstream ``torch.cat`` with cap_feats doesn't inject garbage.
+        """
+        device = next(self.text_encoder.parameters()).device
+        dtype = next(self.text_encoder.parameters()).dtype
+
+        tokens = self.tokenizer(
+            texts,
+            padding="max_length",
+            max_length=self.max_length,
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        input_ids = tokens.input_ids.to(device)
+        attention_mask = tokens.attention_mask.to(device)
+
+        encoder_out = self.text_encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask.float(),
+        )
+        hidden_states = encoder_out[0]
+        feats = self.mapper(hidden_states, attention_mask)
+        feats = feats * attention_mask.unsqueeze(-1).to(dtype=feats.dtype)
+        return feats.to(dtype=dtype)
+
+
+__all__ = ["MingByT5Encoder"]
diff --git a/vllm_omni/diffusion/models/ming_flash_omni/condition_encoder.py b/vllm_omni/diffusion/models/ming_flash_omni/condition_encoder.py
new file mode 100644
index 000000000..6ebca4f5c
--- /dev/null
+++ b/vllm_omni/diffusion/models/ming_flash_omni/condition_encoder.py
@@ -0,0 +1,373 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 The vLLM-Omni team.
+#
+# Adapted from Ming repository (inclusionAI/Ming) — the ``get_condition_embeds
+# _for_image_gen`` path in modeling_bailingmm2.py.
+
+"""Ming-flash-omni-2.0 condition encoder for image generation.
+
+Pipeline (runs inside the imagegen stage):
+
+    thinker hidden states at query-token positions     [B, N, 4096]
+                         │
+                         ▼ proj_in (Linear, bias=True)
+                                                       [B, N, 1536]
+                         │
+                         ▼ Qwen2ForCausalLM connector (is_causal=False)
+                           loaded from <checkpoint>/connector/
+                                                       [B, N, 1536]
+                         │
+                         ▼ proj_out (Linear, bias=True)
+                                                       [B, N, 2560]
+                         │
+                         ▼ F.normalize(dim=-1) × 1000  (text_encoder_norm)
+                                                       [B, N, 2560]
+                         │
+                         ▼
+            cap_feats consumed by ZImageTransformer2DModel
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+
+from vllm_omni.transformers_utils.configs.ming_flash_omni import MingImageGenConfig
+
+logger = logging.getLogger(__name__)
+
+
+class MingConditionEncoder(nn.Module):
+    """Wraps a Qwen2 connector + norm/projection, producing DiT condition embeds.
+
+    The connector is a ``Qwen2ForCausalLM`` loaded from the ``connector/``
+    subfolder of the Ming checkpoint. We run its base model in a non-causal
+    (bidirectional) mode, because the connector is used as an encoder over the
+    pre-baked query-token hidden states, not as an autoregressive decoder.
+
+    Args:
+        image_gen_config: ``MingImageGenConfig`` from ``MingFlashOmniConfig``.
+        thinker_hidden_size: Hidden size of the thinker (BailingMoeV2) model.
+            Used to build a ``proj_in`` layer when the connector embedding
+            dim differs. For the released checkpoint this is 4096.
+        device: Placement for the module.
+        dtype: Parameter dtype (typically bfloat16 / float16).
+    """
+
+    def __init__(
+        self,
+        image_gen_config: MingImageGenConfig,
+        *,
+        thinker_hidden_size: int = 4096,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> None:
+        super().__init__()
+        self.config = image_gen_config
+        self.thinker_hidden_size = thinker_hidden_size
+        self._target_device = torch.device(device) if device is not None else None
+        self._target_dtype = dtype
+
+        # Populated lazily by ``load_from_checkpoint`` to keep this module
+        # cheap to construct (useful for dummy-init paths and unit tests).
+        self.connector: nn.Module | None = None
+        self.connector_hidden_size: int | None = None
+        self.proj_in: nn.Module = nn.Identity()
+        self.proj_out: nn.Module = nn.Identity()
+        self.norm: nn.Module = nn.Identity()
+
+    # ------------------------------------------------------------------
+    # Weight loading
+    # ------------------------------------------------------------------
+
+    def load_from_checkpoint(self, model_path: str | Path) -> None:
+        """Load the Qwen2 connector + optional projection/norm weights.
+
+        This uses HF transformers directly (not vllm's weight loader) because
+        the connector is small (~1.5B params) and only runs once per request
+        as an encoder — vllm's distributed loading machinery is overkill.
+        """
+        from transformers import AutoConfig, Qwen2ForCausalLM
+
+        model_path = Path(model_path)
+        connector_path = model_path / self.config.connector_subfolder
+        logger.info("[MingConditionEncoder] loading connector from %s", connector_path)
+
+        connector_cfg = AutoConfig.from_pretrained(connector_path, trust_remote_code=True, local_files_only=True)
+        # Disable causal masking: the connector is used as a bidirectional
+        # encoder over query-token hidden states in image-gen mode.
+        connector_cfg.is_decoder = False
+        self.connector_hidden_size = int(connector_cfg.hidden_size)
+        logger.info(
+            "[MingConditionEncoder] connector hidden_size=%d, layers=%d",
+            self.connector_hidden_size,
+            getattr(connector_cfg, "num_hidden_layers", -1),
+        )
+
+        connector = Qwen2ForCausalLM.from_pretrained(
+            connector_path,
+            config=connector_cfg,
+            torch_dtype=self._target_dtype,
+            local_files_only=True,
+        )
+        # Force bidirectional attention in every self-attn module. Qwen2 layers
+        # do not branch on ``is_causal`` directly, so we monkey-patch the flag
+        # on each attention block (defensive — some transformers versions read
+        # ``self_attn.is_causal`` in forward).
+        patched = 0
+        for module in connector.modules():
+            if hasattr(module, "is_causal"):
+                module.is_causal = False
+                patched += 1
+        logger.info("[MingConditionEncoder] disabled is_causal on %d sub-modules", patched)
+
+        # We only need the base encoder (no LM head).
+        base = getattr(connector, "model", connector)
+        self.connector = base
+
+        # proj_in: align thinker_hidden_size -> connector_hidden_size.
+        # Observed layout of mlp/ on the released checkpoint
+        # (inclusionAI/Ming-flash-omni-2.0):
+        #   proj_in.{weight,bias}   : Linear(4096 -> 1536)
+        #   proj_out.{weight,bias}  : Linear(1536 -> 2560)
+        #   query_tokens_dict.16x16 : learnable tokens consumed on the thinker
+        #                             side (NOT loaded here).
+        logger.info(
+            "[MingConditionEncoder] adding proj_in: %d -> %d (bias=True)",
+            self.thinker_hidden_size,
+            self.connector_hidden_size,
+        )
+        self.proj_in = nn.Linear(self.thinker_hidden_size, self.connector_hidden_size, bias=True)
+
+        # NOTE: ``text_encoder_norm=True`` in Ming's mlp/config.json refers to
+        # applying ``F.normalize(..., dim=-1)`` (L2 normalization) on the
+        # FINAL cap_feats AFTER proj_out — NOT an intermediate RMSNorm. See
+        # modeling_bailingmm2.py::get_condition_embeds_for_image_gen. We keep
+        # ``self.norm = nn.Identity`` and apply the L2 normalize explicitly
+        # at the end of ``forward``.
+        self.norm = nn.Identity()
+
+        # proj_out: align connector_hidden_size -> diffusion_c_input_dim.
+        c_out = self.config.diffusion_c_input_dim
+        logger.info(
+            "[MingConditionEncoder] adding proj_out: %d -> %d (bias=True)",
+            self.connector_hidden_size,
+            c_out,
+        )
+        self.proj_out = nn.Linear(self.connector_hidden_size, c_out, bias=True)
+
+        mlp_path = model_path / self.config.mlp_subfolder
+        mlp_cfg_path = mlp_path / "config.json"
+        if mlp_cfg_path.exists() and not json.loads(mlp_cfg_path.read_text()).get("use_identity_mlp", False):
+            # MingImagePipeline skips ZImageModel_withMLP, only correct when its inner MLP is Identity.
+            raise NotImplementedError(f"{mlp_cfg_path} has use_identity_mlp=False; ToClipMLP path not implemented.")
+        self._load_optional_mlp_weights(mlp_path)
+
+        if self._target_device is not None:
+            self.to(self._target_device)
+        if self._target_dtype is not None:
+            self.to(dtype=self._target_dtype)
+
+        logger.info(
+            "[MingConditionEncoder] ready: in=%d -> conn=%d -> out=%d",
+            self.thinker_hidden_size,
+            self.connector_hidden_size,
+            c_out,
+        )
+
+    def _load_optional_mlp_weights(self, mlp_path: Path) -> None:
+        """Load proj_in / proj_out (and optional norm) weights from mlp/.
+
+        Expected keys — observed on inclusionAI/Ming-flash-omni-2.0:
+            proj_in.weight           [1536, 4096]
+            proj_in.bias             [1536]
+            proj_out.weight          [2560, 1536]
+            proj_out.bias            [2560]
+            query_tokens_dict.16x16  [256, 4096]   -> thinker-side, skipped here
+            (optional) norm.weight   [1536]
+
+        Any extra keys are logged as warnings so we don't silently miss new
+        fields. Missing proj weights are fatal: they are required for the
+        condition path to be meaningful.
+        """
+        if not mlp_path.exists():
+            logger.warning(
+                "[MingConditionEncoder] mlp/ subfolder missing at %s — proj/norm "
+                "will stay randomly initialized. EXPECT BAD IMAGES until this "
+                "is fixed on real hardware.",
+                mlp_path,
+            )
+            return
+
+        try:
+            from safetensors.torch import load_file  # type: ignore
+        except ImportError:
+            logger.exception("[MingConditionEncoder] safetensors not installed")
+            return
+
+        candidates = sorted(mlp_path.glob("*.safetensors"))
+        if not candidates:
+            candidates = sorted(mlp_path.glob("*.bin"))
+        if not candidates:
+            logger.warning("[MingConditionEncoder] no weight files under %s", mlp_path)
+            return
+
+        state: dict[str, torch.Tensor] = {}
+        for p in candidates:
+            logger.info("[MingConditionEncoder] reading mlp weights: %s", p)
+            if p.suffix == ".safetensors":
+                state.update(load_file(str(p)))
+            else:
+                state.update(torch.load(str(p), map_location="cpu"))
+        logger.info(
+            "[MingConditionEncoder] mlp/ keys: %s",
+            sorted(state.keys()),
+        )
+
+        handled: set[str] = set()
+
+        def _copy(dst: nn.Parameter | torch.Tensor, src_key: str) -> bool:
+            src = state.get(src_key)
+            if src is None:
+                logger.error("[MingConditionEncoder] mlp/ missing key %r", src_key)
+                return False
+            if tuple(src.shape) != tuple(dst.shape):
+                logger.error(
+                    "[MingConditionEncoder] mlp/%s shape mismatch: checkpoint=%s, module=%s",
+                    src_key,
+                    tuple(src.shape),
+                    tuple(dst.shape),
+                )
+                return False
+            with torch.no_grad():
+                dst.copy_(src.to(dtype=dst.dtype, device=dst.device))
+            handled.add(src_key)
+            logger.info(
+                "[MingConditionEncoder] loaded mlp/%s -> %s",
+                src_key,
+                tuple(dst.shape),
+            )
+            return True
+
+        ok_w = _copy(self.proj_in.weight, "proj_in.weight")
+        ok_b = _copy(self.proj_in.bias, "proj_in.bias")
+        ok_ow = _copy(self.proj_out.weight, "proj_out.weight")
+        ok_ob = _copy(self.proj_out.bias, "proj_out.bias")
+        if not (ok_w and ok_b and ok_ow and ok_ob):
+            logger.error(
+                "[MingConditionEncoder] proj_in/proj_out NOT fully loaded; diffusion conditioning will be garbage."
+            )
+
+        # Optional norm weight (Ming uses plain RMSNorm; may or may not ship).
+        if "norm.weight" in state and hasattr(self.norm, "weight"):
+            _copy(self.norm.weight, "norm.weight")
+
+        # query_tokens_dict lives in mlp/ but is consumed on the thinker side
+        # — we neither need nor load it here. Just log the shape so the
+        # thinker-side patch knows what to expect.
+        for k, v in state.items():
+            if k.startswith("query_tokens_dict"):
+                logger.info(
+                    "[MingConditionEncoder] thinker-side parameter %s shape=%s (NOT loaded here; thinker must own it)",
+                    k,
+                    tuple(v.shape),
+                )
+                handled.add(k)
+
+        leftover = set(state.keys()) - handled
+        if leftover:
+            logger.warning("[MingConditionEncoder] mlp/ unhandled keys: %s", sorted(leftover))
+
+    # ------------------------------------------------------------------
+    # Forward
+    # ------------------------------------------------------------------
+
+    def forward(
+        self,
+        thinker_hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Encode thinker hidden states into DiT condition embeddings.
+
+        Args:
+            thinker_hidden_states: ``[B, N, thinker_hidden_size]`` — sliced at
+                the learnable query-token positions by the stage input
+                processor before being passed here.
+            attention_mask: Optional ``[B, N]`` mask. Defaults to all-ones.
+
+        Returns:
+            ``[B, N, diffusion_c_input_dim]`` condition tensor ready for the
+            ZImage transformer's ``cap_feats`` input.
+        """
+        if self.connector is None:
+            raise RuntimeError("MingConditionEncoder.load_from_checkpoint() must be called before forward().")
+        if thinker_hidden_states.dim() != 3:
+            raise ValueError(f"expected [B, N, H], got shape {tuple(thinker_hidden_states.shape)}")
+
+        b, n, _ = thinker_hidden_states.shape
+        logger.debug(
+            "[MingCondEnc] input shape=%s dtype=%s", tuple(thinker_hidden_states.shape), thinker_hidden_states.dtype
+        )
+
+        x = self.proj_in(thinker_hidden_states)  # [B, N, conn_hidden]
+
+        # Ming's ``get_condition_embeds_for_image_gen`` passes a 4D full-ones
+        # attention mask of shape [B, 1, N, N] to the Qwen2 connector,
+        # forcing full bidirectional self-attention over all query positions.
+        if attention_mask is None:
+            attention_mask = torch.ones(
+                (b, 1, n, n),
+                dtype=x.dtype,
+                device=x.device,
+            )
+        elif attention_mask.dim() == 2:
+            mask_2d = attention_mask.to(x.dtype)
+            attention_mask = mask_2d[:, None, None, :].expand(b, 1, n, n)
+
+        out = self.connector(
+            inputs_embeds=x,
+            attention_mask=attention_mask,
+            use_cache=False,
+            output_hidden_states=True,
+            return_dict=True,
+        )
+        hidden = out.hidden_states[-1]  # [B, N, conn_hidden]
+
+        cap_feats = self.proj_out(hidden)  # [B, N, diffusion_c_input_dim]
+
+        # L2 normalize + ×1000 rescale.  Ming's ``get_condition_embeds_for_image_gen``
+        # applies F.normalize; ``ZImageModel_withMLP`` then rescales by 1000
+        # when ``text_encoder_norm=True``.  We fold both into one place.
+        cap_feats = torch.nn.functional.normalize(cap_feats, dim=-1)
+        if self.config.text_encoder_norm:
+            cap_feats = cap_feats * 1000.0
+
+        logger.debug("[MingCondEnc] output shape=%s dtype=%s", tuple(cap_feats.shape), cap_feats.dtype)
+        return cap_feats
+
+    # ------------------------------------------------------------------
+    # Utility
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def zero_negative(
+        self,
+        cap_feats: torch.Tensor,
+    ) -> torch.Tensor:
+        """Return a zero tensor shaped like ``cap_feats`` for CFG negatives."""
+        return torch.zeros_like(cap_feats)
+
+    def extra_repr(self) -> str:
+        return (
+            f"thinker_hidden_size={self.thinker_hidden_size}, "
+            f"connector_hidden_size={self.connector_hidden_size}, "
+            f"diffusion_c_input_dim={self.config.diffusion_c_input_dim}, "
+            f"text_encoder_norm={self.config.text_encoder_norm}"
+        )
+
+
+__all__ = ["MingConditionEncoder"]
diff --git a/vllm_omni/diffusion/models/ming_flash_omni/ming_zimage_transformer.py b/vllm_omni/diffusion/models/ming_flash_omni/ming_zimage_transformer.py
new file mode 100644
index 000000000..4d3925f6d
--- /dev/null
+++ b/vllm_omni/diffusion/models/ming_flash_omni/ming_zimage_transformer.py
@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 The vLLM-Omni team.
+
+"""Ming-specific subclass of ZImageTransformer2DModel that supports ``ref_x``.
+
+Ming's img2img path concatenates a VAE-encoded reference latent along the
+frame axis before patchification, then drops the reference portion from the
+unpatchified output. This is a surgical override — everything else (attention,
+RoPE, final layer) stays on the parent implementation.
+
+"""
+
+from __future__ import annotations
+
+import torch
+
+from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
+from vllm_omni.diffusion.models.z_image.z_image_transformer import ZImageTransformer2DModel
+
+
+class MingZImageTransformer2DModel(ZImageTransformer2DModel):
+    """ZImage DiT with Ming's reference-latent support."""
+
+    def forward(
+        self,
+        x: list[torch.Tensor],
+        t,
+        cap_feats: list[torch.Tensor],
+        patch_size=2,
+        f_patch_size=1,
+    ):
+        ref_latent = get_forward_context().ref_latent if is_forward_context_available() else None
+        if ref_latent is not None:
+            per_item = ref_latent[0].unsqueeze(1).to(dtype=x[0].dtype, device=x[0].device)  # [C, 1, H, W]
+            x = [torch.cat([img, per_item], dim=1) for img in x]
+        return super().forward(x, t, cap_feats, patch_size=patch_size, f_patch_size=f_patch_size)
+
+    def unpatchify(
+        self,
+        x: list[torch.Tensor],
+        size: list[tuple],
+        patch_size,
+        f_patch_size,
+    ) -> list[torch.Tensor]:
+        out = super().unpatchify(x, size, patch_size, f_patch_size)
+        # No-op when F==1 (pure t2i); drops the reference-frame prediction when F==2.
+        return [t[:, :1, :, :] for t in out]
+
+
+__all__ = ["MingZImageTransformer2DModel"]
diff --git a/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py b/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py
new file mode 100644
index 000000000..0da26aff8
--- /dev/null
+++ b/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py
@@ -0,0 +1,520 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 The vLLM-Omni team.
+
+"""Ming-flash-omni-2.0 imagegen (text-to-image / img2img) diffusion pipeline.
+
+Cross-stage data flow:
+
+    Stage 0 (thinker, llm)           Stage 1 (imagegen, diffusion)
+    ────────────────────             ──────────────────────────────
+    forward returns                  thinker2imagegen hook slices
+    multimodal_output[               final_hidden_states at
+      "final_hidden_states"]         <imagePatch> positions,
+       ↓                             returns list[dict] with
+    shared_memory_connector          {"extra": {"thinker_hidden_states"}}
+       ↓                             ──── via OmniMsgpackEncoder ────>
+                                     MingImagePipeline.forward(req):
+                                       hidden = req.prompts[0]["extra"][...]
+                                       cond = condition_encoder(hidden)
+                                       img = ZImagePipeline-style loop
+                                       return DiffusionOutput(output=img)
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn as nn
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl import DistributedAutoencoderKL
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.forward_context import set_forward_context_ref_latent
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
+from vllm_omni.diffusion.models.ming_flash_omni.byte5_encoder import (
+    MingByT5Encoder,
+)
+from vllm_omni.diffusion.models.ming_flash_omni.condition_encoder import (
+    MingConditionEncoder,
+)
+from vllm_omni.diffusion.models.ming_flash_omni.ming_zimage_transformer import (
+    MingZImageTransformer2DModel,
+)
+from vllm_omni.diffusion.models.z_image.pipeline_z_image import ZImagePipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.model_executor.model_loader.weight_utils import (
+    download_weights_from_hf_specific,
+)
+from vllm_omni.transformers_utils.configs.ming_flash_omni import MingImageGenConfig
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class _ZPipelineSamplingParams:
+    """Typed shim satisfying the attributes ZImagePipeline.forward reads
+    from ``req.sampling_params``.  Unlike SimpleNamespace this fails loudly
+    at construction if a required field is missing."""
+
+    height: int
+    width: int
+    num_inference_steps: int
+    guidance_scale: float
+    generator: torch.Generator | None = None
+    strength: float | None = None
+    sigmas: list[float] | None = None
+    max_sequence_length: int = 512
+    guidance_rescale: float | None = None
+    num_outputs_per_prompt: int = 1
+
+
+@dataclass
+class _ZPipelineRequest:
+    """Typed shim satisfying the attributes ZImagePipeline.forward reads
+    from ``req``."""
+
+    request_id: str
+    sampling_params: _ZPipelineSamplingParams
+    prompts: list[dict[str, Any]] = field(default_factory=lambda: [{"prompt": "", "negative_prompt": ""}])
+
+
+class MingImagePipeline(ZImagePipeline):
+    """Ming-flash-omni-2.0 text-to-image diffusion pipeline.
+
+    Ming-specific components added on top of the inherited contract:
+      * ``condition_encoder`` — Qwen2 connector + proj_in/out + F.normalize×1000
+      * ``byte5``             — Optional ByT5 glyph encoder (loaded if checkpoint
+                                ships ``byt5/``)
+    """
+
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",  # noqa: ARG002
+    ) -> None:
+        # Skip ZImagePipeline.__init__ (it would eagerly load the Z-Image text
+        # encoder/tokenizer that Ming replaces with its own condition_encoder).
+        nn.Module.__init__(self)
+
+        model_path = od_config.model
+        if not os.path.exists(model_path):
+            model_path = download_weights_from_hf_specific(model_path, od_config.revision, ["*"])
+
+        dtype = getattr(od_config, "dtype", torch.bfloat16)
+        local_files_only = os.path.exists(model_path)
+
+        self.od_config = od_config
+        self._execution_device = get_local_device()
+        self.device = self._execution_device  # Ming convention alias
+        self._dtype = dtype
+
+        # Ming's per-checkpoint image-gen configuration. We cannot rely on
+        # ``od_config.hf_config.image_gen_config`` because the diffusion
+        # stage is started with ``hf_config_name: thinker_config`` (the
+        # BailingMM2Config), which does not carry a MingImageGenConfig.
+        # Fall back to defaults that match the released checkpoint.
+        self.image_gen_config = MingImageGenConfig()
+        logger.info(
+            "[MingImagePipeline] init: model=%s dtype=%s image_gen_config=%s",
+            model_path,
+            dtype,
+            self.image_gen_config,
+        )
+
+        # ----- weights_sources: DiT transformer + VAE.
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=model_path,
+                subfolder=self.image_gen_config.transformer_subfolder,
+                revision=od_config.revision,
+                prefix="transformer.",
+                fall_back_to_pt=True,
+            ),
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=model_path,
+                subfolder=self.image_gen_config.vae_subfolder,
+                revision=od_config.revision,
+                prefix="vae.",
+            ),
+        ]
+
+        prefetch_subfolders(
+            model_path,
+            [self.image_gen_config.scheduler_subfolder, self.image_gen_config.vae_subfolder],
+            local_files_only=local_files_only,
+        )
+
+        # ----- Scheduler: load config-only from disk + Ming-specific override.
+        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            model_path,
+            subfolder=self.image_gen_config.scheduler_subfolder,
+            local_files_only=local_files_only,
+        )
+        # Ming forces use_dynamic_shifting=True at runtime regardless of what
+        # the checkpoint scheduler_config.json ships.
+        self.scheduler.config["use_dynamic_shifting"] = True
+        logger.info(
+            "[MingImagePipeline] scheduler: %s (use_dynamic_shifting=True)",
+            type(self.scheduler).__name__,
+        )
+
+        # ----- VAE: DistributedAutoencoderKL.
+        vae_config = DistributedAutoencoderKL.load_config(
+            model_path, subfolder=self.image_gen_config.vae_subfolder, local_files_only=local_files_only
+        )
+        self.vae = DistributedAutoencoderKL.from_config(vae_config).to(self._execution_device, dtype=dtype)
+        self.vae.eval()
+
+        # ----- DiT transformer.
+        self.transformer = MingZImageTransformer2DModel(quant_config=None)
+
+        # Ming brings its own conditioning path — no Z-Image text_encoder /
+        # tokenizer.
+        self.text_encoder = None
+        self.tokenizer = None
+
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2, do_convert_rgb=True)
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=getattr(od_config, "enable_diffusion_pipeline_profiler", False)
+        )
+
+        # ----- Condition encoder (Qwen2 connector + proj_in/out + norm×1000).
+        self.condition_encoder = MingConditionEncoder(
+            self.image_gen_config,
+            thinker_hidden_size=self.image_gen_config.thinker_hidden_size,
+            device=self.device,
+            dtype=dtype,
+        )
+        self.condition_encoder.load_from_checkpoint(model_path)
+
+        # Optional ByT5 glyph/text encoder. Only loaded when the checkpoint
+        # ``extra_args.image_gen.byte5_text`` requests will be ignored.
+        byte5_dir = Path(model_path) / "byt5"
+        if byte5_dir.exists():
+            self.byte5 = MingByT5Encoder.from_checkpoint(byte5_dir, device=self.device, dtype=dtype)
+        else:
+            self.byte5 = None
+            logger.info("[MingImagePipeline] no byt5/ subfolder at %s; ByT5 enhancement disabled", byte5_dir)
+
+        logger.info("[MingImagePipeline] ready — vae_scale_factor=%d", self.vae_scale_factor)
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _resolve_byte5_texts(extra: dict, sampling_params) -> list[str]:
+        """Resolve byte5 glyph texts.
+
+        Two sources, in order of priority:
+          1. ``extra["byte5_text"]`` — auto-extracted from the user prompt's
+             quoted spans by ``thinker2imagegen`` (already wrapped as
+             ``'Text "<glyph>". '`` by Ming's ``get_text_from_prompt``).
+          2. ``sampling_params.extra_args["image_gen"]["byte5_text"]`` — an
+             explicit override for programmatic callers; raw strings without
+             the ``Text "..."`` wrapper are auto-wrapped here to match the
+             distribution ByT5 was trained on.
+        """
+        # Source 1: auto-extracted, already wrapped. Return as-is if non-empty.
+        raw = extra.get("byte5_text")
+        if isinstance(raw, str):
+            raw = [raw]
+        if isinstance(raw, list):
+            cleaned = [t for t in raw if isinstance(t, str) and t.strip()]
+            if cleaned:
+                return cleaned
+
+        # Source 2: explicit override — wrap raw strings so the byte5 encoder
+        # sees the same ``Text "<glyph>". `` format Ming used during training.
+        raw = ((getattr(sampling_params, "extra_args", None) or {}).get("image_gen") or {}).get("byte5_text")
+        if isinstance(raw, str):
+            raw = [raw]
+        if isinstance(raw, list):
+            out: list[str] = []
+            for t in raw:
+                if not isinstance(t, str):
+                    continue
+                s = t.strip()
+                if not s:
+                    continue
+                # Don't double-wrap if the caller already supplied ``Text "...". ``.
+                out.append(s if s.startswith('Text "') else f'Text "{s}". ')
+            if out:
+                return out
+        return []
+
+    @torch.inference_mode()
+    def _encode_reference_image(self, ref, height: int, width: int) -> torch.Tensor | None:
+        """Turn a PIL/tensor reference image into a VAE latent for ``ref_x``.
+
+        Applies the same shift/scale Ming uses (``(z - shift_factor) * scaling_factor``)
+        so the concatenated frame lives in the DiT's latent space.
+        """
+        if ref is None:
+            return None
+        if not isinstance(ref, torch.Tensor):
+            ref = self.image_processor.preprocess(ref, height, width)
+        ref = ref.to(device=self.device, dtype=self.vae.dtype)
+        latent = self.vae.encode(ref).latent_dist.mode()
+        return (latent - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+
+    # ------------------------------------------------------------------
+    # Forward
+    # ------------------------------------------------------------------
+
+    @torch.inference_mode()
+    def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
+        """Run one text-to-image generation request.
+
+        Args:
+            req: Diffusion request. The cross-stage thinker hidden states
+                must be present at
+                ``req.prompts[0]["extra"]["thinker_hidden_states"]`` as a
+                ``[N, H]`` (or ``[1, N, H]``) tensor, placed there by
+                ``thinker2imagegen``.
+
+        Returns:
+            DiffusionOutput with ``.output`` set to a ``[B, 3, H, W]``
+            image tensor in ``[-1, 1]``. The vllm-omni diffusion engine's
+            output adapter converts this to PIL/base64 downstream.
+        """
+        first_prompt = req.prompts[0] if req.prompts else None
+        if isinstance(first_prompt, str):
+            prompt_dict: dict[str, Any] = {}
+        elif isinstance(first_prompt, dict):
+            prompt_dict = first_prompt
+        elif first_prompt is not None and hasattr(first_prompt, "_asdict"):
+            prompt_dict = first_prompt._asdict()
+        elif first_prompt is not None and hasattr(first_prompt, "__dict__"):
+            prompt_dict = vars(first_prompt)
+        else:
+            prompt_dict = {}
+
+        extra = prompt_dict.get("extra") or {}
+        hidden = extra.get("thinker_hidden_states")
+        if hidden is None:
+            # Same dual-path convention as glm_image: also check
+            # ``sampling_params.extra_args``.
+            hidden = (req.sampling_params.extra_args or {}).get("thinker_hidden_states")
+        if hidden is None:
+            scale = self.image_gen_config.img_gen_scales[-1]
+            num_query_tokens = scale * scale
+            hidden = torch.zeros(
+                (num_query_tokens, self.image_gen_config.thinker_hidden_size),
+                dtype=self._dtype,
+                device=self.device,
+            )
+            logger.warning(
+                "[MingImagePipeline.forward] 'thinker_hidden_states' missing "
+                "from request; falling back to zero-conditioning %s. This is "
+                "expected during warmup; for real requests verify that "
+                "`custom_process_input_func: thinker2imagegen` is set on the "
+                "diffusion stage in the YAML.",
+                tuple(hidden.shape),
+            )
+
+        if not isinstance(hidden, torch.Tensor):
+            raise TypeError(
+                f"[MingImagePipeline] 'thinker_hidden_states' must be a Tensor, got {type(hidden).__name__}"
+            )
+
+        # Move to the pipeline's device+dtype.
+        target_device = next(self.parameters()).device
+        target_dtype = next(self.parameters()).dtype
+        hidden = hidden.to(device=target_device, dtype=target_dtype)
+        if hidden.dim() == 2:
+            hidden = hidden.unsqueeze(0)  # [N, H] -> [1, N, H]
+        logger.debug(
+            "[MingImagePipeline.forward] thinker_hidden_states=%s on %s (%s)",
+            tuple(hidden.shape),
+            target_device,
+            target_dtype,
+        )
+
+        # ----- Condition encoder → cap_feats
+        cap_feats = self.condition_encoder(hidden)
+        logger.debug("[MingImagePipeline.forward] cap_feats=%s", tuple(cap_feats.shape))
+
+        # Real negative CFG conditioning (opt-in). See expand_cfg_prompts.
+        negative_hidden = extra.get("negative_thinker_hidden_states")
+        negative_cap_feats = None
+        if isinstance(negative_hidden, torch.Tensor):
+            negative_hidden = negative_hidden.to(device=target_device, dtype=target_dtype)
+            if negative_hidden.dim() == 2:
+                negative_hidden = negative_hidden.unsqueeze(0)
+            negative_cap_feats = self.condition_encoder(negative_hidden)
+            logger.debug("[MingImagePipeline.forward] negative_cap_feats=%s", tuple(negative_cap_feats.shape))
+
+        # ByT5 text enhancement (opt-in). Appends glyph-aware features along
+        # the sequence dim; negative side gets zeros for the byte5 portion so
+        # CFG doesn't push away from the rendered text.
+        byte5_texts = self._resolve_byte5_texts(extra, req.sampling_params)
+        if byte5_texts and self.byte5 is not None:
+            byte5_feats = self.byte5(byte5_texts).to(device=target_device, dtype=target_dtype)
+            cap_feats = torch.cat((cap_feats, byte5_feats), dim=1)
+            if negative_cap_feats is not None:
+                negative_cap_feats = torch.cat((negative_cap_feats, torch.zeros_like(byte5_feats)), dim=1)
+            logger.debug("[MingImagePipeline.forward] byte5 cat'd: cap_feats=%s", tuple(cap_feats.shape))
+
+        # Sampling knobs: extra_args.image_gen.* > sampling_params.* > MingImageGenConfig defaults.
+        sp = req.sampling_params
+        cfg = self.image_gen_config
+        ig = (sp.extra_args or {}).get("image_gen") or {}
+        resolved: dict[str, Any] = {}
+        for ig_key, sp_attr, default in (
+            ("height", "height", cfg.default_height),
+            ("width", "width", cfg.default_width),
+            ("steps", "num_inference_steps", cfg.num_inference_steps),
+            ("cfg", "guidance_scale", cfg.guidance_scale),
+            ("seed", "seed", None),
+        ):
+            for v in (ig.get(ig_key), getattr(sp, sp_attr), default):
+                if v is not None:
+                    resolved[ig_key] = v
+                    break
+
+        height = int(resolved["height"])
+        width = int(resolved["width"])
+        num_inference_steps = int(resolved["steps"])
+        guidance_scale = float(resolved["cfg"])
+        seed = resolved.get("seed")
+
+        # Always rebuild the generator from the resolved seed. Reusing
+        # ``sp.generator`` causes two problems:
+        #   (1) if the caller pre-seeded it with sp.seed (e.g. the top-level
+        #       ``seed`` key on OmniDiffusionSamplingParams), any override via
+        #       ``extra_args.image_gen.seed`` would be silently ignored; and
+        #   (2) a persistent generator instance accumulates state across
+        #       requests → same-seed replays produce different outputs.
+        if seed is not None:
+            generator = torch.Generator(device=target_device).manual_seed(int(seed))
+        else:
+            generator = sp.generator
+
+        # Format prompt_embeds / negative_prompt_embeds as list[Tensor]
+        # (one entry per request) — matches ZImagePipeline's contract when
+        # prompt_embeds are pre-computed.
+        prompt_embeds = [cap_feats[i] for i in range(cap_feats.shape[0])]
+        if negative_cap_feats is not None:
+            negative_prompt_embeds = [negative_cap_feats[i] for i in range(negative_cap_feats.shape[0])]
+        else:
+            negative_prompt_embeds = [self.condition_encoder.zero_negative(e) for e in prompt_embeds]
+
+        z_sp = _ZPipelineSamplingParams(
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            generator=generator,
+        )
+        z_req = _ZPipelineRequest(
+            request_id=req.request_id or "ming-imagegen",
+            sampling_params=z_sp,
+        )
+
+        # Reference image (img2img) → VAE-encoded latent published on the
+        # active ForwardContext so MingZImageTransformer2DModel can read it
+        # from request scope inside its forward().
+        ref_latent = self._encode_reference_image(extra.get("reference_image"), height, width)
+        set_forward_context_ref_latent(ref_latent)
+
+        logger.debug(
+            "[MingImagePipeline.forward] running z_pipeline hw=(%d,%d) steps=%d cfg=%.2f seed=%s overrides=%s ref=%s",
+            height,
+            width,
+            num_inference_steps,
+            guidance_scale,
+            seed,
+            ig,
+            None if ref_latent is None else tuple(ref_latent.shape),
+        )
+        try:
+            output = super().forward(
+                z_req,
+                prompt=None,
+                height=height,
+                width=width,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                prompt_embeds=prompt_embeds,
+                negative_prompt_embeds=negative_prompt_embeds,
+                generator=generator,
+                output_type="pt",
+                return_dict=True,
+            )
+        finally:
+            set_forward_context_ref_latent(None)
+
+        if hasattr(output, "output") and output.output is not None:
+            raw = output.output
+        elif hasattr(output, "images"):
+            raw = output.images
+        else:
+            raw = output
+        if not isinstance(raw, torch.Tensor):
+            raise RuntimeError(f"ZImagePipeline returned non-tensor output: {type(raw).__name__}")
+        if logger.isEnabledFor(logging.DEBUG):
+            logger.debug(
+                "[MingImagePipeline.forward] produced image tensor shape=%s range=[%.3f,%.3f]",
+                tuple(raw.shape),
+                raw.float().min().item(),
+                raw.float().max().item(),
+            )
+        return DiffusionOutput(output=raw)
+
+
+# ----------------------------------------------------------------------
+# Helpers
+# ----------------------------------------------------------------------
+
+
+def get_ming_image_post_process_func(od_config: OmniDiffusionConfig):
+    """Return a post-process callable that converts the raw VAE tensor to PIL.
+
+    The diffusion engine calls ``post_process_func(output_data)`` where
+    ``output_data`` is the ``DiffusionOutput.output`` tensor returned by
+    ``MingImagePipeline.forward``. It has shape ``[B, 3, H, W]`` in ``[-1, 1]``
+    (Z-image VAE convention). We run the standard ``VaeImageProcessor``
+    postprocess to convert it to ``list[PIL.Image]`` which vllm-omni's
+    ``OmniRequestOutput.from_diffusion`` then bubbles up as
+    ``omni_outputs.images`` for serving_chat to base64-encode.
+
+    Registered via ``_DIFFUSION_POST_PROCESS_FUNCS["MingImagePipeline"]``
+    in vllm_omni/diffusion/registry.py.
+    """
+    import json
+
+    model_path = od_config.model
+    vae_config_path = os.path.join(model_path, "vae", "config.json")
+    try:
+        with open(vae_config_path) as f:
+            vae_cfg = json.load(f)
+        block_out_channels = vae_cfg.get("block_out_channels", [128, 256, 512, 512])
+        vae_scale_factor = 2 ** (len(block_out_channels) - 1)
+    except Exception:
+        vae_scale_factor = 8  # Ming's Flux-format VAE default
+
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2, do_convert_rgb=True)
+
+    def post_process_func(images: torch.Tensor):
+        # VaeImageProcessor.postprocess with default output_type="pil"
+        # returns ``list[PIL.Image]``.
+        return image_processor.postprocess(images.float())
+
+    return post_process_func
+
+
+__all__ = [
+    "MingImagePipeline",
+    "get_ming_image_post_process_func",
+]
diff --git a/vllm_omni/diffusion/models/ming_flash_omni/t5_block_mapper.py b/vllm_omni/diffusion/models/ming_flash_omni/t5_block_mapper.py
new file mode 100644
index 000000000..75243529e
--- /dev/null
+++ b/vllm_omni/diffusion/models/ming_flash_omni/t5_block_mapper.py
@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 The vLLM-Omni team.
+
+"""T5EncoderBlockByT5Mapper — Ming's per-block T5 stack mapping byte5 features
+onto the DiT condition space.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+from diffusers.models.modeling_utils import ModelMixin
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+from vllm_omni.diffusion.models.t5_encoder.t5_encoder import T5Block
+
+
+class T5EncoderBlockByT5Mapper(ModelMixin):
+    """Stacks ``num_layers`` T5 encoder blocks on top of byte5 features and
+    projects them to ``sdxl_channels`` (= Ming's ``diffusion_c_input_dim``).
+    """
+
+    def __init__(self, byte5_config, num_layers: int, sdxl_channels: int | None = None) -> None:
+        super().__init__()
+        if num_layers > 0:
+            self.blocks = nn.ModuleList(
+                [
+                    T5Block(
+                        byte5_config,
+                        has_relative_attention_bias=(i == 0),
+                        prefix=f"blocks.{i}",
+                    )
+                    for i in range(num_layers)
+                ]
+            )
+        else:
+            self.blocks = None
+        self.layer_norm = RMSNorm(byte5_config.d_model, eps=byte5_config.layer_norm_epsilon)
+        if sdxl_channels is not None:
+            self.channel_mapper = nn.Linear(byte5_config.d_model, sdxl_channels)
+            self.final_layer_norm = RMSNorm(sdxl_channels, eps=byte5_config.layer_norm_epsilon)
+        else:
+            self.channel_mapper = None
+            self.final_layer_norm = None
+
+    def get_extended_attention_mask(self, attention_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+        if attention_mask.dim() == 3:
+            extended = attention_mask[:, None, :, :]
+        elif attention_mask.dim() == 2:
+            extended = attention_mask[:, None, None, :]
+        else:
+            raise ValueError(f"Unexpected attention_mask shape {tuple(attention_mask.shape)}")
+        extended = extended.to(dtype=dtype)
+        return (1.0 - extended) * torch.finfo(dtype).min
+
+    def forward(self, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        extended_mask = self.get_extended_attention_mask(attention_mask, dtype=self.dtype)
+
+        hidden_states = inputs_embeds
+        position_bias = None
+
+        if self.blocks is not None:
+            for block in self.blocks:
+                hidden_states, position_bias = block(
+                    hidden_states,
+                    mask=extended_mask,
+                    position_bias=position_bias,
+                )
+
+        hidden_states = self.layer_norm(hidden_states)
+        if self.channel_mapper is not None:
+            hidden_states = self.channel_mapper(hidden_states)
+            hidden_states = self.final_layer_norm(hidden_states)
+        return hidden_states
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load Ming's HF-format byte5_mapper checkpoint into the fused
+        TP-aware layers.
+
+        Source format (from ``byte5_mapper.pt``):
+            blocks.{i}.layer.0.SelfAttention.{q,k,v,o}.weight
+            blocks.{i}.layer.1.DenseReluDense.{wi_0,wi_1,wo}.weight
+            blocks.{i}.layer.{0,1}.layer_norm.weight
+            {layer_norm, channel_mapper, final_layer_norm}.{weight,bias}
+
+        Target format (after ``T5Block`` from t5_encoder.py):
+            blocks.{i}.layer.0.SelfAttention.qkv_proj.weight  (fused q+k+v)
+            blocks.{i}.layer.1.DenseReluDense.wi.weight       (fused wi_0+wi_1)
+            (others identical)
+        """
+        stacked_params_mapping = [
+            (".qkv_proj.", ".q.", "q"),
+            (".qkv_proj.", ".k.", "k"),
+            (".qkv_proj.", ".v.", "v"),
+            (".DenseReluDense.wi.", ".DenseReluDense.wi_0.", 0),
+            (".DenseReluDense.wi.", ".DenseReluDense.wi_1.", 1),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        for name, loaded_weight in weights:
+            matched = False
+            for target_substr, source_substr, shard_id in stacked_params_mapping:
+                if source_substr not in name:
+                    continue
+                target_name = name.replace(source_substr, target_substr, 1)
+                if target_name in params_dict:
+                    param = params_dict[target_name]
+                    param.weight_loader(param, loaded_weight, shard_id)
+                    loaded_params.add(target_name)
+                    matched = True
+                    break
+
+            if matched:
+                continue
+
+            if name in params_dict:
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+                loaded_params.add(name)
+
+        return loaded_params
+
+
+__all__ = ["T5EncoderBlockByT5Mapper"]
diff --git a/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py b/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py
index 51a93de2c..745baae13 100644
--- a/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py
+++ b/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py
@@ -147,6 +147,8 @@ class OmniVoicePipeline(nn.Module, SupportAudioOutput):
         ref_text = None
         lang = "None"
         instruct = "None"
+        extra = req.sampling_params.extra_args or {}
+        seed = extra.get("seed", None)
 
         voice_name = None
         if isinstance(prompt, dict):
@@ -157,7 +159,6 @@ class OmniVoicePipeline(nn.Module, SupportAudioOutput):
             voice_name = prompt.get("voice_name")
             lang = prompt.get("lang")
             instruct = prompt.get("instruct")
-
             # OmniTextPrompt format (used by offline Omni.generate path):
             # ref_audio comes via multi_modal_data["audio"] and the rest via
             # mm_processor_kwargs. Fall back to those when top-level keys are
@@ -262,7 +263,6 @@ class OmniVoicePipeline(nn.Module, SupportAudioOutput):
         else:
             cond_ids = torch.cat([text_ids, target_ids], dim=1)
         cond_len = cond_ids.shape[1]
-
         uncond_ids = target_ids.clone()
         uncond_len = target_len
         max_len = max(cond_len, uncond_len)
@@ -297,8 +297,8 @@ class OmniVoicePipeline(nn.Module, SupportAudioOutput):
             layer_penalty_factor=self.layer_penalty_factor,
             position_temperature=self.position_temperature,
             class_temperature=self.class_temperature,
+            seed=seed,
         )
-
         # Decode tokens to audio
         audio = self.decoder(tokens)  # [1, 1, samples]
 
diff --git a/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py b/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py
index 8de617dd9..2484f02ed 100644
--- a/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py
+++ b/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py
@@ -168,7 +168,10 @@ class OvisImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMi
         )
 
         self.text_encoder = Qwen3Model.from_pretrained(
-            model, subfolder="text_encoder", local_files_only=local_files_only
+            model,
+            subfolder="text_encoder",
+            local_files_only=local_files_only,
+            torch_dtype=od_config.dtype,
         )
 
         self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py
index 88a66d7f6..abdca8b1a 100644
--- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py
+++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py
@@ -17,7 +17,6 @@ from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous
 from vllm.logger import init_logger
-from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     QKVParallelLinear,
@@ -44,11 +43,16 @@ from vllm_omni.diffusion.distributed.sp_plan import (
 )
 from vllm_omni.diffusion.forward_context import get_forward_context
 from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
+from vllm_omni.diffusion.layers.norm import RMSNorm
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 
 logger = init_logger(__name__)
 
 
+def _join_prefix(prefix: str, suffix: str) -> str:
+    return f"{prefix}.{suffix}" if prefix else suffix
+
+
 class ImageRopePrepare(nn.Module):
     """Prepares image hidden_states and RoPE embeddings for sequence parallel.
 
@@ -471,7 +475,12 @@ class FeedForward(nn.Module):
 
         layers: list[nn.Module] = [
             ColumnParallelApproxGELU(
-                dim, inner_dim, approximate="tanh", bias=bias, quant_config=quant_config, prefix=prefix
+                dim,
+                inner_dim,
+                approximate="tanh",
+                bias=bias,
+                quant_config=quant_config,
+                prefix=_join_prefix(prefix, "net.0.proj"),
             ),
             nn.Identity(),  # placeholder for weight loading
             RowParallelLinear(
@@ -480,7 +489,7 @@ class FeedForward(nn.Module):
                 input_is_parallel=True,
                 return_bias=False,
                 quant_config=quant_config,
-                prefix=prefix,
+                prefix=_join_prefix(prefix, "net.2"),
             ),
         ]
 
@@ -507,6 +516,7 @@ class QwenImageCrossAttention(nn.Module):
         context_pre_only: bool = False,
         out_dim: int | None = None,
         quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         assert dim % num_heads == 0
@@ -523,7 +533,7 @@ class QwenImageCrossAttention(nn.Module):
             head_size=self.head_dim,
             total_num_heads=num_heads,
             quant_config=quant_config,
-            prefix="to_qkv",
+            prefix=_join_prefix(prefix, "to_qkv"),
         )
         self.query_num_heads = self.to_qkv.num_heads
         self.kv_num_heads = self.to_qkv.num_kv_heads
@@ -539,7 +549,7 @@ class QwenImageCrossAttention(nn.Module):
             head_size=head_dim,
             total_num_heads=num_heads,
             quant_config=quant_config,
-            prefix="add_kv_proj",
+            prefix=_join_prefix(prefix, "add_kv_proj"),
         )
         self.add_query_num_heads = self.add_kv_proj.num_heads
         self.add_kv_num_heads = self.add_kv_proj.num_kv_heads
@@ -552,7 +562,7 @@ class QwenImageCrossAttention(nn.Module):
             input_is_parallel=True,
             return_bias=False,
             quant_config=quant_config,
-            prefix="to_add_out",
+            prefix=_join_prefix(prefix, "to_add_out"),
         )
 
         assert not pre_only
@@ -563,7 +573,7 @@ class QwenImageCrossAttention(nn.Module):
             input_is_parallel=True,
             return_bias=False,
             quant_config=quant_config,
-            prefix="to_out",
+            prefix=_join_prefix(prefix, "to_out"),
         )
 
         self.norm_added_q = RMSNorm(head_dim, eps=eps)
@@ -697,6 +707,7 @@ class QwenImageTransformerBlock(nn.Module):
         eps: float = 1e-6,
         zero_cond_t: bool = False,
         quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
     ):
         super().__init__()
 
@@ -705,18 +716,20 @@ class QwenImageTransformerBlock(nn.Module):
         self.attention_head_dim = attention_head_dim
 
         # Image processing modules.
-        # Modulation linear is kept full precision (quant_config=None) — it
+        # Modulation linear is kept unquantized (quant_config=None) — it
         # produces shift/scale/gate values that are precision-sensitive
-        # (see #2728).
+        # (see #2728). Use column TP with gather_output=True so weights are
+        # sharded while downstream modulation still receives full [B, 6 * dim].
         self.img_mod = nn.Sequential(
             nn.SiLU(),
-            ReplicatedLinear(
+            ColumnParallelLinear(
                 dim,
                 6 * dim,
                 bias=True,
+                gather_output=True,
                 return_bias=False,
                 quant_config=None,
-                prefix="img_mod.1",
+                prefix=_join_prefix(prefix, "img_mod.1"),
             ),
         )
         self.img_norm1 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps)
@@ -727,26 +740,38 @@ class QwenImageTransformerBlock(nn.Module):
             context_pre_only=False,
             head_dim=attention_head_dim,
             quant_config=quant_config,
+            prefix=_join_prefix(prefix, "attn"),
         )
         self.img_norm2 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps)
-        self.img_mlp = FeedForward(dim=dim, dim_out=dim, quant_config=quant_config, prefix="img_mlp")
+        self.img_mlp = FeedForward(
+            dim=dim,
+            dim_out=dim,
+            quant_config=quant_config,
+            prefix=_join_prefix(prefix, "img_mlp"),
+        )
 
         # Text processing modules.
         self.txt_mod = nn.Sequential(
             nn.SiLU(),
-            ReplicatedLinear(
+            ColumnParallelLinear(
                 dim,
                 6 * dim,
                 bias=True,
+                gather_output=True,
                 return_bias=False,
                 quant_config=None,
-                prefix="txt_mod.1",
+                prefix=_join_prefix(prefix, "txt_mod.1"),
             ),
         )
         self.txt_norm1 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps)
         # Text doesn't need separate attention - it's handled by img_attn joint computation
         self.txt_norm2 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps)
-        self.txt_mlp = FeedForward(dim=dim, dim_out=dim, quant_config=quant_config, prefix="txt_mlp")
+        self.txt_mlp = FeedForward(
+            dim=dim,
+            dim_out=dim,
+            quant_config=quant_config,
+            prefix=_join_prefix(prefix, "txt_mlp"),
+        )
 
         self.zero_cond_t = zero_cond_t
 
@@ -996,8 +1021,9 @@ class QwenImageTransformer2DModel(CachedTransformer):
                     attention_head_dim=attention_head_dim,
                     zero_cond_t=zero_cond_t,
                     quant_config=quant_config,
+                    prefix=f"transformer_blocks.{i}",
                 )
-                for _ in range(num_layers)
+                for i in range(num_layers)
             ]
         )
 
diff --git a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py
index 906eac974..fe8695b07 100644
--- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py
+++ b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py
@@ -707,7 +707,17 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
         return out.past_key_values, out.hidden_states
 
     def _t2i_predict_v(
-        self, input_embeds, indexes_image, attn_mask, past_key_values, t, z, image_token_num, image_size=None, **_kw
+        self,
+        input_embeds,
+        indexes_image,
+        attn_mask,
+        past_key_values,
+        t,
+        z,
+        image_token_num,
+        image_size=None,
+        cache_dit_skip=False,
+        **_kw,
     ):
         B, L = z.shape[0], z.shape[1]
         outputs = self.language_model(
@@ -719,6 +729,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
             attention_mask=attn_mask,
             past_key_values=past_key_values,
             update_cache=False,
+            cache_dit_skip=cache_dit_skip,
             use_cache=True,
             compute_logits=False,
         )
@@ -1066,6 +1077,13 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
         kv_cond = caches["cond"]
         idx_cond = caches["idx_cond"]
         mask_cond = caches["mask_cond"]
+        is_it2i = "img_cond" in caches
+        if is_it2i:
+            use_cfg = (t > p.cfg_interval[0] and t < p.cfg_interval[1]) or p.cfg_interval[0] == 0
+            needs_cfg = not (p.cfg_scale == 1 and p.img_cfg_scale == 1)
+            has_cached_partner = use_cfg and needs_cfg
+        else:
+            has_cached_partner = t >= p.cfg_interval[0] and t <= p.cfg_interval[1] and p.cfg_scale > 1
 
         out_cond = self._t2i_predict_v(
             image_embeds,
@@ -1076,9 +1094,9 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
             z,
             image_token_num=ns.token_h * ns.token_w,
             image_size=p.image_size,
+            cache_dit_skip=not has_cached_partner,
         )
 
-        is_it2i = "img_cond" in caches
         if is_it2i:
             return self._denoise_step_it2i(out_cond, image_embeds, ns, t, z, caches, p, step_i)
         return self._denoise_step_t2i(out_cond, image_embeds, ns, t, z, caches, p, step_i)
@@ -1153,6 +1171,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
                 caches["img_cond"],
                 t,
                 z,
+                cache_dit_skip=True,
                 **predict_kw,
             )
             out_uncond = self._t2i_predict_v(
diff --git a/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py b/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py
index da43f55df..f9ff8aaf4 100644
--- a/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py
+++ b/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py
@@ -587,6 +587,12 @@ class SenseNovaU1DecoderLayer(nn.Module):
         past_key_values=None,
         **kwargs,
     ):
+        if isinstance(attention_mask, dict):
+            attention_mask = attention_mask.get(
+                self.attention_type,
+                attention_mask.get("full_attention"),
+            )
+
         if exist_und and not exist_gen:
             return self._forward_und(hidden_states, indexes, attention_mask, past_key_values, **kwargs)
         if not exist_und and exist_gen:
@@ -663,7 +669,7 @@ class SenseNovaU1Model(nn.Module):
                 exist_und=exist_und,
                 exist_gen=exist_gen,
                 indexes=indexes,
-                attention_mask=causal_mask_mapping.get(layer.attention_type, causal_mask_mapping.get("full_attention")),
+                attention_mask=causal_mask_mapping,
                 past_key_values=past_key_values,
                 use_cache=use_cache,
                 **kwargs,
diff --git a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py
index 345ccca65..388b6316c 100644
--- a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py
+++ b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py
@@ -832,5 +832,8 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
         loaded_weights = loader.load_weights(weights)
         # Record components loaded by diffusers submodules to satisfy strict checks.
         loaded_weights |= {f"vae.{name}" for name, _ in self.vae.named_parameters()}
-        loaded_weights |= {f"text_encoder.{name}" for name, _ in self.text_encoder.named_parameters()}
+        # downstream pipelines (e.g. MingImagePipeline) may set ``self.text_encoder = None`` when they
+        # bring their own conditioning path.
+        if self.text_encoder is not None:
+            loaded_weights |= {f"text_encoder.{name}" for name, _ in self.text_encoder.named_parameters()}
         return loaded_weights
diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py
index d341c2928..fe9ed8fcf 100644
--- a/vllm_omni/diffusion/registry.py
+++ b/vllm_omni/diffusion/registry.py
@@ -146,6 +146,11 @@ _DIFFUSION_MODELS = {
         "pipeline_lance",
         "LancePipeline",
     ),
+    "MingImagePipeline": (
+        "ming_flash_omni",
+        "pipeline_ming_imagegen",
+        "MingImagePipeline",
+    ),
     "InternVLAA1Pipeline": (
         "internvla_a1",
         "pipeline_internvla_a1",
@@ -266,6 +271,11 @@ _DIFFUSION_MODELS = {
         "pipeline_diffusers_adapter",
         "DiffusersAdapterPipeline",
     ),
+    "HiDreamImagePipeline": (
+        "hidream_image",
+        "pipeline_hidream_image",
+        "HiDreamImagePipeline",
+    ),
 }
 
 
@@ -282,7 +292,6 @@ DiffusionModelRegistry = _ModelRegistry(
 _NO_CACHE_ACCELERATION = {
     # Pipelines that do not support cache acceleration (cache_dit / tea_cache).
     "NextStep11Pipeline",
-    "SenseNovaU1Pipeline",
     "AudioXPipeline",
 }
 
@@ -468,6 +477,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "LongCatImagePipeline": "get_longcat_image_post_process_func",
     "BagelPipeline": "get_bagel_post_process_func",
     "LancePipeline": "get_lance_post_process_func",
+    "MingImagePipeline": "get_ming_image_post_process_func",
     "InternVLAA1Pipeline": "get_internvla_a1_post_process_func",
     "LongCatImageEditPipeline": "get_longcat_image_post_process_func",
     "StableDiffusion3Pipeline": "get_sd3_image_post_process_func",
@@ -488,6 +498,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "OmniVoicePipeline": "get_omnivoice_post_process_func",
     "DreamIDOmniPipeline": "get_dreamid_omni_post_process_func",
     "SenseNovaU1Pipeline": "get_sensenova_u1_post_process_func",
+    "HiDreamImagePipeline": "get_hidream_image_post_process_func",
 }
 
 _DIFFUSION_PRE_PROCESS_FUNCS = {
diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py
index 4d4328d25..79198c3ec 100644
--- a/vllm_omni/diffusion/request.py
+++ b/vllm_omni/diffusion/request.py
@@ -3,10 +3,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import random
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
 
+DUMMY_DIFFUSION_REQUEST_ID = "dummy_req_id"
+
 
 @dataclass
 class OmniDiffusionRequest:
@@ -24,13 +26,14 @@ class OmniDiffusionRequest:
 
     prompts: list[OmniPromptType]  # Actually supporting str-based prompts
     sampling_params: OmniDiffusionSamplingParams
-
-    request_ids: list[str] = field(default_factory=list)
-    request_id: str | None = None
+    request_id: str
     kv_sender_info: dict | None = None
 
     def __post_init__(self):
         """Initialize dependent fields after dataclass initialization."""
+        if not isinstance(self.request_id, str) or not self.request_id:
+            raise ValueError("OmniDiffusionRequest.request_id must be a non-empty string.")
+
         # When neither a generator nor a seed is provided, assign a random seed
         # so that all ranks derive the same generator state.
         if self.sampling_params.generator is None and self.sampling_params.seed is None:
@@ -56,3 +59,10 @@ class OmniDiffusionRequest:
         # so downstream code always has a valid value.
         if self.sampling_params.guidance_scale_2 is None:
             self.sampling_params.guidance_scale_2 = self.sampling_params.guidance_scale
+
+    def is_dummy_run(self) -> bool:
+        return self.is_dummy_run_request_id(self.request_id)
+
+    @classmethod
+    def is_dummy_run_request_id(cls, request_id: str | None) -> bool:
+        return request_id == DUMMY_DIFFUSION_REQUEST_ID
diff --git a/vllm_omni/diffusion/sched/base_scheduler.py b/vllm_omni/diffusion/sched/base_scheduler.py
index 0bd16597f..812873a1c 100644
--- a/vllm_omni/diffusion/sched/base_scheduler.py
+++ b/vllm_omni/diffusion/sched/base_scheduler.py
@@ -46,7 +46,6 @@ class _BaseScheduler(SchedulerInterface):
     def __init__(self) -> None:
         self.od_config: OmniDiffusionConfig | None = None
         self._request_states: dict[str, DiffusionRequestState] = {}
-        self._request_id_to_sched_req_id: dict[str, str] = {}
         self._step_id: int = 0
         self._waiting: deque[str] = deque()
         self._running: list[str] = []
@@ -57,7 +56,6 @@ class _BaseScheduler(SchedulerInterface):
     def initialize(self, od_config: OmniDiffusionConfig) -> None:
         self.od_config = od_config
         self._request_states.clear()
-        self._request_id_to_sched_req_id.clear()
         self._step_id = 0
         self._waiting.clear()
         self._running.clear()
@@ -71,32 +69,31 @@ class _BaseScheduler(SchedulerInterface):
         self._reset_scheduler_state()
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
-        sched_req_id = self._make_sched_req_id(request)
-        return self._add_request_with_sched_req_id(sched_req_id, request)
+        return self._add_request_with_request_id(request.request_id, request)
 
-    def _add_request_with_sched_req_id(self, sched_req_id: str, request: OmniDiffusionRequest) -> str:
-        state = self._make_request_state(sched_req_id, request)
-        request_ids = self._request_ids_for_mapping(request)
-        self._register_request_ids(request_ids, sched_req_id)
-        self._request_states[sched_req_id] = state
-        self._waiting.append(sched_req_id)
-        logger.debug("%s add_request: %s (waiting=%d)", self.__class__.__name__, sched_req_id, len(self._waiting))
-        return sched_req_id
+    def _add_request_with_request_id(self, request_id: str, request: OmniDiffusionRequest) -> str:
+        if request_id in self._request_states:
+            raise ValueError(f"request_id {request_id!r} is already active.")
+        state = self._make_request_state(request_id, request)
+        self._request_states[request_id] = state
+        self._waiting.append(request_id)
+        logger.debug("%s add_request: %s (waiting=%d)", self.__class__.__name__, request_id, len(self._waiting))
+        return request_id
 
     def schedule(self) -> DiffusionSchedulerOutput:
         scheduled_new_reqs: list[NewRequestData] = []
-        scheduled_cached_req_ids: list[str] = []
+        scheduled_cached_request_ids: list[str] = []
 
         # First, schedule the RUNNING request(s)
-        for sched_req_id in self._running:
-            state = self._request_states.get(sched_req_id)
+        for request_id in self._running:
+            state = self._request_states.get(request_id)
             if state is not None:
-                scheduled_cached_req_ids.append(sched_req_id)
+                scheduled_cached_request_ids.append(request_id)
 
         # Second, schedule WAITING requests while capacity remains.
         while self._waiting and len(self._running) < self.max_num_running_reqs:
-            sched_req_id = self._waiting[0]
-            state = self._request_states.get(sched_req_id)
+            request_id = self._waiting[0]
+            state = self._request_states.get(request_id)
             if state is None:
                 self._waiting.popleft()
                 continue
@@ -108,16 +105,16 @@ class _BaseScheduler(SchedulerInterface):
             if not self._running:
                 self._running_sampling_params_key = state.sampling_params_key
             state.status = DiffusionRequestStatus.RUNNING
-            self._running.append(sched_req_id)
+            self._running.append(request_id)
             if was_new_request:
                 scheduled_new_reqs.append(NewRequestData.from_state(state))
             else:
-                scheduled_cached_req_ids.append(sched_req_id)
+                scheduled_cached_request_ids.append(request_id)
 
         scheduler_output = DiffusionSchedulerOutput(
             step_id=self._step_id,
             scheduled_new_reqs=scheduled_new_reqs,
-            scheduled_cached_reqs=CachedRequestData(sched_req_ids=scheduled_cached_req_ids),
+            scheduled_cached_reqs=CachedRequestData(request_ids=scheduled_cached_request_ids),
             finished_req_ids=set(self._finished_req_ids),
             num_running_reqs=len(self._running),
             num_waiting_reqs=len(self._waiting),
@@ -131,40 +128,33 @@ class _BaseScheduler(SchedulerInterface):
     def has_requests(self) -> bool:
         return bool(self._waiting or self._running)
 
-    def get_request_state(self, sched_req_id: str) -> DiffusionRequestState | None:
-        return self._request_states.get(sched_req_id)
+    def get_request_state(self, request_id: str) -> DiffusionRequestState | None:
+        return self._request_states.get(request_id)
 
-    def get_sched_req_id(self, request_id: str) -> str | None:
-        return self._request_id_to_sched_req_id.get(request_id)
+    def pop_request_state(self, request_id: str) -> DiffusionRequestState | None:
+        self._pop_extra_request_state(request_id)
+        return self._request_states.pop(request_id, None)
 
-    def pop_request_state(self, sched_req_id: str) -> DiffusionRequestState | None:
-        self._pop_extra_request_state(sched_req_id)
-        state = self._request_states.pop(sched_req_id, None)
-        if state is not None:
-            self._unregister_request_ids(self._request_ids_for_mapping(state.req), sched_req_id)
-        return state
-
-    def preempt_request(self, sched_req_id: str) -> bool:
-        if sched_req_id not in self._request_states:
+    def preempt_request(self, request_id: str) -> bool:
+        if request_id not in self._request_states:
             return False
-        if sched_req_id in self._running:
-            self._running.remove(sched_req_id)
+        if request_id in self._running:
+            self._running.remove(request_id)
             if not self._running:
                 self._running_sampling_params_key = None
-            self._waiting.appendleft(sched_req_id)
-            self._request_states[sched_req_id].status = DiffusionRequestStatus.PREEMPTED
+            self._waiting.appendleft(request_id)
+            self._request_states[request_id].status = DiffusionRequestStatus.PREEMPTED
             return True
         return False
 
-    def finish_requests(self, sched_req_ids: str | list[str], status: DiffusionRequestStatus) -> None:
+    def finish_requests(self, request_ids: str | list[str], status: DiffusionRequestStatus) -> None:
         assert DiffusionRequestStatus.is_finished(status)
-        if isinstance(sched_req_ids, str):
-            sched_req_ids = [sched_req_ids]
-        self._finish_requests({sched_req_id: status for sched_req_id in sched_req_ids})
+        if isinstance(request_ids, str):
+            request_ids = [request_ids]
+        self._finish_requests({request_id: status for request_id in request_ids})
 
     def close(self) -> None:
         self._request_states.clear()
-        self._request_id_to_sched_req_id.clear()
         self._waiting.clear()
         self._running.clear()
         self._running_sampling_params_key = None
@@ -183,33 +173,31 @@ class _BaseScheduler(SchedulerInterface):
         running_to_remove: set[str] = set()
         waiting_to_remove: set[str] = set()
 
-        for sched_req_id, status in statuses.items():
+        for request_id, status in statuses.items():
             assert DiffusionRequestStatus.is_finished(status)
-            state = self._request_states.get(sched_req_id)
+            state = self._request_states.get(request_id)
             if state is None or state.is_finished():
                 continue
 
-            finished_req_ids.add(sched_req_id)
-            if sched_req_id in self._running:
-                running_to_remove.add(sched_req_id)
-            if sched_req_id in self._waiting:
-                waiting_to_remove.add(sched_req_id)
+            finished_req_ids.add(request_id)
+            if request_id in self._running:
+                running_to_remove.add(request_id)
+            if request_id in self._waiting:
+                waiting_to_remove.add(request_id)
 
         if running_to_remove:
-            self._running = [sched_req_id for sched_req_id in self._running if sched_req_id not in running_to_remove]
+            self._running = [request_id for request_id in self._running if request_id not in running_to_remove]
             if not self._running:
                 self._running_sampling_params_key = None
         if waiting_to_remove:
-            self._waiting = deque(
-                sched_req_id for sched_req_id in self._waiting if sched_req_id not in waiting_to_remove
-            )
+            self._waiting = deque(request_id for request_id in self._waiting if request_id not in waiting_to_remove)
 
-        for sched_req_id in finished_req_ids:
-            state = self._request_states[sched_req_id]
-            status = statuses[sched_req_id]
+        for request_id in finished_req_ids:
+            state = self._request_states[request_id]
+            status = statuses[request_id]
             state.status = status
             if status == DiffusionRequestStatus.FINISHED_ERROR:
-                state.error = None if errors is None else errors.get(sched_req_id)
+                state.error = None if errors is None else errors.get(request_id)
             else:
                 state.error = None
 
@@ -227,7 +215,7 @@ class _BaseScheduler(SchedulerInterface):
         # marked finished at that point, but we still need to surface its id
         # in this update so the engine can observe the terminal state.
         finished_req_ids = {
-            sched_req_id for sched_req_id in sched_output.scheduled_req_ids if sched_req_id in self._finished_req_ids
+            request_id for request_id in sched_output.scheduled_request_ids if request_id in self._finished_req_ids
         }
         finished_req_ids |= self._finish_requests(statuses, errors)
         return finished_req_ids
@@ -235,12 +223,12 @@ class _BaseScheduler(SchedulerInterface):
     def _reset_scheduler_state(self) -> None:
         """Reset subclass-owned state during initialize()/close()."""
 
-    def _pop_extra_request_state(self, sched_req_id: str) -> None:
+    def _pop_extra_request_state(self, request_id: str) -> None:
         """Remove subclass-owned per-request state before popping request state."""
 
-    def _make_request_state(self, sched_req_id: str, request: OmniDiffusionRequest) -> DiffusionRequestState:
+    def _make_request_state(self, request_id: str, request: OmniDiffusionRequest) -> DiffusionRequestState:
         return DiffusionRequestState(
-            sched_req_id=sched_req_id,
+            request_id=request_id,
             req=request,
             sampling_params_key=get_sampling_params_key(request),
         )
@@ -258,24 +246,3 @@ class _BaseScheduler(SchedulerInterface):
         state = self._request_states.get(self._running[0])
         self._running_sampling_params_key = None if state is None else state.sampling_params_key
         return self._running_sampling_params_key
-
-    def _register_request_ids(self, request_ids: list[str], sched_req_id: str) -> None:
-        unique_request_ids = list(dict.fromkeys(request_ids))
-        for request_id in unique_request_ids:
-            existing = self._request_id_to_sched_req_id.get(request_id)
-            if existing is not None and existing != sched_req_id:
-                raise ValueError(f"request_id {request_id!r} is already mapped to active sched_req_id {existing!r}.")
-        for request_id in unique_request_ids:
-            self._request_id_to_sched_req_id[request_id] = sched_req_id
-
-    def _unregister_request_ids(self, request_ids: list[str], sched_req_id: str) -> None:
-        for request_id in request_ids:
-            if self._request_id_to_sched_req_id.get(request_id) == sched_req_id:
-                self._request_id_to_sched_req_id.pop(request_id, None)
-
-    @staticmethod
-    def _request_ids_for_mapping(request: OmniDiffusionRequest) -> list[str]:
-        request_ids = list(request.request_ids)
-        if request_id := getattr(request, "request_id", None):
-            request_ids.append(request_id)
-        return list(dict.fromkeys(request_ids))
diff --git a/vllm_omni/diffusion/sched/interface.py b/vllm_omni/diffusion/sched/interface.py
index 1bfb69458..b50e3555a 100644
--- a/vllm_omni/diffusion/sched/interface.py
+++ b/vllm_omni/diffusion/sched/interface.py
@@ -4,22 +4,17 @@
 from __future__ import annotations
 
 import enum
-import uuid
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from functools import cached_property
 from typing import TYPE_CHECKING
 
-from vllm.logger import init_logger
-
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
 if TYPE_CHECKING:
     from vllm_omni.diffusion.worker.utils import RunnerOutput
 
-logger = init_logger(__name__)
-
 
 class DiffusionRequestStatus(enum.IntEnum):
     """Request status tracked by diffusion scheduler."""
@@ -75,10 +70,7 @@ class SamplingParamsKey:
 class DiffusionRequestState:
     """Scheduler-owned state for one queued OmniDiffusionRequest."""
 
-    # Unique scheduler-owned request ID.
-    # NOTE: This identifies one OmniDiffusionRequest, which may contain multiple request_ids.
-    # TODO: Align this with OmniDiffusionRequest.request_ids once scheduler batching is supported.
-    sched_req_id: str
+    request_id: str
     req: OmniDiffusionRequest
     sampling_params_key: SamplingParamsKey | None = None
     status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
@@ -92,23 +84,23 @@ class DiffusionRequestState:
 class NewRequestData:
     """Full request payload for a newly scheduled diffusion request."""
 
-    sched_req_id: str
+    request_id: str
     req: OmniDiffusionRequest
 
     @classmethod
     def from_state(cls, state: DiffusionRequestState) -> NewRequestData:
-        return cls(sched_req_id=state.sched_req_id, req=state.req)
+        return cls(request_id=state.request_id, req=state.req)
 
 
 @dataclass
 class CachedRequestData:
-    """Cached diffusion requests that only need their scheduler ids resent."""
+    """Cached diffusion requests that only need their request ids resent."""
 
-    sched_req_ids: list[str]
+    request_ids: list[str]
 
     @classmethod
     def make_empty(cls) -> CachedRequestData:
-        return cls(sched_req_ids=[])
+        return cls(request_ids=[])
 
 
 @dataclass
@@ -123,21 +115,18 @@ class DiffusionSchedulerOutput:
     num_waiting_reqs: int
 
     @cached_property
-    def scheduled_req_ids(self) -> list[str]:
+    def scheduled_request_ids(self) -> list[str]:
         """
         All scheduled request ids in this cycle, including both new and cached ones.
-        NOTE:
-            This id is generated and owned by the scheduler,
-            and may be different from the OmniDiffusionRequest.request_ids.
         """
         return [
-            *(req.sched_req_id for req in self.scheduled_new_reqs),
-            *self.scheduled_cached_reqs.sched_req_ids,
+            *(req.request_id for req in self.scheduled_new_reqs),
+            *self.scheduled_cached_reqs.request_ids,
         ]
 
     @property
     def num_scheduled_reqs(self) -> int:
-        return len(self.scheduled_req_ids)
+        return len(self.scheduled_request_ids)
 
     @property
     def is_empty(self) -> bool:
@@ -147,25 +136,6 @@ class DiffusionSchedulerOutput:
 class SchedulerInterface(ABC):
     """Abstract lifecycle contract for diffusion schedulers."""
 
-    def _make_sched_req_id(self, request: OmniDiffusionRequest) -> str:
-        """
-        Generate a unique scheduler request ID for the given request.
-            The default implementation uses the first request_id from the request if available,
-            otherwise generates a random one.
-        """
-        if request.request_ids:
-            base = request.request_ids[0]
-        else:
-            logger.warning("Request has no request_ids, generating a random one. Request: %s", request)
-            base = f"req_{uuid.uuid4().hex[:8]}"
-
-        sched_req_id = base
-        suffix = 1
-        while self.get_request_state(sched_req_id) is not None:
-            sched_req_id = f"{base}#{suffix}"
-            suffix += 1
-        return sched_req_id
-
     @abstractmethod
     def initialize(self, od_config: OmniDiffusionConfig) -> None:
         """Initialize or reset scheduler state."""
@@ -183,7 +153,7 @@ class SchedulerInterface(ABC):
         """Update scheduler state from executor output."""
 
     @abstractmethod
-    def get_request_state(self, sched_req_id: str) -> DiffusionRequestState | None:
+    def get_request_state(self, request_id: str) -> DiffusionRequestState | None:
         """Return request state if present."""
 
     @abstractmethod
@@ -191,19 +161,15 @@ class SchedulerInterface(ABC):
         """Return whether the scheduler still owns runnable requests."""
 
     @abstractmethod
-    def get_sched_req_id(self, request_id: str) -> str | None:
-        """Resolve a public request_id to the active scheduler request id."""
-
-    @abstractmethod
-    def pop_request_state(self, sched_req_id: str) -> DiffusionRequestState | None:
+    def pop_request_state(self, request_id: str) -> DiffusionRequestState | None:
         """Remove and return request state if present."""
 
     @abstractmethod
-    def preempt_request(self, sched_req_id: str) -> bool:
+    def preempt_request(self, request_id: str) -> bool:
         """Preempt a running request back to waiting."""
 
     @abstractmethod
-    def finish_requests(self, sched_req_ids: str | list[str], status: DiffusionRequestStatus) -> None:
+    def finish_requests(self, request_ids: str | list[str], status: DiffusionRequestStatus) -> None:
         """Mark one or more requests finished."""
 
     @abstractmethod
diff --git a/vllm_omni/diffusion/sched/request_scheduler.py b/vllm_omni/diffusion/sched/request_scheduler.py
index 5029e3593..ba9cda616 100644
--- a/vllm_omni/diffusion/sched/request_scheduler.py
+++ b/vllm_omni/diffusion/sched/request_scheduler.py
@@ -26,26 +26,26 @@ class RequestScheduler(_BaseScheduler):
         return super().schedule()
 
     def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
-        scheduled_req_ids = sched_output.scheduled_req_ids
-        if not scheduled_req_ids:
+        scheduled_request_ids = sched_output.scheduled_request_ids
+        if not scheduled_request_ids:
             return set()
 
         terminal_statuses: dict[str, DiffusionRequestStatus] = {}
         terminal_errors: dict[str, str | None] = {}
-        for sched_req_id in scheduled_req_ids:
-            state = self._request_states.get(sched_req_id)
+        for request_id in scheduled_request_ids:
+            state = self._request_states.get(request_id)
             if state is None or state.is_finished():
                 continue
-            req_output = output.get_req_output(sched_req_id)
+            req_output = output.get_request_output(request_id)
             result = req_output.result if req_output is not None else None
             if result is None:
-                terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
-                terminal_errors[sched_req_id] = "No output result"
+                terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR
+                terminal_errors[request_id] = "No output result"
             elif result.error:
-                terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
-                terminal_errors[sched_req_id] = result.error
+                terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR
+                terminal_errors[request_id] = result.error
             else:
-                terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_COMPLETED
-                terminal_errors[sched_req_id] = None
+                terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_COMPLETED
+                terminal_errors[request_id] = None
 
         return self._finalize_update_from_output(sched_output, terminal_statuses, terminal_errors)
diff --git a/vllm_omni/diffusion/sched/step_scheduler.py b/vllm_omni/diffusion/sched/step_scheduler.py
index 58449281b..0bc2d4a3c 100644
--- a/vllm_omni/diffusion/sched/step_scheduler.py
+++ b/vllm_omni/diffusion/sched/step_scheduler.py
@@ -38,84 +38,84 @@ class StepScheduler(_BaseScheduler):
         self._request_progress.clear()
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
-        sched_req_id = self._make_sched_req_id(request)
+        request_id = request.request_id
         total_steps = self._get_total_steps(request)
         if total_steps <= 0:
-            raise ValueError(f"Diffusion request {sched_req_id} must have positive total_steps, got {total_steps}")
+            raise ValueError(f"Diffusion request {request_id} must have positive total_steps, got {total_steps}")
 
         current_step = request.sampling_params.step_index or 0
         if current_step < 0 or current_step >= total_steps:
             raise ValueError(
-                f"Diffusion request {sched_req_id} has invalid initial step_index {current_step} "
+                f"Diffusion request {request_id} has invalid initial step_index {current_step} "
                 f"for total_steps={total_steps}"
             )
 
         request.sampling_params.step_index = current_step
-        sched_req_id = self._add_request_with_sched_req_id(sched_req_id, request)
-        self._request_progress[sched_req_id] = _StepProgress(current_step=current_step, total_steps=total_steps)
+        request_id = self._add_request_with_request_id(request_id, request)
+        self._request_progress[request_id] = _StepProgress(current_step=current_step, total_steps=total_steps)
         logger.debug(
             "StepScheduler add_request: %s (step=%d/%d, waiting=%d)",
-            sched_req_id,
+            request_id,
             current_step,
             total_steps,
             len(self._waiting),
         )
-        return sched_req_id
+        return request_id
 
     def schedule(self) -> DiffusionSchedulerOutput:
         return super().schedule()
 
     def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
-        scheduled_req_ids = sched_output.scheduled_req_ids
-        if not scheduled_req_ids:
+        scheduled_request_ids = sched_output.scheduled_request_ids
+        if not scheduled_request_ids:
             return set()
 
         terminal_statuses: dict[str, DiffusionRequestStatus] = {}
         terminal_errors: dict[str, str | None] = {}
-        for sched_req_id in scheduled_req_ids:
-            state = self._request_states.get(sched_req_id)
-            progress = self._request_progress.get(sched_req_id)
+        for request_id in scheduled_request_ids:
+            state = self._request_states.get(request_id)
+            progress = self._request_progress.get(request_id)
             if state is None or progress is None or state.is_finished():
                 continue
-            req_output = output.get_req_output(sched_req_id)
+            req_output = output.get_request_output(request_id)
             if req_output is None:
                 logger.warning(
                     "No RunnerOutput for request %s, treating as error",
-                    sched_req_id,
+                    request_id,
                 )
-                terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
-                terminal_errors[sched_req_id] = "No output for request"
+                terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR
+                terminal_errors[request_id] = "No output for request"
                 continue
 
             req_result = req_output.result
             output_error = req_result.error if req_result is not None else None
             if output_error is not None:
-                terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
-                terminal_errors[sched_req_id] = output_error
+                terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR
+                terminal_errors[request_id] = output_error
                 continue
 
             if req_output.step_index is None:
                 logger.warning(
                     "Received RunnerOutput with no step_index for request %s, treating as error",
-                    sched_req_id,
+                    request_id,
                 )
-                terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
-                terminal_errors[sched_req_id] = "Missing step_index in RunnerOutput"
+                terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR
+                terminal_errors[request_id] = "Missing step_index in RunnerOutput"
                 continue
 
             # We assume that the decoding stage is executed immediately after the denoising stage completes.
             progress.current_step = req_output.step_index
             state.req.sampling_params.step_index = req_output.step_index
             if req_output.finished:
-                terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_COMPLETED
-                terminal_errors[sched_req_id] = None
+                terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_COMPLETED
+                terminal_errors[request_id] = None
             else:
                 state.error = None
 
         return self._finalize_update_from_output(sched_output, terminal_statuses, terminal_errors)
 
-    def _pop_extra_request_state(self, sched_req_id: str) -> None:
-        self._request_progress.pop(sched_req_id, None)
+    def _pop_extra_request_state(self, request_id: str) -> None:
+        self._request_progress.pop(request_id, None)
 
     def _get_total_steps(self, request: OmniDiffusionRequest) -> int:
         sampling = request.sampling_params
diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py
index 871a29729..07489d4c9 100644
--- a/vllm_omni/diffusion/stage_diffusion_proc.py
+++ b/vllm_omni/diffusion/stage_diffusion_proc.py
@@ -32,6 +32,7 @@ from vllm_omni.distributed.omni_connectors.utils.serialization import (
     OmniMsgpackEncoder,
 )
 from vllm_omni.distributed.omni_coordinator import OmniCoordClientForStage
+from vllm_omni.engine.stage_init_utils import set_death_signal
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -41,6 +42,14 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+_SIGNAL_EXIT_BASE = 128
+
+
+def _signal_exit_code(signum: int) -> int:
+    """Return the conventional process exit code for signal-driven exits."""
+    return _SIGNAL_EXIT_BASE + signum
+
+
 class StageDiffusionProc:
     """Subprocess entry point for diffusion inference.
 
@@ -150,7 +159,6 @@ class StageDiffusionProc:
         request = OmniDiffusionRequest(
             prompts=[prompt],
             sampling_params=sampling_params,
-            request_ids=[request_id],
             request_id=request_id,
             kv_sender_info=kv_sender_info,
         )
@@ -180,7 +188,6 @@ class StageDiffusionProc:
         request = OmniDiffusionRequest(
             prompts=prompts,
             sampling_params=sampling_params,
-            request_ids=[f"{request_id}-{i}" for i in range(len(prompts))],
             request_id=request_id,
             kv_sender_info=kv_sender_info,
         )
@@ -624,11 +631,13 @@ class StageDiffusionProc:
         """
         shutdown_requested = False
 
+        set_death_signal(signal.SIGTERM)
+
         def signal_handler(signum: int, frame: Any) -> None:
             nonlocal shutdown_requested
             if not shutdown_requested:
                 shutdown_requested = True
-                raise SystemExit(128 + signum)
+                raise SystemExit(_signal_exit_code(signum))
 
         signal.signal(signal.SIGTERM, signal_handler)
         signal.signal(signal.SIGINT, signal_handler)
diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py
index ab6dbb79a..636ac7f60 100644
--- a/vllm_omni/diffusion/worker/diffusion_model_runner.py
+++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py
@@ -23,6 +23,10 @@ from vllm.logger import init_logger
 from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes
 
 from vllm_omni.diffusion.cache.cache_dit_backend import cache_summary
+from vllm_omni.diffusion.cache.prompt_embed_cache import (
+    install_prompt_embed_cache,
+    resolve_prompt_embed_cache_config,
+)
 from vllm_omni.diffusion.cache.selector import get_cache_backend
 from vllm_omni.diffusion.compile import regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
@@ -71,6 +75,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         self.pipeline = None
         self.cache_backend = None
         self.offload_backend = None
+        self.prompt_embed_cache = None
 
         # Cache for per-request stepwise state.
         self.state_cache: dict[str, DiffusionRequestState] = {}
@@ -97,7 +102,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         self,
         memory_pool_context_fn: callable | None = None,
         load_format: str | None = None,
-        custom_pipeline_name: str | None = None,
+        custom_pipeline_name: str | type | None = None,
     ) -> None:
         """
         Load the diffusion model, apply compilation and offloading.
@@ -187,12 +192,38 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             else:
                 self.cache_backend.enable(self.pipeline)
 
+        # Install prompt-embedding cache (transparent wrapper around
+        # ``pipeline.encode_prompt``). Enabled via config or env var; a no-op
+        # when the pipeline does not expose ``encode_prompt``.
+        enable_pec, pec_size = resolve_prompt_embed_cache_config(
+            enable=getattr(self.od_config, "enable_prompt_embed_cache", False),
+            max_size=getattr(self.od_config, "prompt_embed_cache_size", 32),
+        )
+        if enable_pec:
+            self.prompt_embed_cache = install_prompt_embed_cache(
+                self.pipeline,
+                max_size=pec_size,
+                enabled=True,
+                model_tag=self.od_config.model_class_name,
+            )
+
         logger.info("Model runner: Initialization complete.")
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load weights into the pipeline."""
         return self.pipeline.load_weights(weights)
 
+    def clear_prompt_embed_cache(self) -> None:
+        """Evict all cached text-encoder outputs (e.g. between training epochs)."""
+        if self.prompt_embed_cache is not None:
+            self.prompt_embed_cache.clear()
+
+    def get_prompt_embed_cache_stats(self) -> dict | None:
+        """Return hit/miss statistics for the prompt-embedding cache, if enabled."""
+        if self.prompt_embed_cache is None:
+            return None
+        return self.prompt_embed_cache.stats()
+
     def _record_peak_memory(self, output: DiffusionOutput) -> None:
         """Record peak GPU memory for the current forward pass into output.
 
@@ -299,6 +330,10 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             if is_primary:
                 self._record_peak_memory(output)
 
+            # Log prompt-embed cache activity (hits/misses accumulate across requests).
+            if is_primary and self.prompt_embed_cache is not None:
+                logger.debug("prompt-embed cache: %s", self.prompt_embed_cache.stats())
+
             # NOTE:
             if (
                 self.cache_backend is not None
@@ -322,50 +357,44 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         self, scheduler_output: DiffusionSchedulerOutput
     ) -> tuple[list[DiffusionRequestState], list[str]]:
         """Step-before update: cleanup finished requests and get/create one running state."""
-        for req_id in scheduler_output.finished_req_ids:
-            self.state_cache.pop(req_id, None)
+        for request_id in scheduler_output.finished_req_ids:
+            self.state_cache.pop(request_id, None)
 
         resolved: list[DiffusionRequestState] = []
-        new_req_id: list[str] = []
+        new_request_ids: list[str] = []
         try:
             # process new requests
             for sched_new_req in scheduler_output.scheduled_new_reqs:
-                # new_req_data = scheduler_output.scheduled_new_reqs[0]
-                req_id = sched_new_req.sched_req_id
+                request_id = sched_new_req.request_id
                 req = sched_new_req.req
-                new_req_id.append(req_id)
-                if req_id in self.state_cache:
-                    raise ValueError(f"Received duplicate new-request payload for cached request {req_id}.")
-                request_ids = req.request_ids or [req_id]
-                if len(request_ids) != len(req.prompts):
-                    raise ValueError(
-                        f"request_ids length ({len(request_ids)}) does not match prompts length ({len(req.prompts)})"
-                    )
+                new_request_ids.append(request_id)
+                if request_id in self.state_cache:
+                    raise ValueError(f"Received duplicate new-request payload for cached request {request_id}.")
                 new_state = DiffusionRequestState(
-                    req_id=req_id,
+                    request_id=request_id,
                     sampling=copy.deepcopy(req.sampling_params),
                     prompts=req.prompts,
                 )
-                self.state_cache[req_id] = new_state
+                self.state_cache[request_id] = new_state
                 resolved.append(new_state)
 
             # process cached requests
-            for req_id in scheduler_output.scheduled_cached_reqs.sched_req_ids:
-                state = self.state_cache.get(req_id)
+            for request_id in scheduler_output.scheduled_cached_reqs.request_ids:
+                state = self.state_cache.get(request_id)
                 if state is None:
-                    raise ValueError(f"Missing cached state for request {req_id}.")
+                    raise ValueError(f"Missing cached state for request {request_id}.")
                 resolved.append(state)
         except Exception:
-            for req_id in new_req_id:
-                self.state_cache.pop(req_id, None)
+            for request_id in new_request_ids:
+                self.state_cache.pop(request_id, None)
             raise
 
-        return resolved, new_req_id
+        return resolved, new_request_ids
 
     def _prepare_batch_inputs(self, states: list[DiffusionRequestState], new_request_ids: list[str]) -> InputBatch:
         # process new reqs
         for state in states:
-            if state.req_id in new_request_ids:
+            if state.request_id in new_request_ids:
                 # set generator
                 if state.sampling.generator is None and state.sampling.seed is not None:
                     if state.sampling.generator_device is not None:
@@ -407,7 +436,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
 
         for state in states:
             if interrupted or state.denoise_completed:
-                self.state_cache.pop(state.req_id, None)
+                self.state_cache.pop(state.request_id, None)
 
     def _prepare_attn_metadata(self, input_batch: InputBatch) -> Any:
         model_state = getattr(self, "model_state", None)
@@ -449,7 +478,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                     for state in states:
                         runner_output_list.append(
                             RunnerOutput(
-                                req_id=state.req_id,
+                                request_id=state.request_id,
                                 step_index=state.step_index,
                                 finished=True,
                                 result=DiffusionOutput(error="stepwise denoise interrupted"),
@@ -470,7 +499,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                             result = None
                         runner_output_list.append(
                             RunnerOutput(
-                                req_id=req.req_id,
+                                request_id=req.request_id,
                                 step_index=req.step_index,
                                 finished=req.denoise_completed,
                                 result=result,
diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py
index cd5fe246d..0e8980a29 100644
--- a/vllm_omni/diffusion/worker/diffusion_worker.py
+++ b/vllm_omni/diffusion/worker/diffusion_worker.py
@@ -11,6 +11,7 @@ to DiffusionModelRunner.
 import gc
 import multiprocessing as mp
 import os
+import signal
 from collections.abc import Iterable, Iterator
 from contextlib import AbstractContextManager, contextmanager, nullcontext
 from dataclasses import dataclass
@@ -46,6 +47,7 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
+from vllm_omni.engine.stage_init_utils import set_death_signal
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.platforms import current_omni_platform
 from vllm_omni.profiler import OmniTorchProfilerWrapper, create_omni_profiler
@@ -187,7 +189,7 @@ class DiffusionWorker:
         self.lora_manager: DiffusionLoRAManager | None = None
         # Worker-side cache of (lora_request, lora_scale) per scheduled
         # request id. Used by step mode to recover LoRA identity for cached
-        # requests, which only carry their sched_req_id in subsequent ticks.
+        # requests, which only carry their request_id in subsequent ticks.
         self._step_lora_state: dict[str, tuple[LoRARequest | None, float]] = {}
         self.stage_id = getattr(od_config, "stage_id", 0)
         self.init_device()
@@ -389,12 +391,12 @@ class DiffusionWorker:
         homogeneity is enforced by the scheduler via ``SamplingParamsKey``,
         so any scheduled request id resolves to the active LoRA identity.
         """
-        for sched_req_id in scheduler_output.finished_req_ids:
-            self._step_lora_state.pop(sched_req_id, None)
+        for request_id in scheduler_output.finished_req_ids:
+            self._step_lora_state.pop(request_id, None)
 
         for new_req in scheduler_output.scheduled_new_reqs:
             sampling = new_req.req.sampling_params
-            self._step_lora_state[new_req.sched_req_id] = (
+            self._step_lora_state[new_req.request_id] = (
                 sampling.lora_request,
                 sampling.lora_scale,
             )
@@ -404,8 +406,8 @@ class DiffusionWorker:
 
         lora_request: LoRARequest | None = None
         lora_scale = 1.0
-        for sched_req_id in scheduler_output.scheduled_req_ids:
-            entry = self._step_lora_state.get(sched_req_id)
+        for request_id in scheduler_output.scheduled_request_ids:
+            entry = self._step_lora_state.get(request_id)
             if entry is not None:
                 lora_request, lora_scale = entry
                 break
@@ -520,11 +522,11 @@ class DiffusionWorker:
             logger.info(f"[Worker {self.rank}] Handshake Received: Task {task.task_id}")
 
             current_omni_platform.synchronize()
-            usage_before = current_omni_platform.get_current_memory_usage(self.device)
+            free_before = current_omni_platform.get_free_memory(self.device)
             self.sleep(level=task.level)
             current_omni_platform.synchronize()
-            usage_after = current_omni_platform.get_current_memory_usage(self.device)
-            real_freed = max(0, usage_before - usage_after)
+            free_after = current_omni_platform.get_free_memory(self.device)
+            real_freed = max(0, free_after - free_before)
             logger.info(f"[Worker {self.rank}] Preparing ACK: freed_bytes={real_freed / GiB_bytes:.2f} GiB.")
 
             # Ensure all ranks have completed sleep before measuring memory and sending ACK
@@ -536,6 +538,11 @@ class DiffusionWorker:
             if self.rank != 0:
                 return None
 
+            try:
+                total_mem = current_omni_platform.get_device_total_memory()
+            except (NotImplementedError, AttributeError):
+                total_mem = torch.cuda.get_device_properties(self.device).total_memory
+            residual_gib = (total_mem - free_after) / GiB_bytes
             ack = OmniACK(
                 task_id=task.task_id,
                 status="SUCCESS",
@@ -546,7 +553,7 @@ class DiffusionWorker:
                 metadata={
                     "source": f"Platform_{current_omni_platform.get_device_name()}",
                     "total_freed_gib": f"{real_freed / GiB_bytes:.2f}",
-                    "rank_residual_gib": f"{usage_after / GiB_bytes:.2f}",
+                    "rank_residual_gib": f"{residual_gib:.2f}",
                 },
             )
             logger.info(f"[Worker {self.rank}] ACK emitted. Freed {real_freed / GiB_bytes:.2f} GiB.")
@@ -574,8 +581,12 @@ class DiffusionWorker:
                 torch.distributed.barrier()
 
             current_omni_platform.synchronize()
-            usage_now = current_omni_platform.get_current_memory_usage(self.device)
-            current_used_gib = usage_now / (1024**3)
+            free_now = current_omni_platform.get_free_memory(self.device)
+            try:
+                total_mem = current_omni_platform.get_device_total_memory()
+            except (NotImplementedError, AttributeError):
+                total_mem = torch.cuda.get_device_properties(self.device).total_memory
+            current_used_gib = (total_mem - free_now) / (1024**3)
 
             if self.rank != 0:
                 return None
@@ -812,11 +823,6 @@ class WorkerProc:
                     continue
 
         logger.info("event loop terminated.")
-        try:
-            self.worker.shutdown()
-        except Exception as exc:
-            logger.warning("Worker %s: Shutdown encountered an error: %s", self.gpu_id, exc)
-        self.context.term()
 
     @staticmethod
     def worker_main(
@@ -831,32 +837,62 @@ class WorkerProc:
         """Worker initialization and execution loops."""
         from vllm_omni.plugins import load_omni_general_plugins
 
+        shutdown_triggered = False
+
+        def signal_handler(signum: int, frame) -> None:
+            nonlocal shutdown_triggered
+            if not shutdown_triggered:
+                shutdown_triggered = True
+                raise SystemExit(128 + signum)
+
+        signal.signal(signal.SIGTERM, signal_handler)
+        signal.signal(signal.SIGINT, signal_handler)
+
+        set_death_signal(signal.SIGTERM)
+
         # Set process title for visibility in nvidia-smi / htop (optional, non-fatal)
         try:
             import setproctitle
 
-            setproctitle.setproctitle(f"vLLM-Omni::StageDiffusionProc-{rank}")
+            setproctitle.setproctitle(f"vLLM-Omni::DiffusionWorker-{rank}")
         except ImportError:
             pass  # setproctitle not installed, skip process title setting
 
         load_omni_general_plugins()
-        worker_proc = WorkerProc(
-            od_config,
-            gpu_id=rank,
-            broadcast_handle=broadcast_handle,
-            wake_event=wake_event,
-            worker_extension_cls=worker_extension_cls,
-            custom_pipeline_args=custom_pipeline_args,
-        )
-        logger.info(f"Worker {rank}: Scheduler loop started.")
-        pipe_writer.send(
-            {
-                "status": "ready",
-                "result_handle": worker_proc.result_mq_handle if rank == 0 else None,
-            }
-        )
-        worker_proc.worker_busy_loop()
-        logger.info(f"Worker {rank}: Shutdown complete.")
+        worker_proc = None
+        try:
+            worker_proc = WorkerProc(
+                od_config,
+                gpu_id=rank,
+                broadcast_handle=broadcast_handle,
+                wake_event=wake_event,
+                worker_extension_cls=worker_extension_cls,
+                custom_pipeline_args=custom_pipeline_args,
+            )
+            logger.info(f"Worker {rank}: Scheduler loop started.")
+            pipe_writer.send(
+                {
+                    "status": "ready",
+                    "result_handle": worker_proc.result_mq_handle if rank == 0 else None,
+                }
+            )
+            worker_proc.worker_busy_loop()
+        except SystemExit:
+            logger.info("Worker %d: Shutdown signal received, starting cleanup.", rank)
+            raise
+        finally:
+            if worker_proc is not None:
+                try:
+                    worker_proc.worker.shutdown()
+                except Exception as exc:
+                    logger.warning("Worker %d: Shutdown encountered an error: %s", rank, exc)
+                worker_proc.context.term()
+            else:
+                # In case of signal interrupting worker_proc initialization
+                # where distributed env is initialized but worker_proc is still None
+                destroy_distributed_env()
+
+        logger.info("Worker %d: Shutdown complete.", rank)
 
 
 class WorkerWrapperBase:
diff --git a/vllm_omni/diffusion/worker/input_batch.py b/vllm_omni/diffusion/worker/input_batch.py
index 48c5bb57b..82033d9dc 100644
--- a/vllm_omni/diffusion/worker/input_batch.py
+++ b/vllm_omni/diffusion/worker/input_batch.py
@@ -81,8 +81,8 @@ def _select_states(
     return selected_states, idx_mapping, idx_mapping.detach().cpu().numpy()
 
 
-def _prepare_req_ids(states: Sequence[DiffusionRequestState]) -> list[str]:
-    return [state.req_id for state in states]
+def _prepare_request_ids(states: Sequence[DiffusionRequestState]) -> list[str]:
+    return [state.request_id for state in states]
 
 
 def _prepare_prompt_field_on_state(
@@ -192,11 +192,13 @@ def _get_request_prompt_seq_lens(
         mask_attr=mask_attr,
     )
     if embeds is None:
-        raise ValueError(f"{embeds_attr} is not initialized on request {state.req_id}.")
+        raise ValueError(f"{embeds_attr} is not initialized on request {state.request_id}.")
 
     if mask is not None:
         if mask.shape[0] != embeds.shape[0]:
-            raise ValueError(f"{mask_attr} batch dimension does not match {embeds_attr} for request {state.req_id}.")
+            raise ValueError(
+                f"{mask_attr} batch dimension does not match {embeds_attr} for request {state.request_id}."
+            )
         return _get_seq_lens_from_mask(mask)
 
     seq_lens = getattr(state, seq_lens_attr)
@@ -220,7 +222,7 @@ def _prepare_request_prompt_field(
         mask_attr=mask_attr,
     )
     if embeds is None:
-        raise ValueError(f"{embeds_attr} is not initialized on request {state.req_id}.")
+        raise ValueError(f"{embeds_attr} is not initialized on request {state.request_id}.")
 
     actual_seq_lens = _get_request_prompt_seq_lens(
         state,
@@ -248,7 +250,7 @@ def _prepare_request_prompt_field(
     if current_seq_len > target_seq_len:
         if max_actual_seq_len > target_seq_len:
             raise ValueError(
-                f"{embeds_attr} for request {state.req_id} requires seq_len "
+                f"{embeds_attr} for request {state.request_id} requires seq_len "
                 f"{max_actual_seq_len}, got target {target_seq_len}."
             )
         return embeds[:, :target_seq_len], None if mask is None else mask[:, :target_seq_len]
@@ -349,7 +351,7 @@ def _require_state_latents(
 ) -> torch.Tensor:
     latents = state.latents
     if latents is None:
-        raise ValueError(f"Request {state.req_id} has no latents while preparing {for_field}.")
+        raise ValueError(f"Request {state.request_id} has no latents while preparing {for_field}.")
     return latents
 
 
@@ -526,12 +528,12 @@ def _prepare_negative_prompt_embeds(
 
 def _same_composition(
     cached_batch: InputBatch | None,
-    req_ids: list[str],
+    request_ids: list[str],
     idx_mapping_np: np.ndarray,
 ) -> bool:
     if cached_batch is None:
         return False
-    if cached_batch.req_ids != req_ids:
+    if cached_batch.request_ids != request_ids:
         return False
     return np.array_equal(cached_batch.idx_mapping_np, idx_mapping_np)
 
@@ -581,7 +583,7 @@ class InputBatch:
     unchanged.
     """
 
-    req_ids: list[str]
+    request_ids: list[str]
     num_reqs: int
     num_reqs_after_padding: int
     idx_mapping: torch.Tensor
@@ -604,9 +606,9 @@ class InputBatch:
     negative_txt_seq_lens: list[int] | None = None
 
     def __post_init__(self) -> None:
-        if len(self.req_ids) != int(self.idx_mapping.numel()):
-            raise ValueError("`req_ids` and `idx_mapping` must have the same length.")
-        if self.num_reqs != len(self.req_ids):
+        if len(self.request_ids) != int(self.idx_mapping.numel()):
+            raise ValueError("`request_ids` and `idx_mapping` must have the same length.")
+        if self.num_reqs != len(self.request_ids):
             raise ValueError("`num_reqs` must match the number of request ids.")
         if self.num_reqs_after_padding < self.num_reqs:
             raise ValueError("`num_reqs_after_padding` must be >= `num_reqs`.")
@@ -653,11 +655,11 @@ class InputBatch:
         selected_states: Sequence[DiffusionRequestState],
         idx_mapping: torch.Tensor,
         idx_mapping_np: np.ndarray,
-        req_ids: list[str],
+        request_ids: list[str],
     ) -> InputBatch:
-        self.req_ids = req_ids
-        self.num_reqs = len(req_ids)
-        self.num_reqs_after_padding = len(req_ids)
+        self.request_ids = request_ids
+        self.num_reqs = len(request_ids)
+        self.num_reqs_after_padding = len(request_ids)
         self.idx_mapping = idx_mapping
         self.idx_mapping_np = idx_mapping_np
         self.latents = _prepare_latents(selected_states, out=self.latents)
@@ -675,9 +677,9 @@ class InputBatch:
     ) -> InputBatch:
         """Build a temporary step-local batch view from request states."""
         selected_states, idx_mapping, idx_mapping_np = _select_states(states, idx_mapping)
-        req_ids = _prepare_req_ids(selected_states)
+        request_ids = _prepare_request_ids(selected_states)
 
-        if _same_composition(cached_batch, req_ids, idx_mapping_np):
+        if _same_composition(cached_batch, request_ids, idx_mapping_np):
             assert cached_batch is not None
             cached_batch._repack_dynamic_fields(selected_states)
             return cached_batch
@@ -687,14 +689,14 @@ class InputBatch:
                 selected_states,
                 idx_mapping,
                 idx_mapping_np,
-                req_ids,
+                request_ids,
             )
 
         prompt_embeds, prompt_embeds_mask = _prepare_prompt_embeds(selected_states)
         negative_prompt_embeds, negative_prompt_embeds_mask = _prepare_negative_prompt_embeds(selected_states)
         do_true_cfg, true_cfg_scale, cfg_normalize = _prepare_cfg_scalars(selected_states)
         return cls(
-            req_ids=req_ids,
+            request_ids=request_ids,
             num_reqs=len(selected_states),
             num_reqs_after_padding=len(selected_states),
             idx_mapping=idx_mapping,
diff --git a/vllm_omni/diffusion/worker/utils.py b/vllm_omni/diffusion/worker/utils.py
index 1e98a7784..a0bfcea9b 100644
--- a/vllm_omni/diffusion/worker/utils.py
+++ b/vllm_omni/diffusion/worker/utils.py
@@ -38,7 +38,7 @@ class DiffusionRequestState:
     """
 
     # ── Identity / request-level inputs ──
-    req_id: str
+    request_id: str
     sampling: OmniDiffusionSamplingParams
     prompts: list[OmniPromptType] | None = None
 
@@ -112,7 +112,7 @@ class DiffusionRequestState:
 
 class BaseRunnerOutput(ABC):
     @abstractmethod
-    def get_req_output(self, sched_req_id: str) -> RunnerOutput | None:
+    def get_request_output(self, request_id: str) -> RunnerOutput | None:
         pass
 
 
@@ -125,13 +125,13 @@ class RunnerOutput(BaseRunnerOutput):
     _request_state_cache.
     """
 
-    req_id: str
+    request_id: str
     step_index: int | None = None
     finished: bool = False
     result: DiffusionOutput | None = None
 
-    def get_req_output(self, sched_req_id: str) -> RunnerOutput | None:
-        return self if self.req_id == sched_req_id else None
+    def get_request_output(self, request_id: str) -> RunnerOutput | None:
+        return self if self.request_id == request_id else None
 
 
 @dataclass
@@ -140,18 +140,18 @@ class BatchRunnerOutput(BaseRunnerOutput):
     _id_to_idx: dict[str, int] = field(init=False, repr=False)
 
     def __post_init__(self) -> None:
-        self._id_to_idx = {out.req_id: i for i, out in enumerate(self.runner_outputs)}
+        self._id_to_idx = {out.request_id: i for i, out in enumerate(self.runner_outputs)}
 
-    def __getitem__(self, req_id: str) -> RunnerOutput | None:
-        """access single RunnerOutput by req_id"""
-        idx = self._id_to_idx.get(req_id)
+    def __getitem__(self, request_id: str) -> RunnerOutput | None:
+        """access single RunnerOutput by request_id"""
+        idx = self._id_to_idx.get(request_id)
         return self.runner_outputs[idx] if idx is not None else None
 
-    def get_req_output(self, sched_req_id: str) -> RunnerOutput | None:
-        return self[sched_req_id]
+    def get_request_output(self, request_id: str) -> RunnerOutput | None:
+        return self[request_id]
 
     @property
-    def req_ids(self) -> list[str]:
+    def request_ids(self) -> list[str]:
         return list(self._id_to_idx.keys())
 
     def __len__(self) -> int:
diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py
index 79a824aa3..db253ebad 100644
--- a/vllm_omni/distributed/omni_connectors/adapter.py
+++ b/vllm_omni/distributed/omni_connectors/adapter.py
@@ -214,3 +214,34 @@ def compute_talker_prompt_ids_length(prompt_ids: list[int]) -> int:
             pass
 
     return sum_user_len + assistant_len
+
+
+def construct_next_stage_streaming_input_prompt(payload_data: dict[str, Any], request: Any) -> None:
+    """Update a downstream streaming request prompt from connector payload ids.
+
+    Async-chunk downstream stages are prewarmed before the real Talker prompt is
+    known. When a Thinker payload carries ``ids.prompt``, this helper:
+
+    * Preserves ``num_computed_tokens`` (the scheduler token watermark).
+    * Moves already-computed output tokens into ``prompt_token_ids``.
+    * Appends a new placeholder prompt slice sized from the upstream ids.
+    * Refreshes block hashes so the scheduler allocates KV slots for the
+      extended prompt without discarding prior computed state.
+    """
+    ids = payload_data.get("ids", {})
+    prompt_token_ids = ids.get("prompt", None)
+    if not prompt_token_ids:
+        return
+    num_computed_tokens = request.num_computed_tokens
+    kept_output_tokens = request._all_token_ids[request.num_prompt_tokens : num_computed_tokens]
+    del request._all_token_ids[num_computed_tokens:]
+    request._output_token_ids.clear()
+    assert request.prompt_token_ids is not None
+    # Extend prompt with kept output tokens.
+    request.prompt_token_ids.extend(kept_output_tokens)
+    next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids))
+    new_prompt = [0] * next_prompt_len
+    request._all_token_ids.extend(new_prompt or ())
+    request.prompt_token_ids.extend(new_prompt or ())
+    request.update_block_hashes()
+    request.num_prompt_tokens = len(request.prompt_token_ids)
diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py
index 372c2cb3d..3b234ad1c 100644
--- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py
+++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py
@@ -12,6 +12,8 @@ from typing import Any
 import torch
 from vllm.logger import init_logger
 
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+
 from .factory import OmniConnectorFactory
 from .utils.config import TRANSFER_ENGINE_CONNECTOR_NAMES, ConnectorSpec
 from .utils.env import expand_env_int
@@ -1014,7 +1016,7 @@ class OmniKVTransferManager:
             return None, 0
 
         # Skip during warmup dummy run — no sender is available.
-        if request_id == "dummy_req_id":
+        if OmniDiffusionRequest.is_dummy_run_request_id(request_id):
             logger.info("Skip receiving KV cache for dummy warmup request")
             return None, 0
 
@@ -1171,12 +1173,7 @@ class OmniKVTransferManager:
     @staticmethod
     def _resolve_request_id(req: Any) -> str | None:
         """Resolve the logical request ID used for KV transfer lookups."""
-        request_id = getattr(req, "request_id", None)
-        if request_id:
-            return request_id
-        if hasattr(req, "request_ids") and req.request_ids:
-            return req.request_ids[0]
-        return None
+        return getattr(req, "request_id", None)
 
     # Legacy compatibility method
     def receive_kv_cache(self, req: Any, target_device: torch.device | None = None) -> bool:
diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py
index bc840c739..0b6e0b572 100644
--- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py
+++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py
@@ -9,8 +9,9 @@ from typing import Any
 import torch
 from vllm.v1.request import Request, RequestStatus
 
-from vllm_omni.data_entry_keys import OmniPayloadStruct, unflatten_payload
+from vllm_omni.data_entry_keys import MetaStruct, OmniPayloadStruct, unflatten_payload
 
+from ..adapter import construct_next_stage_streaming_input_prompt
 from ..factory import OmniConnectorFactory
 from ..utils.config import ConnectorSpec
 from ..utils.logging import get_connector_logger
@@ -54,6 +55,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self.put_req_chunk: dict[str, int] = defaultdict(int)
         self.get_req_chunk: dict[str, int] = defaultdict(int)
         self.finished_requests: set[str] = set()
+        self.segment_finished_requests: set[str] = set()
         self.request_payload = {}
         self.code_prompt_token_ids: dict[str, list[torch.Tensor]] = defaultdict(list)
         self.request_ids_mapping: dict[str, str] = {}
@@ -63,6 +65,22 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self.requests_with_ready_chunks = set()
         self.requests_origin_status = {}
         self.requests_num_chunks_sent: dict[str, int] = defaultdict(int)
+        self._pending_streaming_prefills: dict[str, dict] = {}
+
+    @staticmethod
+    def _is_truthy_scalar(value: Any) -> bool:
+        if isinstance(value, torch.Tensor):
+            return value.numel() == 1 and bool(value.item())
+        return bool(value) if value is not None else False
+
+    @staticmethod
+    def _confirmed_num_computed_tokens(request: Request) -> int:
+        # vLLM async scheduling advances num_computed_tokens with output
+        # placeholders before the corresponding token is committed. Connector
+        # chunk send watermarks must use only committed tokens.
+        num_computed = int(getattr(request, "num_computed_tokens", 0))
+        num_placeholders = int(getattr(request, "num_output_placeholders", 0) or 0)
+        return max(0, num_computed - num_placeholders)
 
     @classmethod
     def create_connector(cls, model_config: Any):
@@ -108,31 +126,45 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self,
         pooling_output: torch.Tensor | None = None,
         request: Request | None = None,
+        is_segment_finished: bool = False,
     ):
         """Build and enqueue one chunk for asynchronous sending.
 
         Payload extraction happens in ``_send_single_request`` on the
         background save_loop thread.
 
+        For streaming input request ``is_segment_finished`` marks the end
+        of the current realtime input segment. It is intentionally separate
+        from ``request.is_finished()``: a resumable `/v1/realtime` session
+        can finish one audio segment and later continue with another segment
+        under the same external request id. For other requests, it is the same
+        as ``request.is_finished()``.
+
         Args:
             pooling_output: Partial pooling output dictionary
             request: Request object
+            is_segment_finished: whether the segment of request is finished
         """
+        is_finished = request.is_finished() and not request.resumable
+
+        confirmed_num_computed_tokens = self._confirmed_num_computed_tokens(request)
 
         # If the request is preempted, skip the already saved chunks.
-        if request.num_computed_tokens < self.requests_num_chunks_sent.get(request.external_req_id, 0):
+        if confirmed_num_computed_tokens < self.requests_num_chunks_sent.get(request.external_req_id, 0):
             logger.warning(
                 f"Enqueue save_async for request {request.external_req_id}, "
                 f"request.num_computed_tokens={request.num_computed_tokens}, "
+                f"request.num_output_placeholders={getattr(request, 'num_output_placeholders', 0)}, "
                 f"previous_chunks_sent={self.requests_num_chunks_sent.get(request.external_req_id, 0)}"
             )
             return
 
-        self.requests_num_chunks_sent[request.external_req_id] = request.num_computed_tokens
+        self.requests_num_chunks_sent[request.external_req_id] = confirmed_num_computed_tokens
         task = {
             "pooling_output": pooling_output,
             "request": request,
-            "is_finished": request.is_finished(),
+            "is_finished": is_finished,
+            "is_segment_finished": is_segment_finished,
         }
         self._pending_save_reqs.append(task)
         with self._save_cond:
@@ -166,24 +198,47 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             self.get_req_chunk[req_id] += 1
 
             meta = payload_data.get("meta", {})
+            payload_finished = self._is_truthy_scalar(meta.get("finished"))
+            payload_segment_finished = self._is_truthy_scalar(meta.get("is_segment_finished"))
             if self.model_mode == "ar":
                 request.additional_information = payload_data
-                if meta.get("finished"):
+                if chunk_id > 0 and request.resumable:
+                    # For new streaming input segment, we should update prompt from payload
+                    construct_next_stage_streaming_input_prompt(payload_data, request)
+
+                if payload_finished:
                     self.finished_requests.add(req_id)
+                    request.resumable = False
+                if payload_segment_finished:
+                    self.segment_finished_requests.add(req_id)
             else:
-                if meta.get("finished"):
+                if payload_finished:
                     self.finished_requests.add(req_id)
+                    request.resumable = False
+                if payload_segment_finished:
+                    self.segment_finished_requests.add(req_id)
 
                 new_ids = payload_data.get("codes", {}).get("audio")
-                if isinstance(new_ids, torch.Tensor):
+                has_tensor_codes = isinstance(new_ids, torch.Tensor)
+                use_tensor_codes = has_tensor_codes and new_ids.ndim >= 2
+                if use_tensor_codes:
+                    request.prompt_token_ids = [0] if new_ids.numel() > 0 else []
+                elif has_tensor_codes:
                     new_ids = new_ids.tolist()
                 elif new_ids is None:
                     new_ids = []
-                request.prompt_token_ids = new_ids
+                    request.prompt_token_ids = new_ids
+                if not use_tensor_codes:
+                    request.prompt_token_ids = new_ids
                 prev_info = getattr(request, "additional_information", None)
                 info = dict(prev_info) if isinstance(prev_info, dict) else {}
                 for key, value in payload_data.items():
                     if key == "codes":
+                        if use_tensor_codes and isinstance(value, dict):
+                            existing_sub = info.get(key)
+                            merged_sub = dict(existing_sub) if isinstance(existing_sub, dict) else {}
+                            merged_sub.update(value)
+                            info[key] = merged_sub
                         continue
                     if isinstance(value, dict):
                         existing_sub = info.get(key)
@@ -199,8 +254,13 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                 request.num_computed_tokens = 0
 
                 # Empty chunk with more data expected: keep polling.
-                if not new_ids and not meta.get("finished"):
-                    return True
+                has_new_ids = bool(new_ids.numel()) if use_tensor_codes else bool(new_ids)
+                if not has_new_ids and not payload_finished:
+                    # The base recv loop treats False as "not ready yet" and
+                    # requeues the request. Do not mark an empty non-terminal
+                    # chunk as ready, otherwise Stage1 can consume before the
+                    # first DAC frame arrives.
+                    return False
 
             # Mark as finished for consumption
             self._finished_load_reqs.add(req_id)
@@ -214,6 +274,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         pooling_output = unflatten_payload(raw_po) if isinstance(raw_po, dict) else raw_po
         request = task["request"]
         is_finished = task["is_finished"]
+        is_segment_finished = task["is_segment_finished"]
         stage_id = self.connector.stage_id
         next_stage_id = stage_id + 1
         external_req_id = request.external_req_id
@@ -227,14 +288,23 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                     transfer_manager=self,
                     pooling_output=pooling_output,
                     request=request,
-                    is_finished=is_finished,
+                    # Existing processors use is_finished as a flush signal.
+                    is_finished=is_segment_finished,
                 )
 
             except Exception as e:
                 logger.error(f"Failed to use custom_process_input_func for payload extraction: {e}")
 
         if payload_data is None:
-            return
+            if not (is_segment_finished or is_finished):
+                return
+            # Segment/request finish markers must still reach downstream even when
+            # the processor has no tensor payload.
+            payload_data = OmniPayloadStruct()
+        if payload_data.meta is None:
+            payload_data.meta = MetaStruct()
+        payload_data.meta.finished = torch.tensor(is_finished, dtype=torch.bool)
+        payload_data.meta.is_segment_finished = torch.tensor(is_segment_finished, dtype=torch.bool)
 
         success, size, metadata = self.connector.put(
             from_stage=str(stage_id),
@@ -263,13 +333,22 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             if is_payload_finished:
                 self.cleanup(request.request_id, external_req_id)
 
-        if is_finished:
+        if is_segment_finished:
             self.code_prompt_token_ids.pop(external_req_id, None)
             self.requests_num_chunks_sent.pop(external_req_id, None)
             cached_ic = getattr(self, "_cached_ic", None)
             if cached_ic is not None:
                 cached_ic.pop(external_req_id, None)
 
+    def is_done_receiving_chunks(self, request_id: str) -> bool:
+        """Return True if the request should stop polling upstream chunks.
+
+        Covers both the whole-request finish marker (``finished_requests``) and
+        the per-segment finish marker (``segment_finished_requests``) used while
+        waiting for the next streaming input slice.
+        """
+        return request_id in self.finished_requests or request_id in self.segment_finished_requests
+
     ########################################################################
     # Cleanup
     ########################################################################
@@ -284,6 +363,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         Idempotent: calling with an already-cleaned or unknown id is safe.
         """
         self.finished_requests.discard(request_id)
+        self.segment_finished_requests.discard(request_id)
         self.get_req_chunk.pop(request_id, None)
         self.requests_with_ready_chunks.discard(request_id)
         self.request_ids_mapping.pop(request_id, None)
@@ -304,6 +384,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self.request_payload.pop(external_req_id, None)
         self.code_prompt_token_ids.pop(external_req_id, None)
         self.requests_num_chunks_sent.pop(external_req_id, None)
+        self._pending_streaming_prefills.pop(external_req_id, None)
 
         cached_ic = getattr(self, "_cached_ic", None)
         if cached_ic is not None:
@@ -354,17 +435,28 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             request.status = RequestStatus.PREEMPTED
             waiting_queue.prepend_requests([request])
 
-    def restore_queues(self, waiting_queue: Any, running_queue: list[Request]) -> None:
+    def restore_queues(
+        self,
+        waiting_queue: Any,
+        running_queue: list[Request],
+        scheduler_requests: dict[str, Request] | None = None,
+    ) -> None:
         """
         Restore requests waiting for chunk to the waiting and running queues.
         """
         # Add request waiting for chunk to the waiting and running queue
         for request in self.waiting_for_chunk_waiting_requests:
-            waiting_queue.add_request(request)
+            if scheduler_requests is None or request.request_id in scheduler_requests:
+                waiting_queue.add_request(request)
         self.waiting_for_chunk_waiting_requests = deque()
 
         if self.waiting_for_chunk_running_requests:
-            running_queue.extend(self.waiting_for_chunk_running_requests)
+            live_running_requests = [
+                request
+                for request in self.waiting_for_chunk_running_requests
+                if scheduler_requests is None or request.request_id in scheduler_requests
+            ]
+            running_queue.extend(live_running_requests)
         self.waiting_for_chunk_running_requests = deque()
 
     def postprocess_scheduler_output(
@@ -413,7 +505,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                     # Requests that have loaded chunk from last round
                     # of schedule, but have not scheduled
                     continue
-                if request.request_id in self.finished_requests:
+                if self.is_done_receiving_chunks(request.request_id):
+                    request.additional_information = None
                     continue
                 # Requests that waiting for chunk
                 self.load_async(request)
@@ -458,3 +551,20 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                 continue
             if req_id in self.requests_origin_status:
                 request.status = self.requests_origin_status.pop(req_id)
+
+        request_ids = set(request_ids)
+
+        self.waiting_for_chunk_waiting_requests = deque(
+            request for request in self.waiting_for_chunk_waiting_requests if request.request_id not in request_ids
+        )
+        self.waiting_for_chunk_running_requests = deque(
+            request for request in self.waiting_for_chunk_running_requests if request.request_id not in request_ids
+        )
+
+        for req_id in request_ids:
+            self.requests_with_ready_chunks.discard(req_id)
+            self.finished_requests.discard(req_id)
+            self._finished_load_reqs.discard(req_id)
+            self._cancelled_load_reqs.add(req_id)
+
+        return []
diff --git a/vllm_omni/distributed/omni_connectors/utils/serialization.py b/vllm_omni/distributed/omni_connectors/utils/serialization.py
index ae9dc4953..c2145dbff 100644
--- a/vllm_omni/distributed/omni_connectors/utils/serialization.py
+++ b/vllm_omni/distributed/omni_connectors/utils/serialization.py
@@ -82,10 +82,12 @@ class OmniMsgpackEncoder:
 
     def _encode_tensor(self, tensor: torch.Tensor) -> dict[str, Any]:
         """Encode torch.Tensor to dict."""
-        t = tensor.detach().contiguous().cpu()
+        t = tensor.detach().cpu()
         # Handle 0-dimensional (scalar) tensors by reshaping to 1D first
         if t.dim() == 0:
             t = t.reshape(1)
+        if not t.is_contiguous():
+            t = t.contiguous()
         t = t.view(torch.uint8)
         return {
             _TENSOR_MARKER: True,
diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py
index 12a5a450a..c46222899 100644
--- a/vllm_omni/engine/arg_utils.py
+++ b/vllm_omni/engine/arg_utils.py
@@ -20,6 +20,7 @@ logger = init_logger(__name__)
 # Used when auto-injecting hf_overrides for models with missing config.json.
 _ARCH_TO_MODEL_TYPE: dict[str, str] = {
     "CosyVoice3Model": "cosyvoice3",
+    "GLMTTSForConditionalGeneration": "glm_tts",
     "OmniVoiceModel": "omnivoice",
     "VoxCPM2TalkerForConditionalGeneration": "voxcpm2",
 }
@@ -27,6 +28,7 @@ _ARCH_TO_MODEL_TYPE: dict[str, str] = {
 # Maps model architecture names to tokenizer subfolder paths within HF repos.
 _TOKENIZER_SUBFOLDER_MAP: dict[str, str] = {
     "CosyVoice3Model": "CosyVoice-BlankEN",
+    "GLMTTSForConditionalGeneration": "vq32k-phoneme-tokenizer",
 }
 
 
@@ -37,6 +39,10 @@ def _register_omni_hf_configs() -> None:
         from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
             Qwen3TTSConfig,
         )
+        from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
+        from vllm_omni.transformers_utils.configs.glm_tts import GLMTTSConfig
+        from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
+        from vllm_omni.transformers_utils.configs.voxcpm2 import VoxCPM2Config
     except Exception as exc:  # pragma: no cover - best-effort optional registration
         logger.warning("Skipping omni HF config registration due to import error: %s", exc)
         return
@@ -51,6 +57,10 @@ def _register_omni_hf_configs() -> None:
 
     for model_type, config_cls in [
         ("qwen3_tts", Qwen3TTSConfig),
+        ("cosyvoice3", CosyVoice3Config),
+        ("glm_tts", GLMTTSConfig),
+        ("omnivoice", OmniVoiceConfig),
+        ("voxcpm2", VoxCPM2Config),
     ]:
         try:
             AutoConfig.register(model_type, config_cls)
@@ -298,6 +308,7 @@ class OmniEngineArgs(EngineArgs):
                             model_path,
                             allow_patterns=[
                                 f"{subfolder}/tokenizer*",
+                                f"{subfolder}/tokenization_*",
                                 f"{subfolder}/special_tokens*",
                                 f"{subfolder}/vocab*",
                                 f"{subfolder}/merges*",
diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py
index 38df28f86..f6eb77e30 100644
--- a/vllm_omni/engine/async_omni_engine.py
+++ b/vllm_omni/engine/async_omni_engine.py
@@ -100,7 +100,7 @@ from vllm_omni.engine.stage_init_utils import (
     load_omni_transfer_config_for_model,
     prepare_engine_environment,
     release_device_locks,
-    setup_stage_devices,
+    stage_runtime_setup,
     terminate_alive_proc,
 )
 from vllm_omni.engine.stage_pool import StagePool, StagePoolClient
@@ -354,6 +354,7 @@ class AsyncOmniEngine:
         self.output_queue: janus.Queue[EngineQueueMessage] = janus.Queue()
         self.rpc_output_queue: janus.Queue[EngineQueueMessage] = janus.Queue()
         self._shutdown_called = False
+        self._stage_init_executor: concurrent.futures.ThreadPoolExecutor | None = None
         self._weak_finalizer: weakref.finalize | None = None
         self._rpc_lock = threading.Lock()
 
@@ -875,50 +876,53 @@ class AsyncOmniEngine:
                     with llm_stage_launch_lock:
                         previous_visible_devices = os.environ.get(device_control_env)
                         try:
-                            setup_stage_devices(plan.metadata.stage_id, plan.metadata.runtime_cfg)
-                            vllm_config = plan.stage_vllm_config
-                            executor_class = plan.executor_class
-                            assert vllm_config is not None
-                            assert executor_class is not None
-                            engine_args_dict = build_engine_args_dict(
-                                stage_cfg,
-                                self.model,
-                                stage_connector_spec=plan.stage_connector_spec,
-                                cli_tokenizer=getattr(self, "tokenizer", None),
-                            )
-                            lock_fds = acquire_device_locks(
+                            with stage_runtime_setup(
                                 plan.metadata.stage_id,
-                                engine_args_dict,
-                                stage_init_timeout,
-                            )
-                            if self.single_stage_mode and self._omni_master_server is not None:
-                                coord_router_addr: str | None = (
-                                    self._coordinator_runtime.router_address
-                                    if self._coordinator_runtime is not None
-                                    else None
+                                plan.metadata.runtime_cfg,
+                            ):
+                                vllm_config = plan.stage_vllm_config
+                                executor_class = plan.executor_class
+                                assert vllm_config is not None
+                                assert executor_class is not None
+                                engine_args_dict = build_engine_args_dict(
+                                    stage_cfg,
+                                    self.model,
+                                    stage_connector_spec=plan.stage_connector_spec,
+                                    cli_tokenizer=getattr(self, "tokenizer", None),
                                 )
-                                engine_manager, coordinator, addresses = launch_stack.enter_context(
-                                    launch_omni_core_engines(
+                                lock_fds = acquire_device_locks(
+                                    plan.metadata.stage_id,
+                                    engine_args_dict,
+                                    stage_init_timeout,
+                                )
+                                if self.single_stage_mode and self._omni_master_server is not None:
+                                    coord_router_addr: str | None = (
+                                        self._coordinator_runtime.router_address
+                                        if self._coordinator_runtime is not None
+                                        else None
+                                    )
+                                    engine_manager, coordinator, addresses = launch_stack.enter_context(
+                                        launch_omni_core_engines(
+                                            vllm_config=vllm_config,
+                                            executor_class=executor_class,
+                                            log_stats=False,
+                                            omni_master_server=self._omni_master_server,
+                                            stage_id=plan.metadata.stage_id,
+                                            stage_config=stage_cfg,
+                                            replica_id=plan.replica_id,
+                                            omni_coordinator_address=coord_router_addr,
+                                        )
+                                    )
+                                else:
+                                    addresses, proc, handshake_address = spawn_stage_core(
                                         vllm_config=vllm_config,
                                         executor_class=executor_class,
                                         log_stats=False,
-                                        omni_master_server=self._omni_master_server,
-                                        stage_id=plan.metadata.stage_id,
-                                        stage_config=stage_cfg,
-                                        replica_id=plan.replica_id,
-                                        omni_coordinator_address=coord_router_addr,
                                     )
+                                logger.info(
+                                    "[AsyncOmniEngine] Stage %s engine launch started",
+                                    plan.metadata.stage_id,
                                 )
-                            else:
-                                addresses, proc, handshake_address = spawn_stage_core(
-                                    vllm_config=vllm_config,
-                                    executor_class=executor_class,
-                                    log_stats=False,
-                                )
-                            logger.info(
-                                "[AsyncOmniEngine] Stage %s engine launch started",
-                                plan.metadata.stage_id,
-                            )
                         finally:
                             if previous_visible_devices is None:
                                 current_omni_platform.unset_device_control_env_var()
@@ -1017,68 +1021,71 @@ class AsyncOmniEngine:
                 with stage_launch_lock:
                     previous_visible_devices = os.environ.get(device_control_env)
                     try:
-                        setup_stage_devices(plan.metadata.stage_id, plan.metadata.runtime_cfg)
-                        omni_conn_cfg, omni_from, omni_to = plan.omni_kv_connector
-                        if omni_conn_cfg:
-                            inject_omni_kv_config(plan.stage_cfg, omni_conn_cfg, omni_from, omni_to)
-                        inject_kv_stage_info(plan.stage_cfg, plan.metadata.stage_id, self.stage_configs)
-                        if self.single_stage_mode:
-                            assert self._omni_master_server is not None
-                            od_config = build_diffusion_config(self.model, plan.stage_cfg, plan.metadata)
-                            lock_fds = acquire_diffusion_device_locks(
-                                plan.metadata.stage_id,
-                                od_config,
-                                stage_init_timeout,
-                            )
-                            handshake_address, request_address, response_address = register_stage_with_omni_master(
-                                omni_master_address=self._omni_master_server.address,
-                                omni_master_port=self._omni_master_server.port,
-                                omni_stage_id=plan.metadata.stage_id,
-                                omni_stage_config=plan.stage_cfg,
-                                return_addresses=True,
-                                replica_id=plan.replica_id,
-                            )
-                            logger.info(
-                                "[AsyncOmniEngine] Stage %s diffusion registration completed",
-                                plan.metadata.stage_id,
-                            )
-                            coord_router_addr: str | None = (
-                                self._coordinator_runtime.router_address
-                                if self._coordinator_runtime is not None
-                                else None
-                            )
-                            proc, _, _, _ = spawn_diffusion_proc(
-                                self.model,
-                                od_config,
-                                handshake_address=handshake_address,
-                                request_address=request_address,
-                                response_address=response_address,
-                                omni_coordinator_address=coord_router_addr,
-                                omni_stage_id=plan.metadata.stage_id,
-                                omni_replica_id=plan.replica_id,
-                            )
-                            complete_diffusion_handshake(proc, handshake_address, stage_init_timeout)
-                            logger.info(
-                                "[AsyncOmniEngine] Stage %s diffusion startup completed",
-                                plan.metadata.stage_id,
-                            )
-                            client = StageDiffusionClient.from_addresses(
-                                plan.metadata,
-                                request_address=request_address,
-                                response_address=response_address,
-                                proc=proc,
-                                batch_size=self.diffusion_batch_size,
-                            )
-                        else:
-                            client = initialize_diffusion_stage(
-                                plan.metadata.stage_id,
-                                self.model,
-                                plan.stage_cfg,
-                                plan.metadata,
-                                stage_init_timeout=stage_init_timeout,
-                                batch_size=self.diffusion_batch_size,
-                                use_inline=self.num_stages == 1 and plan.num_replicas == 1,
-                            )
+                        with stage_runtime_setup(
+                            plan.metadata.stage_id,
+                            plan.metadata.runtime_cfg,
+                        ):
+                            omni_conn_cfg, omni_from, omni_to = plan.omni_kv_connector
+                            if omni_conn_cfg:
+                                inject_omni_kv_config(plan.stage_cfg, omni_conn_cfg, omni_from, omni_to)
+                            inject_kv_stage_info(plan.stage_cfg, plan.metadata.stage_id, self.stage_configs)
+                            if self.single_stage_mode:
+                                assert self._omni_master_server is not None
+                                od_config = build_diffusion_config(self.model, plan.stage_cfg, plan.metadata)
+                                lock_fds = acquire_diffusion_device_locks(
+                                    plan.metadata.stage_id,
+                                    od_config,
+                                    stage_init_timeout,
+                                )
+                                handshake_address, request_address, response_address = register_stage_with_omni_master(
+                                    omni_master_address=self._omni_master_server.address,
+                                    omni_master_port=self._omni_master_server.port,
+                                    omni_stage_id=plan.metadata.stage_id,
+                                    omni_stage_config=plan.stage_cfg,
+                                    return_addresses=True,
+                                    replica_id=plan.replica_id,
+                                )
+                                logger.info(
+                                    "[AsyncOmniEngine] Stage %s diffusion registration completed",
+                                    plan.metadata.stage_id,
+                                )
+                                coord_router_addr: str | None = (
+                                    self._coordinator_runtime.router_address
+                                    if self._coordinator_runtime is not None
+                                    else None
+                                )
+                                proc, _, _, _ = spawn_diffusion_proc(
+                                    self.model,
+                                    od_config,
+                                    handshake_address=handshake_address,
+                                    request_address=request_address,
+                                    response_address=response_address,
+                                    omni_coordinator_address=coord_router_addr,
+                                    omni_stage_id=plan.metadata.stage_id,
+                                    omni_replica_id=plan.replica_id,
+                                )
+                                complete_diffusion_handshake(proc, handshake_address, stage_init_timeout)
+                                logger.info(
+                                    "[AsyncOmniEngine] Stage %s diffusion startup completed",
+                                    plan.metadata.stage_id,
+                                )
+                                client = StageDiffusionClient.from_addresses(
+                                    plan.metadata,
+                                    request_address=request_address,
+                                    response_address=response_address,
+                                    proc=proc,
+                                    batch_size=self.diffusion_batch_size,
+                                )
+                            else:
+                                client = initialize_diffusion_stage(
+                                    plan.metadata.stage_id,
+                                    self.model,
+                                    plan.stage_cfg,
+                                    plan.metadata,
+                                    stage_init_timeout=stage_init_timeout,
+                                    batch_size=self.diffusion_batch_size,
+                                    use_inline=self.num_stages == 1 and plan.num_replicas == 1,
+                                )
                     finally:
                         if previous_visible_devices is None:
                             current_omni_platform.unset_device_control_env_var()
@@ -1161,35 +1168,37 @@ class AsyncOmniEngine:
                 primary_exc = exc
                 break
 
-        # --- 2) LLM replicas: parallel init via a scoped ThreadPoolExecutor. ---
+        # --- 2) LLM replicas: parallel init via a long-lived ThreadPoolExecutor. ---
         if primary_exc is None and llm_replicas:
             future_to_replica: dict[concurrent.futures.Future[StagePoolClient], tuple[int, int]] = {}
-            with concurrent.futures.ThreadPoolExecutor(
-                max_workers=max(1, len(llm_replicas)),
-                thread_name_prefix="stage-init",
-            ) as init_executor:
-                for stage_idx, replica in llm_replicas:
-                    future = init_executor.submit(
-                        self._initialize_replica,
-                        replica,
-                        stage_init_timeout,
-                        stage_launch_lock,
-                    )
-                    future_to_replica[future] = (stage_idx, replica.replica_id)
+            if getattr(self, "_stage_init_executor", None) is None:
+                self._stage_init_executor = concurrent.futures.ThreadPoolExecutor(
+                    max_workers=max(1, len(llm_replicas)),
+                    thread_name_prefix="stage-init",
+                )
+            init_executor = self._stage_init_executor
+            for stage_idx, replica in llm_replicas:
+                future = init_executor.submit(
+                    self._initialize_replica,
+                    replica,
+                    stage_init_timeout,
+                    stage_launch_lock,
+                )
+                future_to_replica[future] = (stage_idx, replica.replica_id)
 
-                for future in concurrent.futures.as_completed(future_to_replica):
-                    stage_idx, replica_id = future_to_replica[future]
-                    try:
-                        initialized_clients_by_stage[stage_idx][replica_id] = future.result()
-                    except concurrent.futures.CancelledError:
-                        continue
-                    except Exception as exc:
-                        if primary_exc is None:
-                            primary_exc = exc
-                            for other_future in future_to_replica:
-                                if other_future is future:
-                                    continue
-                                other_future.cancel()
+            for future in concurrent.futures.as_completed(future_to_replica):
+                stage_idx, replica_id = future_to_replica[future]
+                try:
+                    initialized_clients_by_stage[stage_idx][replica_id] = future.result()
+                except concurrent.futures.CancelledError:
+                    continue
+                except Exception as exc:
+                    if primary_exc is None:
+                        primary_exc = exc
+                        for other_future in future_to_replica:
+                            if other_future is future:
+                                continue
+                            other_future.cancel()
 
         if primary_exc is not None:
             setattr(primary_exc, "_initialized_clients_by_stage", initialized_clients_by_stage)
@@ -1775,6 +1784,12 @@ class AsyncOmniEngine:
             return {
                 "rel_l1_thresh": 0.2,
             }
+        if cache_backend == "mag_cache":
+            return {
+                "mag_threshold": 0.24,
+                "mag_max_skip_steps": 5,
+                "mag_retention_ratio": 0.1,
+            }
         return None
 
     @staticmethod
@@ -1940,6 +1955,8 @@ class AsyncOmniEngine:
                 "mp" if kwargs.get("distributed_executor_backend") is None else kwargs["distributed_executor_backend"]
             ),
             "enable_sleep_mode": kwargs.get("enable_sleep_mode", False),
+            "enable_prompt_embed_cache": kwargs.get("enable_prompt_embed_cache", False),
+            "prompt_embed_cache_size": kwargs.get("prompt_embed_cache_size", 32),
             "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True),
             "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4),
             "quantization": kwargs.get("quantization", None),
@@ -1950,6 +1967,10 @@ class AsyncOmniEngine:
             "force_cutlass_fp8": bool(kwargs.get("force_cutlass_fp8", False)),
             "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False),
             "enable_ar_profiler": kwargs.get("enable_ar_profiler", False),
+            "extras": {
+                "auxiliary_text_encoder": kwargs.get("auxiliary_text_encoder", None),
+                "default_llama_model_id": kwargs.get("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"),
+            },
             **(
                 {
                     "profiler_config": asdict(kwargs["profiler_config"])
@@ -2414,6 +2435,11 @@ class AsyncOmniEngine:
             if self.orchestrator_thread.is_alive():
                 logger.warning("[AsyncOmniEngine] Orchestrator thread did not exit in time")
 
+        stage_init_executor = getattr(self, "_stage_init_executor", None)
+        if stage_init_executor is not None:
+            stage_init_executor.shutdown(wait=False, cancel_futures=True)
+            self._stage_init_executor = None
+
         for q in (self.request_queue, self.output_queue, self.rpc_output_queue):
             try:
                 q.close()
diff --git a/vllm_omni/engine/cfg_companion_tracker.py b/vllm_omni/engine/cfg_companion_tracker.py
index 1105e887c..aeec9e848 100644
--- a/vllm_omni/engine/cfg_companion_tracker.py
+++ b/vllm_omni/engine/cfg_companion_tracker.py
@@ -22,17 +22,18 @@ class CfgCompanionTracker:
         self._companion_to_parent: dict[str, str] = {}  # companion -> parent
         self._done: dict[str, set[str]] = {}  # parent -> completed companion ids
         self._pending_parents: dict[str, dict[str, Any]] = {}  # parent -> deferred result
+        self._companion_outputs: dict[str, Any] = {}  # companion_id -> engine output
 
     def is_companion(self, req_id: str) -> bool:
         return req_id in self._companion_ids
 
-    def has_companions(self, parent_id: str) -> bool:
-        return parent_id in self._companion_map
-
     def get_parent_id(self, req_id: str) -> str | None:
         """Return the parent request id for a companion, or None."""
         return self._companion_to_parent.get(req_id)
 
+    def has_companions(self, parent_id: str) -> bool:
+        return parent_id in self._companion_map
+
     def all_companions_done(self, parent_id: str) -> bool:
         role_map = self._companion_map.get(parent_id, {})
         done_set = self._done.get(parent_id, set())
@@ -52,6 +53,20 @@ class CfgCompanionTracker:
         self._companion_ids.add(companion_id)
         self._companion_to_parent[companion_id] = parent_id
 
+    def set_companion_output(self, companion_id: str, output: Any) -> None:
+        """Stash companion engine output for the parent to bundle at forward time."""
+        self._companion_outputs[companion_id] = output
+
+    def pop_companion_outputs(self, parent_id: str) -> list[Any]:
+        """Pop companion outputs (role-registration order) for bundling into set_engine_outputs."""
+        role_map = self._companion_map.get(parent_id, {})
+        outputs = []
+        for cid in role_map.values():
+            out = self._companion_outputs.pop(cid, None)
+            if out is not None:
+                outputs.append(out)
+        return outputs
+
     def on_companion_completed(self, companion_id: str) -> str | None:
         """Mark done. Returns parent_id only if parent is pending and all companions finished."""
         parent_id = self._companion_to_parent.get(companion_id)
@@ -86,6 +101,7 @@ class CfgCompanionTracker:
         for companion_id in companion_ids:
             self._companion_ids.discard(companion_id)
             self._companion_to_parent.pop(companion_id, None)
+            self._companion_outputs.pop(companion_id, None)
         self._done.pop(parent_id, None)
         self._pending_parents.pop(parent_id, None)
         return companion_ids
diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py
index 27bb44ad7..791b6d5d9 100644
--- a/vllm_omni/engine/orchestrator.py
+++ b/vllm_omni/engine/orchestrator.py
@@ -107,7 +107,7 @@ def build_engine_core_request_from_tokens(
         pooling_params=pooling_params,
         arrival_time=arrival_time,
         lora_request=getattr(params, "lora_request", None),
-        cache_salt=None,
+        cache_salt=prompt.get("cache_salt"),
         data_parallel_rank=None,
         prompt_embeds=prompt_embeds,
         resumable=resumable,
@@ -416,7 +416,7 @@ class Orchestrator:
         stage_id = 0
         request_id = msg.request_id
         request = msg.prompt
-
+        final_stage_id = msg.final_stage_id
         req_state = self.request_states.get(request_id)
         if req_state is None:
             logger.warning(
@@ -449,6 +449,9 @@ class Orchestrator:
             prompt_text=msg.output_prompt_text,
         )
 
+        if self.async_chunk and stage_id == 0 and final_stage_id > 0:
+            await self._prewarm_async_chunk_stages(request_id, request, req_state)
+
     async def _handle_add_companion(self, msg: AddCompanionRequestMessage) -> None:
         """Handle an add_companion_request message: submit companion to stage 0."""
         companion_id = msg.companion_id
@@ -741,7 +744,10 @@ class Orchestrator:
         finished = output.finished
         submit_ts = req_state.stage_submit_ts.get(stage_id)
 
+        # CFG companion: stash output so parent can bundle [parent, *companions]
+        # into source_outputs for the bridge (e.g. thinker2imagegen).
         if finished and self._cfg_tracker.is_companion(req_id):
+            self._cfg_tracker.set_companion_output(req_id, output)
             await self._handle_cfg_companion_ready(req_id)
             await self._cleanup_request_ids([req_id])
             return
@@ -933,12 +939,34 @@ class Orchestrator:
         requires_multimodal_data = getattr(next_client, "requires_multimodal_data", False)
 
         if next_pool.stage_type == "diffusion":
+            companion_outputs = self._cfg_tracker.pop_companion_outputs(req_id)
+            expected = len(self._cfg_tracker.get_companion_request_ids(req_id))
+            if expected > len(companion_outputs):
+                logger.warning(
+                    "[Orchestrator] req=%s: only %d/%d CFG companion outputs arrived; "
+                    "downstream CFG conditioning may degrade",
+                    req_id,
+                    len(companion_outputs),
+                    expected,
+                )
+            diffusion_source_outputs = [output, *companion_outputs]
             if next_client.custom_process_input_func is not None:
                 _t_ar2d = _time.perf_counter()
-                diffusion_prompt = next_client.custom_process_input_func(
-                    source_outputs,
+                _fn = next_client.custom_process_input_func
+                _extra_kwargs: dict[str, Any] = {}
+                # TODO: replace signature probe with explicit kwarg contract.
+                try:
+                    import inspect as _inspect
+
+                    if "sampling_params" in _inspect.signature(_fn).parameters:
+                        _extra_kwargs["sampling_params"] = params
+                except (TypeError, ValueError):
+                    pass
+                diffusion_prompt = _fn(
+                    diffusion_source_outputs,
                     req_state.prompt,
                     requires_multimodal_data,
+                    **_extra_kwargs,
                 )
                 _dt_ar2d = (_time.perf_counter() - _t_ar2d) * 1000
                 req_state.pipeline_timings["ar2diffusion_ms"] = _dt_ar2d
@@ -1135,12 +1163,13 @@ class Orchestrator:
                 base_input["prompt_token_ids"] = [0] * next_prompt_len
                 base_input["multi_modal_data"] = None
                 base_input["mm_processor_kwargs"] = None
-
+                downstream_resumable = bool(getattr(stage0_request, "resumable", req_state.streaming.enabled))
                 request = build_engine_core_request_from_tokens(
                     request_id=request_id,
                     prompt=base_input,
                     params=params,
                     model_config=next_pool.stage_vllm_config.model_config,
+                    resumable=downstream_resumable,
                 )
                 request.external_req_id = request.request_id
                 await next_pool.submit_initial(request_id, req_state, request, prompt_text=None)
diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py
index 84bf72863..b5aecf124 100644
--- a/vllm_omni/engine/stage_engine_core_client.py
+++ b/vllm_omni/engine/stage_engine_core_client.py
@@ -6,11 +6,13 @@ Directly inherits from vLLM's AsyncMPClient to reuse EngineCore architecture.
 
 from __future__ import annotations
 
+import contextlib
 import inspect
 import multiprocessing.connection
 import os
 import socket
 import threading
+import time
 import weakref
 from typing import TYPE_CHECKING, Any
 from urllib.parse import urlparse
@@ -158,6 +160,7 @@ class StageEngineCoreClientBase(StageClientBase):
 
         self.engine_outputs: Any = None
         self._proc = proc
+        self._shutting_down = False
         self.client_addresses = dict(client_addresses or {})
         self._omni_kv_config = getattr(getattr(vllm_config, "model_config", None), "omni_kv_config", None)
         self._kv_sender_host = self._resolve_contact_host()
@@ -233,15 +236,37 @@ class StageEngineCoreClientBase(StageClientBase):
                 multiprocessing.connection.wait([proc.sentinel])
             except Exception:
                 return
+
+            # Give multiprocessing a brief chance to publish a stable exitcode
+            # before emitting the monitor log. Without this, the sentinel can
+            # fire first and the log may misleadingly print ``None``.
+            with contextlib.suppress(Exception):
+                proc.join(timeout=0)
+            exitcode = proc.exitcode
+            if exitcode is None:
+                deadline = time.monotonic() + 0.2
+                while exitcode is None and time.monotonic() < deadline:
+                    time.sleep(0.01)
+                    with contextlib.suppress(Exception):
+                        proc.join(timeout=0)
+                    exitcode = proc.exitcode
+
             resources = resources_ref()
-            if resources is None or resources.engine_dead:
+            if resources is None or resources.engine_dead or self._shutting_down:
+                return
+            if exitcode == 0:
+                logger.info(
+                    "[StageEngineCoreClient] stage-%s [rep-%s] subprocess exited cleanly.",
+                    stage_id,
+                    replica_id,
+                )
                 return
             resources.engine_dead = True
             logger.error(
                 "[StageEngineCoreClient] stage-%s [rep-%s] subprocess died unexpectedly (exit code %s).",
                 stage_id,
                 replica_id,
-                proc.exitcode,
+                exitcode,
             )
 
         t = threading.Thread(
@@ -462,6 +487,7 @@ class StageEngineCoreClientBase(StageClientBase):
 
     def shutdown(self, timeout: float | None = None) -> None:
         """Shutdown managed resources and any externally spawned subprocess."""
+        self._shutting_down = True
         child_procs: list[psutil.Process] = []
         if self._proc is not None and self._proc.pid is not None:
             try:
diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py
index 0b3dadcbd..0bb700369 100644
--- a/vllm_omni/engine/stage_engine_core_proc.py
+++ b/vllm_omni/engine/stage_engine_core_proc.py
@@ -36,6 +36,7 @@ from vllm.v1.engine.utils import (
 from vllm.v1.utils import shutdown
 
 from vllm_omni.distributed.omni_coordinator import OmniCoordClientForStage
+from vllm_omni.engine.stage_init_utils import set_death_signal
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
@@ -44,6 +45,14 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+_SIGNAL_EXIT_BASE = 128
+
+
+def _signal_exit_code(signum: int) -> int:
+    """Return the conventional process exit code for signal-driven exits."""
+    return _SIGNAL_EXIT_BASE + signum
+
+
 class StageEngineCoreProc(EngineCoreProc):
     """Stage-specific engine core process for vLLM-Omni.
 
@@ -89,6 +98,7 @@ class StageEngineCoreProc(EngineCoreProc):
             # like upstream vLLM.
 
             stage_label = f"stage{omni_stage_id}" if omni_stage_id is not None else "noid"
+            set_death_signal(signal.SIGTERM)
             set_process_title(f"StageEngineCoreProc_{stage_label}_replica{omni_replica_id}_DP{dp_rank}")
             decorate_logs()
             os.environ["VLLM_OMNI_REPLICA_ID"] = str(max(int(omni_replica_id), 0))
@@ -142,6 +152,7 @@ class StageEngineCoreProc(EngineCoreProc):
             def signal_handler(signum: int, frame: Any) -> None:
                 engine_core.shutdown_state = EngineShutdownState.REQUESTED
                 signal_callback.trigger()
+                raise SystemExit(_signal_exit_code(signum))
 
             signal.signal(signal.SIGTERM, signal_handler)
             signal.signal(signal.SIGINT, signal_handler)
diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py
index 1a9f98fe5..e650373a9 100644
--- a/vllm_omni/engine/stage_init_utils.py
+++ b/vllm_omni/engine/stage_init_utils.py
@@ -13,9 +13,10 @@ import importlib
 import multiprocessing as mp
 import os
 import time
-from collections.abc import Callable, Sequence
+from collections.abc import Callable, Generator, Sequence
+from contextlib import contextmanager
 from dataclasses import dataclass, replace
-from typing import Any, Literal
+from typing import Any, Literal, cast
 
 from vllm.logger import init_logger
 from vllm.sampling_params import SamplingParams
@@ -181,6 +182,19 @@ def terminate_alive_proc(proc, timeout=5):
             proc.kill()
 
 
+def set_death_signal(sig: int) -> None:
+    """Best-effort parent-death signal for Linux subprocesses."""
+    try:
+        import ctypes
+        import platform
+
+        if platform.system() != "Linux":
+            return
+        ctypes.CDLL("libc.so.6").prctl(1, sig)
+    except Exception:
+        pass
+
+
 def patch_generation_config_if_needed(model_config: Any) -> None:
     """Guard InputProcessor init for models whose config lacks model_type."""
     try:
@@ -635,6 +649,62 @@ def setup_stage_devices(stage_id: int, runtime_cfg: Any) -> None:
         )
 
 
+@contextmanager
+def stage_runtime_setup(stage_id: int, runtime_cfg: Any) -> Generator[None, None, None]:
+    """Apply per-stage ``runtime.env`` and ``runtime.devices`` for the context.
+
+    Restores ``runtime.env`` on exit. Device visibility restore remains the
+    caller's responsibility (e.g. ``AsyncOmniEngine`` saves/restores the
+    platform device-control env var around this block).
+    """
+    with stage_runtime_env(stage_id, runtime_cfg):
+        setup_stage_devices(stage_id, runtime_cfg)
+        yield
+
+
+@contextmanager
+def stage_runtime_env(stage_id: int, runtime_cfg: Any) -> Generator[None, None, None]:
+    """Apply per-stage ``runtime.env`` for the duration of the context."""
+    if runtime_cfg is None:
+        runtime_cfg = {}
+    elif not isinstance(runtime_cfg, dict):
+        runtime_cfg = cast(dict[str, Any], _to_dict(runtime_cfg))
+
+    raw_env = runtime_cfg.get("env")
+    if raw_env is None:
+        yield
+        return
+    if isinstance(raw_env, dict):
+        runtime_env = cast(dict[str, Any], raw_env)
+    else:
+        runtime_env = cast(dict[str, Any], _to_dict(raw_env))
+        if not runtime_env:
+            logger.warning(
+                "[stage_init] Stage-%s ignored runtime.env with unsupported type %s",
+                stage_id,
+                type(raw_env).__name__,
+            )
+            yield
+            return
+
+    previous_env: dict[str, str | None] = {}
+    for key, value in runtime_env.items():
+        env_key = str(key)
+        previous_env[env_key] = os.environ.get(env_key)
+        os.environ[env_key] = str(value)
+
+    if previous_env:
+        logger.info("[stage_init] Stage-%s applied runtime env keys: %s", stage_id, sorted(previous_env))
+    try:
+        yield
+    finally:
+        for key, old_value in previous_env.items():
+            if old_value is None:
+                os.environ.pop(key, None)
+            else:
+                os.environ[key] = old_value
+
+
 def build_engine_args_dict(
     stage_config: Any,
     model: str,
diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py
index 3485f98f7..c07bb2085 100644
--- a/vllm_omni/entrypoints/async_omni.py
+++ b/vllm_omni/entrypoints/async_omni.py
@@ -23,6 +23,7 @@ from vllm.pooling_params import PoolingParams
 from vllm.renderers.inputs.preprocess import extract_prompt_components
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 from vllm.tasks import SupportedTask
+from vllm.utils import random_uuid
 from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask
@@ -215,6 +216,19 @@ class AsyncOmni(EngineClient, OmniBase):
             return None
         return vllm_config.model_config
 
+    @staticmethod
+    def _get_unique_request_id(external_request_id: str):
+        """Get a random new request ID for this request; at the server level,
+        this is usually set by the calling entrypoint, but in direct calls, we
+        need to set it explicitly since we do not allow empty IDs.
+
+        NOTE: in the upstream vLLM, this is done in the InputProcessor's
+        `assign_request_id`.
+        """
+        uuid = random_uuid()
+        prefix = "" if not external_request_id else f"{external_request_id}-"
+        return f"{prefix}{uuid:.8}"
+
     # ==================== Generate Method ====================
 
     async def generate(
@@ -249,7 +263,8 @@ class AsyncOmni(EngineClient, OmniBase):
         Args:
             prompt: A single prompt **or** a list of prompts.  A list
                 triggers batch mode when the diffusion stage is reached.
-            request_id: Unique identifier for this request.
+            request_id: Unique identifier for this request. If one is not provided,
+                a random one will be generated.
             sampling_params_list: List of SamplingParams, one per stage.
                 Must have the same length as the number of stages.
                 If *None*, uses default sampling params for each stage.
@@ -263,11 +278,17 @@ class AsyncOmni(EngineClient, OmniBase):
         Raises:
             ValueError: If sampling_params_list has incorrect length.
         """
+        # Append a random UUID suffix to the request_id to ensure it is unique
+        # and non-empty, similar to vLLM's input processor. The suffix is used
+        # only for internal tracking throughout the request's life.
+        external_request_id = request_id
+        request_id = self._get_unique_request_id(external_request_id)
+
         # Wait until generation is resumed if the engine is paused
         async with self._pause_cond:
             await self._pause_cond.wait_for(lambda: not self._paused)
 
-        logger.debug(f"[AsyncOmni] generate() called for request {request_id}")
+        logger.debug(f"[AsyncOmni] generate() called for request {external_request_id}")
 
         input_stream_task: asyncio.Task | None = None
         try:
@@ -302,7 +323,11 @@ class AsyncOmni(EngineClient, OmniBase):
                 wall_start_ts,
                 final_stage_id_for_e2e,
             )
-            req_state = ClientRequestState(request_id)
+
+            req_state = ClientRequestState(
+                request_id=request_id,
+                external_request_id=external_request_id,
+            )
             req_state.metrics = metrics
             self.request_states[request_id] = req_state
 
@@ -353,11 +378,11 @@ class AsyncOmni(EngineClient, OmniBase):
         except (asyncio.CancelledError, GeneratorExit):
             if input_stream_task is not None and not input_stream_task.done():
                 input_stream_task.cancel()
-            await self.abort(request_id)
+            await self._abort_internal_requests(request_id)
             logger.info(f"[AsyncOmni] Request {request_id} aborted.")
             raise
         except Exception as e:
-            await self.abort(request_id)
+            await self._abort_internal_requests(request_id)
             logger.info(f"[AsyncOmni] Request {request_id} failed (input error): {e}")
             raise
 
@@ -541,6 +566,8 @@ class AsyncOmni(EngineClient, OmniBase):
             )
 
             if output_to_yield:
+                # Set the external request ID back to the user yielded input
+                output_to_yield.request_id = req_state.external_request_id or output_to_yield.request_id
                 logger.debug(
                     "[AsyncOmni] req=%s stage-%s yielding final_output_type=%s",
                     request_id,
@@ -687,9 +714,26 @@ class AsyncOmni(EngineClient, OmniBase):
     async def abort(self, request_id: str | Iterable[str]) -> None:
         """Abort request(s) via the Orchestrator."""
         request_ids = [request_id] if isinstance(request_id, str) else list(request_id)
+        # Map the external user request IDs to internal IDs used by the Orchestrator.
+        # NOTE: If the user request_id matches multiple requests, all of them will be
+        # aborted. This is also what happens in this case in vLLM's output processor.
+        internal_ids = [s.request_id for s in self.request_states.values() if s.external_request_id in request_ids]
+        await self._abort(internal_ids)
+
+    async def _abort_internal_requests(self, request_id: str | Iterable[str]):
+        """Abort request(s) via the Orchestrator given internal request IDs,
+        which take the format <external_request_id>-<UUID>.
+        """
+        request_ids = [request_id] if isinstance(request_id, str) else list(request_id)
+        # Request IDs are already internal, so we just need to get the matching states.
+        internal_req_ids = [rid for rid in request_ids if rid in self.request_states]
+        await self._abort(internal_req_ids)
+
+    async def _abort(self, request_ids: list[str]) -> None:
+        """Submit request IDs to be aborted to the engine."""
         await self.engine.abort_async(request_ids)
-        for req_id in request_ids:
-            self.request_states.pop(req_id, None)
+        for rid in request_ids:
+            self.request_states.pop(rid, None)
         if self.log_stats:
             logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids))
 
diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py
index e4fddc830..a9f3210f0 100644
--- a/vllm_omni/entrypoints/cli/serve.py
+++ b/vllm_omni/entrypoints/cli/serve.py
@@ -507,13 +507,16 @@ class OmniServeCommand(CLISubcommand):
             "--cache-backend",
             type=str,
             default="none",
-            help="Cache backend for diffusion models, options: 'tea_cache', 'cache_dit'",
+            help="Cache backend for diffusion models, options: 'tea_cache', 'cache_dit', 'mag_cache'",
         )
         omni_config_group.add_argument(
             "--cache-config",
             type=str,
             default=None,
-            help="JSON string of cache configuration (e.g., '{\"rel_l1_thresh\": 0.2}').",
+            help="JSON string of cache configuration. "
+            "TeaCache: '{\"rel_l1_thresh\": 0.2}'. "
+            'MagCache: \'{"mag_threshold": 0.24, "mag_max_skip_steps": 5, "mag_retention_ratio": 0.1}\'. '
+            "Calibration mode: add '\"mag_calibrate\": true'",
         )
         omni_config_group.add_argument(
             "--enable-cache-dit-summary",
@@ -650,6 +653,16 @@ class OmniServeCommand(CLISubcommand):
             action="store_true",
             help="Enable AR stage profiler to include AR stage timing in stage_durations.",
         )
+
+        # Supplementary auxiliary text encoder parameters
+        # (e.g., the meta llama/meta llama-3.1-8b-instrument used by hidream)
+        omni_config_group.add_argument(
+            "--auxiliary-text-encoder",
+            type=str,
+            default=None,
+            help="Auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).",
+        )
+
         # Stash via type(self) so the docs hook (which execs this function in a
         # sandboxed globals dict via ``DummySelf``) doesn't fail on a NameError.
         type(self)._parser = serve_parser
diff --git a/vllm_omni/entrypoints/client_request_state.py b/vllm_omni/entrypoints/client_request_state.py
index 1c9103f79..680ab9838 100644
--- a/vllm_omni/entrypoints/client_request_state.py
+++ b/vllm_omni/entrypoints/client_request_state.py
@@ -6,8 +6,14 @@ from vllm_omni.metrics import OrchestratorAggregator
 class ClientRequestState:
     """Tracks the state of an individual request in the orchestrator."""
 
-    def __init__(self, request_id: str, queue: asyncio.Queue | None = None):
+    def __init__(
+        self,
+        request_id: str,
+        external_request_id: str | None = None,
+        queue: asyncio.Queue | None = None,
+    ):
         self.request_id = request_id
+        self.external_request_id = external_request_id
         self.stage_id: int | None = None
         self.queue = queue if queue is not None else asyncio.Queue()
         self.metrics: OrchestratorAggregator | None = None
diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py
index 7bd27be81..de4ef327c 100644
--- a/vllm_omni/entrypoints/omni_base.py
+++ b/vllm_omni/entrypoints/omni_base.py
@@ -277,6 +277,10 @@ class OmniBase(PDDisaggregationMixin):
         try:
             if req_state is None or req_state.metrics is None:
                 return
+            if self.log_stats:
+                # Emit per-request orchestrator timing (including e2e_total_ms)
+                # before dropping request state.
+                req_state.metrics.build_and_log_summary()
         except Exception:
             logger.exception(
                 "[%s] Failed to build/log summary for req=%s",
diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py
index f00e6448a..c13001063 100644
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/vllm_omni/entrypoints/openai/api_server.py
@@ -92,6 +92,7 @@ from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
 from vllm_omni.entrypoints.openai.image_api_utils import (
     SUPPORTED_LAYERED_RESOLUTIONS,
     encode_image_base64,
+    encode_image_base64_with_compression,
     parse_size,
     validate_layered_layers,
 )
@@ -586,6 +587,7 @@ async def omni_init_app_state(
     state.engine_client = engine_client
     state.log_stats = not args.disable_log_stats
     state.args = args
+    state.sleeping_stages = set()
 
     # For omni models
     state.stage_configs = engine_client.stage_configs if hasattr(engine_client, "stage_configs") else None
@@ -949,7 +951,6 @@ async def omni_init_app_state(
 
     state.enable_server_load_tracking = args.enable_server_load_tracking
     state.server_load_metrics = 0
-    state.sleeping_stages = set()
 
 
 def Omnivideo(request: Request) -> OmniOpenAIServingVideo | None:
@@ -1380,22 +1381,6 @@ async def streaming_video_chat(websocket: WebSocket):
 @router.websocket("/v1/realtime")
 async def realtime_websocket(websocket: WebSocket):
     """WebSocket endpoint for OpenAI-style realtime interactions."""
-    engine_client = getattr(websocket.app.state, "engine_client", None)
-    if engine_client is not None and getattr(engine_client, "async_chunk", False):
-        await websocket.accept()
-        await websocket.send_json(
-            {
-                "type": "error",
-                "error": (
-                    "The /v1/realtime API is not supported when async_chunk is enabled on the server. "
-                    "Use a stage configuration with async_chunk disabled and restart the server before using "
-                    "this endpoint."
-                ),
-                "code": "unsupported",
-            }
-        )
-        await websocket.close()
-        return
     serving = getattr(websocket.app.state, "openai_serving_realtime", None)
     if serving is None:
         await websocket.accept()
@@ -1535,6 +1520,12 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
                 # Keep /images validation semantics: invalid LoRA should fail with 400.
                 _parse_lora_request(request.lora)
                 extra_body["lora"] = request.lora
+            if request.bot_task is not None:
+                extra_body["bot_task"] = request.bot_task
+            if request.use_system_prompt is not None:
+                extra_body["use_system_prompt"] = request.use_system_prompt
+            if request.system_prompt is not None:
+                extra_body["system_prompt"] = request.system_prompt
 
             generation_result = await chat_handler.generate_diffusion_images(
                 prompt=request.prompt,
@@ -1546,8 +1537,9 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
                     status_code=generation_result.error.code if generation_result.error else 400,
                     content=generation_result.model_dump(),
                 )
-            flat_images, _, _ = generation_result
+            flat_images, _, _, _ = generation_result
             image_data = [ImageData(b64_json=encode_image_base64(img), revised_prompt=None) for img in flat_images]
+
             return ImageGenerationResponse(created=int(time.time()), data=image_data)
 
         # Build params - pass through user values directly
@@ -1560,6 +1552,8 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
             extra_args["use_system_prompt"] = request.use_system_prompt
         if request.system_prompt is not None:
             extra_args["system_prompt"] = request.system_prompt
+        if request.bot_task is not None:
+            extra_args["bot_task"] = request.bot_task
         if extra_args:
             gen_params.extra_args = extra_args
         # Parse per-request LoRA (compatible with chat's extra_body.lora shape).
@@ -1635,7 +1629,7 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
 
         # Encode images to base64 with the specified format
         image_data = [
-            ImageData(b64_json=_encode_image_base64_with_compression(img, format=output_format), revised_prompt=None)
+            ImageData(b64_json=encode_image_base64_with_compression(img, format=output_format), revised_prompt=None)
             for img in images
         ]
 
@@ -1685,6 +1679,7 @@ async def edit_images(
     output_format: str | None = Form("png"),
     background: str | None = Form("auto"),
     output_compression: Annotated[int, Form(ge=0, le=100)] = 100,
+    stream: bool = Form(False),
     user: str | None = Form(None),  # unused now
     # vllm-omni extensions for image editing
     mask_image: str | UploadFile | None = None,
@@ -1725,8 +1720,14 @@ async def edit_images(
             status_code=HTTPStatus.BAD_REQUEST.value,
             detail="Only response_format 'b64_json' is supported now.",
         )
+    if stream and len(stage_configs) <= 1:
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST.value,
+            detail="stream=true is only supported for multi-stage image editing pipelines.",
+        )
     try:
         # 2. Build prompt & images params
+        cot_output = None
         prompt: OmniTextPrompt = {"prompt": prompt}
         if negative_prompt is not None:
             prompt["negative_prompt"] = negative_prompt
@@ -1931,13 +1932,23 @@ async def edit_images(
                 extra_body=extra_body,
                 reference_images=ref_b64_list,
                 request_id=request_id,
+                stream=stream,
+                model=model_name,
+                output_format=output_format,
+                output_compression=output_compression,
+                size=size_str,
             )
+            if stream and not isinstance(generation_result, ErrorResponse):
+                return StreamingResponse(
+                    content=generation_result,
+                    media_type="text/event-stream",
+                )
             if isinstance(generation_result, ErrorResponse):
                 raise HTTPException(
                     status_code=generation_result.error.code if generation_result.error else 400,
                     detail=generation_result.message,
                 )
-            images, _, _ = generation_result
+            images, _, _, cot_output = generation_result
         else:
             # Single-stage diffusion: use the direct path.
             result = await _generate_with_async_omni(
@@ -1954,7 +1965,7 @@ async def edit_images(
         # Encode images to base64
         image_data = [
             ImageData(
-                b64_json=_encode_image_base64_with_compression(
+                b64_json=encode_image_base64_with_compression(
                     img, format=output_format, output_compression=output_compression
                 ),
                 revised_prompt=None,
@@ -1967,6 +1978,7 @@ async def edit_images(
             data=image_data,
             output_format=output_format,
             size=size_str,
+            cot_output=cot_output,
         )
 
     except (EngineGenerateError, EngineDeadError) as exc:
@@ -2278,48 +2290,6 @@ def _choose_output_format(output_format: str | None, background: str | None) ->
     return "jpeg"
 
 
-def _prepare_image_for_output_format(image: Image.Image, format: str) -> Image.Image:
-    fmt = format.lower()
-    if fmt not in {"jpg", "jpeg"}:
-        return image
-
-    if image.mode == "RGB":
-        return image
-
-    if image.mode in {"RGBA", "LA"} or (image.mode == "P" and "transparency" in image.info):
-        alpha_image = image.convert("RGBA")
-        flattened = Image.new("RGB", alpha_image.size, (255, 255, 255))
-        flattened.paste(alpha_image, mask=alpha_image.getchannel("A"))
-        return flattened
-
-    return image.convert("RGB")
-
-
-def _encode_image_base64_with_compression(
-    image: Image.Image, format: str = "png", output_compression: int = 100
-) -> str:
-    """Encode PIL Image to a base64 image string.
-
-    Args:
-        image: PIL Image object
-        format: Output image format (e.g., "PNG", "JPEG", "WEBP")
-        output_compression: Compression level (0-100%), 100 for best quality
-    Returns:
-        Base64-encoded image as string
-    """
-    buffer = io.BytesIO()
-    image = _prepare_image_for_output_format(image, format)
-    save_kwargs = {}
-    if format in ("jpg", "jpeg", "webp"):
-        save_kwargs["quality"] = output_compression
-    elif format == "png":
-        save_kwargs["compress_level"] = max(0, min(9, 9 - output_compression // 11))  # Map 0-100 to 9-0
-
-    image.save(buffer, format=format, **save_kwargs)
-    buffer.seek(0)
-    return base64.b64encode(buffer.read()).decode("utf-8")
-
-
 def apply_stage_default_sampling_params(
     default_params_json: str | None,
     sampling_params: Any,
diff --git a/vllm_omni/entrypoints/openai/image_api_utils.py b/vllm_omni/entrypoints/openai/image_api_utils.py
index 97847ad11..b958d786a 100644
--- a/vllm_omni/entrypoints/openai/image_api_utils.py
+++ b/vllm_omni/entrypoints/openai/image_api_utils.py
@@ -68,6 +68,50 @@ def encode_image_base64(image: PIL.Image.Image) -> str:
     return base64.b64encode(buffer.read()).decode("utf-8")
 
 
+def prepare_image_for_output_format(image: PIL.Image.Image, format: str) -> PIL.Image.Image:
+    fmt = format.lower()
+    if fmt not in {"jpg", "jpeg"}:
+        return image
+
+    if image.mode == "RGB":
+        return image
+
+    if image.mode in {"RGBA", "LA"} or (image.mode == "P" and "transparency" in image.info):
+        alpha_image = image.convert("RGBA")
+        flattened = PIL.Image.new("RGB", alpha_image.size, (255, 255, 255))
+        flattened.paste(alpha_image, mask=alpha_image.getchannel("A"))
+        return flattened
+
+    return image.convert("RGB")
+
+
+def encode_image_base64_with_compression(
+    image: PIL.Image.Image, format: str = "png", output_compression: int = 100
+) -> str:
+    """Encode a PIL image to a base64 image string.
+
+    Args:
+        image: PIL Image object.
+        format: Output image format, such as "png", "jpeg", or "webp".
+        output_compression: Compression level (0-100), where 100 keeps the
+            best quality for lossy formats and uses the lowest PNG compression.
+
+    Returns:
+        Base64-encoded image string.
+    """
+    buffer = io.BytesIO()
+    image = prepare_image_for_output_format(image, format)
+    save_kwargs = {}
+    if format in ("jpg", "jpeg", "webp"):
+        save_kwargs["quality"] = output_compression
+    elif format == "png":
+        save_kwargs["compress_level"] = max(0, min(9, 9 - output_compression // 11))
+
+    image.save(buffer, format=format, **save_kwargs)
+    buffer.seek(0)
+    return base64.b64encode(buffer.read()).decode("utf-8")
+
+
 def validate_layered_layers(layers: int | None) -> int | None:
     """Validate the Qwen-Image-Layered ``layers`` parameter."""
     if layers is None:
diff --git a/vllm_omni/entrypoints/openai/protocol/__init__.py b/vllm_omni/entrypoints/openai/protocol/__init__.py
index c73203cc4..c68f6f598 100644
--- a/vllm_omni/entrypoints/openai/protocol/__init__.py
+++ b/vllm_omni/entrypoints/openai/protocol/__init__.py
@@ -4,6 +4,10 @@
 from vllm_omni.entrypoints.openai.protocol.chat_completion import OmniChatCompletionStreamResponse
 from vllm_omni.entrypoints.openai.protocol.images import (
     ImageData,
+    ImageEditARDeltaChunk,
+    ImageEditImageChunk,
+    ImageEditStreamError,
+    ImageEditStreamResponse,
     ImageGenerationRequest,
     ImageGenerationResponse,
     ResponseFormat,
@@ -16,6 +20,10 @@ from vllm_omni.entrypoints.openai.protocol.videos import (
 
 __all__ = [
     "ImageData",
+    "ImageEditARDeltaChunk",
+    "ImageEditImageChunk",
+    "ImageEditStreamError",
+    "ImageEditStreamResponse",
     "ImageGenerationRequest",
     "ImageGenerationResponse",
     "ResponseFormat",
diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py
index 3ec2424c0..0325017b4 100644
--- a/vllm_omni/entrypoints/openai/protocol/audio.py
+++ b/vllm_omni/entrypoints/openai/protocol/audio.py
@@ -100,6 +100,50 @@ class OpenAICreateSpeechRequest(BaseModel):
             raise ValueError("'speaker_embedding' values must be finite (no NaN or Inf)")
         return v
 
+    @model_validator(mode="before")
+    @classmethod
+    def reject_higgs_audio_v2_unsupported_aliases(cls, data):
+        """Reject unsupported rich-input aliases for higgs_audio_v2 BEFORE pydantic strips them.
+
+        OpenAICreateSpeechRequest is a permissive schema (it accepts any extra
+        field that the OpenAI Speech API didn't promise to ban) but several
+        of those aliases pull a request out of the v1 higgs_audio_v2 scope.
+        Catching them at parse time lets the API return a deterministic 4xx
+        with a model-specific error message instead of silently dropping the
+        field and proceeding with a degraded request.
+        """
+        # Keys that are silently dropped by pydantic when posted to /v1/audio/speech
+        # (the schema doesn't declare them) but which a model-specific
+        # validator MUST be able to reject. Kept inline to avoid clashing
+        # with pydantic's ModelPrivateAttr discovery for class-level
+        # underscore-prefixed attributes.
+        higgs_audio_v2_reserved_keys = (
+            "messages",  # ChatML rich content (out of scope for v1)
+            "reference_audio",  # voice-cloning alias 1
+            "voice_prompt",  # voice-cloning alias 2
+            "speaker_audio",  # voice-cloning alias 3
+            "speakers",  # multi-speaker dialogue
+        )
+        if not isinstance(data, dict):
+            return data
+        model_id = data.get("model")
+        if not isinstance(model_id, str):
+            return data
+        # Match the "higgs_audio_v2" model_type label, the HF architecture id
+        # in pipeline_registry.hf_architectures, and the hyphenated HF repo id
+        # (e.g. "bosonai/higgs-audio-v2-generation-3B-base") by normalizing
+        # both underscores and hyphens out of the id before substring matching.
+        normalized = model_id.lower().replace("-", "").replace("_", "")
+        if "higgsaudiov2" not in normalized:
+            return data
+        offending = sorted(k for k in higgs_audio_v2_reserved_keys if k in data)
+        if offending:
+            raise ValueError(
+                "higgs_audio_v2 v1 does not support these rich-input fields: "
+                f"{offending}. Supply plain text via the 'input' field instead."
+            )
+        return data
+
     @model_validator(mode="after")
     def validate_embedding_constraints(self) -> "OpenAICreateSpeechRequest":
         if self.speaker_embedding is not None:
diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py
index 0fb22a548..34fee760b 100644
--- a/vllm_omni/entrypoints/openai/protocol/images.py
+++ b/vllm_omni/entrypoints/openai/protocol/images.py
@@ -8,7 +8,7 @@ for text-to-image generation, with vllm-omni specific extensions.
 """
 
 from enum import Enum
-from typing import Any
+from typing import Any, Literal
 
 from pydantic import BaseModel, Field, field_validator
 
@@ -32,6 +32,11 @@ class ImageGenerationRequest(BaseModel):
 
     # Required fields
     prompt: str = Field(..., description="Text description of the desired image(s)")
+    bot_task: str | None = Field(
+        None,
+        description="Task mode for the model (e.g., 'cot' enables chain-of-thought generation). "
+        "Only supported by specific diffusion models.",
+    )
 
     # OpenAI standard fields
     model: str | None = Field(
@@ -165,3 +170,43 @@ class ImageGenerationResponse(BaseModel):
     data: list[ImageData] = Field(..., description="Array of generated images")
     output_format: str = Field(None, description="The output format of the image generation")
     size: str = Field(None, description="The size of the image generated")
+    cot_output: str | None = Field(
+        None,
+        description="Chain-of-thought text output from the AR stage. "
+        "Only present for image editing (IT2I) with CoT-enabled models.",
+    )
+
+
+class ImageEditARDeltaChunk(BaseModel):
+    """Streaming chunk carrying a text delta from the image-edit AR stage."""
+
+    object: Literal["image.edit.chunk"] = "image.edit.chunk"
+    type: Literal["ar_delta"] = "ar_delta"
+    delta: str = Field(..., description="Text delta generated by the AR stage")
+    index: int = Field(default=0, description="Completion index for the AR stage output")
+    created: int = Field(..., description="Unix timestamp of when the stream was created")
+    model: str = Field(..., description="Model used for the image edit request")
+
+
+class ImageEditImageChunk(BaseModel):
+    """Streaming chunk carrying the final image-edit result."""
+
+    object: Literal["image.edit.chunk"] = "image.edit.chunk"
+    type: Literal["image"] = "image"
+    data: list[ImageData] = Field(..., description="Array of generated images")
+    output_format: str = Field(..., description="The output format of the image generation")
+    size: str = Field(..., description="The generated image size")
+    created: int = Field(..., description="Unix timestamp of when the stream was created")
+    model: str = Field(..., description="Model used for the image edit request")
+
+
+class ImageEditStreamError(BaseModel):
+    """Streaming error chunk emitted before the terminal [DONE] event."""
+
+    object: Literal["error"] = "error"
+    created: int = Field(..., description="Unix timestamp of when the stream was created")
+    model: str = Field(..., description="Model used for the image edit request")
+    error: dict[str, Any] = Field(..., description="OpenAI-compatible streaming error payload")
+
+
+ImageEditStreamResponse = ImageEditARDeltaChunk | ImageEditImageChunk | ImageEditStreamError
diff --git a/vllm_omni/entrypoints/openai/realtime_connection.py b/vllm_omni/entrypoints/openai/realtime_connection.py
index f2ce75df2..e04e8cdf9 100644
--- a/vllm_omni/entrypoints/openai/realtime_connection.py
+++ b/vllm_omni/entrypoints/openai/realtime_connection.py
@@ -120,7 +120,6 @@ class RealtimeConnection(VllmRealtimeConnection):
         full_text = ""
         prompt_token_ids_len = 0
         completion_tokens_len = 0
-        last_prompt_token_ids_len = 0  # detect Stage-0 segment rollover
         self._realtime_audio_ref = None
 
         # Coerce cumulative outputs to delta outputs; this ensures
@@ -139,40 +138,26 @@ class RealtimeConnection(VllmRealtimeConnection):
             )
 
             async for output in result_gen:
-                # Handle delta texts; this is very similar to the client from vLLM
-                if output.outputs and len(output.outputs) > 0:
+                stage_id = getattr(output, "stage_id", None)
+                if stage_id == 0 and output.outputs:
                     first_output = output.outputs[0]
                     new_token_ids = list(first_output.token_ids)
-                    new_tokens_len = len(new_token_ids)
-
-                    cur_prompt_token_ids_len = len(output.prompt_token_ids or [])
-                    # Stage-0 segment rollover: buffer_realtime_audio may yield
-                    # multiple TokensPrompt segments and the second segment's
-                    # prompt_token_ids include the first segment's decoded output.
-                    # Clear accumulated full_text at the boundary so the
-                    # carried-over prefix is not duplicated in the final
-                    # TranscriptionDone.text.
-                    stage_id = getattr(output, "stage_id", None)
-                    if stage_id == 0 and cur_prompt_token_ids_len > last_prompt_token_ids_len > 0:
-                        full_text = ""
-                    last_prompt_token_ids_len = cur_prompt_token_ids_len
-
-                    if not prompt_token_ids_len and output.prompt_token_ids:
-                        prompt_token_ids_len = len(output.prompt_token_ids)
-
-                    if new_tokens_len:
+                    if new_token_ids:
                         input_stream.put_nowait(new_token_ids)
 
+                    if output.prompt_token_ids:
+                        prompt_token_ids_len = max(
+                            prompt_token_ids_len,
+                            len(output.prompt_token_ids),
+                        )
+
                     delta_text = first_output.text or ""
                     full_text += delta_text
+                    completion_tokens_len += len(new_token_ids)
 
-                    # append output to input if there was any delta text
                     if delta_text:
                         await self.send(TranscriptionDelta(delta=delta_text))
 
-                    completion_tokens_len += new_tokens_len
-
-                # Handle audio chunking; this is Omni specific
                 audio_chunks, sample_rate = self._extract_audio_chunks(output)
 
                 for chunk in audio_chunks:
diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py
index 66a86e7c9..f98732984 100644
--- a/vllm_omni/entrypoints/openai/serving_chat.py
+++ b/vllm_omni/entrypoints/openai/serving_chat.py
@@ -4,6 +4,7 @@ import json
 import time
 import uuid
 from collections.abc import AsyncGenerator, AsyncIterator, Callable
+from dataclasses import fields, is_dataclass
 from datetime import datetime, timedelta, timezone
 from io import BytesIO
 from typing import Any, Final, cast
@@ -86,9 +87,15 @@ from vllm.utils.collection_utils import as_list
 from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
-from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers
+from vllm_omni.entrypoints.openai.image_api_utils import encode_image_base64_with_compression, validate_layered_layers
 from vllm_omni.entrypoints.openai.protocol import OmniChatCompletionStreamResponse
 from vllm_omni.entrypoints.openai.protocol.audio import AudioResponse, CreateAudio
+from vllm_omni.entrypoints.openai.protocol.images import (
+    ImageData,
+    ImageEditARDeltaChunk,
+    ImageEditImageChunk,
+    ImageEditStreamError,
+)
 from vllm_omni.entrypoints.openai.stage_params import (
     build_stage_sampling_params_list,
     get_default_sampling_params_list,
@@ -404,8 +411,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         engine_prompt_image = None
 
                 # Override the prompts produced by chat-template preprocessing.
+                is_img2img = engine_prompt_image is not None
                 tprompt: OmniTextPrompt = {"prompt": extracted_prompt}
-                if engine_prompt_image:
+                if is_img2img:
                     tprompt["modalities"] = ["img2img"]
                 else:
                     tprompt["modalities"] = ["image"]
@@ -422,6 +430,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     mm_processor_kwargs["target_h"] = height
                 if width is not None:
                     mm_processor_kwargs["target_w"] = width
+                # Pass output modalities so model-specific MM processors can
+                # detect image-generation requests and apply their own prompt
+                # rewrites (e.g. query-token expansion, placeholder injection).
+                mm_processor_kwargs["modalities"] = ["img2img"] if is_img2img else ["image"]
                 tprompt["mm_processor_kwargs"] = mm_processor_kwargs
                 if engine_prompt_image is not None:
                     tprompt["multi_modal_data"] = engine_prompt_image
@@ -436,6 +448,21 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         for k, v in engine_prompt_image.items()
                     }
 
+                # Forward image_gen extra_args from diffusion-stage sampling
+                # params into the prompt dict so that thinker2imagegen can
+                # access them (e.g. byte5_text for ByT5 glyph rendering).
+                if hasattr(request, "sampling_params_list") and request.sampling_params_list:
+                    for sp_raw in request.sampling_params_list:
+                        ea = (
+                            (sp_raw.get("extra_args") or {})
+                            if isinstance(sp_raw, dict)
+                            else (getattr(sp_raw, "extra_args", None) or {})
+                        )
+                        ig = ea.get("image_gen") or {}
+                        if ig:
+                            tprompt["image_gen_extra_args"] = ig
+                            break
+
                 engine_prompts = [tprompt]
                 # Store height/width for applying to diffusion stage sampling params later
                 _image_gen_height = height
@@ -764,13 +791,31 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
         return new_messages
 
-    def _to_sampling_params_list(self, sampling_params_list: list[dict]) -> list[SamplingParams]:
-        final_sampling_params_list = []
-        for sampling_params in sampling_params_list:
+    def _to_sampling_params_list(self, sampling_params_list: list[dict]) -> list[Any]:
+        """Convert request dicts to stage-typed sampling params objects.
+
+        For diffusion stages, build ``OmniDiffusionSamplingParams`` so
+        downstream ``StageDiffusionClient._sampling_params_to_dict`` (which
+        requires a dataclass) works. For LLM stages build ``SamplingParams``.
+        """
+        stage_configs = list(getattr(self.engine_client, "stage_configs", []) or [])
+        final_sampling_params_list: list[Any] = []
+        for idx, sampling_params in enumerate(sampling_params_list):
+            stage_type = get_stage_type(stage_configs[idx]) if idx < len(stage_configs) else "llm"
+            target_cls = OmniDiffusionSamplingParams if stage_type == "diffusion" else SamplingParams
             if isinstance(sampling_params, dict):
-                final_sampling_params_list.append(SamplingParams(**sampling_params))
-            elif isinstance(sampling_params, SamplingParams):
+                final_sampling_params_list.append(target_cls(**sampling_params))
+            elif isinstance(sampling_params, target_cls):
                 final_sampling_params_list.append(sampling_params)
+            elif isinstance(sampling_params, SamplingParams | OmniDiffusionSamplingParams):
+                # Cross-typed (e.g. user passed SamplingParams but this is a
+                # diffusion stage) — rebuild via a dict round-trip so we end
+                # up with the correct target class.
+                as_dict = {
+                    f.name: getattr(sampling_params, f.name)
+                    for f in (fields(sampling_params) if is_dataclass(sampling_params) else [])
+                } or sampling_params.__dict__
+                final_sampling_params_list.append(target_cls(**as_dict))
             else:
                 raise ValueError(f"Invalid sampling params: {sampling_params}")
         return final_sampling_params_list
@@ -2275,7 +2320,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         layers = extra_body.get("layers")
         resolution = extra_body.get("resolution")
         bot_task = extra_body.get("bot_task")
-        sys_type = extra_body.get("sys_type")
+        use_system_prompt = extra_body.get("use_system_prompt") or extra_body.get("sys_type")
         custom_system_prompt = extra_body.get("system_prompt")
 
         engine_prompt_data: dict[str, Any] | None = None
@@ -2289,24 +2334,29 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
         prompt_token_ids: list[int] | None = None
         system_prompt_type: str | None = None
-        if bot_task is not None or sys_type is not None or custom_system_prompt is not None:
+        build_kwargs: dict[str, Any] = {}
+        ar_stop_token_ids: list[int] | None = None
+
+        if bot_task is not None or use_system_prompt is not None or custom_system_prompt is not None:
             from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
                 build_prompt,
                 build_prompt_tokens,
+                resolve_stop_token_ids,
             )
 
             build_kwargs: dict[str, Any] = {
                 "task": "it2i" if reference_images else "t2i",
-                "sys_type": sys_type,
+                "sys_type": use_system_prompt,
                 "custom_system_prompt": custom_system_prompt,
                 "num_images": len(reference_images) if reference_images else 1,
             }
+
             if bot_task is not None:
                 build_kwargs["bot_task"] = bot_task
             elif "bot_task" in extra_body:
                 # Explicit None from the caller is plain-mode; omitted lets
                 # each task fall back to its default trigger.
-                build_kwargs["bot_task"] = None
+                build_kwargs["bot_task"] = extra_body["bot_task"]
             if tokenizer is not None:
                 # Feed segment-tokenized prompt_token_ids so AR matches HF
                 # apply_chat_template byte-for-byte (engine BPE would merge
@@ -2320,6 +2370,19 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 engine_prompt_data = {"image": reference_images[0]}
                 modalities = ["image"]
 
+            ar_task = "it2i" if reference_images else "t2i"
+            # ar_image_size: None -> need_ratio=True (AR predicts ratio);
+            # explicit size -> need_ratio=False (AR stops at terminator).
+            ar_image_size: str | None = None
+            if height is not None and width is not None:
+                ar_image_size = f"{width}x{height}"
+            ar_stop_token_ids = resolve_stop_token_ids(
+                task=ar_task,
+                bot_task=bot_task,
+                tokenizer=tokenizer,
+                image_size=ar_image_size,
+            )
+
         engine_prompt: OmniTextPrompt = {"prompt": prompt}
         if prompt_token_ids is not None:
             engine_prompt["prompt_token_ids"] = prompt_token_ids
@@ -2365,6 +2428,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             stage_type = get_stage_type(stage_cfg)
             default_stage_params = sampling_params_list[idx]
 
+            # AR stop tokens: use stage_type=="llm" instead of comprehension_idx
+            # (None for DictConfig where is_comprehension is nested in engine_args).
+            if stage_type == "llm" and ar_stop_token_ids is not None:
+                default_stage_params.stop_token_ids = ar_stop_token_ids
+
             if (
                 comprehension_idx is not None
                 and idx == comprehension_idx
@@ -2411,17 +2479,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
         return engine_prompt, sampling_params_list
 
-    async def generate_diffusion_images(
+    def _prepare_diffusion_image_request(
         self,
         *,
         prompt: str,
         extra_body: dict[str, Any] | None = None,
         reference_images: list[str] | None = None,
-        request_id: str | None = None,
-    ) -> tuple[list[Image.Image], dict[str, Any], float] | ErrorResponse:
-        """Generate diffusion images and return raw images plus generation stats."""
-        if request_id is None:
-            request_id = f"chatcmpl-{uuid.uuid4().hex[:16]}"
+    ) -> tuple[Any, OmniTextPrompt, OmniDiffusionSamplingParams, list[Image.Image]] | ErrorResponse:
         if extra_body is None:
             extra_body = {}
         if reference_images is None:
@@ -2496,9 +2560,44 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         status_code=400,
                     )
 
+        return engine, gen_prompt, gen_params, pil_images
+
+    async def generate_diffusion_images(
+        self,
+        *,
+        prompt: str,
+        extra_body: dict[str, Any] | None = None,
+        reference_images: list[str] | None = None,
+        request_id: str | None = None,
+        stream: bool = False,
+        model: str | None = None,
+        output_format: str = "png",
+        output_compression: int = 100,
+        size: str = "auto",
+    ) -> tuple[list[Image.Image], dict[str, Any], float, str | None] | ErrorResponse | AsyncIterator[str]:
+        """Generate diffusion images and return raw images plus generation stats."""
+        if request_id is None:
+            request_id = f"chatcmpl-{uuid.uuid4().hex[:16]}"
+
+        prepared = self._prepare_diffusion_image_request(
+            prompt=prompt,
+            extra_body=extra_body,
+            reference_images=reference_images,
+        )
+        if isinstance(prepared, ErrorResponse):
+            return prepared
+        engine, gen_prompt, gen_params, pil_images = prepared
+        if extra_body is None:
+            extra_body = {}
+
         if isinstance(engine, AsyncOmni):
             diffusion_engine = cast(AsyncOmni, engine)
             stage_configs = getattr(diffusion_engine, "stage_configs", None) or []
+            if stream and len(stage_configs) <= 1:
+                return self._create_error_response(
+                    "stream=true is only supported for multi-stage image editing pipelines.",
+                    status_code=400,
+                )
             if len(stage_configs) > 1:
                 # Pull tokenizer from the comprehension (AR) stage so we can
                 # build HF byte-for-byte prompt_token_ids in the helper. If
@@ -2523,15 +2622,31 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 engine_prompt = gen_prompt
                 sampling_params_list = [gen_params]
 
-            result = None
-            async for output in diffusion_engine.generate(
+            sampling_params_list = coerce_param_message_types(sampling_params_list, stream)
+            result_generator = diffusion_engine.generate(
                 prompt=engine_prompt,
                 sampling_params_list=sampling_params_list,
                 request_id=request_id,
-            ):
+            )
+            if stream:
+                return self._stream_diffusion_image_chunks(
+                    result_generator,
+                    model=model or self._diffusion_model_name,
+                    output_format=output_format,
+                    output_compression=output_compression,
+                    size=size,
+                )
+
+            result = None
+            async for output in result_generator:
                 result = output
             if result is None:
                 return self._create_error_response("No output generated from AsyncOmni", status_code=500)
+        elif stream:
+            return self._create_error_response(
+                "Streaming image edits require a multi-stage AsyncOmni engine.",
+                status_code=400,
+            )
         else:
             result = await engine.generate(
                 prompt=gen_prompt,
@@ -2542,15 +2657,103 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         images = getattr(result.request_output, "images", [])
         stage_durations = result.stage_durations
         peak_memory_mb = result.peak_memory_mb
+        cot_output = None
 
+        req_out = getattr(result, "request_output", None)
+        if req_out:
+            prompt_obj = getattr(req_out, "prompt", None)
+            if isinstance(prompt_obj, dict):
+                extra = prompt_obj.get("extra", {})
+                if isinstance(extra, dict):
+                    ar_text = extra.get("ar_generated_text")
+                    if isinstance(ar_text, str) and ar_text.strip():
+                        cot_output = ar_text
+
+        return self._flatten_diffusion_images(images), stage_durations, peak_memory_mb, cot_output
+
+    async def _stream_diffusion_image_chunks(
+        self,
+        result_generator: AsyncIterator[OmniRequestOutput],
+        *,
+        model: str,
+        output_format: str,
+        output_compression: int,
+        size: str,
+    ) -> AsyncIterator[str]:
+        """Yield image edit SSE chunks from multi-stage diffusion outputs."""
+        created = int(time.time())
+        emitted_image = False
+        try:
+            async for output in result_generator:
+                final_output_type = getattr(output, "final_output_type", None)
+                stage_id = getattr(output, "stage_id", None)
+                if final_output_type == "text" and stage_id == 0:
+                    request_output = output.request_output
+                    for completion in request_output.outputs:
+                        text = completion.text or ""
+                        if not text:
+                            continue
+                        chunk = ImageEditARDeltaChunk(
+                            delta=text,
+                            index=completion.index,
+                            created=created,
+                            model=model,
+                        )
+                        yield f"data: {chunk.model_dump_json()}\n\n"
+                elif final_output_type == "image":
+                    images = self._flatten_diffusion_images(getattr(output.request_output, "images", []))
+                    if not images:
+                        raise RuntimeError("Streaming image edit produced an empty final image output.")
+                    image_data = [
+                        ImageData(
+                            b64_json=encode_image_base64_with_compression(
+                                img,
+                                format=output_format,
+                                output_compression=output_compression,
+                            ),
+                            revised_prompt=None,
+                        )
+                        for img in images
+                    ]
+                    chunk = ImageEditImageChunk(
+                        data=image_data,
+                        output_format=output_format,
+                        size=size,
+                        created=created,
+                        model=model,
+                    )
+                    yield f"data: {chunk.model_dump_json()}\n\n"
+                    emitted_image = True
+            if not emitted_image:
+                raise RuntimeError("Streaming image edit completed without a final image output.")
+        except EngineDeadError as exc:
+            logger.error("EngineDeadError during streaming image edit: %s", exc)
+            data = self.create_streaming_error_response(exc)
+            yield f"data: {data}\n\n"
+            terminate_if_errored(engine=self.engine_client)
+        except Exception as exc:
+            logger.exception("Streaming image edit failed: %s", exc)
+            chunk = ImageEditStreamError(
+                created=created,
+                model=model,
+                error={
+                    "message": str(exc),
+                    "type": "server_error",
+                    "code": 500,
+                },
+            )
+            yield f"data: {chunk.model_dump_json()}\n\n"
+        yield "data: [DONE]\n\n"
+
+    @staticmethod
+    def _flatten_diffusion_images(images: Any) -> list[Image.Image]:
         flat_images: list[Image.Image] = []
-        for item in images:
+        for item in images or []:
             if isinstance(item, list):
                 flat_images.extend(item)
             else:
                 flat_images.append(item)
-
-        return flat_images, stage_durations, peak_memory_mb
+        return flat_images
 
     async def _create_diffusion_chat_completion(
         self,
diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py
index 3b7a2e747..e2a4b5af8 100644
--- a/vllm_omni/entrypoints/openai/serving_speech.py
+++ b/vllm_omni/entrypoints/openai/serving_speech.py
@@ -1,5 +1,6 @@
 import asyncio
 import base64
+import hashlib
 import io
 import json
 import math
@@ -68,18 +69,29 @@ _COVO_AUDIO_MODEL_STAGES = {"fused_thinker_talker"}
 _VOXCPM2_TTS_MODEL_STAGES = {"latent_generator"}
 _MING_TTS_MODEL_STAGES = {"ming_tts"}
 _MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"}
+_HIGGS_AUDIO_V2_TTS_MODEL_STAGES = {"higgs_audio_v2"}
+_GLM_TTS_MODEL_STAGES = {"glm_tts"}
 _TTS_MODEL_STAGES: set[str] = (
     _VOXTRAL_TTS_MODEL_STAGES
     | _QWEN3_TTS_MODEL_STAGES
     | _FISH_TTS_MODEL_STAGES
     | _COSYVOICE3_TTS_MODEL_STAGES
     | _OMNIVOICE_TTS_MODEL_STAGES
+    | _HIGGS_AUDIO_V2_TTS_MODEL_STAGES
     | _COVO_AUDIO_MODEL_STAGES
     | _VOXCPM2_TTS_MODEL_STAGES
     | _MING_TTS_MODEL_STAGES
     | _MOSS_TTS_MODEL_STAGES
+    | _GLM_TTS_MODEL_STAGES
 )
-_SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = {"fish_tts", "qwen3_tts", "voxtral_tts", "cosyvoice3", "voxcpm2"}
+_SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = {
+    "fish_tts",
+    "qwen3_tts",
+    "voxtral_tts",
+    "cosyvoice3",
+    "voxcpm2",
+    "higgs_audio_v2",
+}
 _TTS_LANGUAGES: set[str] = {
     "Auto",
     "Chinese",
@@ -184,6 +196,55 @@ def _validate_path_within_directory(file_path: Path, directory: Path) -> bool:
         return False
 
 
+def _conditioning_cache_salt(request, tts_params: dict | None = None) -> str:
+    """Stable hash of the real Stage 0 conditioning for the prefix cache.
+
+    The talker's vLLM prompt is placeholder token ids; the real inputs are
+    rebuilt from text / ref_audio / ref_text into inputs_embeds. vLLM hashes
+    token ids (folded with cache_salt) for prefix caching, so without a salt
+    every request collides and a hit could reuse KV from a semantically
+    different input. Tying the salt to the conditioning keeps a hit safe:
+    identical conditioning may share the prefix, any difference never does.
+
+    Raw request fields alone are not enough for uploaded voices: the request
+    only carries the voice *name* (ref_audio/ref_text/task_type are resolved
+    from stored voice data into ``tts_params``). Delete + re-upload under the
+    same name leaves every raw field identical, so the resolved conditioning
+    must also be folded in. ``voice_created_at`` bumps on every (re-)upload,
+    which uniquely identifies the resolved reference artifact together with
+    the voice name; the decoded ref_audio array itself need not be hashed.
+    """
+    h = hashlib.sha256()
+    for part in (
+        request.input,
+        request.task_type,
+        request.language,
+        request.voice,
+        request.ref_text,
+        request.ref_audio,
+        request.instructions,
+        request.x_vector_only_mode,
+        request.speaker_embedding,
+    ):
+        h.update(b"\x00")
+        if part is not None:
+            h.update(repr(part).encode("utf-8"))
+    # Fold resolved conditioning that is auto-derived for uploaded voices and
+    # absent from the raw request.
+    for key in (
+        "voice_created_at",
+        "task_type",
+        "speaker",
+        "ref_text",
+        "x_vector_only_mode",
+    ):
+        h.update(b"\x00")
+        value = tts_params.get(key) if tts_params is not None else None
+        if value is not None:
+            h.update(repr(value).encode("utf-8"))
+    return h.hexdigest()[:32]
+
+
 class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
     _diffusion_mode: bool = False
     _tts_executor: ThreadPoolExecutor | None = None
@@ -335,6 +396,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         # Determine TTS model type or None
         self._tts_model_type = self._detect_tts_model_type()
 
+        # GLM-TTS lazy-cached resources (populated on first GLM-TTS request)
+        self._glm_tts_text_tokenizer: object | None = None
+        self._glm_tts_text_frontend: object | None = None
+
         # Cache TTS configuration values (computed once, reused per request)
         self._max_instructions_length = self._compute_max_instructions_length()
 
@@ -499,6 +564,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             return "ming_flash_omni_tts"
         if model_stage in _MOSS_TTS_MODEL_STAGES:
             return "moss_tts_nano"
+        if model_stage in _HIGGS_AUDIO_V2_TTS_MODEL_STAGES:
+            return "higgs_audio_v2"
+        if model_stage in _GLM_TTS_MODEL_STAGES:
+            return "glm_tts"
         return None
 
     def _compute_max_instructions_length(self) -> int:
@@ -528,6 +597,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             # is no static speaker list to surface here.
             return set()
         try:
+            if self._tts_model_type == "glm_tts":
+                return set()
             if self._tts_model_type == "voxcpm2":
                 return {"default"}
             if self._tts_model_type == "voxtral_tts":
@@ -589,8 +660,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
     def _estimate_prompt_len(self, tts_params: dict[str, Any]) -> int:
         """Estimate prompt length so the placeholder matches model-side embeddings."""
         try:
-            from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
-                Qwen3TTSTalkerForConditionalGeneration,
+            from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
+                Qwen3TTSPromptEmbedsBuilder,
             )
 
             if self._tts_tokenizer is None:
@@ -605,7 +676,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             hf_config = self.engine_client.model_config.hf_config
             talker_config = hf_config.talker_config
             task_type = (tts_params.get("task_type") or ["CustomVoice"])[0]
-            return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
+            return Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
                 additional_information=tts_params,
                 task_type=task_type,
                 tokenize_prompt=lambda t: self._tts_tokenizer(t, padding=False)["input_ids"],
@@ -785,12 +856,21 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
         voice_lower = request.voice.lower()
         if voice_lower not in self.uploaded_speakers:
-            if self._tts_model_type in ("cosyvoice3", "fish_tts", "omnivoice", "moss_tts_nano"):
+            if self._tts_model_type in (
+                "cosyvoice3",
+                "fish_tts",
+                "omnivoice",
+                "moss_tts_nano",
+                "glm_tts",
+                "higgs_audio_v2",
+            ):
                 label = {
                     "cosyvoice3": "CosyVoice3",
                     "fish_tts": "Fish Speech",
                     "omnivoice": "OmniVoice",
                     "moss_tts_nano": "MOSS-TTS-Nano",
+                    "higgs_audio_v2": "Higgs-Audio V2",
+                    "glm_tts": "GLM-TTS",
                 }.get(self._tts_model_type, self._tts_model_type)
                 return (
                     f"Unknown voice '{request.voice}'. {label} has no "
@@ -1135,6 +1215,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             return self._validate_ming_tts_request(request)
         if self._tts_model_type == "moss_tts_nano":
             return self._validate_moss_tts_request(request)
+        if self._tts_model_type == "glm_tts":
+            return self._validate_glm_tts_request(request)
         return self._validate_qwen_tts_request(request)
 
     def _voxcpm2_encode(self, text: str) -> list[int]:
@@ -1383,6 +1465,167 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         params["prompt_audio_array"] = [[wav_list, sr]]
         return params
 
+    def _validate_higgs_audio_v2_request(self, request: OpenAICreateSpeechRequest) -> str | None:
+        """Validate higgs_audio_v2 request parameters. Returns error message or None.
+
+        Accepted: plain text -> speech, or shallow voice clone via ``ref_audio``
+        + ``ref_text`` (both required together). Still out of scope: preset
+        ``voice``/``speaker`` selection, ``x_vector_only_mode`` /
+        ``speaker_embedding`` helpers, ``task_type``/``language``/
+        ``instructions``/``speed`` overrides, and multi-speaker ``[SPEAKERn]``
+        tags inside the input body.
+        """
+        from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
+            MULTI_SPEAKER_TAG_PATTERN,
+        )
+
+        if not request.input or not request.input.strip():
+            return "higgs_audio_v2: input text cannot be empty"
+
+        # Voice clone: ref_audio and ref_text must come together.
+        if request.ref_audio is not None and not request.ref_text:
+            return (
+                "higgs_audio_v2 voice clone requires both 'ref_audio' and "
+                "'ref_text'; received ref_audio without ref_text"
+            )
+        if request.ref_text and request.ref_audio is None:
+            return (
+                "higgs_audio_v2 voice clone requires both 'ref_audio' and "
+                "'ref_text'; received ref_text without ref_audio"
+            )
+
+        if request.x_vector_only_mode is not None:
+            return "higgs_audio_v2 v1 does not support 'x_vector_only_mode' (voice-cloning helper field)"
+        if request.speaker_embedding is not None:
+            return "higgs_audio_v2 v1 does not support 'speaker_embedding' (voice-cloning helper field)"
+        if request.voice and request.ref_audio is None:
+            # _apply_uploaded_speaker runs before this validator; if voice was
+            # an uploaded speaker, ref_audio is now populated and ref_text is
+            # backfilled from the speaker entry. A bare voice= with no
+            # ref_audio means the name didn't resolve to an uploaded speaker
+            # (and higgs has no built-in preset voices).
+            return (
+                "higgs_audio_v2 v1 does not support 'voice'/'speaker' selection for built-in voices; "
+                f"upload a voice first via POST /v1/audio/voices, or use ref_audio + ref_text. "
+                f"Got voice={request.voice!r}"
+            )
+        if request.instructions:
+            return (
+                "higgs_audio_v2 v1 does not support 'instructions' (voice "
+                "style/emotion control); supply plain text instead"
+            )
+        if request.task_type is not None:
+            return "higgs_audio_v2 v1 does not support 'task_type'; the model is single-mode plain text -> speech"
+        if request.language is not None:
+            return (
+                "higgs_audio_v2 v1 does not accept 'language' overrides; the model infers language from the input text"
+            )
+        if request.speed is not None and request.speed != 1.0:
+            return (
+                "higgs_audio_v2 v1 does not support 'speed' adjustments; the audio is rendered at native rate (24 kHz)"
+            )
+
+        if MULTI_SPEAKER_TAG_PATTERN.search(request.input):
+            return "higgs_audio_v2 v1 does not support multi-speaker [SPEAKERn] tags; remove the tag from the input"
+
+        if request.max_new_tokens is not None:
+            if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN:
+                return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}"
+
+    def _validate_glm_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
+        """Validate GLM-TTS request — requires ref_audio for voice cloning."""
+        if not request.input or not request.input.strip():
+            return "Input text cannot be empty"
+
+        if request.ref_audio is None:
+            return "GLM-TTS requires 'ref_audio' for zero-shot voice cloning"
+        fmt_err = self._validate_ref_audio_format(request.ref_audio)
+        if fmt_err:
+            return fmt_err
+        if not request.ref_text or not request.ref_text.strip():
+            return "GLM-TTS voice cloning requires 'ref_text' (transcript of the reference audio)"
+
+        if request.max_new_tokens is not None:
+            if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN:
+                return f"max_new_tokens must be >= {_TTS_MAX_NEW_TOKENS_MIN}"
+            if request.max_new_tokens > _TTS_MAX_NEW_TOKENS_MAX:
+                return f"max_new_tokens cannot exceed {_TTS_MAX_NEW_TOKENS_MAX}"
+        return None
+
+    async def _build_higgs_audio_v2_params(self, request: OpenAICreateSpeechRequest):
+        """Build prompt_token_ids for higgs_audio_v2 via the upstream processor.
+
+        Plain-text path: runs ``build_plain_text_prompt`` and returns the
+        token-only prompt. Voice-clone path (``ref_audio`` + ``ref_text``):
+        resolves the reference clip via ``_resolve_ref_audio``, runs
+        ``build_voice_clone_prompt`` (which encodes the clip through HF's
+        ``HiggsAudioV2TokenizerModel`` loaded from the k2-fsa/OmniVoice
+        ``audio_tokenizer/`` subdirectory), and attaches the encoded
+        ``audio_input_ids`` + ``audio_input_ids_mask`` tensors via
+        ``additional_information`` so the talker substitutes them at the
+        prompt-side audio placeholders.
+        """
+        from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
+            build_plain_text_prompt,
+            build_voice_clone_prompt,
+            input_ids_to_python_list,
+        )
+
+        processor = await self._resolve_higgs_audio_v2_processor()
+
+        if request.ref_audio is None:
+            inputs = await asyncio.to_thread(build_plain_text_prompt, processor, request.input)
+            prompt_token_ids = input_ids_to_python_list(inputs)
+            return tokens_input(prompt_token_ids=prompt_token_ids)
+
+        wav_list, sr = await self._resolve_ref_audio(request.ref_audio)
+        wav = np.asarray(wav_list, dtype=np.float32)
+        out = await asyncio.to_thread(
+            build_voice_clone_prompt,
+            processor,
+            request.input,
+            wav,
+            int(sr),
+            request.ref_text or "",
+        )
+        prompt = tokens_input(prompt_token_ids=out["prompt_token_ids"])
+        # Pass tensors at the top level of additional_information (NOT list-
+        # wrapped). ``vllm_omni.data_entry_keys.serialize_payload`` routes
+        # bare ``torch.Tensor`` values through ``_serialize_tensor``; a list
+        # containing tensors would fall into the ``list_data`` field which
+        # msgspec cannot serialize and the tensors would be dropped over the
+        # process boundary (silent voice-clone failure).
+        prompt["additional_information"] = {
+            "audio_input_ids": out["audio_input_ids"],
+            "audio_input_ids_mask": out["audio_input_ids_mask"],
+        }
+        return prompt
+
+    async def _resolve_higgs_audio_v2_processor(self):
+        """Lazy-load the AutoProcessor for higgs_audio_v2 (once per serving instance)."""
+        cached = getattr(self, "_higgs_audio_v2_processor", None)
+        if cached is not None:
+            return cached
+
+        from transformers import AutoProcessor
+
+        model_path = None
+        for stage in self.engine_client.stage_configs:
+            model_path = getattr(getattr(stage, "engine_args", None), "model", None)
+            if model_path:
+                break
+        if model_path is None:
+            # Fallback: the orchestrator stores the served model id on the engine
+            # itself (set by AsyncOmniEngine.__init__). Stage-level engine_args
+            # may not surface ``model`` when the deploy yaml doesn't set it per
+            # stage (the CLI-passed model id is the single source of truth).
+            model_path = getattr(self.engine_client, "model", None)
+        if model_path is None:
+            raise RuntimeError("higgs_audio_v2 serving could not resolve the model path from the engine stage configs")
+        processor = AutoProcessor.from_pretrained(model_path)
+        self._higgs_audio_v2_processor = processor
+        return processor
+
     def _validate_fish_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
         """Validate Fish Speech request parameters. Returns error message or None."""
         if not request.input or not request.input.strip():
@@ -1467,6 +1710,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         request_id: str,
         response_format: str = "pcm",
         raw_request: Request | None = None,
+        request_start_s: float | None = None,
     ):
         """Generate audio chunks for streaming response.
 
@@ -1487,6 +1731,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         prev_count = 0
         sample_rate_val = 24000
         first_chunk = True
+        first_audio_chunk_s: float | None = None
+        stream_start_s = request_start_s if request_start_s is not None else time.perf_counter()
 
         try:
             async for res in generator:
@@ -1538,11 +1784,40 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                         stream_format="audio",
                         base64_encode=False,
                     )
+                    if first_audio_chunk_s is None:
+                        first_audio_chunk_s = time.perf_counter()
                     yield self.create_audio(audio_obj).audio_data
+            total_ms = (time.perf_counter() - stream_start_s) * 1000.0
+            if first_audio_chunk_s is not None:
+                first_chunk_ms = (first_audio_chunk_s - stream_start_s) * 1000.0
+                logger.info(
+                    "[SpeechE2E] request_id=%s stream=true status=ok total_ms=%.2f first_chunk_ms=%.2f",
+                    request_id,
+                    total_ms,
+                    first_chunk_ms,
+                )
+            else:
+                logger.info(
+                    "[SpeechE2E] request_id=%s stream=true status=ok total_ms=%.2f first_chunk_ms=NA",
+                    request_id,
+                    total_ms,
+                )
         except asyncio.CancelledError:
+            total_ms = (time.perf_counter() - stream_start_s) * 1000.0
+            logger.info(
+                "[SpeechE2E] request_id=%s stream=true status=cancelled total_ms=%.2f",
+                request_id,
+                total_ms,
+            )
             logger.info("Streaming request %s cancelled by client", request_id)
             raise
         except EngineDeadError as e:
+            total_ms = (time.perf_counter() - stream_start_s) * 1000.0
+            logger.error(
+                "[SpeechE2E] request_id=%s stream=true status=engine_dead total_ms=%.2f",
+                request_id,
+                total_ms,
+            )
             logger.error(
                 "EngineDeadError during streaming speech for %s: %s",
                 request_id,
@@ -1556,6 +1831,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 )
             raise
         except Exception as e:
+            total_ms = (time.perf_counter() - stream_start_s) * 1000.0
+            logger.exception(
+                "[SpeechE2E] request_id=%s stream=true status=error total_ms=%.2f error=%s",
+                request_id,
+                total_ms,
+                e,
+            )
             logger.exception("Streaming speech generation failed for %s: %s", request_id, e)
             raise
 
@@ -1898,6 +2180,105 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         )
         return sampling_params_list
 
+    # ---- GLM-TTS helpers ----
+
+    async def _build_glm_tts_prompt(
+        self,
+        request: OpenAICreateSpeechRequest,
+        *,
+        has_inline_ref_audio: bool = False,
+    ) -> dict[str, Any]:
+        """Build prompt for GLM-TTS.
+
+        Uses the multimodal processor path (same as CosyVoice3):
+        - prompt: synthesis text
+        - multi_modal_data["audio"]: (wav_samples, sr) reference audio
+        - mm_processor_kwargs["prompt_text"]: reference text transcript
+
+        AR preprocess() builds [PromptText | Text | BOA | PromptSpeechTokens + ATS].
+        DiT receives prompt_token, prompt_feat, embedding for conditioning.
+        """
+        # Voice cloning requires ref_audio + ref_text
+        if request.ref_audio is not None and request.ref_text:
+            wav_samples, sr = await self._resolve_ref_audio(request.ref_audio)
+            audio_data = (np.asarray(wav_samples, dtype=np.float32), int(sr))
+
+            mm_kwargs: dict[str, Any] = {
+                "prompt_text": request.ref_text,
+            }
+            if request.voice:
+                voice_lower = request.voice.lower()
+                if voice_lower in self.uploaded_speakers and not has_inline_ref_audio:
+                    mm_kwargs["voice_name"] = voice_lower
+                    mm_kwargs["voice_created_at"] = self._voice_created_at(voice_lower)
+
+            return {
+                "prompt": request.input,
+                "multi_modal_data": {
+                    "audio": audio_data,
+                },
+                "mm_processor_kwargs": mm_kwargs,
+                "additional_information": self._build_glm_tts_prefill_metadata(
+                    request.input,
+                    request.ref_text,
+                ),
+            }
+
+        raise ValueError("GLM-TTS requires ref_audio and ref_text for voice cloning.")
+
+    def _glm_tts_text_tokenizer_and_frontend(self):
+        from vllm_omni.model_executor.models.glm_tts.glm_tts import (
+            load_glm_tts_tokenizer,
+            resolve_glm_tts_tokenizer_path,
+        )
+        from vllm_omni.model_executor.models.glm_tts.text_frontend import GLMTTSTextFrontend
+
+        cached = self._glm_tts_text_tokenizer
+        if cached is None:
+            model_name_or_path = self.engine_client.model_config.model
+            tokenizer_path = getattr(self.engine_client.model_config, "tokenizer", None)
+            if tokenizer_path is None:
+                tokenizer_path = resolve_glm_tts_tokenizer_path(model_name_or_path)
+            cached = load_glm_tts_tokenizer(
+                tokenizer_path,
+                model_name_or_path=model_name_or_path,
+                trust_remote_code=bool(getattr(self.engine_client.model_config, "trust_remote_code", False)),
+            )
+            self._glm_tts_text_tokenizer = cached
+
+        frontend = self._glm_tts_text_frontend
+        if frontend is None:
+            frontend = GLMTTSTextFrontend()
+            self._glm_tts_text_frontend = frontend
+        return cached, frontend
+
+    def _estimate_glm_tts_text_token_len(self, text: str | None, *, add_trailing_space: bool = False) -> int:
+        """Estimate GLM-TTS normalized text length with the model tokenizer."""
+        cached, frontend = self._glm_tts_text_tokenizer_and_frontend()
+        text = text or ""
+        normalized = frontend.text_normalize(text) or text
+        normalized = normalized.strip()
+        if add_trailing_space and normalized:
+            normalized = f"{normalized} "
+        return max(1, len(cached.encode(normalized)))
+
+    def _build_glm_tts_prefill_metadata(self, text: str, prompt_text: str | None) -> dict[str, Any]:
+        """Build GLM-TTS processor length metadata for additional_information.
+
+        The model preprocess hook runs before postprocess can mirror MM kwargs
+        into additional_information, so these scalar fields must originate from
+        the request payload rather than runner-level mm_features.
+        """
+        text_len = self._estimate_glm_tts_text_token_len(text)
+        prompt_text_len = (
+            self._estimate_glm_tts_text_token_len(prompt_text, add_trailing_space=True) if prompt_text else 0
+        )
+        return {
+            "glm_tts_text_token_len": [text_len],
+            "glm_tts_prompt_text_token_len": [prompt_text_len],
+            "input_len": [prompt_text_len + text_len + 1],
+        }
+
     # ---- Ming-flash-omni standalone-talker (TTS) helpers ----
 
     def _build_ming_prompt(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]:
@@ -1963,7 +2344,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         # Resolve uploaded voice for non-Qwen3 models.
         # Qwen3 TTS has its own uploaded voice handling in _build_tts_params().
         has_inline_ref_audio = request.ref_audio is not None
-        if self._tts_model_type in ("fish_tts", "cosyvoice3", "moss_tts_nano"):
+        if self._tts_model_type in ("fish_tts", "cosyvoice3", "moss_tts_nano", "glm_tts"):
             err = self._apply_uploaded_speaker(request)
             if err:
                 raise ValueError(err)
@@ -2005,6 +2386,29 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         elif self._tts_model_type == "covo_audio":
             prompt = self._build_covo_audio_prompt(request)
             tts_params = {}
+        elif self._tts_model_type == "higgs_audio_v2":
+            # Explicit higgs_audio_v2 branch: do NOT fall through to the
+            # Qwen-style generic placeholder path; build prompt_token_ids
+            # directly via the upstream processor + the project's plain-text
+            # helper.
+            #
+            # Resolve uploaded voices first so a voice=<name> request (after
+            # POST /v1/audio/voices) populates ref_audio + ref_text from the
+            # cached speaker entry, mirroring the cosyvoice3 / voxcpm2 flow.
+            err = self._apply_uploaded_speaker(request)
+            if err:
+                raise ValueError(err)
+            validation_error = self._validate_higgs_audio_v2_request(request)
+            if validation_error:
+                raise ValueError(validation_error)
+            prompt = await self._build_higgs_audio_v2_params(request)
+            if request.voice:
+                voice_lower = request.voice.lower()
+                if voice_lower in self.uploaded_speakers and not has_inline_ref_audio:
+                    additional = prompt.setdefault("additional_information", {})
+                    additional["voice_name"] = voice_lower
+                    additional["voice_created_at"] = self._voice_created_at(voice_lower)
+            tts_params = {}
         elif self._tts_model_type == "voxcpm2":
             # voxcpm2 doesn't use `_apply_uploaded_speaker` because the prompt builder needs the
             # raw waveform tuple for prefill-length accounting, not a base64 data URL.
@@ -2041,6 +2445,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             elif self._tts_model_type == "cosyvoice3":
                 prompt = await self._build_cosyvoice3_prompt(request, has_inline_ref_audio=has_inline_ref_audio)
                 tts_params = {}
+            elif self._tts_model_type == "glm_tts":
+                prompt = await self._build_glm_tts_prompt(request, has_inline_ref_audio=has_inline_ref_audio)
+                tts_params = {}
             elif self._tts_model_type == "ming_flash_omni_tts":
                 prompt = self._build_ming_prompt(request)
                 tts_params = {}
@@ -2060,6 +2467,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     tts_params["seed"] = [sampling_params_list[0].seed]
                 prompt = tokens_input(prompt_token_ids=[1])
                 prompt["additional_information"] = tts_params
+                prompt["cache_salt"] = _conditioning_cache_salt(request, tts_params)
             else:
                 tts_params = self._build_tts_params(request)
                 # Resolve ref_audio (explicit or auto-set for uploaded voices)
@@ -2075,6 +2483,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 ph_len = await self._estimate_prompt_len_async(tts_params)
                 prompt = tokens_input(prompt_token_ids=[1] * ph_len)
                 prompt["additional_information"] = tts_params
+                prompt["cache_salt"] = _conditioning_cache_salt(request, tts_params)
         else:
             # Qwen omni models (Qwen3-Omni, Qwen2.5-Omni) use a "talker"
             # stage whose preprocess requires chat-templated tokens.  The
@@ -2112,6 +2521,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             model_type = "ming_flash_omni_tts"
         elif self._tts_model_type == "moss_tts_nano":
             model_type = "moss_tts_nano"
+        elif self._tts_model_type == "higgs_audio_v2":
+            model_type = "higgs_audio_v2"
+        elif self._tts_model_type == "glm_tts":
+            model_type = "glm_tts"
         elif self._is_tts:
             model_type = tts_params.get("task_type", ["unknown"])[0]
         else:
@@ -2129,6 +2542,57 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         if self._tts_model_type == "cosyvoice3" and sampling_params_list:
             sampling_params_list = self._apply_cosyvoice3_dynamic_tokens(sampling_params_list, request)
 
+        # GLM-TTS: set dynamic min/max tokens based on text length.
+        if self._tts_model_type == "glm_tts" and sampling_params_list:
+            import copy
+
+            sampling_params_list = copy.deepcopy(sampling_params_list)
+            glm_metadata = prompt.get("additional_information") if isinstance(prompt, dict) else None
+            text_len_value = None
+            if isinstance(glm_metadata, dict):
+                text_len_value = glm_metadata.get("glm_tts_text_token_len")
+                if isinstance(text_len_value, list) and text_len_value:
+                    text_len_value = text_len_value[0]
+            text_token_len = (
+                int(text_len_value)
+                if text_len_value is not None
+                else self._estimate_glm_tts_text_token_len(request.input)
+            )
+            hf_cfg = self.model_config.hf_config
+            min_ratio = getattr(hf_cfg, "min_token_text_ratio", 2)
+            max_ratio = getattr(hf_cfg, "max_token_text_ratio", 20)
+            stage_min_tokens = getattr(sampling_params_list[0], "min_tokens", None)
+            stage_max_tokens = getattr(sampling_params_list[0], "max_tokens", None)
+            cap_candidates = [int(cap) for cap in (stage_max_tokens, request.max_new_tokens) if cap is not None]
+            hard_cap = min(cap_candidates) if cap_candidates else None
+
+            min_tokens = max(1, int(text_token_len * min_ratio))
+            if stage_min_tokens is not None:
+                min_tokens = max(min_tokens, int(stage_min_tokens))
+            if hard_cap is not None:
+                min_tokens = min(min_tokens, hard_cap)
+
+            max_tokens = max(min_tokens, int(text_token_len * max_ratio))
+            if hard_cap is not None:
+                max_tokens = min(max_tokens, hard_cap)
+            sampling_params_list[0].min_tokens = min_tokens
+            sampling_params_list[0].max_tokens = max_tokens
+            seed = getattr(request, "seed", None)
+            if seed is not None:
+                sampling_params_list[0].seed = seed
+            logger.info(
+                "GLM-TTS dynamic tokens: text_tokens=%d, min_ratio=%s, max_ratio=%s, "
+                "stage_min=%s, stage_max=%s, request_max=%s, min_tokens=%d, max_tokens=%d",
+                text_token_len,
+                min_ratio,
+                max_ratio,
+                stage_min_tokens,
+                stage_max_tokens,
+                request.max_new_tokens,
+                min_tokens,
+                max_tokens,
+            )
+
         # Apply model-specific extra parameters
         if request.extra_params is not None and sampling_params_list:
             if not isinstance(request.extra_params, dict):
@@ -2349,19 +2813,21 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 request.input[:50] + "..." if len(request.input) > 50 else request.input,
                 "ref_audio" in prompt,
             )
-
+            if request.extra_params is not None and not isinstance(request.extra_params, dict):
+                raise ValueError("extra_params must be a JSON object/dict.")
+            extra = dict(request.extra_params or {})
+            if request.seed is not None:
+                extra["seed"] = request.seed
             # Apply extra_params from the request to sampling params
             sampling_params_list = self._diffusion_engine.default_sampling_params_list
-            if request.extra_params is not None:
-                if not isinstance(request.extra_params, dict):
-                    raise ValueError("extra_params must be a JSON object/dict.")
+            if extra:
                 import copy
 
                 sampling_params_list = copy.deepcopy(sampling_params_list)
                 if sampling_params_list[0].extra_args is None:
                     sampling_params_list[0].extra_args = {}
-                sampling_params_list[0].extra_args.update(request.extra_params)
-                logger.info("Applied extra_params to diffusion: %s", request.extra_params)
+                sampling_params_list[0].extra_args.update(extra)
+                logger.info("Applied extra_params to diffusion: %s", extra)
 
             generator = self._diffusion_engine.generate(
                 prompt=prompt,
@@ -2463,6 +2929,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             return error_check_ret
 
         request_id = f"speech-{random_uuid()}"
+        request_start_s = time.perf_counter()
         if raw_request:
             raw_request.state.request_metadata = RequestResponseMetadata(
                 request_id=request_id,
@@ -2495,20 +2962,58 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                         request_id,
                         response_format,
                         raw_request=raw_request,
+                        request_start_s=request_start_s,
                     ),
                     media_type=media_type,
                 )
 
             audio_bytes, media_type = await self._generate_audio_bytes(request, request_id=request_id)
+            total_ms = (time.perf_counter() - request_start_s) * 1000.0
+            logger.info(
+                "[SpeechE2E] request_id=%s stream=false status=ok total_ms=%.2f response_bytes=%d",
+                request_id,
+                total_ms,
+                len(audio_bytes) if isinstance(audio_bytes, (bytes, bytearray)) else len(str(audio_bytes)),
+            )
             return Response(content=audio_bytes, media_type=media_type)
 
         except asyncio.CancelledError:
+            total_ms = (time.perf_counter() - request_start_s) * 1000.0
+            logger.info(
+                "[SpeechE2E] request_id=%s stream=%s status=cancelled total_ms=%.2f",
+                request_id,
+                bool(request.stream),
+                total_ms,
+            )
             return self.create_error_response("Client disconnected")
         except (EngineGenerateError, EngineDeadError):
+            total_ms = (time.perf_counter() - request_start_s) * 1000.0
+            logger.error(
+                "[SpeechE2E] request_id=%s stream=%s status=engine_error total_ms=%.2f",
+                request_id,
+                bool(request.stream),
+                total_ms,
+            )
             raise  # Propagate to the global Omni exception handler
         except ValueError as e:
+            total_ms = (time.perf_counter() - request_start_s) * 1000.0
+            logger.warning(
+                "[SpeechE2E] request_id=%s stream=%s status=bad_request total_ms=%.2f error=%s",
+                request_id,
+                bool(request.stream),
+                total_ms,
+                e,
+            )
             return self.create_error_response(e)
         except Exception as e:
+            total_ms = (time.perf_counter() - request_start_s) * 1000.0
+            logger.exception(
+                "[SpeechE2E] request_id=%s stream=%s status=error total_ms=%.2f error=%s",
+                request_id,
+                bool(request.stream),
+                total_ms,
+                e,
+            )
             logger.exception("Speech generation failed: %s", e)
             return self.create_error_response(f"Speech generation failed: {e}")
 
diff --git a/vllm_omni/entrypoints/pd_utils.py b/vllm_omni/entrypoints/pd_utils.py
index 413d5d6b4..bdc721ab0 100644
--- a/vllm_omni/entrypoints/pd_utils.py
+++ b/vllm_omni/entrypoints/pd_utils.py
@@ -13,8 +13,6 @@ from typing import TYPE_CHECKING, Any
 if TYPE_CHECKING:
     from vllm import SamplingParams
 
-    from vllm_omni.entrypoints.omni_stage import OmniStage
-
 _DEFAULT_MOONCAKE_BOOTSTRAP_PORT = 25201
 
 logger = logging.getLogger(__name__)
@@ -131,7 +129,7 @@ class PDDisaggregationMixin:
         p_stage = self.stage_configs[p_id]
         d_stage = self.stage_configs[d_id]
 
-        def _get_kv_cfg(stage: "OmniStage") -> dict[str, Any]:
+        def _get_kv_cfg(stage: Any) -> dict[str, Any]:
             ea = stage.engine_args
             cfg = getattr(ea, "kv_transfer_config", None)
             if cfg is None:
diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py
index de7eb5f4c..031527895 100644
--- a/vllm_omni/entrypoints/utils.py
+++ b/vllm_omni/entrypoints/utils.py
@@ -183,7 +183,17 @@ def _filter_dict_like_object(obj: dict | Any) -> dict:
     result = {}
     filtered_keys = []
     for k, v in obj.items():
-        if _is_callable_value(v):
+        # Preserve class objects by converting to a fully qualified name string
+        # so callers that resolve via import path (e.g. custom_pipeline_args.pipeline_class)
+        # still work after OmegaConf round-trip.
+        if isinstance(v, type):
+            module = getattr(v, "__module__", None)
+            qualname = getattr(v, "__qualname__", getattr(v, "__name__", None))
+            if module and qualname and module != "builtins":
+                result[k] = f"{module}.{qualname}"
+            else:
+                result[k] = qualname
+        elif _is_callable_value(v):
             filtered_keys.append(str(k))
         else:
             result[k] = _convert_dataclasses_to_dict(v)
@@ -240,6 +250,13 @@ def _convert_dataclasses_to_dict(obj: Any) -> Any:
     # Note: This must come AFTER Counter check since Counter is a dict subclass
     if isinstance(obj, dict):
         return _filter_dict_like_object(obj)
+    # Preserve class objects by converting to a fully qualified name string.
+    if isinstance(obj, type):
+        module = getattr(obj, "__module__", None)
+        qualname = getattr(obj, "__qualname__", getattr(obj, "__name__", None))
+        if module and qualname and module != "builtins":
+            return f"{module}.{qualname}"
+        return qualname
     # Handle callable objects (functions, methods, etc.) - skip them
     # Note: This comes after dict/list checks to avoid misclassifying dict-like objects
     if callable(obj):
@@ -329,11 +346,16 @@ def resolve_model_config_path(model: str) -> str:
             except Exception as e:
                 raise ValueError(f"Failed to read config.json for model: {model}. Error: {e}") from e
         else:
-            raise ValueError(
-                f"Could not determine model_type for model: {model}. "
-                f"Model is not in standard transformers format and does not have model_index.json. "
-                f"Please ensure the model has proper configuration files with 'model_type' field"
-            )
+            # No config.json at repo root (e.g. GLM-TTS stores configs in
+            # subdirectories only).  Try matching against registered deploy
+            # YAML filenames before giving up.
+            model_type = _try_resolve_omni_model_type(model)
+            if model_type is None:
+                raise ValueError(
+                    f"Could not determine model_type for model: {model}. "
+                    f"Model is not in standard transformers format and does not have model_index.json. "
+                    f"Please ensure the model has proper configuration files with 'model_type' field"
+                )
 
     default_config_path = current_omni_platform.get_default_stage_config_path()
     if model_type in _DIFFUSERS_CLASS_TO_CONFIG:
@@ -612,7 +634,7 @@ def load_and_resolve_stage_configs(
         if not stage_configs:
             if default_stage_cfg_factory is not None:
                 default_stage_cfg = default_stage_cfg_factory()
-                stage_configs = create_config(default_stage_cfg)
+                stage_configs = create_config(_convert_dataclasses_to_dict(default_stage_cfg))
             else:
                 stage_configs = []
     elif stage_configs_path is None:
@@ -625,7 +647,7 @@ def load_and_resolve_stage_configs(
         if not stage_configs:
             if default_stage_cfg_factory is not None:
                 default_stage_cfg = default_stage_cfg_factory()
-                stage_configs = create_config(default_stage_cfg)
+                stage_configs = create_config(_convert_dataclasses_to_dict(default_stage_cfg))
             else:
                 stage_configs = []
     else:
diff --git a/vllm_omni/model_executor/models/common/alias_free_activation.py b/vllm_omni/model_executor/models/common/alias_free_activation.py
new file mode 100644
index 000000000..adb7508af
--- /dev/null
+++ b/vllm_omni/model_executor/models/common/alias_free_activation.py
@@ -0,0 +1,243 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Alias-free activation for BigVGAN-style speech decoders.
+
+Provides anti-aliased activation (upsample → activation → downsample) with
+Kaiser-windowed sinc filters.  Includes NPU bf16 workarounds for platforms
+where F.pad(mode='replicate') does not support bfloat16.
+
+Used by: Qwen2.5-Omni, Qwen3-TTS v1, CoVo-Audio, and future BigVGAN vocoders.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from vllm.logger import init_logger
+
+from vllm_omni.platforms import current_omni_platform
+
+logger = init_logger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Helper utilities
+# ---------------------------------------------------------------------------
+
+
+def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor:
+    """Generate a 1-D Kaiser-windowed sinc low-pass filter.
+
+    Args:
+        cutoff: Normalized cutoff frequency (0 to 0.5).
+        half_width: Transition bandwidth.
+        kernel_size: Number of filter taps.
+
+    Returns:
+        Filter tensor of shape ``(1, 1, kernel_size)``.
+    """
+    is_even = kernel_size % 2 == 0
+    half_size = kernel_size // 2
+
+    delta_f = 4 * half_width
+    attenuation = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+
+    if attenuation > 50.0:
+        beta = 0.1102 * (attenuation - 8.7)
+    elif attenuation >= 21.0:
+        beta = 0.5842 * (attenuation - 21) ** 0.4 + 0.07886 * (attenuation - 21.0)
+    else:
+        beta = 0.0
+
+    # torch.kaiser_window is not supported on NPU / XPU — compute on CPU.
+    if current_omni_platform.is_npu():
+        kaiser_window = torch.kaiser_window(
+            kernel_size, beta=beta, periodic=False, dtype=torch.float32, device="cpu"
+        ).to("npu")
+    elif current_omni_platform.is_xpu():
+        kaiser_window = torch.kaiser_window(
+            kernel_size, beta=beta, periodic=False, dtype=torch.float32, device="cpu"
+        ).to("xpu")
+    else:
+        kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32)
+
+    if is_even:
+        time_indices = torch.arange(-half_size, half_size) + 0.5
+    else:
+        time_indices = torch.arange(kernel_size) - half_size
+
+    if cutoff == 0:
+        return torch.zeros((1, 1, kernel_size), dtype=torch.float32)
+
+    sinc_filter = torch.sinc(2 * cutoff * time_indices)
+    normalized_filter = 2 * cutoff * kaiser_window * sinc_filter
+    normalized_filter /= normalized_filter.sum()
+
+    return normalized_filter.view(1, 1, kernel_size)
+
+
+def replication_pad_1d(hidden_states: torch.Tensor, pad_left: int, pad_right: int) -> torch.Tensor:
+    """Manual replicate padding for NPU bf16 compatibility.
+
+    ``F.pad(mode='replicate')`` does not support bfloat16 on NPU (and older
+    CUDA).  This expand+cat fallback produces identical results.
+
+    TODO: Remove when upstream PyTorch fixes replication_pad1d for bf16.
+    """
+    if pad_left == 0 and pad_right == 0:
+        return hidden_states
+
+    segments = []
+    if pad_left > 0:
+        left = hidden_states[..., :1].expand(*hidden_states.shape[:-1], pad_left)
+        segments.append(left)
+
+    segments.append(hidden_states)
+
+    if pad_right > 0:
+        right = hidden_states[..., -1:].expand(*hidden_states.shape[:-1], pad_right)
+        segments.append(right)
+
+    return torch.cat(segments, dim=-1)
+
+
+# ---------------------------------------------------------------------------
+# Up / Down sampling with Kaiser sinc filters
+# ---------------------------------------------------------------------------
+
+
+class UpSample1d(nn.Module):
+    """Kaiser-sinc interpolation upsample (default: 2x, 12-tap filter)."""
+
+    def __init__(self, ratio: int = 2, kernel_size: int | None = None):
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        self.stride = ratio
+        self.pad = self.kernel_size // ratio - 1
+        self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
+        self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+
+        filt = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size)
+        self.register_buffer("filter", filt, persistent=False)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        channels = hidden_states.shape[1]
+        if current_omni_platform.is_npu():
+            input_dtype = hidden_states.dtype
+            hidden_states = replication_pad_1d(hidden_states.to(self.filter.dtype), self.pad, self.pad)
+            hidden_states = self.ratio * F.conv_transpose1d(
+                hidden_states,
+                self.filter.expand(channels, -1, -1),
+                stride=self.stride,
+                groups=channels,
+            ).to(input_dtype)
+        else:
+            input_dtype = hidden_states.dtype
+            hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate").to(self.filter.dtype)
+            hidden_states = self.ratio * F.conv_transpose1d(
+                hidden_states,
+                self.filter.expand(channels, -1, -1),
+                stride=self.stride,
+                groups=channels,
+            ).to(input_dtype)
+        hidden_states = hidden_states[..., self.pad_left : -self.pad_right]
+        return hidden_states
+
+
+class DownSample1d(nn.Module):
+    """Kaiser-sinc anti-aliasing downsample (default: 2x, 12-tap filter)."""
+
+    def __init__(self, ratio: int = 2, kernel_size: int | None = None):
+        super().__init__()
+        if kernel_size is None:
+            kernel_size = int(6 * ratio // 2) * 2
+
+        cutoff = 0.5 / ratio
+        half_width = 0.6 / ratio
+
+        if cutoff < 0.0:
+            raise ValueError("Minimum cutoff must be larger than zero.")
+        if cutoff > 0.5:
+            raise ValueError("A cutoff above 0.5 does not make sense.")
+
+        self.even = kernel_size % 2 == 0
+        self.pad_left = kernel_size // 2 - int(self.even)
+        self.pad_right = kernel_size // 2
+        self.stride = ratio
+
+        filt = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
+        self.register_buffer("filter", filt, persistent=False)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        channels = hidden_states.shape[1]
+        if current_omni_platform.is_npu():
+            input_dtype = hidden_states.dtype
+            hidden_states = replication_pad_1d(hidden_states.to(self.filter.dtype), self.pad_left, self.pad_right)
+            out = F.conv1d(
+                hidden_states,
+                self.filter.to(device=hidden_states.device, dtype=hidden_states.dtype).expand(channels, -1, -1),
+                stride=self.stride,
+                groups=channels,
+            ).to(input_dtype)
+        else:
+            input_dtype = hidden_states.dtype
+            hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate").to(
+                self.filter.dtype
+            )
+            out = F.conv1d(
+                hidden_states,
+                self.filter.expand(channels, -1, -1),
+                stride=self.stride,
+                groups=channels,
+            ).to(input_dtype)
+        return out
+
+
+# ---------------------------------------------------------------------------
+# Alias-free activation: upsample → activation → downsample
+# ---------------------------------------------------------------------------
+
+
+class AliasFreeActivation1d(nn.Module):
+    """Anti-aliased activation following BigVGAN (ICLR 2023).
+
+    Upsamples the signal, applies a pointwise activation (typically SnakeBeta),
+    then downsamples with a low-pass filter to suppress aliasing from the
+    high-frequency components introduced by the activation function.
+
+    The ``activation`` should be an ``nn.Module`` instance — typically
+    ``SnakeBeta`` from ``common.snake_activation``, which already dispatches
+    between Triton and eager PyTorch internally.
+
+    Args:
+        activation: Pointwise activation module (e.g. ``SnakeBeta(channels)``).
+        up_ratio: Upsample factor before activation.
+        down_ratio: Downsample factor after activation.
+        up_kernel_size: Kaiser sinc filter taps for upsampling.
+        down_kernel_size: Kaiser sinc filter taps for downsampling.
+    """
+
+    def __init__(
+        self,
+        activation: nn.Module,
+        up_ratio: int = 2,
+        down_ratio: int = 2,
+        up_kernel_size: int = 12,
+        down_kernel_size: int = 12,
+    ):
+        super().__init__()
+        if not callable(activation):
+            raise TypeError("Activation function must be callable")
+        self.act = activation
+        self.upsample = UpSample1d(up_ratio, up_kernel_size)
+        self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.upsample(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.downsample(hidden_states)
+        return hidden_states
diff --git a/vllm_omni/model_executor/models/common/nucleus_ras_sampling.py b/vllm_omni/model_executor/models/common/nucleus_ras_sampling.py
new file mode 100644
index 000000000..34573ae7d
--- /dev/null
+++ b/vllm_omni/model_executor/models/common/nucleus_ras_sampling.py
@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared TTS sampling primitives: nucleus (top-p/top-k) and RAS.
+
+Used by CosyVoice3 and GLM-TTS (and any future TTS model with RAS).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+
+
+def multinomial_sample(
+    probs: torch.Tensor,
+    generator: torch.Generator | None = None,
+) -> torch.Tensor:
+    """Draw one sample from a categorical distribution."""
+    return torch.multinomial(probs, 1, replacement=True, generator=generator).reshape(())
+
+
+def nucleus_sample_one(
+    weighted_scores: torch.Tensor,
+    *,
+    top_p: float,
+    top_k: int,
+    generator: torch.Generator | None = None,
+) -> int:
+    """Sample one token using nucleus (top-p + top-k) filtering.
+
+    ``weighted_scores`` should be log-softmax-ed logits (for RAS callers)
+    or raw logits (softmax is applied internally).
+
+    Vectorized via ``torch.cumsum`` — no Python loop over the vocabulary.
+    """
+    # Apply top-k before sort.  GLM-TTS and CosyVoice3 usually run with
+    # top_k=25, so sorting the whole audio vocabulary on every AR step is
+    # avoidable work in the hottest path.  Keep probabilities normalized over
+    # the full vocabulary so the top-p cumulative mass matches the old path.
+    if top_k > 0 and int(top_k) < int(weighted_scores.numel()):
+        top_scores, sorted_idx = torch.topk(
+            weighted_scores,
+            k=int(top_k),
+            dim=0,
+            largest=True,
+            sorted=True,
+        )
+        sorted_prob = (top_scores - torch.logsumexp(top_scores, dim=0)).exp()
+    else:
+        probs = weighted_scores.softmax(dim=0)
+        sorted_prob, sorted_idx = probs.sort(descending=True, stable=True)
+
+    # Apply top-p (nucleus) filtering: keep the smallest prefix whose
+    # cumulative probability exceeds top_p, always including the first token.
+    # Use float32 for cumulative sum to match the original Python float64
+    # accumulation loop.  Low-precision (fp16/bf16) cumsum over ~32K vocab
+    # causes rounding drift that shifts the top-p boundary, changing which
+    # tokens are kept and ultimately producing different AR sequences.
+    sorted_prob_f32 = sorted_prob.float()
+    cum_prob = sorted_prob_f32.cumsum(dim=0)
+    # Mask: include tokens where cumsum *before* this token is still < top_p
+    mask = (cum_prob - sorted_prob_f32) < top_p
+    # Safety: ensure top token is always kept.  With sorted descending
+    # probabilities, mask[0] is always True when top_p > 0 (since the
+    # cumsum-before-self of the first element is 0 < top_p).  Setting it
+    # unconditionally avoids a GPU→CPU sync from the old `mask.any()` check.
+    mask[0] = True
+
+    kept_probs = torch.nan_to_num(sorted_prob[mask], nan=0.0, posinf=0.0, neginf=0.0).clamp_min_(0.0)
+    kept_indices = sorted_idx[mask]
+    # ``torch.multinomial`` accepts unnormalized weights, but not an all-zero
+    # vector.  Handle the extreme underflow case on-device by forcing the top
+    # retained token to be selectable without adding another CPU sync.
+    kept_probs[0] = torch.where(
+        kept_probs.sum() > 0,
+        kept_probs[0],
+        torch.ones_like(kept_probs[0]),
+    )
+
+    sample_pos = multinomial_sample(kept_probs, generator=generator)
+    return int(kept_indices[sample_pos].item())
+
+
+def ras_sample_one(
+    weighted_scores: torch.Tensor,
+    decoded_tokens: Sequence[int],
+    *,
+    top_p: float,
+    top_k: int,
+    win_size: int,
+    tau_r: float,
+    generator: torch.Generator | None = None,
+) -> int:
+    """Repetition-Aware Sampling following GLM-TTS/CosyVoice.
+
+    If the nucleus-sampled token appears too often in the recent window,
+    upstream samples once from the original full distribution instead of
+    masking the repeated token.  This matters for GLM-TTS because masking the
+    dominant repeated token can artificially raise EOA probability right after
+    the minimum length guard is lifted.
+    """
+
+    def _random_sample_one() -> int:
+        return int(multinomial_sample(weighted_scores.softmax(dim=0), generator=generator).item())
+
+    top_id = nucleus_sample_one(
+        weighted_scores,
+        top_p=top_p,
+        top_k=top_k,
+        generator=generator,
+    )
+    if win_size > 0 and decoded_tokens:
+        recent = decoded_tokens[-win_size:]
+        rep_num = sum(1 for token in recent if int(token) == top_id)
+        if rep_num >= win_size * tau_r:
+            top_id = _random_sample_one()
+    return top_id
diff --git a/vllm_omni/model_executor/models/common/snake_activation.py b/vllm_omni/model_executor/models/common/snake_activation.py
index 35bb3384a..f50e27674 100644
--- a/vllm_omni/model_executor/models/common/snake_activation.py
+++ b/vllm_omni/model_executor/models/common/snake_activation.py
@@ -1,7 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # Copyright 2025 The Qwen team.
-"""Shared SnakeBeta activation for speech decoders (Qwen3-TTS, Qwen3-Omni Code2Wav)."""
+"""Shared Snake/SnakeBeta activations for speech decoders.
+
+Used by: Qwen3-TTS, Qwen3-Omni Code2Wav, Qwen2.5-Omni, CoVo-Audio, CosyVoice3.
+"""
 
 import torch
 from torch import nn
@@ -70,12 +73,17 @@ class SnakeBeta(nn.Module):
         SnakeBeta._triton_kernel = _kernel
         return True
 
-    def __init__(self, in_features, alpha=1.0):
+    def __init__(self, in_features, alpha=1.0, alpha_logscale=True):
         super().__init__()
         self.in_features = in_features
+        self.alpha_logscale = alpha_logscale
 
-        self.alpha = Parameter(torch.zeros(in_features) * alpha)
-        self.beta = Parameter(torch.zeros(in_features) * alpha)
+        if alpha_logscale:
+            self.alpha = Parameter(torch.zeros(in_features) * alpha)
+            self.beta = Parameter(torch.zeros(in_features) * alpha)
+        else:
+            self.alpha = Parameter(torch.ones(in_features) * alpha)
+            self.beta = Parameter(torch.ones(in_features) * alpha)
 
         self.no_div_by_zero = 0.000000001
 
@@ -86,8 +94,12 @@ class SnakeBeta(nn.Module):
     def precompute_exp_cache(self):
         """Materialize exp(alpha) and 1/(exp(beta)+eps) as frozen buffers."""
         with torch.no_grad():
-            self._exp_alpha = torch.exp(self.alpha).contiguous()
-            self._inv_beta = (1.0 / (torch.exp(self.beta) + self.no_div_by_zero)).contiguous()
+            if self.alpha_logscale:
+                self._exp_alpha = torch.exp(self.alpha).contiguous()
+                self._inv_beta = (1.0 / (torch.exp(self.beta) + self.no_div_by_zero)).contiguous()
+            else:
+                self._exp_alpha = self.alpha.contiguous()
+                self._inv_beta = (1.0 / (self.beta + self.no_div_by_zero)).contiguous()
 
     @property
     def _cached(self):
@@ -104,12 +116,10 @@ class SnakeBeta(nn.Module):
         return self._eager_forward(hidden_states)
 
     def _eager_forward(self, hidden_states):
-        if self._cached:
-            exp_alpha = self._exp_alpha.unsqueeze(0).unsqueeze(-1)
-            inv_beta = self._inv_beta.unsqueeze(0).unsqueeze(-1)
-        else:
-            exp_alpha = torch.exp(self.alpha).unsqueeze(0).unsqueeze(-1)
-            inv_beta = (1.0 / (torch.exp(self.beta) + self.no_div_by_zero)).unsqueeze(0).unsqueeze(-1)
+        if not self._cached:
+            self.precompute_exp_cache()
+        exp_alpha = self._exp_alpha.unsqueeze(0).unsqueeze(-1)
+        inv_beta = self._inv_beta.unsqueeze(0).unsqueeze(-1)
         hidden_states = hidden_states + inv_beta * torch.pow(torch.sin(hidden_states * exp_alpha), 2)
         return hidden_states
 
@@ -134,3 +144,33 @@ class SnakeBeta(nn.Module):
             block_t=block_t,
         )
         return out
+
+
+class Snake(SnakeBeta):
+    """Original Snake activation with a single parameter: x + 1/α * sin²(αx).
+
+    Unlike SnakeBeta which has separate alpha (frequency) and beta (magnitude)
+    parameters, Snake uses alpha for both.  Only ``alpha`` appears in the
+    state_dict — ``beta`` is absent, keeping checkpoint compatibility with
+    CosyVoice3's HiFi-GAN.
+
+    The Triton kernel and precomputed-cache path from SnakeBeta are reused;
+    ``precompute_exp_cache`` derives ``_inv_beta`` from ``alpha`` so the
+    forward path is identical.
+    """
+
+    def __init__(self, in_features, alpha=1.0, alpha_logscale=False):
+        # Initialise SnakeBeta — creates both alpha and beta Parameters.
+        super().__init__(in_features, alpha=alpha, alpha_logscale=alpha_logscale)
+        # Drop beta as a Parameter so it won't appear in state_dict.
+        del self.beta
+
+    def precompute_exp_cache(self):
+        """Derive both exp_alpha and inv_beta from the single alpha parameter."""
+        with torch.no_grad():
+            if self.alpha_logscale:
+                self._exp_alpha = torch.exp(self.alpha).contiguous()
+                self._inv_beta = (1.0 / (torch.exp(self.alpha) + self.no_div_by_zero)).contiguous()
+            else:
+                self._exp_alpha = self.alpha.contiguous()
+                self._inv_beta = (1.0 / (self.alpha + self.no_div_by_zero)).contiguous()
diff --git a/vllm_omni/model_executor/models/common/whisper_vq.py b/vllm_omni/model_executor/models/common/whisper_vq.py
new file mode 100644
index 000000000..9d7af76b5
--- /dev/null
+++ b/vllm_omni/model_executor/models/common/whisper_vq.py
@@ -0,0 +1,291 @@
+# Copyright 2022 The OpenAI Authors and The HuggingFace Inc. team. All rights reserved.
+#               2025 Zhipu AI Inc (authors: CogAudio Group Members)
+# SPDX-License-Identifier: Apache-2.0
+"""WhisperVQEncoder: HF WhisperEncoder + VQ codebook + inter-layer pooling.
+
+Built on standard ``WhisperConfig``.  VQ-specific parameters are patched onto
+the config object after ``from_pretrained`` so that callers never need a
+separate config class.  GLM-TTS (and future VQ-based TTS models) only need::
+
+    from transformers import WhisperConfig
+    from vllm_omni.model_executor.models.common.whisper_vq import WhisperVQEncoder
+
+    cfg = WhisperConfig.from_pretrained(checkpoint_dir)
+    cfg.pooling_kernel_size  = None       # passed through from checkpoint
+    cfg.pooling_type         = "max"      #       or set by caller
+    cfg.pooling_position     = 0
+    cfg.quantize_vocab_size  = 32768
+    cfg.quantize_position    = 16
+    cfg.quantize_encoder_only = True
+    model = WhisperVQEncoder(cfg)
+"""
+
+from __future__ import annotations
+
+import math
+import re
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+from transformers import WhisperConfig
+from transformers.modeling_outputs import BaseModelOutput
+from transformers.models.whisper.modeling_whisper import WhisperEncoder
+
+# ---------------------------------------------------------------------------
+# helpers
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class QuantizedBaseModelOutput(BaseModelOutput):
+    quantized_token_ids: torch.LongTensor | None = None
+
+
+def vector_quantize(inputs: Tensor, codebook: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+    """Nearest-neighbour codebook lookup."""
+    embedding_size = codebook.size(1)
+    inputs_flatten = inputs.reshape(-1, embedding_size)
+    codebook_sqr = torch.sum(codebook**2, dim=1)
+    inputs_sqr = torch.sum(inputs_flatten**2, dim=1, keepdim=True)
+    distances = torch.addmm(
+        codebook_sqr + inputs_sqr,
+        inputs_flatten,
+        codebook.t(),
+        alpha=-2.0,
+        beta=1.0,
+    )
+    _, indices_flatten = torch.min(distances, dim=1)
+    codes_flatten = torch.index_select(codebook, dim=0, index=indices_flatten)
+    codes = codes_flatten.view_as(inputs)
+    return codes, indices_flatten, distances
+
+
+_LEGACY_LAYER_REMAP = (
+    (re.compile(r"^layers\.(\d+)\.attn\.query\."), r"layers.\1.self_attn.q_proj."),
+    (re.compile(r"^layers\.(\d+)\.attn\.key\."), r"layers.\1.self_attn.k_proj."),
+    (re.compile(r"^layers\.(\d+)\.attn\.value\."), r"layers.\1.self_attn.v_proj."),
+    (re.compile(r"^layers\.(\d+)\.attn\.out\."), r"layers.\1.self_attn.out_proj."),
+    (re.compile(r"^layers\.(\d+)\.attn_ln\."), r"layers.\1.self_attn_layer_norm."),
+    (re.compile(r"^layers\.(\d+)\.mlp_ln\."), r"layers.\1.final_layer_norm."),
+    (re.compile(r"^layers\.(\d+)\.mlp\.0\."), r"layers.\1.fc1."),
+    (re.compile(r"^layers\.(\d+)\.mlp\.2\."), r"layers.\1.fc2."),
+)
+
+
+def remap_legacy_whisper_vq_state_dict(state_dict: dict[str, Tensor]) -> dict[str, Tensor]:
+    """Map CogAudio/GLM-TTS fork parameter names onto HF WhisperEncoder names."""
+    remapped: dict[str, Tensor] = {}
+    for name, tensor in state_dict.items():
+        new_name = name
+        if new_name == "embed_positions":
+            new_name = "embed_positions.weight"
+        for pattern, replacement in _LEGACY_LAYER_REMAP:
+            new_name = pattern.sub(replacement, new_name)
+        remapped[new_name] = tensor
+    return remapped
+
+
+# ---------------------------------------------------------------------------
+# main encoder class
+# ---------------------------------------------------------------------------
+
+
+class WhisperVQEncoder(WhisperEncoder):
+    """HF Whisper encoder with optional VQ codebook and pooling.
+
+    Uses a standard ``WhisperConfig`` with the following VQ-specific attrs
+    patched on by the caller (or loaded from the checkpoint's config.json):
+
+    * ``pooling_kernel_size``  -- int | None
+    * ``pooling_type``         -- "max" | "avg"
+    * ``pooling_position``     -- int (0-based layer index)
+    * ``quantize_vocab_size``  -- int | None
+    * ``quantize_position``    -- int (0-based layer index)
+    * ``quantize_encoder_only`` -- bool
+    """
+
+    def __init__(self, config: WhisperConfig):
+        super().__init__(config)
+        embed_dim = config.d_model
+        max_source_positions = config.max_source_positions
+
+        # Truncate layers when quantize_encoder_only is set.
+        qpos = int(getattr(config, "quantize_position", 0) or 0)
+        if getattr(config, "quantize_encoder_only", False):
+            self.layers = nn.ModuleList(list(self.layers[:qpos]))
+            self.layer_norm = None
+        else:
+            qpos = max(qpos, int(config.encoder_layers))  # never exceed total
+
+        # Optional inter-layer pooling.
+        self.pooling_layer: nn.Module | None = None
+        ksize = getattr(config, "pooling_kernel_size", None)
+        if ksize is not None:
+            ptype = str(getattr(config, "pooling_type", "max") or "max")
+            self.pooling_layer = {
+                "max": nn.MaxPool1d(kernel_size=int(ksize)),
+                "avg": nn.AvgPool1d(kernel_size=int(ksize)),
+            }[ptype]
+
+        # Optional vector-quantization codebook.
+        self.codebook: nn.Embedding | None = None
+        self.embed_positions2: nn.Embedding | None = None
+        vq_sz = getattr(config, "quantize_vocab_size", None)
+        if vq_sz is not None:
+            self.codebook = nn.Embedding(int(vq_sz), embed_dim)
+            pos2_len = max_source_positions
+            if ksize is not None:
+                pos2_len = math.ceil(max_source_positions / int(ksize))
+            self.embed_positions2 = nn.Embedding(pos2_len, embed_dim)
+
+        # cached for forward()
+        self._pooling_position = int(getattr(config, "pooling_position", 0) or 0)
+        self._quantize_position = qpos if vq_sz is not None else -1
+
+    # ------------------------------------------------------------------
+    # forward / helpers
+    # ------------------------------------------------------------------
+
+    @property
+    def device(self) -> torch.device:  # type: ignore[override]
+        return self.conv1.weight.device
+
+    @property
+    def dtype(self) -> torch.dtype:  # type: ignore[override]
+        return self.conv1.weight.dtype
+
+    def load_state_dict(
+        self,
+        state_dict: dict[str, Tensor],
+        strict: bool = True,
+        assign: bool = False,
+    ):  # type: ignore[override]
+        remapped = remap_legacy_whisper_vq_state_dict(state_dict)
+        try:
+            return super().load_state_dict(remapped, strict=strict, assign=assign)
+        except TypeError:
+            return super().load_state_dict(remapped, strict=strict)
+
+    def forward(
+        self,
+        input_features: Tensor,
+        attention_mask: Tensor | None = None,
+        **_: Any,
+    ) -> QuantizedBaseModelOutput:
+        batch_size, _, _ = input_features.shape
+        hidden_states = F.gelu(self.conv1(input_features))
+        hidden_states = F.gelu(self.conv2(hidden_states))
+        hidden_states = hidden_states.permute(0, 2, 1)
+        seq_len = int(hidden_states.shape[1])
+
+        pos_embed = self._position_embeddings(seq_len, hidden_states.device, hidden_states.dtype)
+        hidden_states = hidden_states + pos_embed
+
+        valid_mask = self._downsample_attention_mask(attention_mask, batch_size, seq_len, hidden_states.device)
+        hidden_states = hidden_states.masked_fill(~valid_mask.unsqueeze(-1), 0)
+
+        quantized_token_ids = None
+        for idx, layer in enumerate(self.layers):
+            layer_mask = self._make_layer_attention_mask(valid_mask, hidden_states.dtype)
+            layer_outputs = layer(
+                hidden_states,
+                attention_mask=layer_mask,
+                layer_head_mask=None,
+                output_attentions=False,
+            )
+            hidden_states = layer_outputs[0] if isinstance(layer_outputs, tuple) else layer_outputs
+            hidden_states = hidden_states.masked_fill(~valid_mask.unsqueeze(-1), 0)
+
+            if idx + 1 == self._pooling_position and self.pooling_layer is not None:
+                hidden_states, valid_mask = self._apply_pooling(hidden_states, valid_mask)
+
+            if idx + 1 == self._quantize_position and self.codebook is not None:
+                hidden_states, quantized_token_ids = self._apply_vq(hidden_states, valid_mask)
+
+        if self.layer_norm is not None:
+            hidden_states = self.layer_norm(hidden_states)
+            hidden_states = hidden_states.masked_fill(~valid_mask.unsqueeze(-1), 0)
+
+        return QuantizedBaseModelOutput(
+            last_hidden_state=hidden_states,
+            quantized_token_ids=quantized_token_ids,
+        )
+
+    def _position_embeddings(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> Tensor:
+        ep = self.embed_positions
+        if isinstance(ep, nn.Embedding):
+            positions = torch.arange(seq_len, device=device)
+            return ep(positions).to(dtype=dtype).unsqueeze(0)
+        return ep[:seq_len].to(device=device, dtype=dtype).unsqueeze(0)
+
+    def _downsample_attention_mask(
+        self,
+        attention_mask: Tensor | None,
+        batch_size: int,
+        seq_len: int,
+        device: torch.device,
+    ) -> Tensor:
+        if attention_mask is None:
+            return torch.ones(batch_size, seq_len, dtype=torch.bool, device=device)
+        conv_stride = int(self.conv1.stride[0]) * int(self.conv2.stride[0])
+        downsampled = attention_mask[:, ::conv_stride].to(device=device, dtype=torch.bool)
+        if downsampled.shape[1] < seq_len:
+            pad = seq_len - int(downsampled.shape[1])
+            downsampled = F.pad(downsampled, (0, pad), value=False)
+        return downsampled[:, :seq_len]
+
+    @staticmethod
+    def _make_layer_attention_mask(valid_mask: Tensor, dtype: torch.dtype) -> Tensor:
+        min_value = torch.finfo(dtype).min
+        additive = torch.zeros(valid_mask.shape, dtype=dtype, device=valid_mask.device)
+        additive = additive.masked_fill(~valid_mask, min_value)
+        return additive[:, None, None, :]
+
+    def _apply_pooling(self, hidden_states: Tensor, valid_mask: Tensor) -> tuple[Tensor, Tensor]:
+        k = int(self.config.pooling_kernel_size)  # type: ignore[attr-defined]
+        pooled_segments: list[Tensor] = []
+        pooled_masks: list[Tensor] = []
+        for batch_idx in range(int(hidden_states.shape[0])):
+            length = int(valid_mask[batch_idx].sum().item())
+            segment = hidden_states[batch_idx, :length]
+            if segment.numel() == 0:
+                pooled = hidden_states.new_zeros((0, hidden_states.shape[-1]))
+            else:
+                segment_t = segment.unsqueeze(0).transpose(1, 2)
+                if int(segment_t.shape[-1]) % k != 0:
+                    segment_t = F.pad(segment_t, (0, k - int(segment_t.shape[-1]) % k))
+                pooled = self.pooling_layer(segment_t).transpose(1, 2).squeeze(0)  # type: ignore[misc]
+            pooled_segments.append(pooled)
+            pooled_masks.append(torch.ones(pooled.shape[0], dtype=torch.bool, device=hidden_states.device))
+
+        max_len = max((int(seg.shape[0]) for seg in pooled_segments), default=0)
+        output = hidden_states.new_zeros((hidden_states.shape[0], max_len, hidden_states.shape[-1]))
+        output_mask = torch.zeros(hidden_states.shape[0], max_len, dtype=torch.bool, device=hidden_states.device)
+        for i, segment in enumerate(pooled_segments):
+            if segment.shape[0] > 0:
+                output[i, : segment.shape[0]] = segment
+                output_mask[i, : segment.shape[0]] = pooled_masks[i]
+        return output, output_mask
+
+    def _apply_vq(self, hidden_states: Tensor, valid_mask: Tensor) -> tuple[Tensor, Tensor | None]:
+        output = hidden_states.clone()
+        token_ids = torch.full(
+            (hidden_states.shape[0], hidden_states.shape[1]),
+            -1,
+            dtype=torch.long,
+            device=hidden_states.device,
+        )
+        for batch_idx in range(int(hidden_states.shape[0])):
+            length = int(valid_mask[batch_idx].sum().item())
+            if length <= 0:
+                continue
+            segment = hidden_states[batch_idx : batch_idx + 1, :length]
+            quantized, indices, _ = vector_quantize(segment, self.codebook.weight)  # type: ignore[union-attr]
+            if self.embed_positions2 is not None:
+                quantized = quantized + self.embed_positions2.weight[:length].to(dtype=quantized.dtype).unsqueeze(0)
+            output[batch_idx, :length] = quantized.squeeze(0)
+            token_ids[batch_idx, :length] = indices.reshape(-1)
+        return output.masked_fill(~valid_mask.unsqueeze(-1), 0), token_ids
diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py
index bd31443cc..63434d57f 100644
--- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py
+++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py
@@ -11,7 +11,6 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from scipy.signal import get_window
-from torch import pow, sin
 from torch.nn import Conv1d, ConvTranspose1d
 from torch.nn.utils import remove_weight_norm
 
@@ -20,64 +19,8 @@ try:
 except ImportError:
     from torch.nn.utils import weight_norm
 from torch.distributions.uniform import Uniform
-from torch.nn import Parameter
-
-
-# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license.
-#   LICENSE is in incl_licenses directory.
-class Snake(nn.Module):
-    """
-    Implementation of a sine-based periodic activation function
-    Shape:
-        - Input: (B, C, T)
-        - Output: (B, C, T), same shape as the input
-    Parameters:
-        - alpha - trainable parameter
-    References:
-        - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
-        https://arxiv.org/abs/2006.08195
-    Examples:
-        >>> a1 = snake(256)
-        >>> x = torch.randn(256)
-        >>> x = a1(x)
-    """
-
-    def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False):
-        """
-        Initialization.
-        INPUT:
-            - in_features: shape of the input
-            - alpha: trainable parameter
-            alpha is initialized to 1 by default, higher values = higher-frequency.
-            alpha will be trained along with the rest of your model.
-        """
-        super().__init__()
-        self.in_features = in_features
-
-        # initialize alpha
-        self.alpha_logscale = alpha_logscale
-        if self.alpha_logscale:  # log scale alphas initialized to zeros
-            self.alpha = Parameter(torch.zeros(in_features) * alpha)
-        else:  # linear scale alphas initialized to ones
-            self.alpha = Parameter(torch.ones(in_features) * alpha)
-
-        self.alpha.requires_grad = alpha_trainable
-
-        self.no_div_by_zero = 0.000000001
-
-    def forward(self, x):
-        """
-        Forward pass of the function.
-        Applies the function to the input elementwise.
-        Snake ∶= x + 1/a * sin^2 (xa)
-        """
-        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)  # line up with x to [B, C, T]
-        if self.alpha_logscale:
-            alpha = torch.exp(alpha)
-        x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
-
-        return x
 
+from vllm_omni.model_executor.models.common.snake_activation import Snake
 
 """hifigan based generator implementation.
 
diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py
index 5023307ff..2dc81174c 100644
--- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py
+++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py
@@ -784,12 +784,14 @@ class CosyVoice3Model(
                             else:
                                 self._stream_vocoder_cache_by_req[req_id] = new_cache_state
                 else:
+                    token_offset = max(0, meta.talker_prefill_offset or 0) if meta else 0
                     tts_speech = self.code2wav.forward(
                         token=token.unsqueeze(0),
                         prompt_token=speech_token[:1],
                         prompt_feat=speech_feat[:1],
                         embedding=embedding[:1],
                         n_timesteps=10,
+                        token_offset_tokens=token_offset,
                     )
 
                 audio = tts_speech.reshape(-1).to(dtype=torch.float32)
diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py
index 186a258c8..cb3228c13 100644
--- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py
+++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py
@@ -292,6 +292,7 @@ class CosyVoice3Code2Wav(nn.Module):
         prompt_feat: torch.Tensor,
         embedding: torch.Tensor,
         n_timesteps: int = 10,
+        token_offset_tokens: int = 0,
     ) -> torch.Tensor:
         """Generate audio waveform from speech tokens."""
         feat = self._forward_mel(
@@ -300,7 +301,7 @@ class CosyVoice3Code2Wav(nn.Module):
             prompt_feat=prompt_feat,
             embedding=embedding,
             n_timesteps=n_timesteps,
-            token_offset_tokens=0,
+            token_offset_tokens=token_offset_tokens,
             streaming=False,
             finalize=True,
         )
diff --git a/vllm_omni/model_executor/models/cosyvoice3/pipeline.py b/vllm_omni/model_executor/models/cosyvoice3/pipeline.py
index 4480a0dd8..ed35c93bd 100644
--- a/vllm_omni/model_executor/models/cosyvoice3/pipeline.py
+++ b/vllm_omni/model_executor/models/cosyvoice3/pipeline.py
@@ -31,6 +31,7 @@ COSYVOICE3_PIPELINE = PipelineConfig(
             owns_tokenizer=True,
             engine_output_type="latent",
             async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"),
+            custom_process_next_stage_input_func=f"{_PROC}.text2flow_full_payload",
             sampling_constraints={
                 # merged speech stop token (logsumexp of all 200 stop logits)
                 "stop_token_ids": [6562],
@@ -44,7 +45,8 @@ COSYVOICE3_PIPELINE = PipelineConfig(
             final_output=True,
             final_output_type="audio",
             engine_output_type="latent",
-            sync_process_input_func=f"{_PROC}.text2flow",
+            custom_process_input_func=f"{_PROC}.text2flow",
+            sync_process_input_func=f"{_PROC}.text2flow_token_only",
         ),
     ),
 )
diff --git a/vllm_omni/model_executor/models/covo_audio/pipeline.py b/vllm_omni/model_executor/models/covo_audio/pipeline.py
index 5b1a31d6e..97053e328 100644
--- a/vllm_omni/model_executor/models/covo_audio/pipeline.py
+++ b/vllm_omni/model_executor/models/covo_audio/pipeline.py
@@ -29,6 +29,7 @@ COVO_AUDIO_PIPELINE = PipelineConfig(
             owns_tokenizer=True,
             requires_multimodal_data=True,
             engine_output_type="latent",
+            custom_process_next_stage_input_func=f"{_PROC}.llm2code2wav_full_payload",
             sampling_constraints={
                 "detokenize": True,
                 "stop_token_ids": [151645],
@@ -44,6 +45,7 @@ COVO_AUDIO_PIPELINE = PipelineConfig(
             final_output_type="audio",
             engine_output_type="audio",
             custom_process_input_func=f"{_PROC}.llm2code2wav",
+            sync_process_input_func=f"{_PROC}.llm2code2wav_token_only",
             sampling_constraints={"detokenize": False},
         ),
     ),
diff --git a/vllm_omni/model_executor/models/covo_audio/token2wav.py b/vllm_omni/model_executor/models/covo_audio/token2wav.py
index f6bed125c..4a6ce4b7e 100644
--- a/vllm_omni/model_executor/models/covo_audio/token2wav.py
+++ b/vllm_omni/model_executor/models/covo_audio/token2wav.py
@@ -10,10 +10,11 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from einops import rearrange
-from torch import pow, sin
-from torch.nn import Parameter
 from torch.nn.utils import weight_norm
 
+from vllm_omni.model_executor.models.common.alias_free_activation import AliasFreeActivation1d
+from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta
+
 # --------------- Utilities ---------------
 
 
@@ -93,41 +94,7 @@ def apply_rotary_pos_emb(pos, t):
     return t * pos.cos() + rotate_half(t) * pos.sin()
 
 
-if "sinc" in dir(torch):
-    sinc = torch.sinc
-else:
-
-    def sinc(x: torch.Tensor):
-        return torch.where(
-            x == 0, torch.tensor(1.0, device=x.device, dtype=x.dtype), torch.sin(math.pi * x) / math.pi / x
-        )
-
-
-def kaiser_sinc_filter1d(cutoff, half_width, kernel_size):
-    even = kernel_size % 2 == 0
-    half_size = kernel_size // 2
-
-    delta_f = 4 * half_width
-    A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
-    if A > 50.0:
-        beta = 0.1102 * (A - 8.7)
-    elif A >= 21.0:
-        beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0)
-    else:
-        beta = 0.0
-    window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
-
-    if even:
-        time = torch.arange(-half_size, half_size) + 0.5
-    else:
-        time = torch.arange(kernel_size) - half_size
-    if cutoff == 0:
-        filter_ = torch.zeros_like(time)
-    else:
-        filter_ = 2 * cutoff * window * sinc(2 * cutoff * time)
-        filter_ /= filter_.sum()
-
-    return filter_.view(1, 1, kernel_size)
+# sinc, kaiser_sinc_filter1d removed — now in common.alias_free_activation
 
 
 # --------------- Basic Layers ---------------
@@ -320,116 +287,14 @@ class ConvPositionEmbed(nn.Module):
 # --------------- Anti-aliasing (BigVGAN) ---------------
 
 
-class LowPassFilter1d(nn.Module):
-    def __init__(
-        self,
-        cutoff=0.5,
-        half_width=0.6,
-        stride: int = 1,
-        padding: bool = True,
-        padding_mode: str = "replicate",
-        kernel_size: int = 12,
-    ):
-        super().__init__()
-        if cutoff < -0.0:
-            raise ValueError("Minimum cutoff must be larger than zero.")
-        if cutoff > 0.5:
-            raise ValueError("A cutoff above 0.5 does not make sense.")
-        self.kernel_size = kernel_size
-        self.even = kernel_size % 2 == 0
-        self.pad_left = kernel_size // 2 - int(self.even)
-        self.pad_right = kernel_size // 2
-        self.stride = stride
-        self.padding = padding
-        self.padding_mode = padding_mode
-        filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
-        self.register_buffer("filter", filter)
-
-    def forward(self, x):
-        _, C, _ = x.shape
-        if self.padding:
-            x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode)
-        out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C)
-        return out
-
-
-class UpSample1d(nn.Module):
-    def __init__(self, ratio=2, kernel_size=None):
-        super().__init__()
-        self.ratio = ratio
-        self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
-        self.stride = ratio
-        self.pad = self.kernel_size // ratio - 1
-        self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
-        self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
-        filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size)
-        self.register_buffer("filter", filter)
-
-    def forward(self, x):
-        _, C, _ = x.shape
-        x = F.pad(x, (self.pad, self.pad), mode="replicate")
-        x = self.ratio * F.conv_transpose1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C)
-        x = x[..., self.pad_left : -self.pad_right]
-        return x
-
-
-class DownSample1d(nn.Module):
-    def __init__(self, ratio=2, kernel_size=None):
-        super().__init__()
-        self.ratio = ratio
-        self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
-        self.lowpass = LowPassFilter1d(
-            cutoff=0.5 / ratio, half_width=0.6 / ratio, stride=ratio, kernel_size=self.kernel_size
-        )
-
-    def forward(self, x):
-        return self.lowpass(x)
-
-
-class Activation1d(nn.Module):
-    def __init__(
-        self, activation, up_ratio: int = 2, down_ratio: int = 2, up_kernel_size: int = 12, down_kernel_size: int = 12
-    ):
-        super().__init__()
-        self.up_ratio = up_ratio
-        self.down_ratio = down_ratio
-        self.act = activation
-        self.upsample = UpSample1d(up_ratio, up_kernel_size)
-        self.downsample = DownSample1d(down_ratio, down_kernel_size)
-
-    def forward(self, x):
-        x = self.upsample(x)
-        x = self.act(x)
-        x = self.downsample(x)
-        return x
+# LowPassFilter1d, UpSample1d, DownSample1d, Activation1d removed —
+# now provided by common.alias_free_activation.AliasFreeActivation1d
 
 
 # --------------- Activations ---------------
 
 
-class SnakeBeta(nn.Module):
-    def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False):
-        super().__init__()
-        self.in_features = in_features
-        self.alpha_logscale = alpha_logscale
-        if self.alpha_logscale:
-            self.alpha = Parameter(torch.zeros(in_features) * alpha)
-            self.beta = Parameter(torch.zeros(in_features) * alpha)
-        else:
-            self.alpha = Parameter(torch.ones(in_features) * alpha)
-            self.beta = Parameter(torch.ones(in_features) * alpha)
-        self.alpha.requires_grad = alpha_trainable
-        self.beta.requires_grad = alpha_trainable
-        self.no_div_by_zero = 0.000000001
-
-    def forward(self, x):
-        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)
-        beta = self.beta.unsqueeze(0).unsqueeze(-1)
-        if self.alpha_logscale:
-            alpha = torch.exp(alpha)
-            beta = torch.exp(beta)
-        x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
-        return x
+# SnakeBeta removed — now in common.snake_activation (with alpha_logscale support)
 
 
 # --------------- Attention ---------------
@@ -828,7 +693,7 @@ class AMPBlock1(torch.nn.Module):
 
         self.activations = nn.ModuleList(
             [
-                Activation1d(activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale))
+                AliasFreeActivation1d(activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale))
                 for _ in range(self.num_layers)
             ]
         )
@@ -882,7 +747,7 @@ class BigVGANFlowVAE(nn.Module):
                 self.resblocks.append(AMPBlock1(h, ch, k, d, causal=causal))
 
         activation_post = SnakeBeta(ch, alpha_logscale=h.snake_logscale)
-        self.activation_post = Activation1d(activation=activation_post)
+        self.activation_post = AliasFreeActivation1d(activation=activation_post)
 
         self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, causal=causal))
 
diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py
index daef2be18..c6e9064f1 100644
--- a/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py
+++ b/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py
@@ -11,6 +11,7 @@ Install with: ``pip install fish-speech``
 
 from __future__ import annotations
 
+import inspect
 import os
 from collections.abc import Iterable
 from typing import Any
@@ -33,6 +34,41 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput
 logger = init_logger(__name__)
 
 
+def _connector_extra_config(vllm_config: VllmConfig) -> dict[str, Any]:
+    model_config = getattr(vllm_config, "model_config", None)
+    connector_cfg = getattr(model_config, "stage_connector_config", None)
+    if isinstance(connector_cfg, dict):
+        return connector_cfg.get("extra", connector_cfg)
+    extra = getattr(connector_cfg, "extra", None)
+    return extra if isinstance(extra, dict) else {}
+
+
+def _get_int_config(extra_cfg: dict[str, Any], default: int, *names: str) -> int:
+    value = None
+    for name in names:
+        if name in extra_cfg:
+            value = extra_cfg[name]
+            break
+    if value is None:
+        return default
+    try:
+        return max(0, int(value))
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"Invalid Fish Speech DAC integer config {names[0]}={value!r}") from exc
+
+
+def _get_dac_dtype(extra_cfg: dict[str, Any]) -> torch.dtype:
+    value = extra_cfg.get("fish_speech_dac_dtype", "float32")
+    value = str(value).strip().lower()
+    if value in {"fp16", "float16", "half"}:
+        return torch.float16
+    if value in {"bf16", "bfloat16"}:
+        return torch.bfloat16
+    if value in {"fp32", "float32", ""}:
+        return torch.float32
+    raise ValueError(f"Invalid Fish Speech DAC dtype: {value!r}")
+
+
 class FishSpeechDACDecoder(nn.Module):
     """Stage-1 DAC decoder for Fish Speech S2 Pro (GenerationModelRunner).
 
@@ -58,6 +94,43 @@ class FishSpeechDACDecoder(nn.Module):
         self._output_sample_rate: int = DAC_SAMPLE_RATE
         self._hop_length: int = DAC_HOP_LENGTH
         self._logged_codec_stats = False
+        self._codec_decode_takes_lengths: bool | None = None
+        extra_cfg = _connector_extra_config(vllm_config)
+        self._dac_dtype = _get_dac_dtype(extra_cfg)
+        self._decode_batch_max_padded_frames = _get_int_config(
+            extra_cfg,
+            0,
+            "fish_speech_dac_max_padded_frames",
+            "dac_max_padded_frames",
+        )
+        self._decode_batch_max_batch = _get_int_config(
+            extra_cfg,
+            0,
+            "fish_speech_dac_max_batch",
+            "dac_max_batch",
+        )
+
+    def _decode_codes(
+        self,
+        codes_bqf: torch.Tensor,
+        feature_lengths: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        assert self._codec is not None
+        if self._codec_decode_takes_lengths:
+            return self._codec.decode(codes_bqf, feature_lengths)
+
+        if hasattr(self._codec, "from_indices"):
+            wav_batch = self._codec.from_indices(codes_bqf)
+        else:
+            wav_batch = self._codec.decode(codes_bqf)
+        max_frames = max(int(codes_bqf.shape[-1]), 1)
+        scale = wav_batch.shape[-1] / max_frames
+        audio_lengths = torch.clamp(
+            torch.round(feature_lengths.to(torch.float32) * scale).to(torch.long),
+            min=0,
+            max=wav_batch.shape[-1],
+        )
+        return wav_batch, audio_lengths
 
     def _bake_weight_norm(self, codec: nn.Module) -> None:
         baked = 0
@@ -149,14 +222,16 @@ class FishSpeechDACDecoder(nn.Module):
         codec.quantizer.downsample = None
 
         device = self.vllm_config.device_config.device
-        codec = codec.to(device=device, dtype=torch.float32)
+        codec = codec.to(device=device, dtype=self._dac_dtype)
         codec.eval()
         self._codec = codec
+        self._codec_decode_takes_lengths = len(inspect.signature(codec.decode).parameters) >= 2
 
         logger.info(
-            "Fish Speech DAC codec loaded from %s (device=%s, sample_rate=%d)",
+            "Fish Speech DAC codec loaded from %s (device=%s, dtype=%s, sample_rate=%d)",
             codec_path,
             device,
+            self._dac_dtype,
             self._output_sample_rate,
         )
 
@@ -188,6 +263,28 @@ class FishSpeechDACDecoder(nn.Module):
                 return [ids[boundaries[i] : boundaries[i + 1]] for i in range(len(boundaries) - 1)]
         return [ids]
 
+    def _codes_from_runtime_info(
+        self,
+        info: dict[str, Any] | None,
+        device: torch.device,
+    ) -> torch.Tensor | None:
+        if not isinstance(info, dict):
+            return None
+        codes = info.get("codes", {}).get("audio")
+        if not isinstance(codes, torch.Tensor) or codes.numel() == 0:
+            return None
+
+        q = self._num_codebooks
+        codes = codes.to(device=device, dtype=torch.long, non_blocking=True).contiguous()
+        if codes.ndim != 2 or codes.shape[0] != q:
+            logger.warning(
+                "DAC tensor codes must have shape [num_codebooks, frames], got %s for num_codebooks=%d.",
+                tuple(codes.shape),
+                q,
+            )
+            return None
+        return codes
+
     @torch.no_grad()
     def forward(
         self,
@@ -203,14 +300,14 @@ class FishSpeechDACDecoder(nn.Module):
         input_ids layout per request: flat codes [num_codebooks * num_frames].
         Codes are codebook-major: [cb0_f0, cb0_f1, ..., cb0_fN, cb1_f0, ...].
         """
-        self._ensure_codec_loaded()
-        assert self._codec is not None
-
         q = self._num_codebooks
         sr_val = self._output_sample_rate
         sr_tensor = torch.tensor(sr_val, dtype=torch.int32)
         empty = torch.zeros((0,), dtype=torch.float32)
 
+        self._ensure_codec_loaded()
+        assert self._codec is not None
+
         if input_ids is None or input_ids.numel() == 0:
             return OmniOutput(
                 text_hidden_states=None,
@@ -235,9 +332,22 @@ class FishSpeechDACDecoder(nn.Module):
                     left_context_size[i] = meta["left_context_size"]
 
         for i, req_ids in enumerate(request_ids_list):
+            ctx_frames = left_context_size[i]
+            runtime_codes = (
+                self._codes_from_runtime_info(runtime_additional_information[i], ids.device)
+                if runtime_additional_information is not None and i < len(runtime_additional_information)
+                else None
+            )
+            if runtime_codes is not None:
+                frames = runtime_codes.shape[1]
+                parsed_ctx_frames[i] = ctx_frames
+                parsed_total_frames[i] = frames
+                valid_codes_qf.append(runtime_codes)
+                valid_indices.append(i)
+                continue
+
             if req_ids.numel() < 1:
                 continue
-            ctx_frames = left_context_size[i]
             flat = req_ids
             n = flat.numel()
             if n == 0 or n % q != 0:
@@ -279,25 +389,61 @@ class FishSpeechDACDecoder(nn.Module):
             except Exception:
                 pass
 
-        feature_lengths = torch.tensor(
-            [codes_qf.shape[1] for codes_qf in valid_codes_qf],
-            device=valid_codes_qf[0].device,
-            dtype=torch.long,
-        )
-        max_frames = int(feature_lengths.max().item())
-        batch_size = len(valid_codes_qf)
+        wav_tensors: list[torch.Tensor | None] = [None] * len(valid_codes_qf)
 
-        codes_bqf = torch.zeros(
-            (batch_size, q, max_frames),
-            device=valid_codes_qf[0].device,
-            dtype=torch.long,
-        )
-        for i, codes_qf in enumerate(valid_codes_qf):
-            frame_count = int(feature_lengths[i].item())
-            codes_bqf[i, :, :frame_count] = codes_qf
+        def _iter_bounded_groups(group: list[tuple[int, torch.Tensor]]) -> Iterable[list[tuple[int, torch.Tensor]]]:
+            max_padded_frames = int(getattr(self, "_decode_batch_max_padded_frames", 0))
+            max_batch = int(getattr(self, "_decode_batch_max_batch", 0))
+            if max_padded_frames <= 0 and max_batch <= 0:
+                yield group
+                return
 
-        with torch.amp.autocast("cuda", enabled=False):
-            wav_batch, audio_lengths = self._codec.decode(codes_bqf, feature_lengths)
+            current: list[tuple[int, torch.Tensor]] = []
+            current_max_frames = 0
+            for item in sorted(group, key=lambda pair: int(pair[1].shape[1])):
+                frames = int(item[1].shape[1])
+                next_max_frames = max(current_max_frames, frames)
+                next_batch_size = len(current) + 1
+                exceeds_batch = max_batch > 0 and next_batch_size > max_batch
+                exceeds_work = max_padded_frames > 0 and next_batch_size * next_max_frames > max_padded_frames
+                if current and (exceeds_batch or exceeds_work):
+                    yield current
+                    current = [item]
+                    current_max_frames = frames
+                else:
+                    current.append(item)
+                    current_max_frames = next_max_frames
+            if current:
+                yield current
+
+        def _decode_group(group: list[tuple[int, torch.Tensor]]) -> None:
+            actual_frames = [int(codes_qf.shape[1]) for _, codes_qf in group]
+            target_frames = max(actual_frames)
+            batch_size = len(group)
+            first_codes = group[0][1]
+            feature_lengths = torch.tensor(actual_frames, device=first_codes.device, dtype=torch.long)
+            codes_bqf = torch.zeros(
+                (batch_size, q, target_frames),
+                device=first_codes.device,
+                dtype=torch.long,
+            )
+            for row, (_, codes_qf) in enumerate(group):
+                frame_count = int(codes_qf.shape[1])
+                codes_bqf[row, :, :frame_count] = codes_qf
+
+            with torch.amp.autocast("cuda", enabled=False):
+                wav_batch, audio_lengths = self._decode_codes(codes_bqf, feature_lengths)
+            audio_lengths_list = (
+                audio_lengths.detach().to(device="cpu", dtype=torch.long).reshape(-1).tolist()
+                if audio_lengths.numel() > 0
+                else []
+            )
+            for row, (j, _) in enumerate(group):
+                audio_len = int(audio_lengths_list[row]) if len(audio_lengths_list) > row else int(wav_batch.shape[-1])
+                wav_tensors[j] = wav_batch[row, 0, :audio_len]
+
+        for bounded_group in _iter_bounded_groups(list(enumerate(valid_codes_qf))):
+            _decode_group(bounded_group)
 
         audios: list[torch.Tensor] = [empty] * num_req
         srs = [sr_tensor] * num_req
@@ -305,8 +451,8 @@ class FishSpeechDACDecoder(nn.Module):
         for j, idx in enumerate(valid_indices):
             ctx_frames = parsed_ctx_frames[idx]
             total_frames = parsed_total_frames[idx]
-            audio_len = int(audio_lengths[j].item()) if audio_lengths.numel() > j else int(wav_batch.shape[-1])
-            wav = wav_batch[j, 0, :audio_len]
+            wav = wav_tensors[j]
+            assert wav is not None
             # Trim context frames (left overlap for streaming).
             if ctx_frames > 0:
                 # Decode length may deviate from (frames * hop_length) due to model
diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py
index 633652c92..46ceee498 100644
--- a/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py
+++ b/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py
@@ -4,14 +4,13 @@
 autoregressively after each Slow AR step.  Analogous to Qwen3 TTS's
 CodePredictor but with its own embedding table, RoPE, and output head.
 
-Uses re-prefill (no KV cache): each AR step forwards the full growing
-sequence through the 4-layer transformer.  This trades ~O(T^2) attention
-FLOPs (negligible for T=10, 4 layers) for zero KV cache management.
+Uses a tiny per-call KV cache: each AR step forwards one codebook token while
+reusing the previous Fast AR keys/values.
 
 Optimisations:
-  - torch.compile on model forward (kernel fusion)
   - Pre-allocated projection / embedding buffer
   - Pre-allocated position_ids
+  - Pre-allocated per-layer KV cache for the short codebook loop
   - Inline sampling
 """
 
@@ -51,7 +50,8 @@ logger = init_logger(__name__)
 class _FastARAttention(nn.Module):
     """Multi-head attention using F.scaled_dot_product_attention (SDPA).
 
-    Supports fused QKV, RoPE, optional q/k normalization, and native GQA.
+    Supports fused QKV, RoPE, optional q/k normalization, native GQA, and a
+    small one-token decode path for the residual codebook loop.
     Input: [B, seq_len, hidden_size].
     """
 
@@ -124,6 +124,43 @@ class _FastARAttention(nn.Module):
         output, _ = self.o_proj(attn_out)
         return output.view(bsz, seq_len, -1)
 
+    def forward_one(
+        self,
+        hidden_states: torch.Tensor,
+        position_ids: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        cache_pos: int,
+    ) -> torch.Tensor:
+        bsz = int(hidden_states.shape[0])
+
+        qkv, _ = self.qkv_proj(hidden_states.reshape(bsz, -1))
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+
+        if self.q_norm is not None:
+            q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view(q.shape)
+            k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view(k.shape)
+
+        q, k = self.rotary_emb(position_ids.reshape(-1), q, k)
+
+        q = q.view(bsz, self.num_heads, self.head_dim).unsqueeze(2)
+        k = k.view(bsz, self.num_kv_heads, self.head_dim)
+        v = v.view(bsz, self.num_kv_heads, self.head_dim)
+        k_cache[:bsz, :, cache_pos, :] = k
+        v_cache[:bsz, :, cache_pos, :] = v
+
+        attn_out = F.scaled_dot_product_attention(
+            q,
+            k_cache[:bsz, :, : cache_pos + 1, :],
+            v_cache[:bsz, :, : cache_pos + 1, :],
+            scale=self.scaling,
+            is_causal=False,
+            enable_gqa=self._use_gqa,
+        )
+        attn_out = attn_out.transpose(1, 2).reshape(bsz, -1)
+        output, _ = self.o_proj(attn_out)
+        return output
+
 
 class _FastARMLP(nn.Module):
     """SiLU-gated MLP, matching Qwen3/LLaMA MLP structure."""
@@ -175,6 +212,25 @@ class _FastARDecoderLayer(nn.Module):
         hidden_states = residual + hidden_states
         return hidden_states
 
+    def forward_one(
+        self,
+        hidden_states: torch.Tensor,
+        position_ids: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        cache_pos: int,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn.forward_one(hidden_states, position_ids, k_cache, v_cache, cache_pos)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
 
 # ===================================================================
 #  Fast AR Transformer Model
@@ -199,6 +255,25 @@ class FishSpeechFastARModel(nn.Module):
             hidden_states = layer(hidden_states, position_ids)
         return hidden_states
 
+    def forward_one(
+        self,
+        input_embed: torch.Tensor,
+        position_ids: torch.Tensor,
+        k_caches: torch.Tensor,
+        v_caches: torch.Tensor,
+        cache_pos: int,
+    ) -> torch.Tensor:
+        hidden_states = input_embed
+        for layer_idx, layer in enumerate(self.layers):
+            hidden_states = layer.forward_one(
+                hidden_states,
+                position_ids,
+                k_caches[layer_idx],
+                v_caches[layer_idx],
+                cache_pos,
+            )
+        return hidden_states
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         stacked_params_mapping = [
             ("qkv_proj", "q_proj", "q"),
@@ -259,13 +334,13 @@ class FishSpeechFastARModel(nn.Module):
 class FishSpeechFastAR(nn.Module):
     """vLLM-native Fast AR for Fish Speech (residual codebooks).
 
-    Re-prefill approach: each AR step forwards the full growing sequence
-    through the 4-layer transformer.  No KV cache needed.
+    Decode approach: each residual-codebook step forwards one token through
+    the 4-layer transformer with a per-call KV cache.
 
     Optimisations over baseline:
-      1. torch.compile on model forward -- kernel fusion.
-      2. Pre-allocated embedding buffer [B, max_seq, H].
-      3. Pre-allocated position_ids.
+      1. Pre-allocated embedding buffer [B, max_seq, H].
+      2. Pre-allocated position_ids.
+      3. Per-call KV cache for the short residual-codebook loop.
       4. Inline sampling.
     """
 
@@ -307,6 +382,8 @@ class FishSpeechFastAR(nn.Module):
         # Pre-allocated buffers (lazily initialised on first forward).
         self._embed_buf: torch.Tensor | None = None
         self._pos_ids: torch.Tensor | None = None
+        self._k_cache: torch.Tensor | None = None
+        self._v_cache: torch.Tensor | None = None
         self._compiled_model_fwd: object | None = None
         self._compile_attempted = False
         self._compile_failed = False
@@ -316,13 +393,27 @@ class FishSpeechFastAR(nn.Module):
         max_seq = self._num_codebooks + 1  # hidden_state + num_codebooks codes
         if (
             self._embed_buf is not None
+            and self._pos_ids is not None
+            and self._k_cache is not None
+            and self._v_cache is not None
             and self._embed_buf.shape[0] >= bsz
             and self._embed_buf.device == device
             and self._embed_buf.dtype == dtype
+            and self._pos_ids.shape[0] >= bsz
+            and self._pos_ids.shape[1] == max_seq
+            and self._pos_ids.device == device
+            and self._k_cache.shape[1] >= bsz
+            and self._k_cache.device == device
+            and self._k_cache.dtype == dtype
         ):
             return
         self._embed_buf = torch.zeros(bsz, max_seq, self._fast_dim, dtype=dtype, device=device)
-        self._pos_ids = torch.arange(max_seq, dtype=torch.long, device=device)
+        self._pos_ids = torch.arange(max_seq, dtype=torch.long, device=device).unsqueeze(0).expand(bsz, -1).contiguous()
+        num_layers = self.config.num_hidden_layers
+        num_kv_heads = self.config.num_key_value_heads
+        head_dim = self.config.head_dim
+        self._k_cache = torch.empty(num_layers, bsz, num_kv_heads, max_seq, head_dim, dtype=dtype, device=device)
+        self._v_cache = torch.empty_like(self._k_cache)
 
     def _setup_compile(self) -> None:
         if self._compile_attempted:
@@ -375,20 +466,10 @@ class FishSpeechFastAR(nn.Module):
         torch.accelerator.synchronize(device)
 
     @torch.inference_mode()
-    def _run_model(self, step_input: torch.Tensor, step_pos_ids: torch.Tensor, bsz: int) -> torch.Tensor:
-        if self._disable_compile_for_graph:
-            model_fwd = self._compiled_model_fwd or self.model.forward
-        else:
-            model_fwd = self._compiled_model_fwd if bsz == 1 else self.model.forward
-        try:
-            return model_fwd(step_input, step_pos_ids)
-        except Exception as exc:
-            if model_fwd is self.model.forward or self._compile_failed:
-                raise
-            self._compile_failed = True
-            self._compiled_model_fwd = self.model.forward
-            logger.warning("Fish Speech Fast AR torch.compile fallback to eager after runtime failure: %s", exc)
-            return self.model.forward(step_input, step_pos_ids)
+    def _run_model_one(self, input_embed: torch.Tensor, step_pos_ids: torch.Tensor, cache_pos: int) -> torch.Tensor:
+        assert self._k_cache is not None
+        assert self._v_cache is not None
+        return self.model.forward_one(input_embed, step_pos_ids, self._k_cache, self._v_cache, cache_pos)
 
     @torch.inference_mode()
     def forward(
@@ -401,6 +482,7 @@ class FishSpeechFastAR(nn.Module):
         top_k: int = 30,
         top_p: float = 0.9,
         seed: int | None = None,
+        generator: torch.Generator | None = None,
     ) -> torch.Tensor:
         """Predict residual codebook codes 0..num_codebooks-1 autoregressively.
 
@@ -428,7 +510,6 @@ class FishSpeechFastAR(nn.Module):
         all_codes[:, 0] = semantic_code
 
         self._ensure_buffers(bsz, device, dtype)
-        self._setup_compile()
 
         embed_buf = self._embed_buf
         pos_ids = self._pos_ids
@@ -445,8 +526,7 @@ class FishSpeechFastAR(nn.Module):
         inv_temperature = 1.0 / max(temperature, 1e-6) if use_sampling else 0.0
 
         # Create a seeded generator for deterministic residual codebook sampling.
-        generator = None
-        if seed is not None and use_sampling:
+        if generator is None and seed is not None and use_sampling:
             generator = torch.Generator(device=device)
             generator.manual_seed(seed)
 
@@ -455,15 +535,11 @@ class FishSpeechFastAR(nn.Module):
         # codebooks only have 1024 entries.  Truncate logits for steps > 0.
         residual_codebook_size = 1024
 
-        for step in range(1, num_cb):
-            seq_len = step + 1
-            step_input = embed_buf[:bsz, :seq_len, :]
-            # Use a dense 2D position tensor for every batch size; stride-0
-            # views from expand() were fragile under compiled execution.
-            step_pos_ids = pos_ids[:seq_len].unsqueeze(0).repeat(bsz, 1)
+        self._run_model_one(embed_buf[:bsz, 0, :], pos_ids[:bsz, 0], 0)
 
-            hidden_out = self._run_model(step_input, step_pos_ids, bsz)
-            logits = self.fast_output(self.fast_norm(hidden_out[:, -1, :]))
+        for step in range(1, num_cb):
+            hidden_out = self._run_model_one(embed_buf[:bsz, step, :], pos_ids[:bsz, step], step)
+            logits = self.fast_output(self.fast_norm(hidden_out))
 
             # Residual codebooks (step >= 1) only have 1024 entries.
             if step >= 1:
diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py
index 24182b93e..1317dc044 100644
--- a/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py
+++ b/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py
@@ -199,6 +199,12 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module):
 
         # Qwen3 transformer backbone.
         self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
+        from vllm_omni.attention.fish_kvcache_attn import is_fish_kvcache_attn_enabled
+
+        if is_fish_kvcache_attn_enabled():
+            from vllm_omni.attention.fish_kvcache_backend import install_fish_kvcache_attn_backend
+
+            install_fish_kvcache_attn_backend(self.model)
 
         # Fish Speech uses interleaved (GPT-J) RoPE, not NeoX style.
         # vLLM's Qwen3Attention defaults to NeoX (is_neox_style=True).
@@ -672,14 +678,19 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module):
         past_hidden = last_talker_hidden.reshape(bsz, -1).to(dtype=torch.bfloat16, device=dev)
 
         # Run Fast AR to predict all num_codebooks codes.
+        do_sample = kwargs.get("do_sample")
+        temperature = kwargs.get("temperature")
+        top_k = kwargs.get("top_k")
+        top_p = kwargs.get("top_p")
         audio_codes = self.fast_ar(
             slow_ar_hidden=past_hidden,
             semantic_token_id=input_ids.reshape(bsz),
-            do_sample=True,
-            temperature=0.8,
-            top_k=30,
-            top_p=0.9,
+            do_sample=True if do_sample is None else bool(do_sample),
+            temperature=0.8 if temperature is None else float(temperature),
+            top_k=30 if top_k is None else int(top_k),
+            top_p=0.9 if top_p is None else float(top_p),
             seed=seed,
+            generator=kwargs.get("generator"),
         )  # [B, num_codebooks]
 
         # Add codebook embeddings to the input embedding (from preprocess).
diff --git a/vllm_omni/model_executor/models/glm_tts/__init__.py b/vllm_omni/model_executor/models/glm_tts/__init__.py
new file mode 100644
index 000000000..9618227ca
--- /dev/null
+++ b/vllm_omni/model_executor/models/glm_tts/__init__.py
@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GLM-TTS model support for vLLM-Omni.
+
+GLM-TTS is a two-stage text-to-speech system:
+  - Stage 0 (LLM): Llama-based AR model generates speech tokens from text
+  - Stage 1 (DiT): Flow matching model converts speech tokens to mel-spectrogram
+
+Reference: https://github.com/zai-org/GLM-TTS
+"""
+
+from vllm_omni.model_executor.models.glm_tts.glm_tts import (
+    GLMTTSForConditionalGeneration,
+)
+
+__all__ = [
+    "GLMTTSForConditionalGeneration",
+]
diff --git a/vllm_omni/model_executor/models/glm_tts/glm_tts.py b/vllm_omni/model_executor/models/glm_tts/glm_tts.py
new file mode 100644
index 000000000..6e14840e2
--- /dev/null
+++ b/vllm_omni/model_executor/models/glm_tts/glm_tts.py
@@ -0,0 +1,1423 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GLM-TTS AR Model (Stage 0): Text → Speech Tokens.
+
+Based on Llama architecture, generates speech token sequences from input text.
+Analogous to Fish Speech Slow AR and Qwen3-TTS Talker models.
+"""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Iterable, Mapping, Sequence
+from dataclasses import replace
+from typing import Any
+
+import torch
+import torch.nn as nn
+from transformers import AutoTokenizer
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.utils.hub import cached_file
+from vllm.config import VllmConfig
+from vllm.config.multimodal import BaseDummyOptions
+from vllm.distributed import get_pp_group
+from vllm.inputs import MultiModalDataDict
+from vllm.logger import init_logger
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.interfaces import SupportsMultiModal
+from vllm.model_executor.models.llama import LlamaModel
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    WeightsMapper,
+    maybe_prefix,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
+from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
+from vllm.multimodal.processing import (
+    BaseDummyInputsBuilder,
+    BaseMultiModalProcessor,
+    BaseProcessingInfo,
+    ProcessorInputs,
+    PromptIndexTargets,
+    PromptInsertion,
+    PromptUpdate,
+)
+from vllm.sequence import IntermediateTensors
+from vllm.v1.outputs import SamplerOutput
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.sample.sampler import Sampler
+
+from vllm_omni.model_executor.models.common.nucleus_ras_sampling import ras_sample_one as _ras_sample_one
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.transformers_utils.configs.glm_tts import GLMTTSConfig
+from vllm_omni.utils.speaker_cache import get_speaker_cache
+
+from .text_frontend import GLMTTSTextFrontend
+from .voice_clone import (
+    extract_prompt_feat,
+    extract_prompt_speech_token,
+    extract_spk_embedding,
+    load_voice_clone_frontend,
+)
+
+logger = init_logger(__name__)
+
+_GLM_TTS_DEFAULT_REPO_ID = "zai-org/GLM-TTS"
+_GLM_TTS_TOKENIZER_SUBDIR = "vq32k-phoneme-tokenizer"
+_GLM_TTS_MAX_PROMPT_SPEECH_TOKENS = 1024
+
+
+def _req_float(param: torch.Tensor | None, req_idx: int, default: float) -> float:
+    if param is None or param.numel() == 0:
+        return default
+    index = min(req_idx, int(param.numel()) - 1)
+    return float(param.reshape(-1)[index].item())
+
+
+def _tensor_param_values(
+    param: torch.Tensor | None,
+    count: int,
+    default: float,
+) -> list[float]:
+    if param is None or param.numel() == 0:
+        return [default] * count
+    values = param.reshape(-1).detach().cpu().tolist()
+    if not values:
+        return [default] * count
+    if len(values) < count:
+        values.extend([values[-1]] * (count - len(values)))
+    return [float(value) for value in values[:count]]
+
+
+def resolve_glm_tts_tokenizer_path(model_name_or_path: Any) -> str:
+    """Resolve vq32k-phoneme-tokenizer directory to a local path.
+
+    For local model dirs, scans for the tokenizer subfolder on disk.
+    For HF hub names, downloads the tokenizer subfolder via
+    ``snapshot_download`` (same pattern as CosyVoice3's
+    ``_ensure_cached_runtime_components`` and
+    ``arg_utils._TOKENIZER_SUBFOLDER_MAP``).
+    """
+    model_path = os.fspath(model_name_or_path)
+    if os.path.exists(model_path):
+        for candidate in [
+            os.path.join(model_path, _GLM_TTS_TOKENIZER_SUBDIR),
+            os.path.join(os.path.dirname(model_path), _GLM_TTS_TOKENIZER_SUBDIR),
+        ]:
+            if os.path.isdir(candidate):
+                return candidate
+        return model_path
+
+    from huggingface_hub import snapshot_download
+
+    local_dir = snapshot_download(
+        model_path,
+        allow_patterns=[
+            f"{_GLM_TTS_TOKENIZER_SUBDIR}/tokenizer*",
+            f"{_GLM_TTS_TOKENIZER_SUBDIR}/tokenization_*",
+            f"{_GLM_TTS_TOKENIZER_SUBDIR}/special_tokens*",
+            f"{_GLM_TTS_TOKENIZER_SUBDIR}/vocab*",
+            f"{_GLM_TTS_TOKENIZER_SUBDIR}/merges*",
+            f"{_GLM_TTS_TOKENIZER_SUBDIR}/added_tokens*",
+        ],
+    )
+    candidate = os.path.join(local_dir, _GLM_TTS_TOKENIZER_SUBDIR)
+    if os.path.isdir(candidate):
+        return candidate
+    return model_path
+
+
+def resolve_glm_tts_model_dir(
+    model_name_or_path: Any,
+    *,
+    tokenizer_path: Any | None = None,
+    required_files: Sequence[str] = (),
+    optional_files: Sequence[str] = (),
+) -> str:
+    """Resolve GLM-TTS repo root: local dir → tokenizer parent → cached_file → snapshot."""
+    model_path = os.fspath(model_name_or_path)
+    if os.path.isdir(model_path):
+        return os.path.abspath(model_path)
+
+    def _has_files(root: str) -> bool:
+        has_req = all(os.path.isfile(os.path.join(root, *f.split("/"))) for f in required_files)
+        has_opt = not optional_files or any(os.path.isfile(os.path.join(root, *f.split("/"))) for f in optional_files)
+        return has_req and has_opt
+
+    # Try tokenizer parent directory
+    if tokenizer_path:
+        tp = os.path.abspath(os.path.normpath(os.fspath(tokenizer_path)))
+        root = os.path.dirname(tp) if os.path.basename(tp) == _GLM_TTS_TOKENIZER_SUBDIR else tp
+        if os.path.isdir(root) and _has_files(root):
+            return root
+
+    # Try cached_file to locate individual files without full snapshot
+    for filename in list(required_files) + list(optional_files):
+        try:
+            resolved = cached_file(model_name_or_path, filename)
+            if resolved:
+                root = os.path.abspath(resolved)
+                for _ in filename.split("/"):
+                    root = os.path.dirname(root)
+                if _has_files(root):
+                    return root
+        except Exception:
+            pass
+
+    from huggingface_hub import snapshot_download
+
+    return snapshot_download(model_name_or_path)
+
+
+def _first_glm_tts_value(value: Any) -> Any:
+    return value[0] if isinstance(value, list) and value else value
+
+
+def _glm_tts_int_value(value: Any) -> int | None:
+    """Extract a scalar integer from vLLM multimodal/runtime wrappers."""
+    value = _first_glm_tts_value(value)
+    if isinstance(value, torch.Tensor):
+        if value.numel() == 0:
+            return None
+        return int(value.reshape(-1)[0].item())
+    if isinstance(value, (int, float)):
+        return int(value)
+    return None
+
+
+def _decode_glm_tts_audio_data(value: Any) -> tuple[torch.Tensor | None, int | None]:
+    """Extract (waveform_tensor, sample_rate) from multimodal input."""
+    if isinstance(value, list) and len(value) == 1:
+        value = value[0]
+    if isinstance(value, torch.Tensor):
+        return value, None
+    if isinstance(value, (tuple, list)) and len(value) == 2 and isinstance(value[1], (int, float)):
+        return torch.as_tensor(value[0]), int(value[1])
+    if hasattr(value, "shape"):  # numpy array
+        return torch.as_tensor(value), None
+    return None, None
+
+
+_glm_tts_tokenizer_cache: dict[str, Any] = {}
+
+
+def load_glm_tts_tokenizer(
+    tokenizer_path: Any,
+    *,
+    model_name_or_path: Any | None = None,
+    trust_remote_code: bool = False,
+    **kwargs: Any,
+) -> Any:
+    """Load GLM-TTS phoneme tokenizer (ChatGLM4Tokenizer, slow-only).
+
+    Expects a resolved local directory (use ``resolve_glm_tts_tokenizer_path``
+    to handle HF hub names first).  Results are cached by *tokenizer_path*.
+    """
+    cache_key = str(tokenizer_path)
+    cached = _glm_tts_tokenizer_cache.get(cache_key)
+    if cached is not None:
+        return cached
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        tokenizer_path,
+        use_fast=False,
+        trust_remote_code=trust_remote_code,
+        **kwargs,
+    )
+    _glm_tts_tokenizer_cache[cache_key] = tokenizer
+    return tokenizer
+
+
+def get_glm_tts_special_token_ids(tokenizer: Any) -> dict[str, int]:
+    """Resolve GLM-TTS special IDs from the phoneme/audio tokenizer."""
+    special_tokens = {
+        "ats": "<|audio_0|>",
+        "ate": "<|audio_32767|>",
+        "boa": "<|begin_of_audio|>",
+        "eoa": "<|user|>",
+        "pad": "<|endoftext|>",
+    }
+
+    result: dict[str, int] = {}
+    for key, token_str in special_tokens.items():
+        token_ids = tokenizer.encode(token_str, add_special_tokens=False)
+        if len(token_ids) != 1:
+            raise ValueError(f"Token '{key}' ({token_str}) should encode to single ID, got: {token_ids}")
+        result[key] = int(token_ids[0])
+    return result
+
+
+def resolve_glm_tts_campplus_path(model_dir: str) -> str | None:
+    """Resolve ``campplus.onnx`` for GLM-TTS voice cloning."""
+    local = os.path.join(model_dir, "frontend", "campplus.onnx")
+    if os.path.isfile(local):
+        return local
+
+    try:
+        resolved = cached_file("FunAudioLLM/CosyVoice-300M", "campplus.onnx")
+        if resolved is not None:
+            logger.info("Resolved campplus.onnx from FunAudioLLM/CosyVoice-300M: %s", resolved)
+            return resolved
+    except Exception:
+        logger.debug("cached_file could not fetch campplus.onnx", exc_info=True)
+
+    logger.warning(
+        "campplus.onnx not found locally or could not be downloaded. "
+        "Voice cloning speaker embedding will not be available.",
+    )
+    return None
+
+
+def _normalize_glm_tts_processor_text(
+    frontend: GLMTTSTextFrontend,
+    value: str | None,
+    *,
+    add_trailing_space: bool = False,
+) -> str:
+    if value is None:
+        return ""
+    normalized = frontend.text_normalize(value)
+    normalized = (normalized or value).strip()
+    if add_trailing_space and normalized:
+        normalized = f"{normalized} "
+    return normalized
+
+
+_glm_tts_text_frontend: GLMTTSTextFrontend | None = None
+
+
+def _get_text_frontend() -> GLMTTSTextFrontend:
+    global _glm_tts_text_frontend
+    if _glm_tts_text_frontend is None:
+        _glm_tts_text_frontend = GLMTTSTextFrontend()
+    return _glm_tts_text_frontend
+
+
+def build_glm_tts_prefill_metadata(
+    model_name_or_path: Any,
+    text: str,
+    prompt_text: str | None,
+    *,
+    tokenizer_path: Any | None = None,
+    trust_remote_code: bool = False,
+) -> dict[str, list[int]]:
+    """Build request-local GLM-TTS length metadata for AR prefill.
+
+    These scalars are mirrored into ``model_intermediate_buffer`` so the model
+    preprocess hook can initialize request state from the metadata alone.
+    """
+    if tokenizer_path is None:
+        tokenizer_path = resolve_glm_tts_tokenizer_path(model_name_or_path)
+    tokenizer = load_glm_tts_tokenizer(
+        tokenizer_path,
+        model_name_or_path=model_name_or_path,
+        trust_remote_code=trust_remote_code,
+    )
+    frontend = _get_text_frontend()
+    text_len = max(1, len(tokenizer.encode(_normalize_glm_tts_processor_text(frontend, text))))
+    prompt_text_len = 0
+    if prompt_text:
+        prompt_text_len = max(
+            1,
+            len(
+                tokenizer.encode(
+                    _normalize_glm_tts_processor_text(
+                        frontend,
+                        prompt_text,
+                        add_trailing_space=True,
+                    )
+                )
+            ),
+        )
+    return {
+        "glm_tts_text_token_len": [text_len],
+        "glm_tts_prompt_text_token_len": [prompt_text_len],
+        "input_len": [prompt_text_len + text_len + 1],
+    }
+
+
+class GLMTTSMultiModalProcessingInfo(BaseProcessingInfo):
+    def get_hf_config(self):
+        return self.ctx.get_hf_config(GLMTTSConfig)
+
+    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+        return {"audio": 1}
+
+    def get_mm_max_tokens_per_item(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+    ) -> Mapping[str, int] | None:
+        if mm_counts.get("audio", 0) <= 0:
+            return {}
+        return {"audio": min(seq_len, _GLM_TTS_MAX_PROMPT_SPEECH_TOKENS)}
+
+    def get_data_parser(self):
+        return MultiModalDataParser(
+            target_sr=24000,
+            expected_hidden_size=self._get_expected_hidden_size(),
+        )
+
+
+class GLMTTSMultiModalProcessor(BaseMultiModalProcessor[GLMTTSMultiModalProcessingInfo]):
+    """GLM-TTS voice-clone processor.
+
+    Unlike CosyVoice3, GLM-TTS prompt speech tokens are normal Llama vocab IDs
+    (``<|audio_N|>``).  The processor exposes them as multimodal embeddings for
+    the AR prompt and also carries WhisperVQ/CampPlus outputs to the AR->DiT
+    handoff.
+    """
+
+    def _ensure_cached_runtime_components(
+        self,
+        model_dir: str,
+        config: GLMTTSConfig,
+        *,
+        load_voice_clone: bool = True,
+    ) -> None:
+        # Skip if already loaded for this model_dir + mode
+        cache_key = (model_dir, load_voice_clone)
+        if getattr(self, "_loaded_cache_key", None) == cache_key:
+            return
+
+        tokenizer_path = getattr(self.info.ctx.model_config, "tokenizer", None)
+        required = ("speech_tokenizer/config.json", "speech_tokenizer/model.safetensors") if load_voice_clone else ()
+        model_dir = resolve_glm_tts_model_dir(
+            model_dir,
+            tokenizer_path=tokenizer_path,
+            required_files=required,
+        )
+
+        # Load text tokenizer + frontend (always needed)
+        if getattr(self, "_text_model_dir", None) != model_dir:
+            tok_path = tokenizer_path
+            if not tok_path or not os.path.isdir(os.fspath(tok_path)):
+                tok_path = os.path.join(model_dir, _GLM_TTS_TOKENIZER_SUBDIR)
+                if not os.path.isdir(tok_path):
+                    tok_path = model_dir
+            trust_remote_code = bool(getattr(self.info.ctx.model_config, "trust_remote_code", False))
+            self.tokenizer = load_glm_tts_tokenizer(
+                tok_path,
+                model_name_or_path=self.info.ctx.model_config.model,
+                trust_remote_code=trust_remote_code,
+            )
+            self.special_ids = get_glm_tts_special_token_ids(self.tokenizer)
+            self.text_frontend = GLMTTSTextFrontend()
+            self._text_model_dir = model_dir
+
+        # Load voice clone components (speech tokenizer + CampPlus)
+        if load_voice_clone:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            self.processor_device = device
+            campplus_path = getattr(self, "_campplus_path", None) or resolve_glm_tts_campplus_path(model_dir)
+            self._campplus_path = campplus_path
+            self.speech_tokenizer, self.campplus_session = load_voice_clone_frontend(
+                model_dir,
+                device,
+                speech_tokenizer_cache=getattr(self, "speech_tokenizer", None),
+                campplus_cache=getattr(self, "campplus_session", None),
+                campplus_path=campplus_path,
+            )
+
+        if not hasattr(self, "_speaker_cache"):
+            self._speaker_cache = get_speaker_cache()
+
+        self._loaded_cache_key = cache_key
+
+    def _encode_text(self, text: str) -> torch.Tensor:
+        token_ids = self.tokenizer.encode(text)
+        return torch.tensor([token_ids], dtype=torch.long)
+
+    def _get_audio(self, mm_data: Mapping[str, object]) -> tuple[torch.Tensor | None, int | None]:
+        audio = mm_data.get("audio")
+        if audio is None:
+            audios = mm_data.get("audios")
+            if isinstance(audios, (list, tuple)) and audios:
+                # Match CosyVoice3's compatibility behavior and accept the
+                # first item from the legacy plural field name.
+                audio = audios[0]
+            else:
+                audio = audios
+        wav, sr = _decode_glm_tts_audio_data(audio)
+        if wav is not None and wav.ndim > 1:
+            wav = wav.float()
+            wav = wav.mean(dim=0) if wav.shape[0] <= wav.shape[-1] else wav.mean(dim=-1)
+        return wav, int(sr or 24000) if wav is not None else sr
+
+    def _call_hf_processor(
+        self,
+        prompt: str,
+        mm_data: Mapping[str, object],
+        mm_kwargs: Mapping[str, object],
+        tok_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        config = self.info.ctx.get_hf_config()
+        model_dir = self.info.ctx.model_config.model
+        wav, sr = self._get_audio(mm_data)
+        if wav is None:
+            audio_supplied = mm_data.get("audio") is not None or mm_data.get("audios") is not None
+            assert not audio_supplied, (
+                "GLM-TTS received an audio multimodal payload that could not be decoded. "
+                "Runtime voice cloning must provide decodable multi_modal_data['audio']; "
+                "the text-only processor path is reserved for profiling/cache."
+            )
+            self._ensure_cached_runtime_components(model_dir, config, load_voice_clone=False)
+            normalized_text = _normalize_glm_tts_processor_text(self.text_frontend, prompt)
+            text_ids = self._encode_text(normalized_text)
+            dummy_audio_len = min(
+                _GLM_TTS_MAX_PROMPT_SPEECH_TOKENS,
+                max(1, int(getattr(config, "max_position_embeddings", _GLM_TTS_MAX_PROMPT_SPEECH_TOKENS))),
+            )
+            boa_tensor = torch.tensor([[int(self.special_ids["boa"])]], dtype=torch.long)
+            input_ids = torch.cat([text_ids, boa_tensor], dim=1)
+            # vLLM profiling/cache setup may call the processor without an
+            # audio item. Do not fabricate a dummy wav here: for GLM-TTS that
+            # would exercise the real voice-clone frontend during profiling.
+            # This shape-only branch is not a supported runtime inference mode.
+            return BatchFeature(
+                {
+                    "input_ids": input_ids,
+                    "prompt_speech_token": torch.zeros((1, dummy_audio_len), dtype=torch.long),
+                    "prompt_speech_token_len": [torch.tensor([dummy_audio_len], dtype=torch.long)],
+                    "glm_tts_prompt_text_token_len": [torch.tensor([0], dtype=torch.long)],
+                    "glm_tts_text_token_len": [torch.tensor([int(text_ids.shape[1])], dtype=torch.long)],
+                }
+            )
+
+        self._ensure_cached_runtime_components(model_dir, config)
+        normalized_text = _normalize_glm_tts_processor_text(self.text_frontend, prompt)
+        text_ids = self._encode_text(normalized_text)
+        prompt_text = mm_kwargs.get("prompt_text")
+        if not isinstance(prompt_text, str) or not prompt_text.strip():
+            raise ValueError("GLM-TTS voice cloning requires mm_processor_kwargs['prompt_text'].")
+
+        normalized_prompt_text = _normalize_glm_tts_processor_text(
+            self.text_frontend,
+            prompt_text,
+            add_trailing_space=True,
+        )
+        prompt_text_ids = self._encode_text(normalized_prompt_text)
+
+        # Speaker cache: skip WhisperVQ + mel + CampPlus on hit
+        voice_name = mm_kwargs.get("voice_name")
+        cache_key = None
+        if voice_name and isinstance(voice_name, str):
+            cache_key = self._speaker_cache.make_cache_key(
+                voice_name,
+                model_type="glm_tts",
+                created_at=int(mm_kwargs.get("voice_created_at") or 0),
+            )
+            cached = self._speaker_cache.get(cache_key)
+            if cached is not None:
+                boa_tensor = torch.tensor([[int(self.special_ids["boa"])]], dtype=torch.long)
+                input_ids = torch.cat([prompt_text_ids, text_ids, boa_tensor], dim=1)
+                return BatchFeature(
+                    {
+                        "input_ids": input_ids,
+                        "prompt_speech_token": cached["prompt_speech_token"].clone(),
+                        "prompt_speech_token_len": [cached["prompt_speech_token_len"].clone()],
+                        "glm_tts_prompt_text_token_len": [
+                            torch.tensor([int(prompt_text_ids.shape[1])], dtype=torch.long)
+                        ],
+                        "prompt_feat": cached["prompt_feat"].clone(),
+                        "embedding": cached["embedding"].clone(),
+                        "glm_tts_text_token_len": [torch.tensor([int(text_ids.shape[1])], dtype=torch.long)],
+                    }
+                )
+
+        prompt_speech_token = extract_prompt_speech_token(wav, int(sr or 24000), self.speech_tokenizer)
+        if not prompt_speech_token:
+            raise RuntimeError("GLM-TTS failed to extract WhisperVQ prompt speech tokens from ref_audio.")
+
+        prompt_speech_token_tensor = torch.tensor([prompt_speech_token], dtype=torch.long)
+        boa_tensor = torch.tensor([[int(self.special_ids["boa"])]], dtype=torch.long)
+        input_ids = torch.cat([prompt_text_ids, text_ids, boa_tensor], dim=1)
+        logger.info(
+            "GLM-TTS processor prompt: prompt_text_tokens=%d text_tokens=%d "
+            "prompt_speech_tokens=%d input_tokens_before_audio=%d expected_prefill_tokens=%d",
+            int(prompt_text_ids.shape[1]),
+            int(text_ids.shape[1]),
+            len(prompt_speech_token),
+            int(input_ids.shape[1]),
+            int(input_ids.shape[1]) + len(prompt_speech_token),
+        )
+
+        prompt_feat = extract_prompt_feat(wav, int(sr or 24000), self.processor_device)
+        if prompt_feat is None:
+            raise RuntimeError("GLM-TTS failed to extract prompt mel features from ref_audio.")
+        embedding = extract_spk_embedding(wav, int(sr or 24000), self.campplus_session)
+        if embedding is None:
+            raise RuntimeError("GLM-TTS failed to extract CampPlus speaker embedding from ref_audio.")
+
+        prompt_speech_token_len = torch.tensor([len(prompt_speech_token)], dtype=torch.long)
+        prompt_feat_cpu = prompt_feat.detach().to("cpu").unsqueeze(0).contiguous()
+        embedding_tensor = torch.tensor([embedding], dtype=torch.float32)
+
+        if cache_key is not None:
+            self._speaker_cache.put(
+                cache_key,
+                {
+                    "prompt_speech_token": prompt_speech_token_tensor.detach().cpu(),
+                    "prompt_speech_token_len": prompt_speech_token_len.detach().cpu(),
+                    "prompt_feat": prompt_feat_cpu.detach().cpu(),
+                    "embedding": embedding_tensor.detach().cpu(),
+                },
+            )
+
+        return BatchFeature(
+            {
+                "input_ids": input_ids,
+                "prompt_speech_token": prompt_speech_token_tensor,
+                "prompt_speech_token_len": [prompt_speech_token_len],
+                "glm_tts_prompt_text_token_len": [torch.tensor([int(prompt_text_ids.shape[1])], dtype=torch.long)],
+                "prompt_feat": prompt_feat_cpu,
+                "embedding": embedding_tensor,
+                "glm_tts_text_token_len": [torch.tensor([int(text_ids.shape[1])], dtype=torch.long)],
+            }
+        )
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs: BatchFeature,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        return {
+            key: MultiModalFieldConfig.batched("audio")
+            for key in (
+                "prompt_speech_token",
+                "prompt_speech_token_len",
+                "glm_tts_prompt_text_token_len",
+                "prompt_feat",
+                "embedding",
+                "glm_tts_text_token_len",
+            )
+            if key in hf_inputs
+        }
+
+    def _hf_processor_applies_updates(
+        self,
+        prompt_text: str,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        tokenization_kwargs: Mapping[str, object],
+    ) -> bool:
+        return False
+
+    def _cached_apply_hf_processor(self, inputs: ProcessorInputs, timing_ctx: Any):
+        # GLM-TTS builds the actual AR prompt from both the request text and
+        # mm_processor_kwargs["prompt_text"]. The base cache path separates text
+        # processing from audio processing, which drops the reference transcript
+        # and BOA token from the final prompt.
+        return self._apply_hf_processor(inputs, timing_ctx)
+
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        out_mm_kwargs: MultiModalKwargsItems,
+    ) -> Sequence[PromptUpdate]:
+        def insertion_end(item_idx: int) -> list[int]:
+            audio_items = out_mm_kwargs["audio"]
+            item = audio_items[item_idx] if item_idx < len(audio_items) else audio_items[0]
+            token_len = item["prompt_speech_token_len"].data[0].item()
+            return [1] * int(token_len)
+
+        return [
+            PromptInsertion(
+                modality="audio",
+                target=PromptIndexTargets.end(),
+                insertion=insertion_end,
+            )
+        ]
+
+
+class GLMTTSDummyInputsBuilder(BaseDummyInputsBuilder[GLMTTSMultiModalProcessingInfo]):
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        return "This is a test of the GLM-TTS voice cloning system."
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, BaseDummyOptions] | None = None,
+    ) -> MultiModalDataDict:
+        return {}
+
+    def get_dummy_processor_inputs(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, BaseDummyOptions] | None = None,
+    ) -> ProcessorInputs:
+        inputs = super().get_dummy_processor_inputs(seq_len, mm_counts, mm_options)
+        inputs.hf_processor_mm_kwargs = {"prompt_text": "This is the reference voice."}
+        return inputs
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    GLMTTSMultiModalProcessor,
+    info=GLMTTSMultiModalProcessingInfo,
+    dummy_inputs=GLMTTSDummyInputsBuilder,
+)
+class GLMTTSForConditionalGeneration(nn.Module, SupportsMultiModal):
+    """vLLM model for GLM-TTS.
+
+    Handles both stages via model_stage branching:
+      - ``glm_tts`` (Stage 0): Text → Speech tokens (LLM AR, Llama backbone).
+      - ``glm_tts_dit`` (Stage 1): Speech tokens → Audio (DiT flow-matching + vocoder).
+
+    Attributes:
+        have_multimodal_outputs: Signals scheduler to collect multimodal outputs.
+        has_preprocess: Model has preprocess hook for input preparation (stage 0 only).
+        has_postprocess: Model has postprocess hook for hidden state caching (stage 0 only).
+    """
+
+    supports_multimodal_raw_input_only = True
+    supports_multimodal = True
+    requires_raw_input_tokens = True
+    prefer_model_sampler = True
+    _sampling_eps = 1e-5
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            "llama_embedding.": "model.embed_tokens.",
+            "llama.model.": "model.",
+            "llama.": "model.",
+        }
+    )
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.vllm_config = vllm_config
+        self.model_path = vllm_config.model_config.model
+        self.model_dir = self.model_path
+        self.model_stage = getattr(vllm_config.model_config, "model_stage", "glm_tts")
+
+        # Load configuration
+        config: GLMTTSConfig = vllm_config.model_config.hf_config  # type: ignore[assignment]
+        self.config = config
+
+        # ---- Stage 1 (DiT): delegate to GLMTTSDiTForGeneration ----
+        if self.model_stage == "glm_tts_dit":
+            from .glm_tts_dit_wrapper import GLMTTSDiTForGeneration
+
+            self._dit_gen = GLMTTSDiTForGeneration(vllm_config=vllm_config, prefix=prefix)
+            # Expose the DiT as self.model for weight loading routing
+            self.model = self._dit_gen
+            self.have_multimodal_outputs = True
+            self.enable_update_additional_information = True
+            # DiT stage does not use preprocess/postprocess/sample
+            self.has_preprocess = False
+            self.has_postprocess = False
+            # DiT loads all weights internally from flow/flow.pt etc.
+            # Point DefaultModelLoader to flow/flow.pt (a .pt file) so it can
+            # find and load it without triggering the vLLM#39699 bug where
+            # subdirectory safetensors patterns cause pt_weights_iterator
+            # to be used on .safetensors files.  load_weights() ignores the
+            # yielded weights and loads flow/flow.pt itself.
+            self.allow_patterns_overrides = ["flow/flow.pt"]
+            self.fall_back_to_pt_during_load = False
+            return
+        sample_method = str(getattr(config, "sample_method", "ras")).lower()
+        if sample_method != "ras":
+            raise ValueError(f"Unsupported GLM-TTS sample_method={sample_method!r}; expected 'ras'.")
+
+        # Resolve repo root to local path (CosyVoice3 pattern).
+        # Model weights are in llm/ subdirectory; tokenizer and other resources
+        # are siblings of llm/ under the repo root.
+        self.model_dir = resolve_glm_tts_model_dir(
+            self.model_dir,
+            tokenizer_path=getattr(vllm_config.model_config, "tokenizer", None),
+        )
+
+        # Stage 0 weights live under llm/model-*.safetensors and are loaded
+        # explicitly by _iter_llm_safetensors() during load_weights().
+        #
+        # We still point DefaultModelLoader at flow/flow.pt as a bootstrap
+        # sentinel because the current upstream loader path does not reliably
+        # target subdirectory safetensors patterns here without hitting the
+        # vLLM#39699 iterator bug. Keep the real stage-0 source of truth in
+        # _iter_llm_safetensors(); this override is only to get through the
+        # generic loader bootstrap step.
+        self.allow_patterns_overrides = ["flow/flow.pt"]
+        self.fall_back_to_pt_during_load = False
+
+        # Load tokenizer for special token ID resolution.
+        # Prefer the tokenizer path auto-detected by arg_utils
+        # (_TOKENIZER_SUBFOLDER_MAP).  Fall back to vq32k-phoneme-tokenizer/
+        # under model_dir.
+        tokenizer_path = getattr(vllm_config.model_config, "tokenizer", None)
+        if not (tokenizer_path and os.path.isdir(tokenizer_path)):
+            tokenizer_path = os.path.join(self.model_dir, _GLM_TTS_TOKENIZER_SUBDIR)
+            if not os.path.exists(tokenizer_path):
+                tokenizer_path = self.model_dir
+        trust_remote_code = bool(getattr(vllm_config.model_config, "trust_remote_code", False))
+        self._tokenizer = load_glm_tts_tokenizer(
+            tokenizer_path,
+            model_name_or_path=self.model_path,
+            trust_remote_code=trust_remote_code,
+        )
+        special_ids = get_glm_tts_special_token_ids(self._tokenizer)
+        self._ats = special_ids["ats"]
+        self._ate = special_ids["ate"]
+        self._boa = special_ids["boa"]
+        self._eoa = special_ids["eoa"]
+        self._pad = special_ids["pad"]
+        self._bos = self._tokenizer.bos_token_id or self._pad
+
+        logger.info(
+            "GLM-TTS token IDs: ATS=%d, ATE=%d, BOA=%d, EOA=%d, PAD=%d, vocab_size=%d",
+            self._ats,
+            self._ate,
+            self._boa,
+            self._eoa,
+            self._pad,
+            config.vocab_size,
+        )
+
+        # Validate special token sanity with runtime exceptions (asserts can be
+        # stripped under python -O).
+        if self._ate - self._ats != 32767:
+            raise ValueError(f"Audio token range should be 32768, got {self._ate - self._ats + 1}")
+        if self._boa >= self._ats:
+            raise ValueError(f"BOA={self._boa} should be < ATS={self._ats} (BOA is text token)")
+
+        # Validate vocab_size covers all special tokens
+        max_token = max(self._ats, self._ate, self._boa, self._eoa, self._pad)
+        if max_token >= config.vocab_size:
+            raise ValueError(
+                f"vocab_size ({config.vocab_size}) must be > max token ID ({max_token}). "
+                f"Check model's config.json has correct vocab_size."
+            )
+
+        # Update config with dynamic token IDs so vLLM uses correct eos_token
+        # This enables proper stop detection when EOA is sampled
+        config.eos_token_id = self._eoa
+        config.eoa_token_id = self._eoa
+        config.audio_token_start = self._ats
+        config.audio_token_end = self._ate
+        config.boa_token_id = self._boa
+        config.pad_token_id = self._pad
+        config.bos_token_id = self._bos
+
+        # Model flags for AR scheduler
+        self.have_multimodal_outputs = True
+        self.has_preprocess = True
+        self.has_postprocess = True
+        # GLM-TTS async chunk transfer consumes speech_tokens plus voice-clone
+        self.gpu_resident_buffer_keys: set[tuple[str, str]] = {("last_hidden", "last")}
+
+        self.model = LlamaModel(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
+
+        # LM head for speech token prediction
+        if get_pp_group().is_last_rank:
+            self.lm_head = ParallelLMHead(
+                config.vocab_size,
+                config.hidden_size,
+                quant_config=vllm_config.quant_config,
+                prefix=maybe_prefix(prefix, "lm_head"),
+            )
+        else:
+            self.lm_head = PPMissingLayer()
+
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+        self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors
+
+        # Runtime validation: dynamic EOA must match the hardcoded pipeline
+        # stop_token_ids constant (59253).  A mismatch means the upstream
+        # tokenizer vocabulary has changed and pipeline.py needs updating.
+        _PIPELINE_EOA = 59253
+        if self._eoa != _PIPELINE_EOA:
+            logger.warning(
+                "GLM-TTS EOA token mismatch: tokenizer resolved %d but "
+                "pipeline.py hardcodes stop_token_ids=[%d]. Update "
+                "pipeline.py to match the current checkpoint.",
+                self._eoa,
+                _PIPELINE_EOA,
+            )
+
+        # RAS sampling config — stored as attributes for sample()
+        self._ras_win_size = int(getattr(config, "ras_win_size", 10))
+        self._ras_tau_r = float(getattr(config, "ras_tau_r", 0.1))
+        self._ras_top_p = float(getattr(config, "ras_top_p", 0.8))
+        self._ras_top_k = int(getattr(config, "ras_top_k", 25))
+
+    def _model_dtype(self) -> torch.dtype:
+        """Return the active parameter dtype for locally-created embeddings."""
+        try:
+            return next(self.model.parameters()).dtype
+        except StopIteration:
+            dtype = getattr(self.vllm_config.model_config, "dtype", None)
+            if isinstance(dtype, torch.dtype):
+                return dtype
+            if isinstance(dtype, str):
+                _aliases = {"bf16": "bfloat16", "fp16": "float16", "half": "float16", "fp32": "float32"}
+                resolved = getattr(torch, _aliases.get(dtype.lower(), dtype.lower()), None)
+                if resolved is not None:
+                    return resolved
+            return torch.get_default_dtype()
+
+    def _resolve_prefill_target_text_len(
+        self,
+        span_len: int,
+        info_dict: Mapping[str, Any],
+    ) -> int | None:
+        """Resolve the target-text token count before the first AR sample.
+
+        The request builder mirrors the processor's length metadata into
+        model_intermediate_buffer before the first prefill.  Prefer lengths
+        inferred from that metadata, then fall back to the explicit target
+        scalar.  Never silently use a prompt-text length as the target length.
+        """
+
+        provided_text_len = _glm_tts_int_value(info_dict.get("glm_tts_text_token_len"))
+        prompt_text_len = _glm_tts_int_value(info_dict.get("glm_tts_prompt_text_token_len"))
+        prompt_speech_len = _glm_tts_int_value(info_dict.get("prompt_speech_token_len"))
+        input_len = _glm_tts_int_value(info_dict.get("input_len"))
+
+        inferred_text_len: int | None = None
+        if input_len is not None and prompt_text_len is not None:
+            inferred_text_len = input_len - prompt_text_len - 1  # strip prompt text and BOA.
+        elif prompt_text_len is not None and prompt_speech_len is not None:
+            inferred_text_len = span_len - prompt_text_len - prompt_speech_len - 1
+
+        if inferred_text_len is not None and inferred_text_len > 0:
+            if provided_text_len is not None and provided_text_len != inferred_text_len:
+                logger.warning(
+                    "GLM-TTS target text token length mismatch: processor=%d inferred=%d; using inferred length.",
+                    provided_text_len,
+                    inferred_text_len,
+                )
+            return inferred_text_len
+
+        return provided_text_len if provided_text_len is not None and provided_text_len > 0 else None
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: Any | None = None,
+        is_multimodal: Any | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        if self.model_stage == "glm_tts_dit":
+            return self._dit_gen.embed_input_ids(input_ids, **kwargs)
+        embed_tokens = self.model.embed_tokens(input_ids)
+        if multimodal_embeddings is None or is_multimodal is None:
+            return embed_tokens
+
+        mask = torch.as_tensor(is_multimodal, device=embed_tokens.device, dtype=torch.bool).reshape(-1)
+        if not bool(mask.any()):
+            return embed_tokens
+
+        if isinstance(multimodal_embeddings, (list, tuple)):
+            tensors = [
+                torch.as_tensor(item, device=embed_tokens.device, dtype=embed_tokens.dtype)
+                for item in multimodal_embeddings
+                if item is not None
+            ]
+            if not tensors:
+                return embed_tokens
+            mm_embeds = torch.cat([item.reshape(-1, item.shape[-1]) for item in tensors], dim=0)
+        else:
+            mm_embeds = multimodal_embeddings
+        if mm_embeds is None:
+            return embed_tokens
+        mm_embeds = torch.as_tensor(mm_embeds, device=embed_tokens.device, dtype=embed_tokens.dtype)
+        if mm_embeds.ndim == 3 and int(mm_embeds.shape[0]) == 1:
+            mm_embeds = mm_embeds.squeeze(0)
+        if mm_embeds.ndim != 2:
+            shape = tuple(mm_embeds.shape)
+            raise ValueError(f"GLM-TTS multimodal embeddings should be 2D, got shape={shape}")
+
+        flat_tokens = embed_tokens.reshape(-1, embed_tokens.shape[-1])
+        if int(flat_tokens.shape[0]) != int(mask.numel()):
+            raise ValueError(
+                "GLM-TTS multimodal mask/token length mismatch: "
+                f"mask={int(mask.numel())}, tokens={int(flat_tokens.shape[0])}"
+            )
+        mm_len = int(mask.sum().item())
+        if int(mm_embeds.shape[0]) < mm_len:
+            raise ValueError(
+                "GLM-TTS multimodal embedding length mismatch: "
+                f"embeddings={int(mm_embeds.shape[0])}, placeholders={mm_len}"
+            )
+
+        flat_tokens = flat_tokens.clone()
+        flat_tokens[mask] = mm_embeds[:mm_len]
+        return flat_tokens.reshape_as(embed_tokens)
+
+    def embed_multimodal(self, **kwargs: Any) -> list[torch.Tensor] | None:
+        if self.model_stage != "glm_tts":
+            return None
+        prompt_speech_token = kwargs.get("prompt_speech_token")
+        if prompt_speech_token is None:
+            return None
+
+        def embed_one(value: Any) -> torch.Tensor | None:
+            if value is None:
+                return None
+            speech_token = torch.as_tensor(value, device=next(self.model.parameters()).device)
+            if speech_token.numel() == 0:
+                return None
+            speech_token = speech_token.to(dtype=torch.long)
+            if speech_token.ndim == 0:
+                speech_token = speech_token.reshape(1, 1)
+            elif speech_token.ndim == 1:
+                speech_token = speech_token.unsqueeze(0)
+            elif speech_token.ndim > 2:
+                speech_token = speech_token.reshape(int(speech_token.shape[0]), -1)
+            if int(speech_token.min().item()) >= self._ats:
+                speech_ids = speech_token
+            else:
+                speech_ids = speech_token + self._ats
+            speech_embeds = self.model.embed_tokens(speech_ids)
+            return speech_embeds.reshape(-1, speech_embeds.shape[-1])
+
+        def is_flat_token_sequence(value: list[Any] | tuple[Any, ...]) -> bool:
+            if not value:
+                return False
+            try:
+                return all(torch.as_tensor(item).ndim == 0 for item in value)
+            except (TypeError, ValueError):
+                return False
+
+        if isinstance(prompt_speech_token, (list, tuple)):
+            if is_flat_token_sequence(prompt_speech_token):
+                one_embed = embed_one(prompt_speech_token)
+                return [one_embed] if one_embed is not None else None
+            embeds = [embed_one(item) for item in prompt_speech_token]
+            embeds = [item for item in embeds if item is not None]
+            if not embeds:
+                return None
+            return embeds
+
+        one_embed = embed_one(prompt_speech_token)
+        return [one_embed] if one_embed is not None else None
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor | IntermediateTensors | OmniOutput:
+        if self.model_stage == "glm_tts_dit":
+            return self._dit_gen.forward(input_ids, positions, **kwargs)
+        return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
+
+    def compute_logits(
+        self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None
+    ) -> torch.Tensor | None:
+        if self.model_stage == "glm_tts_dit":
+            return self._dit_gen.compute_logits(hidden_states, sampling_metadata)
+        if isinstance(hidden_states, OmniOutput):
+            hidden_states = hidden_states.text_hidden_states
+        if hidden_states is None:
+            return None
+        logits = self.logits_processor(self.lm_head, hidden_states)
+        if logits is None:
+            return None
+
+        return logits
+
+    def _glm_tts_ras_enabled(self, sampling_metadata: SamplingMetadata) -> bool:
+        if self.model_stage != "glm_tts":
+            return False
+        if sampling_metadata.max_num_logprobs is not None:
+            return False
+        if sampling_metadata.temperature is None:
+            return False
+        if bool(sampling_metadata.bad_words_token_ids):
+            return False
+        if torch.any(sampling_metadata.frequency_penalties != 0):
+            return False
+        if torch.any(sampling_metadata.presence_penalties != 0):
+            return False
+        return True
+
+    def _apply_allowed_mask(self, logits: torch.Tensor) -> torch.Tensor:
+        """Restrict logits to valid audio tokens [ATS..ATE] ∪ {EOA}."""
+        allowed_mask = getattr(self, "_ar_allowed_mask", None)
+        if allowed_mask is None or allowed_mask.device != logits.device or allowed_mask.shape[0] != logits.shape[-1]:
+            allowed_mask = torch.full((logits.shape[-1],), float("-inf"), device=logits.device)
+            allowed_mask[self._ats : self._ate + 1] = 0.0
+            allowed_mask[self._eoa] = 0.0
+            self._ar_allowed_mask = allowed_mask
+        return logits + allowed_mask
+
+    def _select_allowed_audio_logits(self, logits: torch.Tensor) -> torch.Tensor:
+        """Return compact logits for [audio tokens] plus EOA.
+
+        The RAS path samples only from GLM-TTS audio tokens and EOA.  Keeping a
+        compact view avoids top-k/logsumexp over the full LLM vocabulary on
+        every AR step; this is equivalent to sampling after ``_apply_allowed_mask``.
+        """
+        audio_logits = logits[:, self._ats : self._ate + 1]
+        if self._ats <= self._eoa <= self._ate:
+            return audio_logits
+        return torch.cat((audio_logits, logits[:, self._eoa : self._eoa + 1]), dim=-1)
+
+    def _allowed_local_to_token_id(self, local_id: int) -> int:
+        audio_vocab = self._ate - self._ats + 1
+        return self._ats + local_id if local_id < audio_vocab else self._eoa
+
+    def _recent_allowed_local_ids(self, decoded_tokens: Sequence[int], win_size: int) -> list[int]:
+        if win_size <= 0 or not decoded_tokens:
+            return []
+        audio_vocab = self._ate - self._ats + 1
+        recent: list[int] = []
+        for token in decoded_tokens[-win_size:]:
+            token_id = int(token)
+            if self._ats <= token_id <= self._ate:
+                recent.append(token_id - self._ats)
+            elif token_id == self._eoa:
+                recent.append(audio_vocab)
+            else:
+                recent.append(-1)
+        return recent
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> SamplerOutput | None:
+        """RAS sampler following CosyVoice3 pattern.
+
+        Uses vLLM Sampler for logits processing (logit_bias_state handles
+        min_tokens/max_tokens/stop_token_ids).  When RAS is enabled, applies
+        per-request nucleus+repetition-aware sampling; otherwise falls back
+        to standard vLLM sampling.
+        """
+        if logits is None or logits.numel() == 0:
+            return None
+        if self.model_stage != "glm_tts":
+            return None
+
+        sampler = getattr(self, "_ar_sampler", None)
+        if sampler is None:
+            sampler = Sampler()
+            self._ar_sampler = sampler
+
+        if not self._glm_tts_ras_enabled(sampling_metadata):
+            logits = self._apply_allowed_mask(logits)
+            return sampler(logits=logits, sampling_metadata=sampling_metadata)
+
+        logits = logits.to(torch.float32)
+        sampling_for_processors = replace(sampling_metadata, no_penalties=True)
+        logits = sampler.apply_logits_processors(logits, sampling_for_processors, predict_bonus_token=False)
+        allowed_logits = self._select_allowed_audio_logits(logits)
+
+        default_top_p = self._ras_top_p
+        default_top_k = self._ras_top_k
+        ws = self._ras_win_size
+        tr = self._ras_tau_r
+
+        num_reqs = int(allowed_logits.shape[0])
+        t_data = sampling_metadata.temperature
+        tp_data = sampling_metadata.top_p
+        tk_data = sampling_metadata.top_k
+        cache_key = (
+            t_data.data_ptr() if t_data is not None else 0,
+            tp_data.data_ptr() if tp_data is not None else 0,
+            tk_data.data_ptr() if tk_data is not None else 0,
+            num_reqs,
+        )
+        _sp_cache = getattr(self, "_sampling_param_cache", None)
+        if _sp_cache is not None and _sp_cache[0] == cache_key:
+            temperatures, top_ps, top_ks = _sp_cache[1]
+        else:
+            temperatures = _tensor_param_values(t_data, num_reqs, 1.0)
+            top_ps = _tensor_param_values(tp_data, num_reqs, default_top_p)
+            top_ks = _tensor_param_values(tk_data, num_reqs, default_top_k)
+            self._sampling_param_cache = (cache_key, (temperatures, top_ps, top_ks))
+
+        sampled_ids: list[int] = []
+        for req_idx in range(num_reqs):
+            row_logits = allowed_logits[req_idx]
+            temperature = temperatures[req_idx]
+            if temperature < self._sampling_eps:
+                local_id = int(torch.argmax(row_logits).item())
+                sampled_ids.append(self._allowed_local_to_token_id(local_id))
+                continue
+
+            top_p = top_ps[req_idx]
+            top_k = int(top_ks[req_idx])
+            generator = sampling_metadata.generators.get(req_idx)
+            weighted_scores = row_logits / max(temperature, self._sampling_eps)
+            decoded_tokens = (
+                sampling_metadata.output_token_ids[req_idx] if req_idx < len(sampling_metadata.output_token_ids) else []
+            )
+            local_id = _ras_sample_one(
+                weighted_scores,
+                self._recent_allowed_local_ids(decoded_tokens, ws),
+                top_p=top_p,
+                top_k=top_k,
+                win_size=ws,
+                tau_r=tr,
+                generator=generator,
+            )
+            sampled_ids.append(self._allowed_local_to_token_id(local_id))
+
+        sampled = torch.tensor(sampled_ids, device=logits.device, dtype=torch.int32)
+        return SamplerOutput(sampled_token_ids=sampled.unsqueeze(-1), logprobs_tensors=None)
+
+    def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
+        """Package hidden states, speech tokens, and voice clone data into OmniOutput.
+
+        Streaming contract: **delta**.  Each decode step emits exactly one
+        speech token (or a prefill placeholder).  The engine's output
+        processor concatenates per-step deltas into the final tensor.
+        """
+        if isinstance(model_outputs, OmniOutput):
+            return model_outputs
+
+        hidden = model_outputs
+        info_dicts = kwargs.get("model_intermediate_buffer") or []
+        if isinstance(info_dicts, dict):
+            info_dicts = [info_dicts]
+
+        speech_tokens_list: list[torch.Tensor] = []
+        multimodal_extras: dict[str, Any] = {}
+        mm_conditioning_keys = (
+            "prompt_speech_token",
+            "prompt_speech_token_len",
+            "glm_tts_prompt_text_token_len",
+            "prompt_feat",
+            "embedding",
+            "glm_tts_text_token_len",
+        )
+        info_for_flags = next((info for info in info_dicts if isinstance(info, dict)), None)
+        if info_for_flags is None or not info_for_flags.get("_voice_clone_emitted"):
+            copied_from_kwargs = False
+            for key in mm_conditioning_keys:
+                val = kwargs.get(key)
+                if val is not None:
+                    multimodal_extras[key] = val
+                    copied_from_kwargs = True
+            if copied_from_kwargs and info_for_flags is not None:
+                info_for_flags["_voice_clone_emitted"] = True
+
+        for info in info_dicts:
+            if not isinstance(info, dict):
+                continue
+            tokens = info.get("speech_tokens")
+            if isinstance(tokens, torch.Tensor) and tokens.numel() > 0:
+                speech_tokens_list.append(tokens)
+            # Propagate voice cloning features from preprocess info_update.
+            # IMPORTANT: Only emit once — the output processor accumulates
+            # every decode step and concatenates tensors, which would corrupt
+            # constant data (e.g. embedding [192] → [N*192]).
+            if not info.get("_voice_clone_emitted"):
+                for key in mm_conditioning_keys:
+                    val = info.get(key)
+                    if val is not None and key not in multimodal_extras:
+                        multimodal_extras[key] = val
+                if any(info.get(k) is not None for k in mm_conditioning_keys):
+                    info["_voice_clone_emitted"] = True
+
+        if not speech_tokens_list:
+            return OmniOutput(text_hidden_states=hidden, multimodal_outputs=multimodal_extras)
+
+        speech_tokens = torch.cat(speech_tokens_list, dim=0)
+        span_len = int(speech_tokens.shape[0])
+        hidden = hidden[:span_len]
+        multimodal_extras["speech_tokens"] = speech_tokens
+        return OmniOutput(
+            text_hidden_states=hidden,
+            multimodal_outputs=multimodal_extras,
+        )
+
+    def preprocess(
+        self,
+        input_ids: torch.Tensor,
+        input_embeds: torch.Tensor | None,
+        **info_dict: Any,
+    ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
+        """Prepare inputs for GLM-TTS AR model.
+
+        GLM-TTS only supports the multimodal processor path:
+        text prompt + ``multi_modal_data["audio"]`` + ``mm_processor_kwargs["prompt_text"]``.
+        Legacy placeholder prompts via ``additional_information`` are rejected.
+        """
+        if info_dict.get("additional_information") is not None:
+            raise ValueError(
+                "GLM-TTS no longer accepts legacy additional_information prompts; "
+                "use prompt + multi_modal_data['audio'] + mm_processor_kwargs['prompt_text']."
+            )
+
+        span_len = int(input_ids.shape[0])
+        logger.debug("preprocess: span_len=%d, input_ids.shape=%s", span_len, input_ids.shape)
+        if span_len <= 0:
+            return input_ids, input_embeds if input_embeds is not None else self.embed_input_ids(input_ids), {}
+
+        device = input_ids.device
+
+        if isinstance(info_dict.get("text"), list) and info_dict["text"]:
+            raise ValueError(
+                "GLM-TTS no longer accepts legacy text/additional_information prefill payloads; "
+                "use the multimodal processor path."
+            )
+
+        mm_prefill_done = bool(info_dict.get("glm_tts_mm_prefill_done"))
+        sampled_token: int | None = None
+        if input_embeds is not None and span_len == 1:
+            sampled_token = int(input_ids[0].item())
+        one_token_prefill_tail = (
+            input_embeds is not None
+            and span_len == 1
+            and mm_prefill_done
+            and sampled_token is not None
+            and sampled_token != self._eoa
+            and not (self._ats <= sampled_token <= self._ate)
+        )
+        is_prefill_span = span_len > 1 or not mm_prefill_done or one_token_prefill_tail
+        if input_embeds is None and is_prefill_span:
+            raise ValueError("Missing GLM-TTS multimodal input embeddings.")
+        if input_embeds is not None and is_prefill_span:
+            input_ids_out = input_ids.clone()
+            input_ids_out[:] = self._pad
+            info_update: dict[str, Any] = {
+                "glm_tts_mm_prefill_done": True,
+                "speech_tokens": torch.full((span_len, 1), -1, device=device, dtype=torch.long),
+            }
+            # One-token prefill tails re-enter the prefill branch but
+            # text_token_len was already resolved during the initial prefill.
+            if not one_token_prefill_tail:
+                text_token_len = self._resolve_prefill_target_text_len(span_len, info_dict)
+                if text_token_len is not None:
+                    info_update["glm_tts_text_token_len"] = torch.tensor(
+                        [text_token_len], device=device, dtype=torch.long
+                    )
+                prompt_text_len = info_dict.get("glm_tts_prompt_text_token_len")
+                if prompt_text_len is not None:
+                    info_update["glm_tts_prompt_text_token_len"] = prompt_text_len
+                prompt_speech_len = info_dict.get("prompt_speech_token_len")
+                if prompt_speech_len is None:
+                    input_len = _glm_tts_int_value(info_dict.get("input_len"))
+                    if input_len is not None:
+                        prompt_speech_len = torch.tensor(
+                            [max(0, span_len - input_len)], device=device, dtype=torch.long
+                        )
+                if prompt_speech_len is not None:
+                    info_update["prompt_speech_token_len"] = prompt_speech_len
+            return input_ids_out, input_embeds.to(dtype=self._model_dtype()), info_update
+
+        # Decode: span_len == 1
+        # Standard autoregressive decode - use input_ids directly
+        if input_embeds is not None and int(input_embeds.shape[0]) == 1:
+            inputs_embeds_out = input_embeds.reshape(1, -1).to(dtype=self._model_dtype())
+        else:
+            inputs_embeds_out = self.embed_input_ids(input_ids.reshape(1, 1).to(torch.long))
+            inputs_embeds_out = inputs_embeds_out.reshape(1, -1).to(dtype=self._model_dtype())
+
+        # Convert sampled token to speech token (relative to ATS)
+        # -1 = invalid/EOA, valid range = [0, ATE-ATS]
+        if sampled_token is None:
+            sampled_token = int(input_ids[0].item())
+        if self._ats <= sampled_token <= self._ate:
+            speech_token = sampled_token - self._ats
+        else:
+            # EOA or other non-audio token → mark as invalid (-1)
+            speech_token = -1
+            logger.debug("GLM-TTS decode: non-audio token %d (EOA=%d)", sampled_token, self._eoa)
+        speech_tokens = torch.tensor([[speech_token]], device=device, dtype=torch.long)
+
+        info_update = {"speech_tokens": speech_tokens}
+        return input_ids, inputs_embeds_out, info_update
+
+    def postprocess(self, hidden_states: torch.Tensor, **kwargs: Any) -> dict[str, Any]:
+        """Cache last hidden state for next decode step."""
+        if hidden_states.numel() == 0:
+            return {}
+        last = hidden_states[-1, :].detach()
+        update: dict[str, Any] = {"last_hidden": {"last": last}}
+
+        multimodal_outputs = kwargs.get("multimodal_outputs")
+        if isinstance(multimodal_outputs, dict):
+            copied_conditioning = False
+            for key in ("prompt_speech_token", "prompt_speech_token_len", "prompt_feat", "embedding"):
+                val = multimodal_outputs.get(key)
+                if val is not None:
+                    update[key] = _first_glm_tts_value(val)
+                    copied_conditioning = True
+
+            prompt_text_len = multimodal_outputs.get("glm_tts_prompt_text_token_len")
+            if prompt_text_len is not None:
+                update["glm_tts_prompt_text_token_len"] = _first_glm_tts_value(prompt_text_len)
+
+            text_token_len = _glm_tts_int_value(multimodal_outputs.get("glm_tts_text_token_len"))
+            if text_token_len is not None:
+                update["glm_tts_text_token_len"] = torch.tensor([int(text_token_len)], dtype=torch.long)
+
+            if copied_conditioning:
+                update["_voice_clone_emitted"] = True
+
+        return update
+
+    def _iter_llm_safetensors(self) -> Iterable[tuple[str, torch.Tensor]]:
+        """Yield (name, tensor) pairs from llm/model-*.safetensors."""
+        import glob as glob_module
+
+        from safetensors.torch import load_file
+
+        llm_dir = os.path.join(self.model_dir, "llm")
+        if os.path.isdir(llm_dir):
+            sf_files = sorted(glob_module.glob(os.path.join(llm_dir, "model-*.safetensors")))
+            if sf_files:
+                for sf_path in sf_files:
+                    yield from load_file(sf_path, device="cpu").items()
+                return
+
+        from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+
+        model_root = download_weights_from_hf_specific(
+            self.model_path,
+            self.vllm_config.load_config.download_dir,
+            allow_patterns=["llm/model-*.safetensors"],
+        )
+        llm_dir = os.path.join(model_root, "llm")
+        sf_files = sorted(glob_module.glob(os.path.join(llm_dir, "model-*.safetensors")))
+        if not sf_files:
+            raise RuntimeError(f"No LLM safetensors found under {model_root}. Expected llm/model-*.safetensors.")
+        for sf_path in sf_files:
+            yield from load_file(sf_path, device="cpu").items()
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load weights from checkpoint.
+
+        Stage 0 (glm_tts): HuggingFace Llama-format checkpoint from llm/ subdir.
+        Stage 1 (glm_tts_dit): DiT flow.pt + vocoder.
+        """
+        if self.model_stage == "glm_tts_dit":
+            return self._dit_gen.load_weights(weights)
+
+        def _glm_tts_weights() -> Iterable[tuple[str, torch.Tensor]]:
+            for name, loaded_weight in self._iter_llm_safetensors():
+                if "rotary_emb.inv_freq" in name:
+                    continue
+                if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
+                    continue
+                yield name, loaded_weight
+
+        loader = AutoWeightsLoader(self)
+        loaded_params = loader.load_weights(_glm_tts_weights(), mapper=self.hf_to_vllm_mapper)
+
+        params_dict = dict(self.named_parameters())
+        # Handle tie_word_embeddings: copy embed_tokens to lm_head if not loaded.
+        lm_head_key = "lm_head.weight"
+        embed_key = "model.embed_tokens.weight"
+        if lm_head_key not in loaded_params and embed_key in loaded_params:
+            if lm_head_key in params_dict and embed_key in params_dict:
+                lm_head_param = params_dict[lm_head_key]
+                embed_param = params_dict[embed_key]
+                weight_loader = getattr(lm_head_param, "weight_loader", default_weight_loader)
+                weight_loader(lm_head_param, embed_param.data)
+                loaded_params.add(lm_head_key)
+                logger.info("Tied lm_head.weight to embed_tokens.weight")
+
+        logger.info("Loaded %d weights for GLMTTSForConditionalGeneration", len(loaded_params))
+        return loaded_params
diff --git a/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py b/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py
new file mode 100644
index 000000000..5cda6ec45
--- /dev/null
+++ b/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py
@@ -0,0 +1,650 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GLM-TTS DiT (Diffusion Transformer) Model.
+
+Flow Matching model that converts speech tokens to mel-spectrogram.
+Ported from: https://github.com/zai-org/GLM-TTS/blob/main/flow/dit.py
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb
+
+from vllm_omni.diffusion.attention.layer import Attention as DiffusionAttention
+
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor: float = 1.0):
+    """Precompute Rotary Positional Embeddings (RoPE)."""
+    theta *= theta_rescale_factor ** (dim / (dim - 2))
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)
+    freqs = torch.outer(t, freqs).float()
+    freqs_cos = torch.cos(freqs)
+    freqs_sin = torch.sin(freqs)
+    return torch.cat([freqs_cos, freqs_sin], dim=-1)
+
+
+def get_pos_embed_indices(start, length, max_pos, scale=1.0):
+    """Get position embedding indices."""
+    scale = scale * torch.ones_like(start, dtype=torch.float32)
+    pos = (
+        start.unsqueeze(1)
+        + (torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) * scale.unsqueeze(1)).long()
+    )
+    pos = torch.where(pos < max_pos, pos, max_pos - 1)
+    return pos
+
+
+class SinusPositionEmbedding(nn.Module):
+    """Sinusoidal position embedding for timesteps."""
+
+    def __init__(self, dim: int):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x: torch.Tensor, scale: float = 1000) -> torch.Tensor:
+        device = x.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
+        emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+class TimestepEmbedding(nn.Module):
+    """Timestep embedding with MLP."""
+
+    def __init__(self, dim: int, freq_embed_dim: int = 256):
+        super().__init__()
+        self.time_embed = SinusPositionEmbedding(freq_embed_dim)
+        self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+
+    def forward(self, timestep: torch.Tensor) -> torch.Tensor:
+        time_hidden = self.time_embed(timestep)
+        # Cast to model weight dtype (bfloat16) not input dtype (may be float32)
+        weight_dtype = self.time_mlp[0].weight.dtype
+        time_hidden = time_hidden.to(weight_dtype)
+        return self.time_mlp(time_hidden)
+
+
+class GRN(nn.Module):
+    """Global Response Normalization layer."""
+
+    def __init__(self, dim: int):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gx = torch.norm(x, p=2, dim=1, keepdim=True)
+        nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-6)
+        return self.gamma * (x * nx) + self.beta + x
+
+
+class ConvNeXtV2Block(nn.Module):
+    """ConvNeXt-V2 Block for text embedding."""
+
+    def __init__(self, dim: int, intermediate_dim: int, dilation: int = 1):
+        super().__init__()
+        padding = (dilation * (7 - 1)) // 2
+        self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation)
+        self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, intermediate_dim)
+        self.act = nn.GELU()
+        self.grn = GRN(intermediate_dim)
+        self.pwconv2 = nn.Linear(intermediate_dim, dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = x.transpose(1, 2)
+        x = self.dwconv(x)
+        x = x.transpose(1, 2)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv2(x)
+        return residual + x
+
+
+class ConvPositionEmbedding(nn.Module):
+    """Convolutional position embedding."""
+
+    def __init__(self, dim: int, kernel_size: int = 31, groups: int = 16):
+        super().__init__()
+        if kernel_size % 2 == 0:
+            raise ValueError(f"kernel_size must be odd, got {kernel_size}")
+        self.conv1d = nn.Sequential(
+            nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),
+            nn.Mish(),
+            nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),
+            nn.Mish(),
+        )
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
+        if mask is not None:
+            mask = mask[..., None]
+            x = x.masked_fill(~mask, 0.0)
+
+        x = rearrange(x, "b n d -> b d n")
+        x = self.conv1d(x)
+        out = rearrange(x, "b d n -> b n d")
+
+        if mask is not None:
+            out = out.masked_fill(~mask, 0.0)
+        return out
+
+
+class TextEmbedding(nn.Module):
+    """Text/speech token embedding with optional ConvNeXt modeling."""
+
+    def __init__(
+        self,
+        text_num_embeds: int,
+        output_dim: int,
+        conv_layers: int = 0,
+        conv_mult: int = 2,
+        length_align: str = "fill",
+    ):
+        super().__init__()
+        self.pad_id = int(text_num_embeds)
+        self.text_embed = nn.Embedding(text_num_embeds + 1, output_dim, padding_idx=self.pad_id)
+        self.length_align = length_align
+
+        if conv_layers > 0:
+            self.extra_modeling = True
+            self.precompute_max_pos = 4096
+            self.register_buffer(
+                "freqs_cis", precompute_freqs_cis(output_dim, self.precompute_max_pos), persistent=False
+            )
+            self.text_blocks = nn.Sequential(
+                *[ConvNeXtV2Block(output_dim, output_dim * conv_mult) for _ in range(conv_layers)]
+            )
+        else:
+            self.extra_modeling = False
+
+    def forward(
+        self,
+        text: torch.Tensor,
+        aim_seq_len: int,
+        text_lens: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        batch, text_len = text.shape[0], text.shape[1]
+        if self.length_align == "fill":
+            text = text[:, :aim_seq_len]
+            text = F.pad(text, (0, aim_seq_len - min(text_len, aim_seq_len)), value=self.pad_id)
+        elif self.length_align == "interpolate_token":
+            if text_lens is None:
+                text = (
+                    F.interpolate(
+                        text.unsqueeze(1).float(),
+                        size=aim_seq_len,
+                        mode="nearest",
+                    )
+                    .squeeze(1)
+                    .long()
+                )
+            else:
+                text_lens = text_lens.to(device=text.device, dtype=torch.long)
+                text_lens = torch.clamp(text_lens, min=1, max=text_len)
+                positions = torch.arange(aim_seq_len, device=text.device, dtype=torch.long)
+                src_idx = positions.unsqueeze(0) * text_lens.unsqueeze(1) // int(aim_seq_len)
+                src_idx = torch.minimum(src_idx, (text_lens - 1).unsqueeze(1))
+                text = torch.gather(text, dim=1, index=src_idx)
+        elif self.length_align != "interpolate_feature":
+            raise ValueError(f"Unsupported length_align={self.length_align!r}")
+
+        hidden = self.text_embed(text)
+
+        if self.length_align == "interpolate_feature":
+            hidden = F.interpolate(
+                hidden.permute(0, 2, 1),
+                size=aim_seq_len,
+                mode="nearest",
+            ).permute(0, 2, 1)
+
+        if self.extra_modeling:
+            # Get model dtype from text_blocks weights
+            model_dtype = self.text_blocks[0].dwconv.weight.dtype
+            hidden = hidden.to(model_dtype)
+
+            batch_start = torch.zeros((batch,), dtype=torch.long, device=text.device)
+            pos_idx = get_pos_embed_indices(batch_start, aim_seq_len, max_pos=self.precompute_max_pos)
+            text_pos_embed = self.freqs_cis[pos_idx].to(model_dtype)
+            hidden = hidden + text_pos_embed
+            hidden = self.text_blocks(hidden)
+
+        return hidden
+
+
+class EmbeddingConcater(nn.Module):
+    """Concatenates and projects noisy audio, condition, and text embeddings."""
+
+    def __init__(self, in_dim: int, out_dim: int):
+        super().__init__()
+        self.proj = nn.Linear(in_dim, out_dim)
+        self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond: torch.Tensor,
+        text_embed: torch.Tensor,
+        drop_audio_cond: bool = False,
+    ) -> torch.Tensor:
+        if drop_audio_cond:
+            cond = torch.zeros_like(cond)
+        x = self.proj(torch.cat((x, cond, text_embed), dim=-1))
+        x = self.conv_pos_embed(x) + x
+        return x
+
+
+class AdaLayerNormZero(nn.Module):
+    """Adaptive LayerNorm with modulation from time/condition embedding."""
+
+    def __init__(self, dim: int, additional_dim: int = 0):
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(dim + additional_dim, dim * 6)
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(self, x: torch.Tensor, emb: torch.Tensor):
+        emb = self.linear(self.silu(emb))
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)
+        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+
+
+class AdaLayerNormZeroFinal(nn.Module):
+    """Adaptive LayerNorm for final layer."""
+
+    def __init__(self, dim: int, additional_dim: int = 0):
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(dim + additional_dim, dim * 2)
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
+        emb = self.linear(self.silu(emb))
+        scale, shift = torch.chunk(emb, 2, dim=1)
+        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+        return x
+
+
+class FeedForward(nn.Module):
+    """Feed-forward network with GELU activation.
+
+    Matches original GLM-TTS structure with nested Sequential for checkpoint compatibility.
+    Weight paths: ff.ff.0.0.* (first Linear), ff.ff.2.* (second Linear)
+    """
+
+    def __init__(self, dim: int, dim_out: int | None = None, mult: int = 4, dropout: float = 0.0):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+
+        # Nested structure to match original checkpoint paths
+        project_in = nn.Sequential(
+            nn.Linear(dim, inner_dim),
+            nn.GELU(approximate="tanh"),
+        )
+        self.ff = nn.Sequential(
+            project_in,
+            nn.Dropout(dropout),
+            nn.Linear(inner_dim, dim_out),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.ff(x)
+
+
+class DiTAttention(nn.Module):
+    """Attention module using diffusion infrastructure."""
+
+    def __init__(self, dim: int, heads: int = 8, dim_head: int = 64, dropout: float = 0.0):
+        super().__init__()
+        self.dim = dim
+        self.heads = heads
+        self.dim_head = dim_head
+        self.inner_dim = dim_head * heads
+        self.scale = 1.0 / math.sqrt(dim_head)
+
+        self.to_q = nn.Linear(dim, self.inner_dim)
+        self.to_k = nn.Linear(dim, self.inner_dim)
+        self.to_v = nn.Linear(dim, self.inner_dim)
+        self.to_out = nn.Sequential(nn.Linear(self.inner_dim, dim), nn.Dropout(dropout))
+
+        self.attn = DiffusionAttention(
+            num_heads=heads,
+            head_size=dim_head,
+            softmax_scale=self.scale,
+            causal=False,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        padding_mask: torch.Tensor | None = None,
+        rope=None,
+        attn_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        batch_size, seq_len = x.shape[0], x.shape[1]
+
+        query = self.to_q(x)
+        key = self.to_k(x)
+        value = self.to_v(x)
+
+        # Apply RoPE
+        if rope is not None:
+            freqs, xpos_scale = rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+            query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
+            key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
+
+        # Reshape for attention
+        query = query.view(batch_size, seq_len, self.heads, self.dim_head)
+        key = key.view(batch_size, seq_len, self.heads, self.dim_head)
+        value = value.view(batch_size, seq_len, self.heads, self.dim_head)
+
+        # Block-causal mask requires arbitrary 2D attention mask which
+        # FlashAttention does not support.  Fall back to PyTorch SDPA
+        # (matches official GLM-TTS flow/modules.py:272).
+        if attn_mask is not None:
+            q = query.permute(0, 2, 1, 3)  # [B, H, T, D]
+            k = key.permute(0, 2, 1, 3)
+            v = value.permute(0, 2, 1, 3)
+            out = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                attn_mask=attn_mask,
+                dropout_p=0.0,
+                is_causal=False,
+            )
+            out = out.permute(0, 2, 1, 3)  # [B, T, H, D]
+        else:
+            out = self.attn(query, key, value, attn_metadata=None)
+        out = out.view(batch_size, seq_len, self.inner_dim)
+        out = out.to(query.dtype)
+        out = self.to_out(out)
+
+        # Apply padding mask
+        if padding_mask is not None:
+            padding_mask = rearrange(padding_mask, "b 1 n -> b n 1").bool()
+            out = out.masked_fill(~padding_mask, 0.0)
+
+        return out
+
+
+class DiTBlock(nn.Module):
+    """DiT block with AdaLayerNorm modulation."""
+
+    def __init__(
+        self,
+        dim: int,
+        heads: int,
+        dim_head: int,
+        ff_mult: int = 4,
+        dropout: float = 0.1,
+        additional_condition_dim: int = 0,
+    ):
+        super().__init__()
+        self.attn_norm = AdaLayerNormZero(dim, additional_condition_dim)
+        self.attn = DiTAttention(dim=dim, heads=heads, dim_head=dim_head, dropout=dropout)
+        self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        t: torch.Tensor,
+        padding_mask: torch.Tensor | None = None,
+        rope=None,
+        attn_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # Pre-norm & modulation for attention
+        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t)
+
+        # Attention
+        attn_output = self.attn(x=norm, padding_mask=padding_mask, rope=rope, attn_mask=attn_mask)
+        x = x + gate_msa.unsqueeze(1) * attn_output
+
+        # FeedForward
+        norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        ff_output = self.ff(norm)
+        x = x + gate_mlp.unsqueeze(1) * ff_output
+
+        return x
+
+
+class GLMTTSDiT(nn.Module):
+    """GLM-TTS Diffusion Transformer.
+
+    Converts speech tokens to mel-spectrogram via flow matching.
+    """
+
+    _repeated_blocks = ["DiTBlock"]
+    _layerwise_offload_blocks_attrs = ["transformer_blocks"]
+
+    def __init__(
+        self,
+        *,
+        trans_dim: int = 768,
+        depth: int = 18,
+        heads: int = 12,
+        dim_head: int = 64,
+        ff_mult: int = 2,
+        dropout: float = 0.1,
+        mel_dim: int = 80,
+        text_vocab_size: int = 100000,
+        text_emb_dim: int = 512,
+        conv_layers: int = 4,
+        condition_dim: int = 160,  # mel_dim + spk_embed_dim
+        spkr_emb_adaln: bool = False,
+        spkr_dim: int = 192,
+        use_wavlm_emb: bool = False,
+        long_skip_connection: bool = False,
+    ):
+        super().__init__()
+
+        self.time_embed = TimestepEmbedding(trans_dim)
+        self.text_emb_layer = TextEmbedding(
+            text_vocab_size,
+            text_emb_dim,
+            conv_layers=conv_layers,
+            length_align="interpolate_token",
+        )
+
+        in_dim = mel_dim + text_emb_dim + condition_dim
+        self.emb_concator = EmbeddingConcater(in_dim, trans_dim)
+
+        self.rotary_embed = RotaryEmbedding(dim_head)
+
+        self.dim = trans_dim
+        self.depth = depth
+        self.num_heads = heads
+        self.spkr_emb_adaln = spkr_emb_adaln
+
+        # Frame rates for block-causal mask token→mel conversion
+        self._input_frame_rate: float = 25.0
+        self._mel_framerate: float = 50.0
+
+        # Speaker embedding dimension for AdaLN
+        if spkr_emb_adaln:
+            additional_dim = spkr_dim + 256 if use_wavlm_emb else spkr_dim
+        else:
+            additional_dim = 0
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                DiTBlock(
+                    dim=trans_dim,
+                    heads=heads,
+                    dim_head=dim_head,
+                    ff_mult=ff_mult,
+                    dropout=dropout,
+                    additional_condition_dim=additional_dim,
+                )
+                for _ in range(depth)
+            ]
+        )
+
+        self.long_skip_connection = nn.Linear(trans_dim * 2, trans_dim, bias=False) if long_skip_connection else None
+
+        self.norm_out = AdaLayerNormZeroFinal(trans_dim, additional_dim=additional_dim)
+        self.proj_out = nn.Linear(trans_dim, mel_dim)
+
+    # ------------------------------------------------------------------
+    # Block-causal attention mask (ported from GLM-TTS block_mask_util)
+    # ------------------------------------------------------------------
+
+    def _token_to_mel_size(self, token_count: int) -> int:
+        """Convert token block size to mel-frame block size for streaming mask.
+
+        Matches upstream GLM-TTS ``flow/dit.py`` block-mask conversion.
+        This is intentionally separate from normal ``feat_len`` calculation,
+        which still follows ``input_frame_rate`` and ``mel_framerate``.
+        """
+        return int(token_count / 12.5 * 22050 / 256)
+
+    def _create_block_causal_mask(
+        self,
+        block_pattern: list[int],
+        seq_len: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Build block-causal attention mask [1, 1, T, T].
+
+        Ported from GLM-TTS utils/block_mask_util.py ``create()``.
+        Each position can attend to its own block plus a look-ahead into
+        the next block (50% for small blocks, 66% for large blocks).
+
+        Vectorized via ``torch.searchsorted`` — no Python loop over T.
+        """
+        import itertools
+
+        mel_blocks = [self._token_to_mel_size(b) for b in block_pattern]
+        # Extend with copies of the last block (enough for any seq_len)
+        extended = mel_blocks + [mel_blocks[-1]] * ((seq_len // max(mel_blocks[-1], 1)) + 2)
+        accum = list(itertools.accumulate(extended))
+
+        # Vectorized: map each position to its block index via searchsorted
+        accum_t = torch.tensor(accum, dtype=torch.long)
+        extended_t = torch.tensor(extended, dtype=torch.long)
+        positions = torch.arange(seq_len, dtype=torch.long)
+
+        # searchsorted(accum, positions, right=True) gives the block index
+        # for each position (accum is strictly increasing).
+        block_idx = torch.searchsorted(accum_t, positions, right=True)
+
+        # Per-position: block start, block size, position within block
+        block_start = torch.where(block_idx > 0, accum_t[block_idx - 1], torch.zeros_like(block_idx))
+        cur_block_size = extended_t[block_idx]
+        idx_in_block = positions - block_start
+
+        # Look-ahead: positions near the end of a block can see into the next
+        look_future = cur_block_size - 1 - idx_in_block
+        min_future = torch.where(
+            cur_block_size <= 172,
+            cur_block_size // 2,
+            (cur_block_size.float() / 1.5).long(),
+        )
+        delta = torch.clamp(min_future - look_future, min=0)
+        aim_len = torch.clamp(block_start + cur_block_size + delta, max=seq_len)
+
+        # Build mask: position i can attend to columns [0, aim_len[i])
+        col_idx = torch.arange(seq_len, device=device)
+        mask = col_idx.unsqueeze(0) < aim_len.unsqueeze(1).to(device)
+
+        # [1, 1, T, T] for broadcast across batch and heads
+        return mask.unsqueeze(0).unsqueeze(0)
+
+    def forward(
+        self,
+        middle_point: torch.Tensor,
+        condition: torch.Tensor,
+        text: torch.Tensor,
+        time_step: torch.Tensor,
+        padding_mask: torch.Tensor,
+        spkr_emb: torch.Tensor | None = None,
+        text_lens: torch.Tensor | None = None,
+        is_causal: bool = False,
+        attn_mask: torch.Tensor | None = None,
+        block_pattern: list[int] | None = None,
+    ) -> torch.Tensor:
+        """Forward pass.
+
+        Args:
+            middle_point: Noised audio tensor [B, T, mel_dim]
+            condition: Masked audio condition [B, T, condition_dim]
+            text: Speech token tensor [B, T_text]
+            time_step: Timestep tensor [B]
+            padding_mask: Padding mask [B, T]
+            spkr_emb: Speaker embedding [B, spkr_dim]
+            is_causal: Use block-causal attention
+            attn_mask: Custom attention mask (overrides is_causal)
+            block_pattern: Token-level block sizes for block-causal mask
+
+        Returns:
+            Predicted mel-spectrogram [B, T, mel_dim]
+        """
+        # Cast floating inputs to model dtype (weights may be bfloat16)
+        model_dtype = self.proj_out.weight.dtype
+        middle_point = middle_point.to(model_dtype)
+        condition = condition.to(model_dtype)
+        time_step = time_step.to(model_dtype)
+        if spkr_emb is not None:
+            spkr_emb = spkr_emb.to(model_dtype)
+
+        _, seq_len = middle_point.shape[0], middle_point.shape[1]
+
+        # Time embedding
+        time_emb = self.time_embed(time_step)
+
+        # Concatenate speaker embedding if using AdaLN
+        if self.spkr_emb_adaln and spkr_emb is not None:
+            time_emb = torch.cat([time_emb, spkr_emb], dim=-1)
+
+        # Text embedding - ensure cast to model dtype
+        text_embed = self.text_emb_layer(text, seq_len, text_lens=text_lens).to(model_dtype)
+
+        # Input projection
+        x = self.emb_concator(middle_point, condition, text_embed, drop_audio_cond=False)
+
+        # Rotary embeddings
+        rope = self.rotary_embed.forward_from_seq_len(seq_len)
+
+        # Build block-causal mask when is_causal and no explicit attn_mask
+        if attn_mask is None and is_causal and block_pattern is not None:
+            attn_mask = self._create_block_causal_mask(
+                block_pattern,
+                seq_len,
+                middle_point.device,
+                model_dtype,
+            )
+
+        # Long skip connection
+        if self.long_skip_connection is not None:
+            residual = x
+
+        # Transformer blocks
+        for block in self.transformer_blocks:
+            x = block(x, time_emb, padding_mask=padding_mask.unsqueeze(1), rope=rope, attn_mask=attn_mask)
+
+        # Long skip connection
+        if self.long_skip_connection is not None:
+            x = self.long_skip_connection(torch.cat((x, residual), dim=-1))
+
+        # Output projection
+        x = self.norm_out(x, time_emb)
+        output = self.proj_out(x)
+
+        return output
diff --git a/vllm_omni/model_executor/models/glm_tts/glm_tts_dit_wrapper.py b/vllm_omni/model_executor/models/glm_tts/glm_tts_dit_wrapper.py
new file mode 100644
index 000000000..a7c26af30
--- /dev/null
+++ b/vllm_omni/model_executor/models/glm_tts/glm_tts_dit_wrapper.py
@@ -0,0 +1,1301 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GLM-TTS DiT wrapper: LLM_GENERATION-compatible flow-matching + vocoder."""
+
+from __future__ import annotations
+
+import dataclasses
+import json
+import math
+import os
+import threading
+from collections.abc import Iterable
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.cuda import CUDAGraph
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+from vllm_omni.model_executor.models.glm_tts.glm_tts import (
+    resolve_glm_tts_model_dir,
+)
+from vllm_omni.model_executor.models.glm_tts.vocoder import (
+    load_vocoder,
+    mel_to_audio,
+)
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+logger = init_logger(__name__)
+
+
+@dataclasses.dataclass
+class _StreamRequestState:
+    """Per-request streaming state (replaces 5 separate dicts)."""
+
+    diff_cache: dict[int | str, Any] = dataclasses.field(default_factory=dict)
+    prompt_payload: dict[str, torch.Tensor] = dataclasses.field(default_factory=dict)
+    audio_tail: torch.Tensor | None = None
+    wav_pointer: int = 0
+    fade_tail: torch.Tensor | None = None
+
+
+_GLM_TTS_OFFICIAL_FLOW_FALLBACK: dict[str, Any] = {
+    "spkr_emb_adaLN": True,
+    "speech_token_cfg": False,
+    "remove_spkr_concat_condition": True,
+    "mel_dim": 80,
+    "mel_framerate": 50,
+    "input_frame_rate": 25,
+}
+
+_GLM_TTS_RUNTIME_FLOW_DEFAULTS: dict[str, Any] = {
+    "sample_rate": 24000,
+    "n_timesteps": 10,
+    "t_scheduler": "cosine",
+    "inference_cfg_rate": 0.7,
+    "spk_embed_dim": 80,
+    "trans_dim": 768,
+    "depth": 18,
+    "heads": 12,
+}
+
+
+def _load_glm_tts_flow_config(flow_dir: str) -> dict[str, Any]:
+    """Load GLM-TTS flow config from JSON/YAML, falling back to official defaults."""
+    config: dict[str, Any] = dict(_GLM_TTS_OFFICIAL_FLOW_FALLBACK)
+
+    # Try config.json first (standard HuggingFace layout)
+    json_path = os.path.join(flow_dir, "config.json")
+    if os.path.exists(json_path):
+        try:
+            with open(json_path) as f:
+                raw = json.load(f)
+            flow = raw.get("flow", raw) if isinstance(raw, dict) else {}
+            if isinstance(flow, dict):
+                config.update(flow)
+        except Exception as e:
+            logger.warning("Failed to parse GLM-TTS flow config.json: %s", e)
+
+    # Try config.yaml (HyperPyYAML — strip tags, only extract kwargs)
+    yaml_path = os.path.join(flow_dir, "config.yaml")
+    if os.path.exists(yaml_path):
+        try:
+            import yaml
+
+            class _TagStripper(yaml.SafeLoader):
+                pass
+
+            _TagStripper.add_multi_constructor(
+                "!",
+                lambda loader, _tag, node: (
+                    loader.construct_mapping(node, deep=True)
+                    if isinstance(node, yaml.MappingNode)
+                    else loader.construct_sequence(node, deep=True)
+                    if isinstance(node, yaml.SequenceNode)
+                    else loader.construct_scalar(node)
+                ),
+            )
+            with open(yaml_path) as f:
+                raw = yaml.load(f, Loader=_TagStripper)
+            flow = raw.get("flow", raw) if isinstance(raw, dict) else {}
+            if isinstance(flow, dict):
+                config.update(flow)
+        except Exception as e:
+            logger.warning("Failed to parse GLM-TTS flow config.yaml: %s", e)
+
+    return config
+
+
+def make_pad_mask(lengths: torch.Tensor, max_len: int | None = None) -> torch.Tensor:
+    if max_len is None:
+        max_len = int(lengths.max().item())
+    seq_range = torch.arange(0, max_len, device=lengths.device)
+    return seq_range.unsqueeze(0).expand(lengths.shape[0], max_len) >= lengths.unsqueeze(1)
+
+
+def _unwrap(value: object) -> object:
+    """Unwrap single-element lists (common runtime_info wrapping)."""
+    return value[0] if isinstance(value, list) and value else value
+
+
+def nested_get(info: dict[str, Any], section: str, key: str) -> Any:
+    nested = info.get(section)
+    return nested[key] if isinstance(nested, dict) and key in nested else info.get(key)
+
+
+def _glm_tts_private_info(info: dict[str, Any]) -> dict[str, Any]:
+    private = info.get("glm_tts")
+    if isinstance(private, dict):
+        return private
+    kv_metadata = info.get("kv_metadata")
+    if isinstance(kv_metadata, dict):
+        private = kv_metadata.get("glm_tts")
+        if isinstance(private, dict):
+            return private
+    return {}
+
+
+def as_tensor(value: object) -> torch.Tensor | None:
+    value = _unwrap(value)
+    return value if isinstance(value, torch.Tensor) else None
+
+
+def as_bool(value: object) -> bool:
+    value = _unwrap(value)
+    if isinstance(value, torch.Tensor) and value.numel() > 0:
+        return bool(value.reshape(-1)[0].item())
+    return bool(value) if value is not None else False
+
+
+def as_str(value: object) -> str | None:
+    value = _unwrap(value)
+    return str(value) if value is not None else None
+
+
+def _as_scalar(value: object, typ: type, default: int | float = 0):
+    """Extract scalar int/float from runtime_info value."""
+    value = _unwrap(value)
+    if isinstance(value, torch.Tensor) and value.numel() > 0:
+        return typ(value.reshape(-1)[0].item())
+    if value is None:
+        return default
+    try:
+        return typ(value)
+    except (TypeError, ValueError):
+        return default
+
+
+def as_int(value: object, default: int = 0) -> int:
+    return _as_scalar(value, int, default)
+
+
+def as_float(value: object, default: float = 0.0) -> float:
+    return _as_scalar(value, float, default)
+
+
+def as_int_list(value: object) -> list[int] | None:
+    if value is None:
+        return None
+    if isinstance(value, torch.Tensor):
+        return [int(x) for x in value.reshape(-1).tolist()] if value.numel() else []
+    if isinstance(value, list):
+        return [int(x.item() if hasattr(x, "item") else x) for x in value]
+    return None
+
+
+def valid_speech_tokens(
+    value: object,
+    *,
+    device: torch.device,
+    fallback: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Return GLM-TTS speech tokens, treating an explicit empty list as empty."""
+    if value is None:
+        token = fallback if fallback is not None else torch.empty(0, device=device, dtype=torch.long)
+    elif isinstance(value, torch.Tensor):
+        token = value.to(device=device, dtype=torch.long)
+    elif isinstance(value, np.ndarray):
+        token = torch.as_tensor(value, device=device, dtype=torch.long)
+    elif isinstance(value, list):
+        if len(value) == 1 and isinstance(value[0], torch.Tensor):
+            token = value[0].to(device=device, dtype=torch.long)
+        elif value and all(isinstance(item, torch.Tensor) for item in value):
+            token = torch.cat([item.reshape(-1).to(device=device, dtype=torch.long) for item in value])
+        else:
+            token = torch.tensor(value, device=device, dtype=torch.long)
+    else:
+        token = fallback if fallback is not None else torch.empty(0, device=device, dtype=torch.long)
+
+    token = token.reshape(-1).to(device=device, dtype=torch.long)
+    return token[(token >= 0) & (token < 32768)]
+
+
+def split_request_ids(ids: torch.Tensor, seq_token_counts: list[int] | None = None) -> list[torch.Tensor]:
+    if seq_token_counts is not None:
+        boundaries, total = [0], ids.numel()
+        for count in seq_token_counts:
+            boundaries.append(boundaries[-1] + int(count))
+        return [ids[boundaries[i] : min(boundaries[i + 1], total)] for i in range(len(seq_token_counts))]
+    try:
+        from vllm.forward_context import get_forward_context, is_forward_context_available
+
+        if is_forward_context_available():
+            slices = get_forward_context().ubatch_slices
+            if slices is not None and len(slices) > 1 and not any(hasattr(s, "token_slice") for s in slices):
+                boundaries = [0]
+                for s in slices:
+                    boundaries.append(boundaries[-1] + int(s))
+                return [ids[boundaries[i] : boundaries[i + 1]] for i in range(len(boundaries) - 1)]
+    except ImportError:
+        pass
+    return [ids]
+
+
+class GLMTTSDiTForGeneration(nn.Module):
+    """GLM-TTS DiT flow-matching stage wrapped as LLM_GENERATION."""
+
+    have_multimodal_outputs = True
+    enable_update_additional_information = True
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        self.config = config
+        self.model_stage = "glm_tts_dit"
+        self.max_num_seqs = int(getattr(vllm_config.scheduler_config, "max_num_seqs", 1))
+        self._use_dit_cuda_graphs = bool(getattr(config, "use_dit_cuda_graphs", False))
+        self._dit_cudagraph: CUDAGraphGLMTTSDiTWrapper | None = None
+        self._codec_chunk_frames, self._codec_left_context_frames = self._connector_chunk_config(vllm_config)
+
+        model_path = vllm_config.model_config.model
+        model_path = resolve_glm_tts_model_dir(
+            model_path,
+            tokenizer_path=getattr(vllm_config.model_config, "tokenizer", None),
+            required_files=("flow/flow.pt",),
+            optional_files=("hift/hift.pt", "vocos2d/generator_jit.ckpt"),
+        )
+        self.model_dir = model_path
+        flow_dir = os.path.join(self.model_dir, "flow")
+        self.hf_flow_config = _load_glm_tts_flow_config(flow_dir)
+        logger.info("GLM-TTS flow config: %s", self.hf_flow_config)
+
+        self.mel_dim = self.hf_flow_config.get("mel_dim", 80)
+        self.input_frame_rate = self.hf_flow_config.get("input_frame_rate", 25.0)
+        self.mel_framerate = self.hf_flow_config.get("mel_framerate", 50)
+        self.sample_rate = self.hf_flow_config.get("sample_rate", _GLM_TTS_RUNTIME_FLOW_DEFAULTS["sample_rate"])
+
+        self.n_timesteps = self.hf_flow_config.get("n_timesteps", _GLM_TTS_RUNTIME_FLOW_DEFAULTS["n_timesteps"])
+        config_n_ts = getattr(config, "n_timesteps", None)
+        if config_n_ts is not None:
+            self.n_timesteps = int(config_n_ts)
+        self.t_scheduler = self.hf_flow_config.get("t_scheduler", _GLM_TTS_RUNTIME_FLOW_DEFAULTS["t_scheduler"])
+        self.inference_cfg_rate = self.hf_flow_config.get(
+            "inference_cfg_rate", _GLM_TTS_RUNTIME_FLOW_DEFAULTS["inference_cfg_rate"]
+        )
+        self.speech_token_cfg = self.hf_flow_config.get("speech_token_cfg", False)
+
+        self.spk_embed_dim = self.hf_flow_config.get("spk_embed_dim", _GLM_TTS_RUNTIME_FLOW_DEFAULTS["spk_embed_dim"])
+        self.spkr_emb_adaln = self.hf_flow_config.get("spkr_emb_adaLN", True)
+        self.remove_spkr_concat_condition = self.hf_flow_config.get("remove_spkr_concat_condition", True)
+
+        from vllm_omni.model_executor.models.glm_tts.glm_tts_dit import GLMTTSDiT
+
+        trans_dim = self.hf_flow_config.get("trans_dim", _GLM_TTS_RUNTIME_FLOW_DEFAULTS["trans_dim"])
+        depth = self.hf_flow_config.get("depth", _GLM_TTS_RUNTIME_FLOW_DEFAULTS["depth"])
+        heads = self.hf_flow_config.get("heads", _GLM_TTS_RUNTIME_FLOW_DEFAULTS["heads"])
+        condition_dim = self.mel_dim if self.remove_spkr_concat_condition else self.mel_dim + self.spk_embed_dim
+
+        self.dit = GLMTTSDiT(
+            trans_dim=trans_dim,
+            depth=depth,
+            heads=heads,
+            dim_head=64,
+            ff_mult=2,
+            dropout=0.0,
+            mel_dim=self.mel_dim,
+            text_vocab_size=100000,
+            text_emb_dim=512,
+            conv_layers=4,
+            condition_dim=condition_dim,
+            spkr_emb_adaln=self.spkr_emb_adaln,
+            spkr_dim=192,
+            use_wavlm_emb=False,
+        )
+        self._model_dtype = self._resolve_model_dtype(vllm_config)
+        if self._model_dtype is not None:
+            self.dit.to(dtype=self._model_dtype)
+
+        if not self.remove_spkr_concat_condition:
+            self.spk_embed_affine_layer = nn.Linear(192, self.spk_embed_dim)
+            if self._model_dtype is not None:
+                self.spk_embed_affine_layer.to(dtype=self._model_dtype)
+
+        self._vocoder = None
+        self._vocoder_sample_rate: int | None = None
+
+        self._stream_cache_lock = threading.Lock()
+        self._stream_state_by_req: dict[str, _StreamRequestState] = {}
+        # Track recently-seen req_ids so stale entries (aborted requests)
+        # can be evicted without modifying the shared model runner.
+        self._stream_last_seen: dict[str, int] = {}
+        self._stream_forward_count: int = 0
+
+        self._flow_weights_path = os.path.join(self.model_dir, "flow", "flow.pt")
+        self._hift_weights_path = os.path.join(self.model_dir, "hift", "hift.pt")
+        self._vocos_jit_path = os.path.join(self.model_dir, "vocos2d", "generator_jit.ckpt")
+
+    @staticmethod
+    def _connector_chunk_config(vllm_config: VllmConfig) -> tuple[int, int]:
+        cc = getattr(vllm_config.model_config, "stage_connector_config", None)
+        extra = (cc or {}).get("extra") if isinstance(cc, dict) else getattr(cc, "extra", None)
+        if not isinstance(extra, dict):
+            return 25, 25
+        cf = extra.get("codec_chunk_frames", 25)
+        return (int(cf[0]) if isinstance(cf, list) else int(cf)), int(extra.get("codec_left_context_frames", 25))
+
+    @staticmethod
+    def _resolve_model_dtype(vllm_config: VllmConfig) -> torch.dtype | None:
+        dtype = getattr(vllm_config.model_config, "dtype", None)
+        if isinstance(dtype, torch.dtype):
+            return dtype
+        if isinstance(dtype, str):
+            _aliases = {"bf16": "bfloat16", "fp16": "float16", "half": "float16", "fp32": "float32"}
+            return getattr(torch, _aliases.get(dtype.lower(), dtype.lower()), None)
+        return None
+
+    def _ensure_vocoder_loaded(self) -> None:
+        if self._vocoder is not None:
+            return
+
+        device = next(self.dit.parameters()).device
+        self._vocoder, self._vocoder_sample_rate = load_vocoder(self.model_dir, device, self.sample_rate)
+        if self._vocoder is None:
+            raise RuntimeError(
+                "GLM-TTS stage 1 could not load a vocoder. "
+                f"Expected either {self._hift_weights_path} or {self._vocos_jit_path}."
+            )
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        **_: Any,
+    ) -> torch.Tensor:
+        """Return zero embeddings (DiT does not use token embeddings)."""
+        if input_ids.dim() != 1:
+            raise ValueError(f"GLM-TTS DiT input_ids must be 1D, got shape={tuple(input_ids.shape)}")
+        hidden = int(self.config.hidden_size)
+        return torch.zeros(
+            (input_ids.shape[0], hidden),
+            device=input_ids.device,
+            dtype=self._model_dtype or torch.float32,
+        )
+
+    def compute_logits(self, hidden_states: Any, sampling_metadata: Any = None) -> None:
+        """DiT does not produce logits for sampling."""
+        return None
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        **kwargs: Any,
+    ) -> OmniOutput:
+        """Run DiT flow-matching + vocoder on speech tokens."""
+        # Evict stale streaming caches (aborted requests whose stream_finished
+        # was never sent).  Entries not seen for 64 forward calls are dropped.
+        self._evict_stale_stream_caches()
+
+        runtime_info = kwargs.get("model_intermediate_buffer") or []
+
+        seq_token_counts = kwargs.get("seq_token_counts")
+        flat_ids = input_ids.reshape(-1).to(dtype=torch.long)
+        if seq_token_counts is None and self.max_num_seqs > 1:
+            raise RuntimeError(
+                "GLM-TTS DiT stage requires seq_token_counts when max_num_seqs > 1; "
+                "otherwise concatenated speech tokens cannot be split per request."
+            )
+        request_ids_list = split_request_ids(flat_ids, seq_token_counts)
+
+        num_reqs = max(1, len(request_ids_list))
+        sample_rate = torch.tensor(int(self.sample_rate), dtype=torch.int32)
+        empty_audio = torch.zeros((0,), dtype=torch.float32, device=input_ids.device)
+        audios: list[torch.Tensor] = [empty_audio] * num_reqs
+        srs: list[torch.Tensor] = [sample_rate] * num_reqs
+
+        if not isinstance(runtime_info, list):
+            runtime_info = []
+
+        for idx, req_ids in enumerate(request_ids_list):
+            info = runtime_info[idx] if idx < len(runtime_info) and isinstance(runtime_info[idx], dict) else {}
+            req_id = as_str(nested_get(info, "meta", "req_id")) if info else None
+            if req_id is not None:
+                with self._stream_cache_lock:
+                    self._stream_last_seen[req_id] = self._stream_forward_count
+            stream_finished_raw = nested_get(info, "meta", "stream_finished") if info else None
+            if stream_finished_raw is None and info:
+                stream_finished_raw = info.get("stream_finished")
+            stream_finished = as_bool(stream_finished_raw) if info else False
+            prompt_token = as_tensor(nested_get(info, "embed", "speech_token")) if info else None
+            if prompt_token is None and info:
+                prompt_token = as_tensor(info.get("prompt_token"))  # backward compat
+            if prompt_token is None and info:
+                prompt_token = as_tensor(info.get("prompt_speech_token"))  # backward compat
+            prompt_feat = as_tensor(nested_get(info, "embed", "speech_feat")) if info else None
+            if prompt_feat is None and info:
+                prompt_feat = as_tensor(info.get("prompt_feat"))  # backward compat
+            embedding = as_tensor(nested_get(info, "embed", "embedding")) if info else None
+            if embedding is None and info:
+                embedding = as_tensor(info.get("embedding"))  # backward compat
+
+            speech_tokens_raw = info.get("speech_tokens") if info else None
+            has_meta = isinstance(info.get("meta"), dict) if info else False
+            uses_streaming = bool(info) and (
+                has_meta
+                or "stream_finished" in info
+                or "token_offset" in info
+                or "left_context_size" in info
+                or "glm_tts" in info
+                or "kv_metadata" in info
+            )
+            if uses_streaming and req_id is not None:
+                with self._stream_cache_lock:
+                    state = self._stream_state_by_req.setdefault(req_id, _StreamRequestState())
+                    prompt_payload = state.prompt_payload
+                if prompt_token is not None:
+                    prompt_payload["prompt_speech_token"] = prompt_token
+                elif "prompt_speech_token" in prompt_payload:
+                    prompt_token = prompt_payload["prompt_speech_token"]
+                if prompt_feat is not None:
+                    prompt_payload["prompt_feat"] = prompt_feat
+                elif "prompt_feat" in prompt_payload:
+                    prompt_feat = prompt_payload["prompt_feat"]
+                if embedding is not None:
+                    prompt_payload["embedding"] = embedding
+                elif "embedding" in prompt_payload:
+                    embedding = prompt_payload["embedding"]
+
+            fallback_token = req_ids.reshape(-1).to(device=input_ids.device, dtype=torch.long)
+            if uses_streaming:
+                token = valid_speech_tokens(None, device=input_ids.device, fallback=fallback_token)
+            else:
+                token = valid_speech_tokens(
+                    speech_tokens_raw,
+                    device=input_ids.device,
+                    fallback=fallback_token,
+                )
+
+            if token.numel() == 0:
+                if stream_finished and req_id is not None:
+                    with self._stream_cache_lock:
+                        self._stream_state_by_req.pop(req_id, None)
+                        self._stream_last_seen.pop(req_id, None)
+                audios[idx] = empty_audio
+                continue
+
+            if uses_streaming:
+                private_info = _glm_tts_private_info(info)
+                token_offset_raw = nested_get(info, "meta", "left_context_size")
+                if token_offset_raw is None:
+                    token_offset_raw = info.get("left_context_size")
+                if token_offset_raw is None:
+                    token_offset_raw = nested_get(info, "meta", "token_offset")
+                if token_offset_raw is None:
+                    token_offset_raw = info.get("token_offset")
+                token_offset = max(
+                    0,
+                    as_int(token_offset_raw, 0),
+                )
+
+                chunk_sizes_raw = private_info.get("chunk_sizes_history")
+                if chunk_sizes_raw is None:
+                    chunk_sizes_raw = nested_get(info, "meta", "chunk_sizes_history")
+                if chunk_sizes_raw is None:
+                    chunk_sizes_raw = info.get("chunk_sizes_history")
+                block_pattern_raw = private_info.get("block_pattern")
+                if block_pattern_raw is None:
+                    block_pattern_raw = nested_get(info, "meta", "block_pattern")
+                if block_pattern_raw is None:
+                    block_pattern_raw = info.get("block_pattern")
+                crossfade_raw = private_info.get("crossfade_sec")
+                if crossfade_raw is None:
+                    crossfade_raw = nested_get(info, "meta", "crossfade_sec")
+                if crossfade_raw is None:
+                    crossfade_raw = info.get("crossfade_sec")
+                chunk_sizes_history = as_int_list(chunk_sizes_raw)
+                block_pattern = as_int_list(block_pattern_raw)
+                crossfade_sec = as_float(crossfade_raw, 0.1)
+
+                audio = self._forward_full(
+                    speech_tokens=token,
+                    prompt_token=prompt_token,
+                    prompt_feat=prompt_feat,
+                    embedding=embedding,
+                    token_offset=token_offset,
+                    req_id=req_id,
+                    stream_finished=stream_finished,
+                    chunk_sizes_history=chunk_sizes_history,
+                    block_pattern=block_pattern,
+                    crossfade_sec=crossfade_sec,
+                )
+                if stream_finished and req_id is not None:
+                    with self._stream_cache_lock:
+                        self._stream_state_by_req.pop(req_id, None)
+                        self._stream_last_seen.pop(req_id, None)
+            else:
+                audio = self._forward_full(
+                    speech_tokens=token,
+                    prompt_token=prompt_token,
+                    prompt_feat=prompt_feat,
+                    embedding=embedding,
+                )
+
+            audios[idx] = audio.reshape(-1).to(dtype=torch.float32)
+
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={"audio": audios, "sr": srs},
+        )
+
+    @staticmethod
+    def _to_2d(val: torch.Tensor | np.ndarray | None, device: torch.device, dtype: torch.dtype) -> torch.Tensor | None:
+        """Convert numpy/tensor to 2D+ tensor on device, or return None."""
+        if val is None:
+            return None
+        if isinstance(val, np.ndarray):
+            val = torch.from_numpy(val)
+        val = val.to(device=device, dtype=dtype)
+        if val.ndim == 1:
+            val = val.unsqueeze(0)
+        return val
+
+    @staticmethod
+    def _validate_voice_clone(prompt_token, prompt_feat, embedding) -> None:
+        """Raise if voice clone payload is incomplete."""
+        fields = {"prompt_speech_token": prompt_token, "prompt_feat": prompt_feat, "embedding": embedding}
+        present = [k for k, v in fields.items() if v is not None]
+        missing = [k for k, v in fields.items() if v is None]
+        if present and missing:
+            raise RuntimeError(f"GLM-TTS voice clone payload incomplete: present={present}, missing={missing}.")
+
+    @torch.inference_mode()
+    def _forward_full(
+        self,
+        speech_tokens: torch.Tensor,
+        prompt_token: torch.Tensor | None = None,
+        prompt_feat: torch.Tensor | None = None,
+        embedding: torch.Tensor | None = None,
+        token_offset: int = 0,
+        req_id: str | None = None,
+        stream_finished: bool = False,
+        chunk_sizes_history: list[int] | None = None,
+        block_pattern: list[int] | None = None,
+        crossfade_sec: float = 0.1,
+    ) -> torch.Tensor:
+        """DiT flow-matching + vocoder pipeline. Returns audio waveform."""
+        device = speech_tokens.device
+        self._validate_voice_clone(prompt_token, prompt_feat, embedding)
+
+        if speech_tokens.ndim == 1:
+            speech_tokens = speech_tokens.unsqueeze(0)
+        generated_token_len = int(speech_tokens.shape[1])
+
+        has_voice_clone = prompt_token is not None or embedding is not None
+        prompt_token = self._to_2d(prompt_token, device, torch.long)
+        if prompt_token is not None:
+            prompt_token_len = prompt_token.shape[1]
+            speech_tokens = torch.cat([prompt_token, speech_tokens], dim=1)
+        else:
+            prompt_token_len = 0
+        has_prompt_audio_context = prompt_token_len > 0
+
+        token_len = speech_tokens.shape[1]
+        if token_len == 0:
+            return torch.zeros(0, device=device)
+
+        text_lens = torch.tensor([token_len], device=device, dtype=torch.long)
+        feat_len = max(1, int(token_len / self.input_frame_rate * self.mel_framerate))
+
+        padding_mask = ~make_pad_mask(torch.tensor([feat_len], device=device), max_len=feat_len)
+        mel_cond = torch.zeros([1, feat_len, self.mel_dim], device=device)
+
+        prompt_feat_len = 0
+        prompt_feat = self._to_2d(prompt_feat, device, mel_cond.dtype)
+        if prompt_feat is not None:
+            if prompt_feat.ndim == 2:
+                prompt_feat = prompt_feat.unsqueeze(0)
+            copy_len = min(prompt_feat.shape[1], feat_len)
+            mel_cond[:, :copy_len, :] = prompt_feat[:, -copy_len:, :]
+            prompt_feat_len = copy_len
+        has_prompt_audio_context = has_prompt_audio_context and prompt_feat_len > 0
+
+        embedding = self._to_2d(embedding, device, mel_cond.dtype)
+        spkr_embedding = None
+        if not self.remove_spkr_concat_condition:
+            if has_voice_clone and embedding is not None:
+                spk_embedding_normed = torch.nn.functional.normalize(embedding, dim=1)
+                spk_embedding = self.spk_embed_affine_layer(spk_embedding_normed)
+                spk_embedding_expanded = spk_embedding.unsqueeze(1).expand(-1, feat_len, -1)
+                condition = torch.cat([mel_cond, spk_embedding_expanded], dim=-1)
+                spkr_embedding = spk_embedding_normed
+            else:
+                spk_zeros = torch.zeros([1, feat_len, self.spk_embed_dim], device=device)
+                condition = torch.cat([mel_cond, spk_zeros], dim=-1)
+        else:
+            condition = mel_cond
+
+        if self.spkr_emb_adaln:
+            if has_voice_clone and embedding is not None:
+                if spkr_embedding is None:
+                    spkr_embedding = torch.nn.functional.normalize(embedding, dim=1)
+            else:
+                spkr_embedding = torch.zeros([1, 192], device=device, dtype=mel_cond.dtype)
+
+        with self._stream_cache_lock:
+            _rs = self._stream_state_by_req.get(req_id) if req_id is not None else None
+            last_step_cache = _rs.diff_cache if _rs is not None else None
+
+        dit_block_pattern: list[int] | None = None
+        is_streaming = req_id is not None and token_offset >= 0
+        if is_streaming:
+            pattern_source = block_pattern or chunk_sizes_history
+            if pattern_source is not None and len(pattern_source) > 0:
+                if prompt_token_len > 0:
+                    dit_block_pattern = [prompt_token_len] + pattern_source
+                else:
+                    dit_block_pattern = list(pattern_source)
+
+        lookahead_mel_len = 0
+        if is_streaming and chunk_sizes_history is not None and len(chunk_sizes_history) > 0:
+            last_chunk_tokens = chunk_sizes_history[-1]
+            lookahead_tokens = last_chunk_tokens // 2
+            lookahead_mel_len = int(lookahead_tokens / self.input_frame_rate * self.mel_framerate)
+
+        # Compute cache_len: only cache the stable region for streaming
+        stream_cache_len: int | None = None
+        if req_id is not None:
+            stable_mel_len = int(generated_token_len / self.input_frame_rate * self.mel_framerate)
+            override_mel = max(0, stable_mel_len - lookahead_mel_len)
+            stream_cache_len = min(feat_len, prompt_feat_len + override_mel + 1)
+
+        mel, current_step_cache = self._do_sample(
+            speech_tokens=speech_tokens,
+            mel_cond=mel_cond,
+            condition=condition,
+            padding_mask=padding_mask,
+            spkr_embedding=spkr_embedding,
+            n_timesteps=self.n_timesteps,
+            last_step_cache=last_step_cache,
+            block_pattern=dit_block_pattern,
+            cache_len=stream_cache_len,
+            text_lens=text_lens,
+        )
+
+        if req_id is not None:
+            current_step_cache["override_len"] = int(stream_cache_len or 0)
+            with self._stream_cache_lock:
+                rs = self._stream_state_by_req.setdefault(req_id, _StreamRequestState())
+                rs.diff_cache = current_step_cache
+
+        if has_prompt_audio_context:
+            crop_len = min(prompt_feat_len, max(0, int(mel.shape[1]) - 1))
+            if crop_len > 0:
+                mel = mel[:, crop_len:, :]
+
+        self._ensure_vocoder_loaded()
+
+        audio = mel_to_audio(self._vocoder, mel)
+
+        effective_sr = self._vocoder_sample_rate or self.sample_rate
+        if effective_sr != self.sample_rate and audio is not None and audio.numel() > 0:
+            try:
+                import torchaudio
+
+                audio = torchaudio.functional.resample(audio, effective_sr, self.sample_rate)
+            except Exception as e:
+                logger.warning(
+                    "Failed to resample audio from %d to %d: %s",
+                    effective_sr,
+                    self.sample_rate,
+                    e,
+                )
+
+        audio = audio.reshape(-1)
+        if is_streaming and req_id is not None:
+            audio = self._slice_streaming_audio(
+                audio,
+                req_id=req_id,
+                stream_finished=stream_finished,
+                lookahead_mel_len=lookahead_mel_len,
+                crossfade_sec=crossfade_sec,
+            )
+
+        return audio
+
+    def _slice_streaming_audio(
+        self,
+        audio: torch.Tensor,
+        *,
+        req_id: str,
+        stream_finished: bool,
+        lookahead_mel_len: int,
+        crossfade_sec: float,
+    ) -> torch.Tensor:
+        """Slice cumulative streaming audio using the official wav pointer.
+
+        Official GLM-TTS runs the flow/vocoder on the cumulative token prefix
+        every chunk, then drops the waveform samples that were already emitted.
+        Non-final chunks additionally keep a lookahead plus fade tail for the
+        next chunk; the final full-prefix pass flushes that tail.
+        """
+        look_back_len = max(0, int(lookahead_mel_len * self.sample_rate / self.mel_framerate))
+        fade_len = max(0, int(crossfade_sec * self.sample_rate))
+        overlap_len = look_back_len + fade_len
+
+        with self._stream_cache_lock:
+            _rs = self._stream_state_by_req.get(req_id)
+            wav_pointer = _rs.wav_pointer if _rs is not None else 0
+            fade_tail = _rs.fade_tail if _rs is not None else None
+
+        current = audio[min(wav_pointer, int(audio.numel())) :].clone()
+        if fade_tail is not None and fade_tail.numel() > 0 and current.numel() > 0:
+            blend_len = min(int(fade_tail.numel()), int(current.numel()))
+            fade_in = torch.linspace(0.0, 1.0, blend_len, dtype=current.dtype, device=current.device)
+            fade_out = 1.0 - fade_in
+            current[:blend_len] = (
+                fade_tail[:blend_len].to(device=current.device, dtype=current.dtype) * fade_out
+                + current[:blend_len] * fade_in
+            )
+
+        if stream_finished:
+            with self._stream_cache_lock:
+                _rs = self._stream_state_by_req.get(req_id)
+                if _rs is not None:
+                    _rs.wav_pointer = 0
+                    _rs.fade_tail = None
+            return current
+
+        if overlap_len <= 0 or current.numel() <= overlap_len:
+            emit = torch.zeros(0, dtype=current.dtype, device=current.device)
+            next_tail = current.detach().clone()
+        else:
+            emit = current[:-overlap_len]
+            tail_end = -look_back_len if look_back_len > 0 else None
+            next_tail = current[-overlap_len:tail_end].detach().clone()
+
+        with self._stream_cache_lock:
+            rs = self._stream_state_by_req.setdefault(req_id, _StreamRequestState())
+            rs.wav_pointer = wav_pointer + int(emit.numel())
+            rs.fade_tail = next_tail
+        return emit
+
+    # Maximum number of forward calls before a cache entry is considered stale.
+    _STALE_THRESHOLD = 64
+
+    def _evict_stale_stream_caches(self) -> None:
+        """Evict streaming cache entries for requests not seen recently.
+
+        This is a self-contained fallback: if ``stream_finished`` is never
+        sent (e.g. client disconnect), the entry would otherwise leak forever.
+        Called at the top of every ``forward`` so no changes to the shared
+        model runner are needed.
+        """
+        self._stream_forward_count += 1
+        # Only check every 16 forward calls to amortize the lock overhead.
+        if self._stream_forward_count % 16 != 0:
+            return
+        with self._stream_cache_lock:
+            if not self._stream_last_seen:
+                return
+            stale_ids = [
+                rid
+                for rid, last_seen in self._stream_last_seen.items()
+                if self._stream_forward_count - last_seen > self._STALE_THRESHOLD
+            ]
+            for rid in stale_ids:
+                self._stream_state_by_req.pop(rid, None)
+                self._stream_last_seen.pop(rid, None)
+
+    def _do_sample(
+        self,
+        speech_tokens: torch.Tensor,
+        mel_cond: torch.Tensor,
+        condition: torch.Tensor,
+        padding_mask: torch.Tensor,
+        spkr_embedding: torch.Tensor | None,
+        n_timesteps: int,
+        last_step_cache: dict[int | str, Any] | None = None,
+        block_pattern: list[int] | None = None,
+        cache_len: int | None = None,
+        text_lens: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, dict[int | str, Any]]:
+        """Euler ODE sampling for flow matching."""
+        # CUDA graph path: the captured graph does NOT pass is_causal /
+        # block_pattern to the DiT, so it always uses bidirectional attention.
+        # That is correct for non-streaming (full-sequence) inference but
+        # wrong for streaming where block-causal attention is required.
+        # Skip the graph when block_pattern is set so we fall through to the
+        # eager path that correctly forwards these arguments.
+        use_cudagraph = block_pattern is None
+        graph_result = (
+            self._try_cudagraph_sample(
+                speech_tokens=speech_tokens,
+                mel_cond=mel_cond,
+                condition=condition,
+                padding_mask=padding_mask,
+                spkr_embedding=spkr_embedding,
+                n_timesteps=n_timesteps,
+                last_step_cache=last_step_cache,
+                cache_len=cache_len,
+                text_lens=text_lens,
+            )
+            if use_cudagraph
+            else None
+        )
+        if graph_result is not None:
+            return graph_result
+
+        device = speech_tokens.device
+        # Always accumulate the Euler ODE state in float32 regardless of the
+        # DiT model dtype.  Half-precision rounding errors across 10 steps (×2
+        # CFG passes each) manifest as audible hiss on H20/H100.
+        x = torch.randn(mel_cond.shape, device=device, dtype=torch.float32)
+        # Cast conditioning tensors to float32 for the ODE loop.
+        condition_f32 = condition.to(torch.float32)
+        spkr_embedding_f32 = spkr_embedding.to(torch.float32) if spkr_embedding is not None else None
+        # Pre-cast invariant conditioning tensors to DiT dtype once (avoids
+        # repeated .to() calls inside the inner Euler loop).
+        condition_dit = condition_f32.to(mel_cond.dtype)
+        spkr_embedding_dit = spkr_embedding_f32.to(mel_cond.dtype) if spkr_embedding_f32 is not None else None
+        current_step_cache: dict[int | str, Any] = {}
+        effective_cache_len = min(int(cache_len or 0), int(mel_cond.shape[1]))
+        should_cache_steps = effective_cache_len > 0
+        t_span = torch.linspace(0, 1, n_timesteps + 1, device=device, dtype=torch.float32)
+        if self.t_scheduler == "cosine":
+            t_span = 1 - torch.cos(t_span * 0.5 * torch.pi)
+
+        for step in range(n_timesteps):
+            cache_step = step + 1
+            if last_step_cache is not None and cache_step in last_step_cache:
+                x_cache = last_step_cache[cache_step].get("x")
+                if isinstance(x_cache, torch.Tensor):
+                    override_len = int(last_step_cache.get("override_len", x_cache.shape[1]))
+                    safe_len = min(int(x.shape[1]), int(x_cache.shape[1]), max(0, override_len))
+                    if safe_len > 0:
+                        x[:, :safe_len, :] = x_cache[:, :safe_len, :].to(device=device, dtype=torch.float32)
+            if should_cache_steps:
+                current_step_cache[cache_step] = {"x": x[:, :effective_cache_len, :].detach().clone()}
+
+            t_current = t_span[step]
+            dt = t_span[step + 1] - t_current
+
+            use_causal = block_pattern is not None
+            x_dit = x.to(mel_cond.dtype)
+
+            if self.inference_cfg_rate > 0:
+                # Batch conditional + unconditional into a single forward
+                middle_b = x_dit.expand(2, -1, -1)
+                condition_b = torch.cat([condition_dit, torch.zeros_like(condition_dit)], dim=0)
+                if self.speech_token_cfg:
+                    text_b = torch.cat([speech_tokens, torch.zeros_like(speech_tokens)], dim=0)
+                else:
+                    text_b = speech_tokens.expand(2, -1)
+                time_b = t_current.unsqueeze(0).expand(2)
+                padding_b = padding_mask.expand(2, -1)
+                if spkr_embedding_dit is not None and self.spkr_emb_adaln:
+                    spkr_b = torch.cat([spkr_embedding_dit, torch.zeros_like(spkr_embedding_dit)], dim=0)
+                else:
+                    spkr_b = None
+                text_lens_b = text_lens.expand(2) if text_lens is not None else None
+                pred_b = self.dit(
+                    middle_point=middle_b,
+                    condition=condition_b,
+                    text=text_b,
+                    time_step=time_b,
+                    padding_mask=padding_b,
+                    spkr_emb=spkr_b,
+                    is_causal=use_causal,
+                    block_pattern=block_pattern,
+                    text_lens=text_lens_b,
+                ).to(torch.float32)
+                dphi_dt, cfg_dphi_dt = pred_b.chunk(2, dim=0)
+                dphi_dt = (1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt
+            else:
+                dphi_dt = self.dit(
+                    middle_point=x_dit,
+                    condition=condition_dit,
+                    text=speech_tokens,
+                    time_step=t_current.unsqueeze(0),
+                    padding_mask=padding_mask,
+                    spkr_emb=spkr_embedding_dit if spkr_embedding_dit is not None and self.spkr_emb_adaln else None,
+                    is_causal=use_causal,
+                    block_pattern=block_pattern,
+                    text_lens=text_lens,
+                ).to(torch.float32)
+
+            x = x + dt * dphi_dt
+
+        return x.to(mel_cond.dtype), current_step_cache
+
+    def _try_cudagraph_sample(
+        self,
+        *,
+        speech_tokens: torch.Tensor,
+        mel_cond: torch.Tensor,
+        condition: torch.Tensor,
+        padding_mask: torch.Tensor,
+        spkr_embedding: torch.Tensor | None,
+        n_timesteps: int,
+        last_step_cache: dict[int | str, Any] | None,
+        cache_len: int | None = None,
+        text_lens: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, dict[int | str, Any]] | None:
+        if not self._use_dit_cuda_graphs or mel_cond.device.type != "cuda":
+            return None
+        if self._dit_cudagraph is None:
+            condition_dim = self.mel_dim if self.remove_spkr_concat_condition else self.mel_dim + self.spk_embed_dim
+            self._dit_cudagraph = CUDAGraphGLMTTSDiTWrapper(
+                self.dit,
+                mel_dim=self.mel_dim,
+                condition_dim=condition_dim,
+                input_frame_rate=self.input_frame_rate,
+                mel_framerate=self.mel_framerate,
+                inference_cfg_rate=self.inference_cfg_rate,
+                speech_token_cfg=self.speech_token_cfg,
+                spkr_emb_adaln=self.spkr_emb_adaln,
+                t_scheduler=self.t_scheduler,
+                enabled=True,
+            )
+            try:
+                self._dit_cudagraph.warmup(
+                    device=mel_cond.device,
+                    dtype=mel_cond.dtype,
+                    n_timesteps=n_timesteps,
+                    codec_chunk_frames=self._codec_chunk_frames,
+                    codec_left_context_frames=self._codec_left_context_frames,
+                )
+            except Exception:
+                logger.warning("Disabling GLM-TTS DiT CUDA graphs after warmup failure", exc_info=True)
+                self._dit_cudagraph.enabled = False
+                return None
+
+        try:
+            return self._dit_cudagraph.sample(
+                speech_tokens=speech_tokens,
+                mel_cond=mel_cond,
+                condition=condition,
+                padding_mask=padding_mask,
+                spkr_embedding=spkr_embedding,
+                n_timesteps=n_timesteps,
+                last_step_cache=last_step_cache,
+                cache_len=cache_len,
+                text_lens=text_lens,
+            )
+        except Exception:
+            logger.warning("GLM-TTS DiT CUDA graph replay failed; falling back to eager", exc_info=True)
+            return None
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load DiT + vocoder weights from {model_dir}/flow/flow.pt."""
+        for _ in weights:
+            pass
+
+        weight_dict: dict[str, torch.Tensor] = {}
+        if not os.path.exists(self._flow_weights_path):
+            raise FileNotFoundError(f"GLM-TTS flow weights not found at {self._flow_weights_path}")
+
+        logger.info("Loading GLM-TTS flow weights from %s", self._flow_weights_path)
+        checkpoint = torch.load(self._flow_weights_path, map_location="cpu", weights_only=True)
+        if "model" in checkpoint:
+            weight_dict = checkpoint["model"]
+        else:
+            weight_dict = checkpoint
+
+        dit_state = {
+            (name.removeprefix("estimator.") if name.startswith("estimator.") else name.removeprefix("dit.")): tensor
+            for name, tensor in weight_dict.items()
+            if (name.startswith("estimator.") or name.startswith("dit."))
+            and "inv_freq" not in name
+            and "rotary_embed" not in name
+        }
+        if not dit_state:
+            raise RuntimeError(
+                "GLM-TTS flow checkpoint did not contain any DiT weights under 'estimator.*' or 'dit.*'."
+            )
+        missing, unexpected = self.dit.load_state_dict(dit_state, strict=False)
+        missing = [k for k in missing if "inv_freq" not in k]
+        if missing:
+            logger.warning("Missing DiT weights: %s", missing[:10])
+        if unexpected:
+            logger.warning("Unexpected DiT weights: %s", unexpected[:10])
+
+        loaded_spk: set[str] = set()
+        if hasattr(self, "spk_embed_affine_layer"):
+            for suffix in ("weight", "bias"):
+                key = f"spk_embed_affine_layer.{suffix}"
+                if key in weight_dict:
+                    param = getattr(self.spk_embed_affine_layer, suffix)
+                    param.data.copy_(weight_dict[key].to(device=param.device, dtype=param.dtype))
+                    loaded_spk.add(key)
+            self.spk_embed_affine_layer.eval()
+        self.dit.eval()
+        self._ensure_vocoder_loaded()
+
+        # Return both prefixed and unprefixed param names for nested module compat
+        loaded_params: set[str] = set()
+        for k in dit_state:
+            loaded_params.update((f"dit.{k}", f"_dit_gen.dit.{k}"))
+        for k in loaded_spk:
+            loaded_params.update((k, f"_dit_gen.{k}"))
+        logger.info("Loaded GLMTTSDiTForGeneration: %d params", len(loaded_params))
+        return loaded_params
+
+
+class CUDAGraphGLMTTSDiTWrapper:
+    """Capture GLM-TTS DiT Euler sampling for fixed mel-frame buckets.
+
+    The wrapper owns static buffers for a single-request DiT stage. Dynamic
+    requests are right-padded to the nearest captured mel-frame bucket and
+    trimmed back to the actual mel length after graph replay.
+    """
+
+    def __init__(
+        self,
+        dit: torch.nn.Module,
+        *,
+        mel_dim: int,
+        condition_dim: int,
+        input_frame_rate: float,
+        mel_framerate: float,
+        inference_cfg_rate: float,
+        speech_token_cfg: bool,
+        spkr_emb_adaln: bool,
+        t_scheduler: str,
+        enabled: bool = True,
+        capture_sizes: list[int] | None = None,
+    ) -> None:
+        self.dit = dit
+        self.mel_dim = int(mel_dim)
+        self.condition_dim = int(condition_dim)
+        self.input_frame_rate = float(input_frame_rate)
+        self.mel_framerate = float(mel_framerate)
+        self.inference_cfg_rate = float(inference_cfg_rate)
+        self.speech_token_cfg = bool(speech_token_cfg)
+        self.spkr_emb_adaln = bool(spkr_emb_adaln)
+        self.t_scheduler = str(t_scheduler)
+        self.enabled = bool(enabled)
+        self.capture_sizes = sorted(set(capture_sizes or []))
+
+        self.graphs: dict[int, CUDAGraph] = {}
+        self.static_x: dict[int, torch.Tensor] = {}
+        self.static_condition: dict[int, torch.Tensor] = {}
+        self.static_text: dict[int, torch.Tensor] = {}
+        self.static_text_lens: dict[int, torch.Tensor] = {}
+        self.static_padding_mask: dict[int, torch.Tensor] = {}
+        self.static_spkr_emb: dict[int, torch.Tensor] = {}
+        self.static_output: dict[int, torch.Tensor] = {}
+        self.static_step_cache: dict[int, list[torch.Tensor]] = {}
+
+        self._warmed_up = False
+        self._n_timesteps: int | None = None
+
+    @staticmethod
+    def compute_capture_sizes(
+        *,
+        codec_chunk_frames: int = 25,
+        codec_left_context_frames: int = 25,
+        input_frame_rate: float = 25.0,
+        mel_framerate: float = 50.0,
+        max_bucket: int = 1024,
+    ) -> list[int]:
+        sizes: set[int] = set()
+
+        def token_to_mel(tokens: int) -> int:
+            return max(1, int(math.ceil(float(tokens) / input_frame_rate * mel_framerate)))
+
+        if codec_chunk_frames > 0:
+            sizes.add(token_to_mel(codec_chunk_frames))
+            if codec_left_context_frames > 0:
+                sizes.add(token_to_mel(codec_chunk_frames + codec_left_context_frames))
+
+        for size in (32, 64, 128, 256, 512, 1024):
+            if size <= max_bucket:
+                sizes.add(size)
+        return sorted(sizes)
+
+    def _get_padded_size(self, actual_feat_len: int) -> int | None:
+        for size in self.capture_sizes:
+            if actual_feat_len <= size:
+                return size
+        return None
+
+    def _token_bucket_len(self, bucket_feat_len: int) -> int:
+        return max(1, int(math.ceil(float(bucket_feat_len) * self.input_frame_rate / self.mel_framerate)))
+
+    def warmup(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        n_timesteps: int,
+        codec_chunk_frames: int = 25,
+        codec_left_context_frames: int = 25,
+    ) -> None:
+        if not self.enabled or self._warmed_up or device.type != "cuda" or torch.cuda.is_current_stream_capturing():
+            return
+        self._n_timesteps = int(n_timesteps)
+        if not self.capture_sizes:
+            self.capture_sizes = self.compute_capture_sizes(
+                codec_chunk_frames=codec_chunk_frames,
+                codec_left_context_frames=codec_left_context_frames,
+                input_frame_rate=self.input_frame_rate,
+                mel_framerate=self.mel_framerate,
+            )
+
+        self.dit.eval()
+        logger.info("Starting GLM-TTS DiT CUDA graph warmup for sizes: %s", self.capture_sizes)
+        for size in self.capture_sizes:
+            try:
+                self._allocate_static_buffers(size, device, dtype)
+                self._capture(size, device)
+                logger.info("Captured GLM-TTS DiT CUDA graph for mel frames=%d", size)
+            except Exception:
+                logger.warning("Failed to capture GLM-TTS DiT CUDA graph for size=%d", size, exc_info=True)
+        self._warmed_up = True
+
+    def _allocate_static_buffers(self, size: int, device: torch.device, dtype: torch.dtype) -> None:
+        token_len = self._token_bucket_len(size)
+        self.static_x[size] = torch.zeros(1, size, self.mel_dim, device=device, dtype=dtype)
+        self.static_condition[size] = torch.zeros(1, size, self.condition_dim, device=device, dtype=dtype)
+        self.static_text[size] = torch.zeros(1, token_len, device=device, dtype=torch.long)
+        self.static_text_lens[size] = torch.full((1,), token_len, device=device, dtype=torch.long)
+        self.static_padding_mask[size] = torch.zeros(1, size, device=device, dtype=torch.bool)
+        self.static_spkr_emb[size] = torch.zeros(1, 192, device=device, dtype=dtype)
+        self.static_output[size] = torch.zeros(1, size, self.mel_dim, device=device, dtype=dtype)
+        self.static_step_cache[size] = [
+            torch.zeros(1, size, self.mel_dim, device=device, dtype=dtype) for _ in range(int(self._n_timesteps or 0))
+        ]
+
+    def _time_schedule(self, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+        n_timesteps = int(self._n_timesteps or 0)
+        t_span = torch.linspace(0, 1, n_timesteps + 1, device=device, dtype=dtype)
+        if self.t_scheduler == "cosine":
+            t_span = 1 - torch.cos(t_span * 0.5 * torch.pi)
+        return t_span
+
+    def _run_static(self, size: int) -> torch.Tensor:
+        x = self.static_x[size]
+        condition = self.static_condition[size]
+        text = self.static_text[size]
+        text_lens = self.static_text_lens[size]
+        padding_mask = self.static_padding_mask[size]
+        spkr_emb = self.static_spkr_emb[size] if self.spkr_emb_adaln else None
+        t_span = self._time_schedule(x.device, x.dtype)
+
+        for step in range(int(self._n_timesteps or 0)):
+            self.static_step_cache[size][step].copy_(x)
+            t_current = t_span[step]
+            dt = t_span[step + 1] - t_current
+            if self.inference_cfg_rate > 0:
+                middle_b = x.expand(2, -1, -1)
+                condition_b = torch.cat([condition, torch.zeros_like(condition)], dim=0)
+                if self.speech_token_cfg:
+                    text_b = torch.cat([text, torch.zeros_like(text)], dim=0)
+                else:
+                    text_b = text.expand(2, -1)
+                time_b = t_current.unsqueeze(0).expand(2)
+                padding_b = padding_mask.expand(2, -1)
+                text_lens_b = text_lens.expand(2)
+                if spkr_emb is not None:
+                    spkr_b = torch.cat([spkr_emb, torch.zeros_like(spkr_emb)], dim=0)
+                else:
+                    spkr_b = None
+                pred_b = self.dit(
+                    middle_point=middle_b,
+                    condition=condition_b,
+                    text=text_b,
+                    text_lens=text_lens_b,
+                    time_step=time_b,
+                    padding_mask=padding_b,
+                    spkr_emb=spkr_b,
+                )
+                dphi_dt, cfg_dphi_dt = pred_b.chunk(2, dim=0)
+                dphi_dt = (1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt
+            else:
+                dphi_dt = self.dit(
+                    middle_point=x,
+                    condition=condition,
+                    text=text,
+                    text_lens=text_lens,
+                    time_step=t_current.unsqueeze(0),
+                    padding_mask=padding_mask,
+                    spkr_emb=spkr_emb,
+                )
+            x = x + dt * dphi_dt
+        self.static_output[size].copy_(x)
+        return self.static_output[size]
+
+    def _capture(self, size: int, device: torch.device) -> None:
+        with torch.no_grad():
+            self.static_x[size].normal_()
+            _ = self._run_static(size)
+        torch.accelerator.synchronize(device)
+
+        graph = CUDAGraph()
+        with torch.no_grad():
+            with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()):
+                self._run_static(size)
+        self.graphs[size] = graph
+
+    def sample(
+        self,
+        *,
+        speech_tokens: torch.Tensor,
+        mel_cond: torch.Tensor,
+        condition: torch.Tensor,
+        padding_mask: torch.Tensor,
+        spkr_embedding: torch.Tensor | None,
+        n_timesteps: int,
+        last_step_cache: dict[int | str, Any] | None = None,
+        cache_len: int | None = None,
+        text_lens: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, dict[int | str, Any]] | None:
+        if (
+            not self.enabled
+            or not self._warmed_up
+            or int(n_timesteps) != int(self._n_timesteps or -1)
+            or speech_tokens.shape[0] != 1
+            or mel_cond.shape[0] != 1
+            or mel_cond.device.type != "cuda"
+            or torch.cuda.is_current_stream_capturing()
+        ):
+            return None
+
+        actual_feat_len = int(mel_cond.shape[1])
+        size = self._get_padded_size(actual_feat_len)
+        if size is None or size not in self.graphs:
+            return None
+
+        token_bucket_len = int(self.static_text[size].shape[1])
+        if int(speech_tokens.shape[1]) > token_bucket_len:
+            return None
+
+        self.static_x[size].normal_()
+        self.static_condition[size].zero_()
+        self.static_text[size].zero_()
+        self.static_padding_mask[size].zero_()
+        self.static_spkr_emb[size].zero_()
+
+        # Warm-start: when streaming, pre-fill the stable region of the
+        # initial noise from the step-1 cache of the previous chunk.
+        # The graph's Euler loop then evolves from this warm-started state.
+        # New (lookahead) region keeps random noise — no per-step overlay.
+        if last_step_cache is not None and 1 in last_step_cache:
+            cached_x = last_step_cache[1]["x"]
+            override_len = int(last_step_cache.get("override_len", cached_x.shape[1]))
+            safe = min(size, int(cached_x.shape[1]), max(0, override_len))
+            if safe > 0:
+                self.static_x[size][:, :safe, :].copy_(cached_x[:, :safe, :].to(self.static_x[size]))
+
+        self.static_condition[size][:, :actual_feat_len, : condition.shape[-1]] = condition
+        self.static_text[size][:, : speech_tokens.shape[1]] = speech_tokens
+        actual_text_len = int(speech_tokens.shape[1])
+        if text_lens is not None:
+            self.static_text_lens[size].copy_(text_lens)
+        else:
+            self.static_text_lens[size].fill_(actual_text_len)
+        self.static_padding_mask[size][:, :actual_feat_len] = padding_mask
+        if spkr_embedding is not None:
+            self.static_spkr_emb[size].copy_(spkr_embedding.to(self.static_spkr_emb[size]))
+
+        self.graphs[size].replay()
+        mel = self.static_output[size][:, :actual_feat_len, :].clone()
+        current_step_cache: dict[int | str, Any] = {}
+        cache_len = min(int(cache_len or 0), actual_feat_len)
+        if cache_len > 0:
+            for idx, cache in enumerate(self.static_step_cache[size], start=1):
+                current_step_cache[idx] = {"x": cache[:, :cache_len, :].detach().clone()}
+        return mel, current_step_cache
diff --git a/vllm_omni/model_executor/models/glm_tts/pipeline.py b/vllm_omni/model_executor/models/glm_tts/pipeline.py
new file mode 100644
index 000000000..f016a993e
--- /dev/null
+++ b/vllm_omni/model_executor/models/glm_tts/pipeline.py
@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GLM-TTS pipeline: Stage 0 (AR) → Stage 1 (DiT)."""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+_PROC = "vllm_omni.model_executor.stage_input_processors.glm_tts"
+
+GLM_TTS_PIPELINE = PipelineConfig(
+    model_type="glm_tts",
+    model_arch="GLMTTSForConditionalGeneration",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="glm_tts",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            owns_tokenizer=True,
+            engine_output_type="latent",
+            async_chunk_process_next_stage_input_func=(f"{_PROC}.ar_to_dit_async_chunk"),
+            sampling_constraints={
+                # GLM-TTS uses 👂 (token string "👂", ID 59253) as
+                # end-of-audio marker.  The ID is resolved dynamically by
+                # GLMTTSForConditionalGeneration.__init__ from the tokenizer and
+                # validated at runtime against this hardcoded value.  If the
+                # upstream checkpoint changes the mapping, the model will log a
+                # warning so the constant here can be updated.
+                "stop_token_ids": [59253],
+            },
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="glm_tts_dit",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="audio",
+            engine_output_type="latent",
+            sync_process_input_func=f"{_PROC}.ar_to_dit",
+        ),
+    ),
+)
diff --git a/vllm_omni/model_executor/models/glm_tts/text_frontend.py b/vllm_omni/model_executor/models/glm_tts/text_frontend.py
new file mode 100644
index 000000000..8cfc0e5c8
--- /dev/null
+++ b/vllm_omni/model_executor/models/glm_tts/text_frontend.py
@@ -0,0 +1,393 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Lightweight text frontend aligned with the official GLM-TTS preprocessing."""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+
+try:
+    import emoji
+except ImportError:
+    emoji = None  # type: ignore[assignment]
+
+try:
+    import inflect
+except ImportError:
+    inflect = None  # type: ignore[assignment]
+
+try:
+    from tn.chinese.normalizer import Normalizer as ZhNormalizer
+except ImportError:
+    ZhNormalizer = None
+
+try:
+    from tn.english.normalizer import Normalizer as EnNormalizer
+except ImportError:
+    EnNormalizer = None
+
+
+CHINESE_CHAR_PATTERN = re.compile(r"[\u4e00-\u9fff]+")
+PUNCTUATION_CHARS = r"。？！；：、.?!;:，,"
+
+
+def contains_chinese(text: str) -> bool:
+    return bool(CHINESE_CHAR_PATTERN.search(text))
+
+
+def markdown_norm(markdown_text: str) -> str:
+    markdown_text = re.sub(r"^(\d+)\. ", r"\1。", markdown_text)
+    return markdown_text.replace("\\n", "\n")
+
+
+def multi_line_process(plain_text: str) -> str:
+    lines: list[str] = []
+    for line in plain_text.split("\n"):
+        line = line.strip()
+        if not line:
+            continue
+        if line[-1] not in {".", "!", "?", ";", ":", "：", "。", "！", "？", "；", "，"}:
+            line = f"{line}{'。' if contains_chinese(line) else '. '}"
+        lines.append(line)
+    return "".join(lines) if contains_chinese("".join(lines)) else " ".join(lines)
+
+
+def remove_bracket(text: str, lang: str = "zh") -> str:
+    brackets_to_remove = [
+        ("(", ")"),
+        ("（", "）"),
+        ("【", "】"),
+        ("「", "」"),
+        ("`", "`"),
+        ("《", "》"),
+        ("『", "』"),
+        ("{", "}"),
+        ("[", "]"),
+    ]
+    if lang != "en":
+        brackets_to_remove.append(("‘", "’"))
+    for left, right in brackets_to_remove:
+        text = text.replace(left, "").replace(right, "")
+    return text
+
+
+def number_to_chinese(number: int | float | str) -> str:
+    units = ["", "十", "百", "千", "万", "十万", "百万", "千万", "亿", "十亿", "百亿", "千亿"]
+    nums = "零一二三四五六七八九"
+    if isinstance(number, (float, str)):
+        str_num = str(float(number))
+        if "." in str_num:
+            int_part, decimal_part = str_num.split(".")
+            if decimal_part.strip("0") == "":
+                return number_to_chinese(int(int_part))
+            chinese_int = number_to_chinese(int(int_part)) if int(int_part) != 0 else "零"
+            chinese_decimal = "".join(nums[int(d)] for d in decimal_part)
+            return f"{chinese_int}点{chinese_decimal}"
+
+    number = int(number)
+    if number == 0:
+        return "零"
+    if number < 0:
+        return "负" + number_to_chinese(abs(number))
+
+    result: list[str] = []
+    unit_position = 0
+    zero_flag = False
+    last_unit = ""
+    while number > 0:
+        num = number % 10
+        current_unit = units[unit_position] if unit_position < len(units) else ""
+        if num == 0:
+            if not zero_flag and result and last_unit not in ["万", "亿"]:
+                result.append(nums[num])
+                zero_flag = True
+        else:
+            if current_unit:
+                result.append(current_unit)
+            result.append(nums[num])
+            zero_flag = False
+            last_unit = current_unit
+        unit_position += 1
+        number //= 10
+
+    result.reverse()
+    text = "".join(result)
+    text = re.sub(r"^一十", "十", text)
+    text = re.sub(r"零+", "零", text)
+    text = re.sub(r"零$", "", text)
+    return text
+
+
+def tn_scientific_notation(text: str) -> str:
+    pattern = re.compile(r"(.*?)(-?\d+(?:\.\d+)?)[*x×](\d+(?:\.\d+)?)\^(-?\d+(?:\.\d+)?)(.*?)$")
+    match = pattern.match(text)
+    if not match:
+        return text
+    prefix, base, multiplier, exponent, suffix = match.groups()
+    description = f"{number_to_chinese(base)}乘{number_to_chinese(multiplier)}的{number_to_chinese(exponent)}次方"
+    return f"{prefix}{description}{suffix}"
+
+
+def replace_asterisk_with_multiply(text: str, lang: str) -> str:
+    if lang == "en":
+        replace_name = "multiply"
+        replace_number = "0-9"
+    else:
+        replace_name = "乘"
+        replace_number = "一二三四五六七八九十百千万亿"
+
+    rule0 = rf"(?<=[{replace_number}a-zA-Z])\s*\*\s*(?=[{replace_number}a-zA-Z])"
+    rule1 = rf"(?<=[{replace_number}a-zA-Z]\))\s*\*\s*(?=[{replace_number}a-zA-Z])"
+    rule2 = rf"(?<=[{replace_number}a-zA-Z])\s*\*\s*(?=\([{replace_number}a-zA-Z])"
+    rule3 = rf"(?<=[{replace_number}a-zA-Z]\))\s*\*\s*(?=\([{replace_number}a-zA-Z])"
+    return re.sub(f"{rule0}|{rule1}|{rule2}|{rule3}", replace_name, text).replace("*", "")
+
+
+def special_replace(text: str) -> str:
+    text = text.replace("\\", "")
+    text = f" {text} "
+    text = re.sub(r"(?<=\W)额(?=\W)", "呃", text)
+    text = re.sub(r"[~～]+", "。", text)
+    text = replace_asterisk_with_multiply(text, "zh")
+    return text.strip()
+
+
+def spell_out_number(text: str, inflect_parser: object | None) -> str:
+    if inflect_parser is None:
+        return text
+    new_text: list[str] = []
+    start = None
+    for idx, char in enumerate(text):
+        if not char.isdigit():
+            if start is not None:
+                new_text.append(inflect_parser.number_to_words(text[start:idx]))
+                start = None
+            new_text.append(char)
+        elif start is None:
+            start = idx
+    if start is not None and start < len(text):
+        new_text.append(inflect_parser.number_to_words(text[start:]))
+    return "".join(new_text)
+
+
+def replace_space(text: str) -> str:
+    alphanumeric_pattern = r"[a-zA-Z0-9]"
+    punctuation_pattern = r"[.,!?;:]"
+    text = re.sub(r" +", " ", text)
+    result = ""
+    idx = 0
+    while idx < len(text):
+        current_char = text[idx]
+        if current_char != " ":
+            result += current_char
+            idx += 1
+            continue
+        prev_char = text[idx - 1] if idx > 0 else ""
+        next_char = text[idx + 1] if idx + 1 < len(text) else ""
+        if re.match(alphanumeric_pattern, prev_char) and re.match(alphanumeric_pattern, next_char):
+            result += " "
+        elif re.match(punctuation_pattern, prev_char) and re.match(alphanumeric_pattern, next_char):
+            result += " "
+        idx += 1
+    return result
+
+
+def normalize_punctuation(text: str, punctuation_chars: str) -> str:
+    text = replace_space(text)
+    text = re.sub(rf"([{punctuation_chars}])\1+", r"\1", text)
+    text = re.sub(rf"([{punctuation_chars}])(?=[{punctuation_chars}])", "", text)
+    text = text.replace("#", "")
+    text = text.replace("！", "。")
+    text = text.replace("!", ".")
+    return text
+
+
+def ensure_proper_ending(text: str) -> str:
+    if not text:
+        return text
+    if text[-1] in "？?":
+        return text
+    if text[-1] in PUNCTUATION_CHARS:
+        return f"{text[:-1]}{'。' if contains_chinese(text) else '.'}"
+    return f"{text}{'。' if contains_chinese(text) else '.'}"
+
+
+class GLMTTSTextFrontend:
+    """Subset of the official GLM-TTS text frontend used by the adapter."""
+
+    def __init__(self) -> None:
+        self.inflect_parser = inflect.engine() if inflect is not None else None
+        self.zh_tn_model = (
+            ZhNormalizer(
+                remove_erhua=False,
+                full_to_half=True,
+                remove_interjections=False,
+                overwrite_cache=True,
+            )
+            if ZhNormalizer is not None
+            else None
+        )
+        self.en_tn_model = EnNormalizer() if EnNormalizer is not None else None
+        self._custom_replacements = self._load_custom_replacements()
+
+    @staticmethod
+    def _load_custom_replacements() -> list[tuple[str, str]]:
+        custom_replace_path = os.path.join(os.path.dirname(__file__), "configs", "custom_replace.jsonl")
+        if not os.path.exists(custom_replace_path):
+            return []
+        replacements: list[tuple[str, str]] = []
+        with open(custom_replace_path, encoding="utf-8") as f:
+            for line in f:
+                if not line.strip():
+                    continue
+                item = json.loads(line)
+                replacements.append((item["origin"], item["new"]))
+        return replacements
+
+    def pre_replace(self, sentence: str) -> str:
+        sentence = tn_scientific_notation(sentence)
+        sentence = re.sub(r"(?<=\d)\s*-\s*(?=\d)", "减", sentence)
+        sentence = sentence.replace("-", "")
+        sentence = re.sub(r"咯([" + re.escape(PUNCTUATION_CHARS) + r"])", r"喽\1", sentence)
+        for origin, new in self._custom_replacements:
+            sentence = sentence.replace(origin, new)
+        return sentence
+
+    def post_replace(self, sentence: str) -> str:
+        sentence = remove_bracket(sentence)
+        sentence = sentence.replace(" - ", "，")
+        sentence = sentence.replace("——", "，")
+        sentence = re.sub(r"[,:：;；、]+", "，", sentence)
+        sentence = re.sub(r"[.…]+", "。", sentence)
+        sentence = re.sub(r"[_·]+", "", sentence)
+        sentence = re.sub(r"""['"‘’“”|]+""", "", sentence)
+
+        symbol_map = {
+            "†": "，",
+            "²": "平方",
+            "³": "立方",
+            "/": "每",
+            "~": "到",
+            "～": "到",
+            "①": "一",
+            "②": "二",
+            "③": "三",
+            "④": "四",
+            "⑤": "五",
+            "⑥": "六",
+            "⑦": "七",
+            "⑧": "八",
+            "⑨": "九",
+            "⑩": "十",
+            "α": "阿尔法",
+            "β": "贝塔",
+            "γ": "伽玛",
+            "Γ": "伽玛",
+            "δ": "德尔塔",
+            "Δ": "德尔塔",
+            "△": "德尔塔",
+            "ε": "艾普西龙",
+            "ζ": "捷塔",
+            "η": "依塔",
+            "θ": "西塔",
+            "Θ": "西塔",
+            "ι": "艾欧塔",
+            "κ": "喀帕",
+            "λ": "拉姆达",
+            "Λ": "拉姆达",
+            "μ": "缪",
+            "ν": "拗",
+            "ξ": "克西",
+            "Ξ": "克西",
+            "ο": "欧米克伦",
+            "π": "派",
+            "Π": "派",
+            "ρ": "肉",
+            "ς": "西格玛",
+            "Σ": "西格玛",
+            "σ": "西格玛",
+            "τ": "套",
+            "υ": "宇普西龙",
+            "φ": "服艾",
+            "Φ": "服艾",
+            "χ": "器",
+            "ψ": "普赛",
+            "Ψ": "普赛",
+            "ω": "欧米伽",
+            "Ω": "欧米伽",
+            "□": "方框",
+            ">": "大于",
+            "<": "小于",
+            "∈": "属于",
+            "∉": "不属于",
+            "∪": "并",
+            "∩": "交",
+            "⊥": "垂直",
+            "∥": "平行",
+            "≠": "不等于",
+            "∵": "因为",
+            "∴": "所以",
+            "∅": "空集",
+            "⊂": "真包含于",
+            "⊃": "包含",
+            "⊆": "包含于",
+            "⊇": "真包含",
+            "⊄": "不属于",
+            "⊅": "非超集",
+            "⊈": "不属于",
+            "⊉": "非超集",
+        }
+        for old, new in symbol_map.items():
+            sentence = sentence.replace(old, new)
+        return sentence
+
+    def text_normalize(self, text: str | None) -> str | None:
+        if text is None:
+            return None
+        text = self._preprocess_text(text)
+        if contains_chinese(text):
+            text = self._normalize_chinese_text(text).lower()
+        else:
+            text = self._normalize_english_text(text)
+        text = normalize_punctuation(text, PUNCTUATION_CHARS)
+        return ensure_proper_ending(text)
+
+    def _preprocess_text(self, text: str) -> str:
+        text = markdown_norm(text)
+        text = multi_line_process(text)
+        if emoji is not None:
+            text = emoji.replace_emoji(text, replace="")
+        return re.sub(r"(?<=[a-zA-Z])-(?=[a-zA-Z])", " ", text)
+
+    def _normalize_chinese_text(self, text: str) -> str:
+        text = self.pre_replace(text)
+        if self.zh_tn_model is not None:
+            text = self.zh_tn_model.normalize(text)
+        text = special_replace(text)
+        return self.post_replace(text).strip()
+
+    def _normalize_english_text(self, text: str) -> str:
+        text = text.replace("'", "’")
+        if self.en_tn_model is not None:
+            text = self.en_tn_model.normalize(text)
+        text = remove_bracket(text, "en")
+        text = replace_asterisk_with_multiply(text, "en")
+        text = text.replace("—", " ")
+        text = text.replace("’", "'")
+        text = spell_out_number(text, self.inflect_parser)
+        text = re.sub(r"\s+", " ", text)
+        keep_punctuation = r"\.,!\?'\:;"
+        pattern = rf"[^\w\s{keep_punctuation}]"
+        text = re.sub(pattern, "", text)
+        text = text.lower()
+        text = re.sub(r"\.+", ".", text)
+        text = re.sub(r"\,+", ",", text)
+        text = re.sub(r"!+", "!", text)
+        text = re.sub(r"\?+", "?", text)
+        text = re.sub(r"'+", "'", text)
+        text = re.sub(r":+", ":", text)
+        text = re.sub(r";+", ";", text)
+        text = re.sub(r"\s*([.,?!':;])\s*", r"\1 ", text)
+        return text.strip()
diff --git a/vllm_omni/model_executor/models/glm_tts/vocoder.py b/vllm_omni/model_executor/models/glm_tts/vocoder.py
new file mode 100644
index 000000000..77c61c0de
--- /dev/null
+++ b/vllm_omni/model_executor/models/glm_tts/vocoder.py
@@ -0,0 +1,228 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Vocoder loading and mel-to-audio conversion for GLM-TTS.
+
+Supports:
+  - HiFT (24kHz) — reuses CosyVoice3's vendored HiFTGenerator
+  - Vocos2D JIT (32kHz) from TorchScript checkpoint
+
+Extracted from glm_tts_dit_wrapper.py to keep file sizes under 800 lines.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import torch
+import torch.nn as nn
+from vllm.logger import init_logger
+
+try:
+    from torch.nn.utils.parametrizations import weight_norm
+except ImportError:
+    from torch.nn.utils import weight_norm
+
+logger = init_logger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# HiFT vocoder (24kHz) — reuses CosyVoice3's vendored HiFTGenerator
+# ---------------------------------------------------------------------------
+
+
+class ConvRNNF0Predictor(nn.Module):
+    """Non-causal F0 predictor for GLM-TTS HiFT vocoder.
+
+    GLM-TTS ships a non-causal HiFT checkpoint whose F0 predictor uses
+    standard ``nn.Conv1d`` layers (kernel_size=3, padding=1).  CosyVoice3's
+    vendored hifigan only includes the *causal* variant
+    (``CausalConvRNNF0Predictor``), so we provide the non-causal version
+    here.  The network structure and state_dict keys are identical to the
+    original ``cosyvoice.hifigan_cosy2.f0_predictor.ConvRNNF0Predictor``.
+    """
+
+    def __init__(
+        self,
+        num_class: int = 1,
+        in_channels: int = 80,
+        cond_channels: int = 512,
+    ):
+        super().__init__()
+        self.num_class = num_class
+        self.condnet = nn.Sequential(
+            weight_norm(nn.Conv1d(in_channels, cond_channels, kernel_size=3, padding=1)),
+            nn.ELU(),
+            weight_norm(nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1)),
+            nn.ELU(),
+            weight_norm(nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1)),
+            nn.ELU(),
+            weight_norm(nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1)),
+            nn.ELU(),
+            weight_norm(nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1)),
+            nn.ELU(),
+        )
+        self.classifier = nn.Linear(in_features=cond_channels, out_features=self.num_class)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.condnet(x)
+        x = x.transpose(1, 2)
+        return torch.abs(self.classifier(x).squeeze(-1))
+
+
+class HiFTWrapper:
+    """Thin wrapper around CosyVoice3's vendored HiFTGenerator."""
+
+    def __init__(self, state_dict: dict[str, Any], device: torch.device):
+        self.device = device
+        self.sample_rate = 24000
+        from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
+            HiFTGenerator,
+        )
+
+        f0_predictor = ConvRNNF0Predictor(num_class=1, in_channels=80, cond_channels=512)
+        self.model = HiFTGenerator(
+            in_channels=80,
+            base_channels=512,
+            nb_harmonics=8,
+            sampling_rate=24000,
+            nsf_alpha=0.1,
+            nsf_sigma=0.003,
+            nsf_voiced_threshold=10,
+            upsample_rates=[8, 5, 3],
+            upsample_kernel_sizes=[16, 11, 7],
+            istft_params={"n_fft": 16, "hop_len": 4},
+            resblock_kernel_sizes=[3, 7, 11],
+            resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+            source_resblock_kernel_sizes=[7, 7, 11],
+            source_resblock_dilation_sizes=[
+                [1, 3, 5],
+                [1, 3, 5],
+                [1, 3, 5],
+            ],
+            lrelu_slope=0.1,
+            audio_limit=0.99,
+            f0_predictor=f0_predictor,
+        ).to(device)
+        self.model.load_state_dict(state_dict)
+        self.model.eval()
+        self.model.float()
+
+    def __call__(self, mel: torch.Tensor) -> torch.Tensor:
+        mel = mel.to(self.device)
+        with torch.no_grad():
+            audio, _ = self.model.inference(mel)
+        return audio
+
+
+def load_hift(ckpt_path: str, device: torch.device) -> HiFTWrapper:
+    """Load HiFT vocoder from checkpoint."""
+    state_dict = torch.load(ckpt_path, map_location=device, weights_only=True)
+    if isinstance(state_dict, dict):
+        for key in ("state_dict", "model", "generator"):
+            nested = state_dict.get(key)
+            if isinstance(nested, dict):
+                state_dict = nested
+                break
+    return HiFTWrapper(state_dict, device)
+
+
+# ---------------------------------------------------------------------------
+# Vocos2D JIT vocoder (32kHz)
+# ---------------------------------------------------------------------------
+
+_MEL_LOGDIFF = -7.847762537473608
+
+
+class Vocos2DWrapper:
+    """Thin wrapper around Vocos2D TorchScript model."""
+
+    def __init__(self, ckpt_path: str, device: torch.device):
+        self.device = device
+        self.gen_model = torch.jit.load(ckpt_path, map_location=device)
+        self.gen_model.eval()
+        self._mel_logdiff = _MEL_LOGDIFF
+
+    def __call__(self, mel: torch.Tensor) -> torch.Tensor:
+        if mel.size(-1) == 1:
+            mel = torch.cat([mel, mel], dim=-1)
+        xs_mel = mel.transpose(-1, -2) + self._mel_logdiff
+        xs = self.gen_model(xs_mel.to(self.device))
+        if xs.ndim == 2:
+            xs = xs.unsqueeze(1)
+        return xs
+
+
+def load_vocos2d_jit(ckpt_path: str, device: torch.device) -> Vocos2DWrapper:
+    """Load Vocos2D JIT vocoder from TorchScript checkpoint."""
+    return Vocos2DWrapper(ckpt_path, device)
+
+
+# ---------------------------------------------------------------------------
+# Unified loader
+# ---------------------------------------------------------------------------
+
+
+def load_vocoder(
+    model_root: str,
+    device: torch.device,
+    sample_rate: int = 24000,
+) -> tuple[Any | None, int]:
+    """Try to load the best available vocoder.
+
+    Returns:
+        (vocoder, actual_sample_rate) or (None, sample_rate) on failure.
+    """
+    # Try HiFT first for 24kHz
+    if sample_rate == 24000:
+        hift_path = os.path.join(model_root, "hift", "hift.pt")
+        if os.path.isfile(hift_path):
+            try:
+                vocoder = load_hift(hift_path, device)
+                logger.info("Loaded HiFT vocoder from %s", hift_path)
+                return vocoder, 24000
+            except Exception as e:
+                logger.warning("Failed to load HiFT: %s", e)
+
+    # Try Vocos2D JIT
+    vocos_jit_path = os.path.join(model_root, "vocos2d", "generator_jit.ckpt")
+    if os.path.isfile(vocos_jit_path):
+        try:
+            vocoder = load_vocos2d_jit(vocos_jit_path, device)
+            logger.info("Loaded Vocos2D JIT vocoder (32kHz)")
+            return vocoder, 32000
+        except Exception as e:
+            logger.warning("Failed to load Vocos2D JIT: %s", e)
+
+    logger.warning("No vocoder available")
+    return None, sample_rate
+
+
+def mel_to_audio(
+    vocoder: Any | None,
+    mel: torch.Tensor,
+) -> torch.Tensor:
+    """Convert mel-spectrogram to audio waveform.
+
+    Args:
+        vocoder: Vocoder instance (HiFTWrapper, Vocos2DWrapper, or None).
+        mel: Mel-spectrogram [B, T, mel_dim].
+
+    Returns:
+        Audio waveform [B, 1, samples].
+    """
+    # Transpose to [B, mel_dim, T] for vocoder.
+    # HiFT / Vocos2D are precision-sensitive; always feed float32.
+    mel = mel.to(torch.float32).transpose(1, 2)
+
+    if vocoder is not None:
+        if callable(vocoder) and not hasattr(vocoder, "decode"):
+            audio = vocoder(mel)
+        else:
+            audio = vocoder.decode(mel)
+        if audio.ndim == 2:
+            audio = audio.unsqueeze(1)
+    else:
+        raise RuntimeError("GLM-TTS vocoder is not loaded.")
+
+    return audio
diff --git a/vllm_omni/model_executor/models/glm_tts/voice_clone.py b/vllm_omni/model_executor/models/glm_tts/voice_clone.py
new file mode 100644
index 000000000..d4d32f4f6
--- /dev/null
+++ b/vllm_omni/model_executor/models/glm_tts/voice_clone.py
@@ -0,0 +1,299 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GLM-TTS voice cloning frontend: speech tokenizer, speaker embedding, mel features.
+
+Lazy-loads external models (WhisperVQ, CampPlus ONNX) and provides extraction
+helpers consumed by the GLM-TTS multimodal processor.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import torchaudio.compliance.kaldi as kaldi
+import torchaudio.transforms as T
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_MEL_BASIS_CACHE: dict[tuple[int, int, int, int, int, str], torch.Tensor] = {}
+_HANN_WINDOW_CACHE: dict[tuple[int, str], torch.Tensor] = {}
+
+
+def _spectral_normalize_torch(magnitudes: torch.Tensor) -> torch.Tensor:
+    return torch.log(torch.clamp(magnitudes, min=1e-5))
+
+
+def _get_glm_tts_mel_basis(
+    *,
+    device: torch.device,
+    sample_rate: int,
+    n_fft: int,
+    n_mels: int,
+    f_min: int,
+    f_max: int,
+) -> torch.Tensor:
+    """Compute or retrieve cached GLM-TTS mel filter bank.
+
+    Uses torchaudio melscale_fbanks with Slaney parameters (norm="slaney",
+    mel_scale="slaney") to match the official GLM-TTS ``utils/audio.py``
+    librosa.filters.mel output.  The previous HTK / norm=None config
+    produced a numerically different filter bank that shifted the
+    prompt_feat distribution and degraded voice-clone conditioning.
+    """
+    cache_key = (sample_rate, n_fft, n_mels, f_min, f_max, str(device))
+    mel_basis = _MEL_BASIS_CACHE.get(cache_key)
+    if mel_basis is not None:
+        return mel_basis
+
+    import torchaudio.functional as AF
+
+    mel_basis = (
+        AF.melscale_fbanks(
+            n_freqs=(n_fft // 2) + 1,
+            f_min=float(f_min),
+            f_max=float(f_max),
+            n_mels=n_mels,
+            sample_rate=sample_rate,
+            norm="slaney",
+            mel_scale="slaney",
+        )
+        .transpose(0, 1)
+        .contiguous()
+        .float()
+        .to(device)
+    )
+
+    _MEL_BASIS_CACHE[cache_key] = mel_basis
+    return mel_basis
+
+
+def _extract_glm_tts_mel_feature(
+    audio: torch.Tensor,
+    *,
+    sample_rate: int,
+    n_fft: int,
+    win_length: int,
+    hop_length: int,
+    n_mels: int,
+    f_min: int,
+    f_max: int,
+) -> torch.Tensor:
+    """Extract time-major log-mel spectrogram features from audio waveform."""
+    device = audio.device
+    mel_basis = _get_glm_tts_mel_basis(
+        device=device,
+        sample_rate=sample_rate,
+        n_fft=n_fft,
+        n_mels=n_mels,
+        f_min=f_min,
+        f_max=f_max,
+    )
+
+    hann_key = (win_length, str(device))
+    window = _HANN_WINDOW_CACHE.get(hann_key)
+    if window is None:
+        window = torch.hann_window(win_length, device=device)
+        _HANN_WINDOW_CACHE[hann_key] = window
+
+    if audio.ndim == 1:
+        audio = audio.unsqueeze(0)
+
+    pad = int((n_fft - hop_length) / 2)
+    audio = torch.nn.functional.pad(
+        audio.unsqueeze(1),
+        (pad, pad),
+        mode="reflect",
+    ).squeeze(1)
+
+    spec = torch.stft(
+        audio,
+        n_fft,
+        hop_length=hop_length,
+        win_length=win_length,
+        window=window,
+        center=False,
+        pad_mode="reflect",
+        normalized=False,
+        onesided=True,
+        return_complex=True,
+    )
+    spec = torch.sqrt(spec.real.pow(2) + spec.imag.pow(2) + 1e-9)
+    spec = torch.matmul(mel_basis, spec)
+    spec = _spectral_normalize_torch(spec)
+    return spec.squeeze(0).transpose(0, 1).contiguous()
+
+
+def load_voice_clone_frontend(
+    model_root: str,
+    model_device: torch.device,
+    *,
+    speech_tokenizer_cache: tuple | None,
+    campplus_cache: object | None,
+    campplus_path: str | None = None,
+) -> tuple[tuple | None, object | None]:
+    """Lazy-load voice cloning frontend models.
+
+    Args:
+        campplus_path: Pre-resolved path to ``campplus.onnx``.  When provided
+            this takes priority over the default ``model_root/frontend/``
+            lookup.  The caller (``GLMTTSForConditionalGeneration``) resolves
+            the path via ``transformers.utils.hub.cached_file`` during
+            ``__init__`` so the file is already downloaded.
+
+    Returns:
+        ``(speech_tokenizer, campplus_session)`` — guaranteed non-None
+        on success.  Raises ``RuntimeError`` when a required component
+        cannot be loaded.
+    """
+    speech_tokenizer = speech_tokenizer_cache
+    campplus_session = campplus_cache
+
+    if speech_tokenizer is None:
+        speech_tokenizer_paths = (
+            os.path.join(model_root, "speech_tokenizer"),
+            os.path.join(model_root, "ckpt", "speech_tokenizer"),
+        )
+        speech_tokenizer_path = next((path for path in speech_tokenizer_paths if os.path.isdir(path)), None)
+        if speech_tokenizer_path is None:
+            raise RuntimeError(
+                f"GLM-TTS speech tokenizer directory not found at {speech_tokenizer_paths[0]} "
+                f"or {speech_tokenizer_paths[1]}.  Voice cloning requires the WhisperVQ "
+                f"speech tokenizer from the model snapshot."
+            )
+        try:
+            from safetensors.torch import load_file
+            from transformers import AutoFeatureExtractor, WhisperConfig
+
+            from vllm_omni.model_executor.models.common.whisper_vq import (
+                WhisperVQEncoder,
+            )
+
+            _config = WhisperConfig.from_pretrained(speech_tokenizer_path)
+            _config.quantize_encoder_only = True
+            _model = WhisperVQEncoder(_config)
+            _sd = load_file(os.path.join(speech_tokenizer_path, "model.safetensors"))
+            _prefix, _cleaned = "model.encoder.", {}
+            for k, v in _sd.items():
+                _cleaned[k[len(_prefix) :] if k.startswith(_prefix) else k] = v
+            _model.load_state_dict(_cleaned, strict=False)
+            _model = _model.to(model_device).eval()
+            _fe = AutoFeatureExtractor.from_pretrained(speech_tokenizer_path)
+            speech_tokenizer = (_model, _fe)
+            logger.info("Loaded GLM-TTS speech tokenizer from %s", speech_tokenizer_path)
+        except Exception as exc:
+            raise RuntimeError(f"Failed to load GLM-TTS speech tokenizer from {speech_tokenizer_path}") from exc
+
+    if campplus_session is None:
+        cp = campplus_path or os.path.join(model_root, "frontend", "campplus.onnx")
+        if not os.path.isfile(cp):
+            raise RuntimeError(
+                f"GLM-TTS campplus.onnx not found at {cp}.  "
+                f"Voice cloning requires the CampPlus speaker embedding model."
+            )
+        try:
+            import onnxruntime
+
+            cuda_avail = "CUDAExecutionProvider" in onnxruntime.get_available_providers()
+            cuda_ok = model_device.type == "cuda" and cuda_avail
+            providers = (["CUDAExecutionProvider"] if cuda_ok else []) + ["CPUExecutionProvider"]
+            campplus_session = onnxruntime.InferenceSession(cp, providers=providers)
+            logger.info("Loaded GLM-TTS campplus ONNX from %s", cp)
+        except Exception as exc:
+            raise RuntimeError(f"Failed to load GLM-TTS campplus ONNX from {cp}") from exc
+
+    return speech_tokenizer, campplus_session
+
+
+def extract_prompt_speech_token(
+    ref_audio_wav: torch.Tensor,
+    ref_audio_sr: int,
+    speech_tokenizer: tuple,
+) -> list[int] | None:
+    """Extract prompt speech tokens from reference audio using WhisperVQ."""
+
+    model, feature_extractor = speech_tokenizer
+    device = model.device
+
+    audio = ref_audio_wav.float().to(device)
+    if ref_audio_sr != 16000:
+        resampler = T.Resample(orig_freq=ref_audio_sr, new_freq=16000).to(device)
+        audio = resampler(audio)
+
+    if audio.ndim > 1:
+        audio = audio[0]
+
+    audio_np = audio.cpu().numpy()
+    pooling_kernel_size = getattr(model.config, "pooling_kernel_size", 1)
+    stride = model.conv1.stride[0] * model.conv2.stride[0] * pooling_kernel_size * feature_extractor.hop_length
+
+    all_tokens: list[int] = []
+    time_step = 0
+    while time_step * 16000 < audio_np.shape[0]:
+        segment = audio_np[time_step * 16000 : (time_step + 30) * 16000]
+        features = feature_extractor(
+            [segment],
+            sampling_rate=16000,
+            return_attention_mask=True,
+            return_tensors="pt",
+            padding="longest",
+            pad_to_multiple_of=stride,
+        ).to(device)
+        with torch.no_grad():
+            outputs = model(**features)
+        tokens = outputs.quantized_token_ids
+        attn = features.attention_mask[:, :: model.conv1.stride[0] * model.conv2.stride[0]][:, ::pooling_kernel_size]
+        all_tokens.extend(tokens[0][attn[0].bool()].tolist())
+        time_step += 30
+
+    return all_tokens if all_tokens else None
+
+
+def extract_spk_embedding(
+    ref_audio_wav: torch.Tensor,
+    ref_audio_sr: int,
+    campplus_session: object,
+) -> list[float] | None:
+    """Extract speaker embedding from reference audio using CampPlus ONNX."""
+
+    audio = ref_audio_wav.float().cpu()
+    if ref_audio_sr != 16000:
+        resampler = T.Resample(orig_freq=ref_audio_sr, new_freq=16000)
+        audio = resampler(audio)
+    if audio.ndim > 1:
+        audio = audio[0]
+
+    feat = kaldi.fbank(audio.unsqueeze(0), num_mel_bins=80, dither=0, sample_frequency=16000)
+    feat = feat - feat.mean(dim=0, keepdim=True)
+
+    input_name = campplus_session.get_inputs()[0].name
+    embedding = campplus_session.run(None, {input_name: feat.unsqueeze(dim=0).cpu().numpy()})[0].flatten().tolist()
+    return embedding
+
+
+def extract_prompt_feat(
+    ref_audio_wav: torch.Tensor,
+    ref_audio_sr: int,
+    model_device: torch.device,
+) -> torch.Tensor | None:
+    """Extract mel features from reference audio for DiT conditioning."""
+    audio = ref_audio_wav.float().to(model_device)
+    if ref_audio_sr != 24000:
+        resampler = T.Resample(orig_freq=ref_audio_sr, new_freq=24000).to(model_device)
+        audio = resampler(audio)
+    if audio.ndim > 1:
+        audio = audio[0]
+
+    with torch.no_grad():
+        feat = _extract_glm_tts_mel_feature(
+            audio,
+            sample_rate=24000,
+            n_fft=1920,
+            win_length=1920,
+            hop_length=480,
+            n_mels=80,
+            f_min=0,
+            f_max=8000,
+        )
+    return feat  # [T_mel, mel_dim]
diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/__init__.py b/vllm_omni/model_executor/models/higgs_audio_v2/__init__.py
new file mode 100644
index 000000000..1a7c23eba
--- /dev/null
+++ b/vllm_omni/model_executor/models/higgs_audio_v2/__init__.py
@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""vllm-omni integration for boson-ai's higgs-audio v2 (two-stage TTS)."""
+
+from vllm_omni.model_executor.models.higgs_audio_v2.configuration_higgs_audio_v2 import (
+    HiggsAudioV2Config,
+)
+
+__all__ = ["HiggsAudioV2Config"]
diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py b/vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py
new file mode 100644
index 000000000..dc5eb12b4
--- /dev/null
+++ b/vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py
@@ -0,0 +1,169 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Configuration class for higgs-audio v2 in vllm-omni.
+
+Mirrors `bosonai/higgs-audio-v2-generation-3B-base/config.json` and exposes the
+Stage-1 codec decoder reference (sample rate, codebook layout, stream-special
+token ids) so the rest of the integration can read a single typed object.
+
+The text-LM backbone constants match `Llama-3.2-3B` (hidden=3072, 28 layers,
+GQA 24Q/8KV, head_dim=128, vocab=128256, max_position_embeddings=2048, RoPE
+``llama3`` scaling with factor=32). The audio side adds 8 codebooks of size
+1026 (real codes ``[0..1023]``, stream specials ``audio_stream_bos_id=1024``
+and ``audio_stream_eos_id=1025``) plus the DualFFN audio expert.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from transformers import PretrainedConfig
+
+__all__ = ["HiggsAudioV2Config"]
+
+
+class HiggsAudioV2Config(PretrainedConfig):
+    """Typed wrapper around the upstream `higgs_audio_v2` config.
+
+    The field set here is a strict subset of the upstream
+    ``transformers.models.higgs_audio_v2.HiggsAudioV2Config``: only the
+    knobs that vllm-omni reads directly are surfaced. Anything extra coming
+    from the HF config dict is preserved via ``PretrainedConfig`` so
+    ``AutoConfig.from_pretrained(...)`` round-trips correctly.
+    """
+
+    model_type: str = "higgs_audio_v2"
+    keys_to_ignore_at_inference = ("past_key_values",)
+
+    # Architecture identifier persisted in config.json and used by both
+    # vllm_omni.config.pipeline_registry and vllm_omni.model_executor.models.registry.
+    DEFAULT_ARCHITECTURES = ("HiggsAudioV2ForConditionalGeneration",)
+
+    def __init__(
+        self,
+        # --- Text LM (Llama-3.2-3B) ---
+        vocab_size: int = 128256,
+        hidden_size: int = 3072,
+        intermediate_size: int = 8192,
+        num_hidden_layers: int = 28,
+        num_attention_heads: int = 24,
+        num_key_value_heads: int = 8,
+        head_dim: int = 128,
+        hidden_act: str = "silu",
+        max_position_embeddings: int = 2048,
+        rms_norm_eps: float = 1e-5,
+        rope_parameters: dict[str, Any] | None = None,
+        initializer_range: float = 0.02,
+        attention_bias: bool = False,
+        attention_dropout: float = 0.0,
+        mlp_bias: bool = False,
+        tie_word_embeddings: bool = False,
+        pretraining_tp: int = 1,
+        # --- Audio side ---
+        num_codebooks: int = 8,
+        codebook_size: int = 1026,
+        audio_token_id: int = 128016,
+        audio_bos_token_id: int = 128013,
+        audio_eos_token_id: int = 128012,
+        audio_delay_token_id: int = 128014,
+        audio_stream_bos_id: int = 1024,
+        audio_stream_eos_id: int = 1025,
+        sample_rate: int = 24000,
+        frame_rate: int = 25,
+        # --- Standard LM specials ---
+        bos_token_id: int = 1,
+        eos_token_id: int = 128009,
+        pad_token_id: int = 128001,
+        # --- DualFFN ---
+        # Upstream `HiggsAudioV2DecoderLayer` always carries a parallel `audio_mlp`
+        # expert; this flag is kept for forward-compat in case a future config
+        # disables it (no-op for the boson-ai checkpoint).
+        use_audio_dual_ffn: bool = True,
+        # --- Stage-1 tokenizer / codec decoder ---
+        # Path to the audio-tokenizer subdir (relative to the model directory).
+        # Empty string means the model_dir itself contains config.json +
+        # model.safetensors (the layout used by the standalone
+        # ``bosonai/higgs-audio-v2-tokenizer`` HF repo). Non-empty means the
+        # tokenizer files live in ``<model_dir>/<subdir>/`` (the layout used by
+        # the 3B Stage-0 checkpoint which bundles its tokenizer under
+        # ``audio_tokenizer/``).
+        audio_tokenizer_subdir: str = "",
+        # HF repo id (or local path) of a STANDALONE audio tokenizer. When set,
+        # Stage-1 resolves this via ``huggingface_hub.snapshot_download`` and
+        # loads the codec from there, ignoring ``audio_tokenizer_subdir``. This
+        # matches the boson-ai release layout where the codec lives in a
+        # separate repo (``bosonai/higgs-audio-v2-tokenizer``).
+        audio_tokenizer_id: str | None = "bosonai/higgs-audio-v2-tokenizer",
+        # --- Misc ---
+        use_cache: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        if rope_parameters is None:
+            rope_parameters = {
+                "factor": 32.0,
+                "high_freq_factor": 0.5,
+                "low_freq_factor": 0.125,
+                "original_max_position_embeddings": 1024,
+                "rope_theta": 500000.0,
+                "rope_type": "llama3",
+            }
+
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_parameters = dict(rope_parameters)
+        self.initializer_range = initializer_range
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.mlp_bias = mlp_bias
+        self.pretraining_tp = pretraining_tp
+
+        self.num_codebooks = num_codebooks
+        self.codebook_size = codebook_size
+        self.audio_token_id = audio_token_id
+        self.audio_bos_token_id = audio_bos_token_id
+        self.audio_eos_token_id = audio_eos_token_id
+        self.audio_delay_token_id = audio_delay_token_id
+        self.audio_stream_bos_id = audio_stream_bos_id
+        self.audio_stream_eos_id = audio_stream_eos_id
+        self.sample_rate = sample_rate
+        self.frame_rate = frame_rate
+
+        self.use_audio_dual_ffn = use_audio_dual_ffn
+        self.audio_tokenizer_subdir = audio_tokenizer_subdir
+        self.audio_tokenizer_id = audio_tokenizer_id
+        self.use_cache = use_cache
+
+    # Codes [0..codebook_size-1] include 2 stream specials at indices
+    # `audio_stream_bos_id` and `audio_stream_eos_id`. Real codes ready for the
+    # codec decoder are the strict prefix [0..audio_stream_bos_id).
+    @property
+    def num_real_codes(self) -> int:
+        return self.audio_stream_bos_id  # i.e. 1024 for the boson-ai checkpoint
+
+    # The full vocab size of each per-codebook output head.
+    @property
+    def codebook_output_size(self) -> int:
+        return self.codebook_size  # i.e. 1026 for the boson-ai checkpoint
+
+    # Canonical MusicGen-style delay pattern used by
+    # `HiggsAudioV2DelayPatternLogitsProcessor`: codebook k starts emitting
+    # real codes only after k frames. Kept as a function so subclasses can
+    # override if upstream ever publishes a different convention.
+    def default_delay_pattern(self) -> list[int]:
+        return list(range(self.num_codebooks))
diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py
new file mode 100644
index 000000000..74cdbb2db
--- /dev/null
+++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py
@@ -0,0 +1,471 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""HiggsAudio codec decoder kernel for higgs-audio v2.
+
+This module hosts the parameter-side building blocks for the higgs-audio-v2
+audio tokenizer's decoder path:
+
+  audio_codes [B, 8, T]
+    -> RVQ codebook lookup + project_out -> sum -> [B, hidden_size, T]
+    -> fc2 Linear(hidden_size, 256) -> [B, 256, T]
+    -> DAC acoustic decoder (conv-transpose upsampling) -> [B, 1, T*960]
+    -> 24 kHz waveform (25 fps x 960 samples/frame)
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+__all__ = [
+    "HiggsAudioVQLayer",
+    "HiggsAudioRVQ",
+    "adjust_conv_transpose_output_padding",
+    "build_higgs_audio_acoustic_decoder",
+    "BosonDacDecoder",
+    "build_boson_dac_decoder",
+    "load_higgs_audio_codec",
+    "_remap_boson_model_pth_state_dict",  # exported for unit-testing the mapper
+]
+
+
+class HiggsAudioVQLayer(nn.Module):
+    """Single VQ layer: codebook lookup + project_out."""
+
+    def __init__(self, codebook_size: int = 1024, codebook_dim: int = 64, hidden_size: int = 1024):
+        super().__init__()
+        self.codebook = nn.Embedding(codebook_size, codebook_dim)
+        self.project_out = nn.Linear(codebook_dim, hidden_size)
+
+    def decode(self, indices: torch.Tensor) -> torch.Tensor:
+        """indices: [B, T] -> [B, hidden_size, T]."""
+        quantized = self.codebook(indices)
+        quantized = self.project_out(quantized)
+        return quantized.permute(0, 2, 1)
+
+
+class HiggsAudioRVQ(nn.Module):
+    """Residual Vector Quantizer with ``num_quantizers`` codebook layers."""
+
+    def __init__(
+        self,
+        num_quantizers: int = 8,
+        codebook_size: int = 1024,
+        codebook_dim: int = 64,
+        hidden_size: int = 1024,
+    ):
+        super().__init__()
+        self.quantizers = nn.ModuleList(
+            [HiggsAudioVQLayer(codebook_size, codebook_dim, hidden_size) for _ in range(num_quantizers)]
+        )
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """codes: [num_quantizers, B, T] -> [B, hidden_size, T]."""
+        result = torch.zeros(
+            codes.shape[1],
+            self.quantizers[0].project_out.out_features,
+            codes.shape[2],
+            device=codes.device,
+            dtype=torch.float32,
+        )
+        for i, quantizer in enumerate(self.quantizers):
+            result = result + quantizer.decode(codes[i])
+        return result
+
+
+@torch.jit.script
+def _snake(x: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor:
+    shape = x.shape
+    x = x.reshape(shape[0], shape[1], -1)
+    x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+    return x.reshape(shape)
+
+
+class _Snake1d(nn.Module):
+    """Snake activation (per-channel learnable alpha). Mirrors the boson-ai DAC
+    Snake1d implementation byte-for-byte so the boson ``model.pth`` state-dict
+    loads directly into ``self.alpha``.
+    """
+
+    def __init__(self, channels: int):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return _snake(x, self.alpha)
+
+
+def _wn_conv1d(*args: Any, **kwargs: Any) -> nn.Module:
+    from torch.nn.utils import weight_norm
+
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def _wn_conv_transpose1d(*args: Any, **kwargs: Any) -> nn.Module:
+    from torch.nn.utils import weight_norm
+
+    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+class _BosonResidualUnit(nn.Module):
+    def __init__(self, dim: int, dilation: int = 1):
+        super().__init__()
+        pad = ((7 - 1) * dilation) // 2
+        self.block = nn.Sequential(
+            _Snake1d(dim),
+            _wn_conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
+            _Snake1d(dim),
+            _wn_conv1d(dim, dim, kernel_size=1),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = self.block(x)
+        pad = (x.shape[-1] - y.shape[-1]) // 2
+        if pad > 0:
+            x = x[..., pad:-pad]
+        return x + y
+
+
+class _BosonDecoderBlock(nn.Module):
+    def __init__(self, input_dim: int, output_dim: int, stride: int):
+        super().__init__()
+        import math
+
+        self.block = nn.Sequential(
+            _Snake1d(input_dim),
+            _wn_conv_transpose1d(
+                input_dim,
+                output_dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+                output_padding=stride % 2,
+            ),
+            _BosonResidualUnit(output_dim, dilation=1),
+            _BosonResidualUnit(output_dim, dilation=3),
+            _BosonResidualUnit(output_dim, dilation=9),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.block(x)
+
+
+class BosonDacDecoder(nn.Module):
+    """Vendored copy of ``boson_multimodal``'s standalone DAC decoder, with the
+    exact ``decoder_2.model.<i>.{...}`` key layout the boson-ai
+    ``bosonai/higgs-audio-v2-tokenizer`` checkpoint expects.
+
+    See ``boson_multimodal/audio_processing/descriptaudiocodec/dac/model/dac.py``
+    (``class Decoder``) in the upstream package for the reference implementation
+    we mirror.
+
+    Architecture for the v2 24 kHz tokenizer:
+      input_channel=256, channels=1024, rates=[8, 5, 4, 2, 3], d_out=1
+    Cumulative upsample factor is ``prod(rates) = 960`` (25 fps -> 24 kHz).
+    """
+
+    def __init__(
+        self,
+        input_channel: int = 256,
+        channels: int = 1024,
+        rates: tuple[int, ...] = (8, 5, 4, 2, 3),
+        d_out: int = 1,
+    ):
+        super().__init__()
+        layers: list[nn.Module] = [_wn_conv1d(input_channel, channels, kernel_size=7, padding=3)]
+        last_dim = channels
+        for i, stride in enumerate(rates):
+            input_dim = channels // 2**i
+            output_dim = channels // 2 ** (i + 1)
+            layers.append(_BosonDecoderBlock(input_dim, output_dim, stride))
+            last_dim = output_dim
+        layers.append(_Snake1d(last_dim))
+        layers.append(_wn_conv1d(last_dim, d_out, kernel_size=7, padding=3))
+        self.model = nn.Sequential(*layers)
+        # Cache hop length (= prod(rates)) so callers don't need to recompute it.
+        hop = 1
+        for s in rates:
+            hop *= int(s)
+        self._hop_length = hop
+
+    @property
+    def hop_length(self) -> int:
+        return self._hop_length
+
+    # Expose conv1 with a ``.in_channels`` attribute so the existing fc2
+    # fallback shape inference (acoustic_decoder.conv1.in_channels) works
+    # unchanged for the boson decoder too.
+    @property
+    def conv1(self) -> nn.Module:
+        return self.model[0]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)
+
+
+def build_boson_dac_decoder(device: torch.device) -> BosonDacDecoder:
+    """Construct the vendored boson DAC decoder on ``device``.
+
+    The architectural constants (input_channel=256, channels=1024,
+    rates=[8,5,4,2,3], d_out=1) are baked in because they're invariant for the
+    v2 24 kHz tokenizer ``bosonai/higgs-audio-v2-tokenizer`` checkpoint we
+    target. Other channel counts are not supported until we encounter a
+    different boson checkpoint in the wild.
+    """
+    return BosonDacDecoder().to(device)
+
+
+def adjust_conv_transpose_output_padding(decoder: nn.Module) -> None:
+    """Set ConvTranspose1d output_padding = stride % 2 (HiggsAudioV2 modification).
+
+    The vanilla DAC decoder ships with the default output_padding (0); the
+    boson-ai checkpoint expects ``stride % 2`` instead. This is a no-op for
+    even strides and adds a single sample for odd strides.
+    """
+    for module in decoder.modules():
+        if isinstance(module, nn.ConvTranspose1d):
+            stride = module.stride[0] if isinstance(module.stride, tuple) else module.stride
+            module.output_padding = (stride % 2,)
+
+
+def build_higgs_audio_acoustic_decoder(
+    tokenizer_config: dict[str, Any],
+    device: torch.device,
+) -> nn.Module:
+    """Build the DAC acoustic decoder used by the HiggsAudioV2 tokenizer.
+
+    Returns the decoder sub-module of ``transformers.DacModel`` with the
+    HiggsAudioV2 output-padding fix already applied. The tanh activation is
+    replaced with ``Identity`` so the network matches the boson-ai checkpoint.
+    Weights are NOT loaded here; the caller is responsible for copying them in.
+    """
+    from transformers import DacConfig, DacModel
+
+    dac_cfg = DacConfig(**tokenizer_config["acoustic_model_config"])
+    dac_model = DacModel(dac_cfg)
+    decoder = dac_model.decoder.to(device)
+    adjust_conv_transpose_output_padding(decoder)
+    if hasattr(decoder, "tanh"):
+        decoder.tanh = nn.Identity()
+    return decoder
+
+
+def _load_higgs_audio_state_dict(audio_tokenizer_dir: str, device: torch.device) -> dict[str, torch.Tensor]:
+    """Load the codec state dict from either layout used in the wild.
+
+    Tries layouts in order:
+    1. ``<dir>/model.safetensors`` (OmniVoice-bundled layout used by
+       ``k2-fsa/OmniVoice/audio_tokenizer/``).
+    2. ``<dir>/model.pth`` (boson-ai standalone ``bosonai/higgs-audio-v2-tokenizer``
+       layout). The state-dict keys differ structurally
+       (``quantizer.vq.layers.<i>._codebook.embed`` etc.) and are remapped to
+       the OmniVoice-style names this kernel expects via
+       :func:`_remap_boson_model_pth_state_dict`.
+    """
+    safetensors_path = os.path.join(audio_tokenizer_dir, "model.safetensors")
+    pth_path = os.path.join(audio_tokenizer_dir, "model.pth")
+    # Preference order: try model.pth first when present (boson-ai standalone
+    # ships the codec there; its model.safetensors is a different artefact).
+    # Fall back to model.safetensors (the OmniVoice-bundled layout).
+    sd: dict[str, torch.Tensor]
+    if os.path.exists(pth_path):
+        sd = torch.load(pth_path, map_location=device, weights_only=False)
+    elif os.path.exists(safetensors_path):
+        from safetensors.torch import load_file
+
+        sd = load_file(safetensors_path, device=str(device))
+    else:
+        raise FileNotFoundError(f"Audio tokenizer weights not found at {pth_path} or {safetensors_path}")
+    # Run the boson-layout remap if needed; for the OmniVoice-bundled layout
+    # it is a no-op (no ``quantizer.vq.layers.*`` keys to rewrite).
+    if any(k.startswith("quantizer.vq.layers.") for k in sd):
+        sd = _remap_boson_model_pth_state_dict(sd)
+    return sd
+
+
+def _remap_boson_model_pth_state_dict(sd: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """Translate boson-ai's standalone ``model.pth`` keys into OmniVoice-style names
+    that the shared kernel's RVQ + fc2 + DAC sites consume.
+
+    Best-effort: only the RVQ-side keys map cleanly. The boson-ai decoder
+    uses ``decoder_2.model.<i>.weight_g/weight_v`` (DAC with weight-norm +
+    Snake activations) rather than the OmniVoice DAC layout; that side
+    requires either vendoring the upstream decoder module or rewriting the
+    DAC builder to consume weight-normed tensors. This function returns the
+    mapped RVQ keys plus any acoustic_decoder.* / fc2.* / fc.* tensors that
+    happen to share the OmniVoice names, and leaves the decoder-2 keys
+    untouched (the caller will see them as MISSING when copying into the
+    OmniVoice DAC decoder).
+
+    Mapping (RVQ side only):
+        quantizer.vq.layers.<i>._codebook.embed -> quantizer.quantizers.<i>.codebook.embed
+        quantizer.vq.layers.<i>.project_out.weight -> quantizer.quantizers.<i>.project_out.weight
+        quantizer.vq.layers.<i>.project_out.bias   -> quantizer.quantizers.<i>.project_out.bias
+    """
+    if not isinstance(sd, dict):
+        raise TypeError(f"expected a state dict, got {type(sd)!r}")
+    remapped: dict[str, torch.Tensor] = {}
+    for key, tensor in sd.items():
+        if not isinstance(tensor, torch.Tensor):
+            continue
+        # Quantizer rewrite: vq.layers.<i>._codebook.embed -> quantizers.<i>.codebook.embed
+        if key.startswith("quantizer.vq.layers."):
+            parts = key.split(".")
+            # parts[3] is the layer index, parts[4]+ is the tail
+            if len(parts) >= 5:
+                idx = parts[3]
+                tail = ".".join(parts[4:])
+                # The actual codebook centroids the decode path uses live at
+                # ``_codebook.embed`` (a [codebook_size, codebook_dim] buffer).
+                # ``_codebook.embed_avg`` is the un-normalised EMA accumulator
+                # used during training (~cluster_size x bigger); MATCHING IT
+                # HERE produces a 4-5x amplified codebook and unintelligible
+                # PCM downstream. Use an EXACT-match check, not startswith.
+                if tail == "_codebook.embed":
+                    new_key = f"quantizer.quantizers.{idx}.codebook.embed"
+                    remapped[new_key] = tensor
+                    continue
+                if tail.startswith("project_out."):
+                    new_key = f"quantizer.quantizers.{idx}.{tail}"
+                    remapped[new_key] = tensor
+                    continue
+                # project_in / _codebook.cluster_size / _codebook.embed_avg / inited
+                # are encoder-side or training-only state that the decode path
+                # does not need; drop them.
+                continue
+        # boson decoder + fc_post2 -> our vendored kernel names.
+        if key.startswith("decoder_2."):
+            # decoder_2.model.<i>.{...} -> acoustic_decoder.model.<i>.{...}
+            new_key = "acoustic_decoder." + key[len("decoder_2.") :]
+            remapped[new_key] = tensor
+            continue
+        if key.startswith("fc_post2."):
+            # fc_post2.{weight,bias} -> fc2.{weight,bias}
+            new_key = "fc2." + key[len("fc_post2.") :]
+            remapped[new_key] = tensor
+            continue
+        # Anything else (acoustic_decoder.*, fc2.*, fc.*, fc_post*.*, decoder_2.*,
+        # decoder_semantic.*, semantic_model.*) passes through unchanged so
+        # the caller's lookup is unambiguous.
+        remapped[key] = tensor
+    return remapped
+
+
+def load_higgs_audio_codec(
+    audio_tokenizer_dir: str,
+    device: torch.device,
+) -> tuple[HiggsAudioRVQ, nn.Linear, nn.Module, dict[str, Any]]:
+    """Load the HiggsAudioV2 RVQ + fc2 + DAC decoder from a checkpoint folder.
+
+    Accepts both layouts: ``<dir>/model.safetensors`` (OmniVoice-bundled) and
+    ``<dir>/model.pth`` (boson-ai standalone). The standalone path remaps
+    quantizer keys before consumption; structural decoder differences
+    (boson uses Snake + weight-norm) still leave some DAC parameters missing
+    when loading the standalone path -- those entries log warnings but the
+    RVQ side completes successfully.
+
+    Args:
+        audio_tokenizer_dir: Path to a directory containing ``config.json``
+            and EITHER ``model.safetensors`` (OmniVoice layout) OR
+            ``model.pth`` (boson-ai standalone layout).
+        device: Device to place the loaded modules and state dict on.
+
+    Returns:
+        (quantizer, fc2, acoustic_decoder, tokenizer_config)
+        - quantizer: ``HiggsAudioRVQ`` with ``num_quantizers`` discovered from
+            the state dict (defaults to 8 for boson-ai checkpoints).
+        - fc2: ``nn.Linear`` projecting RVQ output (1024) into the DAC's
+            hidden dimension (typically 256). May be uninitialized when loaded
+            from the boson-ai standalone layout (no ``fc2.*`` keys present).
+        - acoustic_decoder: DAC decoder with HiggsAudioV2 output-padding fix
+            and tanh-replaced-by-Identity. Fully initialized for the OmniVoice
+            layout; partially initialized for boson-ai standalone (missing
+            ``decoder_2.*`` -> ``acoustic_decoder.*`` mapping is not yet
+            implemented; full upstream-decoder vendoring is the next step).
+        - tokenizer_config: The loaded ``config.json`` dict; useful for
+            callers that need ``sample_rate`` or other tokenizer metadata.
+    """
+    config_path = os.path.join(audio_tokenizer_dir, "config.json")
+
+    with open(config_path) as f:
+        tokenizer_config: dict[str, Any] = json.load(f)
+
+    state_dict = _load_higgs_audio_state_dict(audio_tokenizer_dir, device)
+
+    codebook_dim = tokenizer_config.get("codebook_dim", 64)
+    codebook_size = tokenizer_config.get("codebook_size", 1024)
+    # Discover hidden_size and num_quantizers from the (possibly remapped) state dict.
+    if "quantizer.quantizers.0.project_out.weight" not in state_dict:
+        raise KeyError(
+            "Codec state dict is missing 'quantizer.quantizers.0.project_out.weight'. "
+            "If you loaded a boson-ai standalone tokenizer, ensure the model.pth "
+            "remap fired (see _remap_boson_model_pth_state_dict)."
+        )
+    hidden_size = state_dict["quantizer.quantizers.0.project_out.weight"].shape[0]
+    num_quantizers = sum(
+        1 for k in state_dict if k.startswith("quantizer.quantizers.") and k.endswith(".codebook.embed")
+    )
+
+    quantizer = HiggsAudioRVQ(
+        num_quantizers=num_quantizers,
+        codebook_size=codebook_size,
+        codebook_dim=codebook_dim,
+        hidden_size=hidden_size,
+    ).to(device)
+    for i in range(num_quantizers):
+        prefix = f"quantizer.quantizers.{i}"
+        embed_key = f"{prefix}.codebook.embed"
+        if embed_key in state_dict:
+            quantizer.quantizers[i].codebook.weight.data.copy_(state_dict[embed_key])
+        proj_out_w = f"{prefix}.project_out.weight"
+        proj_out_b = f"{prefix}.project_out.bias"
+        if proj_out_w in state_dict:
+            quantizer.quantizers[i].project_out.weight.data.copy_(state_dict[proj_out_w])
+        if proj_out_b in state_dict:
+            quantizer.quantizers[i].project_out.bias.data.copy_(state_dict[proj_out_b])
+
+    # Decide which acoustic decoder we need. A boson-layout state dict (after
+    # _remap_boson_model_pth_state_dict) carries ``acoustic_decoder.model.0.weight_v``
+    # (the WNConv first conv) which the OmniVoice DacModel doesn't have; the
+    # OmniVoice layout carries ``acoustic_decoder.block.0.*`` instead.
+    is_boson_layout = "acoustic_decoder.model.0.weight_v" in state_dict
+    if is_boson_layout:
+        acoustic_decoder = build_boson_dac_decoder(device)
+        boson_sd = {
+            k[len("acoustic_decoder.") :]: v for k, v in state_dict.items() if k.startswith("acoustic_decoder.")
+        }
+        load_report = acoustic_decoder.load_state_dict(boson_sd, strict=False)
+        if load_report.missing_keys:
+            raise RuntimeError(
+                "Boson DAC decoder is missing keys after load: "
+                f"{load_report.missing_keys[:5]}... ({len(load_report.missing_keys)} total)"
+            )
+        # ``unexpected_keys`` is informational only.
+    else:
+        acoustic_decoder = build_higgs_audio_acoustic_decoder(tokenizer_config, device)
+        for name, param in acoustic_decoder.named_parameters():
+            higgs_name = f"acoustic_decoder.{name}"
+            if higgs_name in state_dict:
+                param.data.copy_(state_dict[higgs_name])
+    acoustic_decoder.eval()
+
+    if "fc2.weight" in state_dict and "fc2.bias" in state_dict:
+        fc2_w = state_dict["fc2.weight"]
+        fc2_b = state_dict["fc2.bias"]
+        fc2 = nn.Linear(fc2_w.shape[1], fc2_w.shape[0]).to(device)
+        fc2.weight.data.copy_(fc2_w)
+        fc2.bias.data.copy_(fc2_b)
+    else:
+        # Last-resort fallback: shape fc2 from the decoder's first conv input
+        # channels so the matmul produces compatible shapes. Weights stay
+        # uninitialised — PCM quality will be noise on this path.
+        first_conv = getattr(acoustic_decoder, "conv1", None)
+        if first_conv is None or not hasattr(first_conv, "in_channels"):
+            raise RuntimeError("acoustic_decoder has no conv1.in_channels; cannot size fc2 fallback")
+        fc2 = nn.Linear(hidden_size, int(first_conv.in_channels)).to(device)
+
+    return quantizer, fc2, acoustic_decoder, tokenizer_config
diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py
new file mode 100644
index 000000000..6a14707ee
--- /dev/null
+++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py
@@ -0,0 +1,413 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Stage 1 (codec decoder) for higgs-audio v2.
+
+Two surfaces exposed:
+
+1. **Direct decode API** (kept for offline / unit tests):
+   ``decode_codes(audio_codes: [B, num_codebooks=8, T])`` -> ``[B, 1, T*hop]`` PCM.
+   ``forward_chunk(audio_codes, *, left_context_size, hop_length=960)``
+   trims overlap from streamed PCM. ``forward(audio_codes)`` retains the same
+   single-tensor signature for backward compatibility with code that imports
+   the class directly.
+
+2. **vLLM stage runtime API** (used by the engine):
+   ``__init__(*, vllm_config, prefix)``, ``embed_input_ids``,
+   ``compute_logits=None``, plus a runtime ``forward(input_ids=..., positions=...,
+   runtime_additional_information=...)`` that takes flat codebook-major
+   ``input_ids`` (``[Q * num_frames]``) per request and returns an
+   :class:`vllm_omni.model_executor.models.output_templates.OmniOutput`.
+
+The kernel is the HiggsAudio codec helper in
+``vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py``; both
+surfaces share the same RVQ + DAC weights.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import torch
+import torch.nn as nn
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+
+from vllm_omni.model_executor.models.higgs_audio_v2.configuration_higgs_audio_v2 import (
+    HiggsAudioV2Config,
+)
+from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import (
+    HiggsAudioRVQ,
+    load_higgs_audio_codec,
+)
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+__all__ = [
+    "HiggsAudioV2Code2Wav",
+    "HiggsAudioV2Code2WavForConditionalGeneration",
+]
+
+logger = init_logger(__name__)
+
+
+class HiggsAudioV2Code2Wav(nn.Module):
+    """Stage-1 codec decoder for higgs-audio v2.
+
+    Constructor accepts either the vLLM-runtime signature (``*, vllm_config,
+    prefix``) or the direct config-object form (``HiggsAudioV2Code2Wav(config)``)
+    so both engine-side and unit-test callers work. When the runtime form is
+    used we read the higgs_audio_v2 config out of ``vllm_config.model_config.hf_config``
+    and the model_path off ``vllm_config.model_config.model`` so the codec can
+    be loaded directly from disk.
+    """
+
+    input_modalities = "audio"
+
+    def __init__(
+        self,
+        config: HiggsAudioV2Config | None = None,
+        *,
+        vllm_config: VllmConfig | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+
+        # Resolve config from either positional (unit-test) or kwarg (engine) form.
+        if vllm_config is not None:
+            hf_config = vllm_config.model_config.hf_config
+            if isinstance(hf_config, HiggsAudioV2Config):
+                self.config = hf_config
+            else:
+                self.config = HiggsAudioV2Config(**hf_config.to_dict())
+            self._model_path: str | None = vllm_config.model_config.model
+            self.vllm_config: VllmConfig | None = vllm_config
+        else:
+            if config is None:
+                raise TypeError(
+                    "HiggsAudioV2Code2Wav: provide either positional `config` "
+                    "(HiggsAudioV2Config) or keyword `vllm_config` (VllmConfig)."
+                )
+            self.config = config
+            self._model_path = None
+            self.vllm_config = None
+
+        self.sample_rate: int = int(self.config.sample_rate)
+        self.num_codebooks: int = int(self.config.num_codebooks)
+        self.num_real_codes: int = int(self.config.num_real_codes)
+        # Each codec frame upsamples to 960 24 kHz samples (= 25 fps * 960 = 24000).
+        self.hop_length: int = 960
+
+        # Engine-runner hooks (Stage 1 has no token sampling).
+        self.have_multimodal_outputs = True
+        self.has_preprocess = False
+        self.has_postprocess = False
+        self.enable_update_additional_information = True
+        self.requires_raw_input_tokens = True
+
+        # Populated by load_weights().
+        self.quantizer: HiggsAudioRVQ | None = None
+        self.fc2: nn.Linear | None = None
+        self.acoustic_decoder: nn.Module | None = None
+        self._loaded: bool = False
+
+        # When constructed via the engine path, eagerly load codec weights from
+        # the model directory. The engine's subsequent ``load_weights(weights)``
+        # call (with the Stage-0 safetensors iterator) is then a cheap no-op.
+        if self._model_path is not None:
+            try:
+                self.load_codec_from_disk(self._model_path)
+            except FileNotFoundError as exc:
+                logger.warning("HiggsAudioV2Code2Wav: eager codec load deferred (%s)", exc)
+
+    # ------------------------------------------------------------- engine hooks
+    def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
+        """Stage 1 ignores embeddings; vLLM's runner still needs a stable shape."""
+        if input_ids.numel() == 0:
+            return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32)
+        return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32)
+
+    def compute_logits(self, hidden_states: Any, sampling_metadata: Any = None) -> None:
+        return None
+
+    # ------------------------------------------------------------------ load
+    def load_weights(self, weights_or_model_dir, device: torch.device | None = None):
+        """Dual-signature load:
+
+        - ``load_weights(weights_iterator)`` -- the vLLM engine path. The iterator
+          yields ``(name, tensor)`` pairs from the Stage-0 safetensors which DO NOT
+          contain codec weights, so we ignore them and eagerly load the codec from
+          ``self._model_path``. Returns the set of all our params as "initialized"
+          so the engine's loader passes its sanity check.
+        - ``load_weights(model_dir: str)`` -- the legacy file-based call used by
+          offline unit tests and ``end2end.py --mode stage1_only``. Returns ``None``.
+        """
+        if not isinstance(weights_or_model_dir, (str, bytes, os.PathLike)):
+            # Engine call: consume the iterator (we don't need it) and eagerly
+            # load the codec from disk.
+            try:
+                for _ in weights_or_model_dir:
+                    pass
+            except TypeError:
+                pass
+            if self._model_path is None:
+                raise RuntimeError(
+                    "HiggsAudioV2Code2Wav.load_weights called via the engine path but "
+                    "no model directory is known. Construct with vllm_config so the "
+                    "constructor caches vllm_config.model_config.model."
+                )
+            self.load_codec_from_disk(self._model_path, device=device)
+            return {name for name, _ in self.named_parameters()}
+
+        # Legacy file-based call.
+        self.load_codec_from_disk(weights_or_model_dir, device=device)
+        return None
+
+    def load_codec_from_disk(self, model_dir: str, device: torch.device | None = None) -> None:
+        """Load codec weights from disk (RVQ + fc2 + DAC decoder).
+
+        Resolution order for the codec directory:
+        1. If :attr:`HiggsAudioV2Config.audio_tokenizer_id` is set, treat it as
+           either a local path (if it exists on disk) or an HF repo id that we
+           resolve via ``huggingface_hub.snapshot_download``. This is the
+           default for the boson-ai release (codec in a separate repo).
+        2. Otherwise fall back to ``<model_dir>/<audio_tokenizer_subdir>``.
+        """
+        if device is None:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        audio_tokenizer_path = self._resolve_audio_tokenizer_path(model_dir)
+        quantizer, fc2, acoustic_decoder, _tokenizer_config = load_higgs_audio_codec(audio_tokenizer_path, device)
+        if len(quantizer.quantizers) != self.num_codebooks:
+            raise ValueError(
+                f"checkpoint has {len(quantizer.quantizers)} quantizers but config.num_codebooks={self.num_codebooks}"
+            )
+        self.quantizer = quantizer
+        self.fc2 = fc2
+        self.acoustic_decoder = acoustic_decoder
+        self._loaded = True
+        logger.info(
+            "Loaded HiggsAudioV2Code2Wav: %d quantizers, fc2(%d->%d), sample_rate=%d",
+            len(self.quantizer.quantizers),
+            self.fc2.in_features,
+            self.fc2.out_features,
+            self.sample_rate,
+        )
+
+    def _resolve_audio_tokenizer_path(self, model_dir: str) -> str:
+        """Return a local filesystem path containing ``config.json`` + codec
+        weights. See :meth:`load_codec_from_disk` for the resolution rules."""
+        tokenizer_id = getattr(self.config, "audio_tokenizer_id", None)
+        if tokenizer_id:
+            if os.path.isdir(tokenizer_id):
+                return tokenizer_id
+            return self._resolve_hf_id_to_local(tokenizer_id)
+        subdir = self.config.audio_tokenizer_subdir or ""
+        if os.path.isdir(model_dir):
+            return os.path.join(model_dir, subdir) if subdir else model_dir
+        local = self._resolve_hf_id_to_local(model_dir)
+        return os.path.join(local, subdir) if subdir else local
+
+    @staticmethod
+    def _resolve_hf_id_to_local(repo_id: str) -> str:
+        """Resolve a HF repo id to a local snapshot directory.
+
+        Prefers a read-only cache lookup (no write locks) so a quota-constrained
+        Lustre cache can still serve already-downloaded weights. Falls back to
+        ``snapshot_download`` only when nothing is cached locally.
+        """
+        from huggingface_hub import try_to_load_from_cache
+
+        cached = try_to_load_from_cache(repo_id=repo_id, filename="config.json")
+        if isinstance(cached, str) and os.path.isfile(cached):
+            return os.path.dirname(cached)
+
+        from huggingface_hub.constants import HF_HUB_CACHE
+
+        safe = repo_id.replace("/", "--")
+        snapshots_dir = os.path.join(HF_HUB_CACHE, f"models--{safe}", "snapshots")
+        if os.path.isdir(snapshots_dir):
+            for rev in os.listdir(snapshots_dir):
+                candidate = os.path.join(snapshots_dir, rev)
+                if os.path.isfile(os.path.join(candidate, "config.json")):
+                    return candidate
+
+        from huggingface_hub import snapshot_download
+
+        return snapshot_download(repo_id=repo_id)
+
+    # ------------------------------------------------------ direct decode API
+    @torch.inference_mode()
+    def decode_codes(self, audio_codes: torch.Tensor) -> torch.Tensor:
+        """Decode a ``[B, num_codebooks=8, T]`` code tensor to PCM ``[B, 1, T*hop]``."""
+        if not self._loaded:
+            raise RuntimeError("HiggsAudioV2Code2Wav not loaded. Call load_weights() first.")
+
+        codes = self._validate_codes(audio_codes)
+        device = codes.device
+
+        rvq_codes = codes.transpose(0, 1).long()  # [num_codebooks, B, T]
+        quantized = self.quantizer.decode(rvq_codes)  # [B, hidden, T]
+        # Match fc2's parameter dtype — vLLM may auto-cast the codec to the
+        # stage's model dtype (e.g. bf16), while the quantizer's RVQ codebooks
+        # often stay in float32 because they're embedded lookups. A defensive
+        # cast keeps the matmul homogeneous.
+        quantized = quantized.to(dtype=self.fc2.weight.dtype)
+        quantized = self.fc2(quantized.transpose(1, 2)).transpose(1, 2)
+        # Acoustic decoder may have parameters in yet another dtype — match
+        # its first parameter so the linear/conv ops stay consistent.
+        first_param = next(self.acoustic_decoder.parameters(), None)
+        if first_param is not None and quantized.dtype != first_param.dtype:
+            quantized = quantized.to(dtype=first_param.dtype)
+        audio = self.acoustic_decoder(quantized)
+        if audio.dim() == 2:
+            audio = audio.unsqueeze(1)
+        return audio.to(device)
+
+    @torch.inference_mode()
+    def forward_chunk(
+        self,
+        audio_codes: torch.Tensor,
+        *,
+        left_context_size: int = 0,
+        hop_length: int | None = None,
+    ) -> torch.Tensor:
+        """Chunked decode that trims ``left_context_size * hop_length`` samples
+        off the leading edge of the upsampled PCM. Mirrors the qwen3_tts
+        ``talker2code2wav_async_chunk`` overlap contract.
+        """
+        if left_context_size < 0:
+            raise ValueError(f"left_context_size must be >= 0; got {left_context_size}")
+        hop = int(hop_length) if hop_length is not None else self.hop_length
+        if hop <= 0:
+            raise ValueError(f"hop_length must be > 0; got {hop}")
+        pcm = self.decode_codes(audio_codes)
+        if left_context_size == 0:
+            return pcm
+        trim = left_context_size * hop
+        if pcm.shape[-1] <= trim:
+            return pcm[..., :0]
+        return pcm[..., trim:]
+
+    # ---------------------------------------------------- vLLM runtime forward
+    @torch.no_grad()
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        positions: torch.Tensor | None = None,
+        intermediate_tensors: Any = None,
+        inputs_embeds: torch.Tensor | None = None,
+        runtime_additional_information: list[dict[str, Any]] | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor | OmniOutput:
+        """Dual-mode forward.
+
+        - When ``input_ids`` is a 3-D ``[B, num_codebooks, T]`` tensor we treat
+          this as a direct decode call (legacy API) and return raw PCM.
+        - When ``input_ids`` is a flat 1-D / 2-D tensor we treat it as the
+          engine runtime payload (codebook-major flat per request) and return
+          an :class:`OmniOutput` carrying multimodal audio.
+        """
+        # Legacy direct-decode signature: caller passed a 3-D code tensor as input_ids.
+        if isinstance(input_ids, torch.Tensor) and input_ids.ndim == 3:
+            return self.decode_codes(input_ids)
+
+        sr_val = int(self.sample_rate)
+        sr_tensor = torch.tensor(sr_val, dtype=torch.int32)
+        empty = torch.zeros((0,), dtype=torch.float32)
+
+        if input_ids is None or input_ids.numel() == 0:
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={"model_outputs": [empty], "sr": [sr_tensor]},
+            )
+
+        ids = input_ids.reshape(-1).to(dtype=torch.long)
+        request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts"))
+
+        left_context_size = [0] * len(request_ids_list)
+        if runtime_additional_information is not None:
+            for i, info in enumerate(runtime_additional_information):
+                if i >= len(left_context_size):
+                    break
+                meta = info.get("meta", {}) if isinstance(info, dict) else {}
+                if "left_context_size" in meta:
+                    left_context_size[i] = int(meta["left_context_size"])
+
+        wavs: list[torch.Tensor] = []
+        for i, req_ids in enumerate(request_ids_list):
+            n = int(req_ids.numel())
+            if n == 0:
+                wavs.append(empty)
+                continue
+            if n % self.num_codebooks != 0:
+                logger.warning(
+                    "HiggsAudioV2Code2Wav: flat code length %d not divisible by num_codebooks=%d; dropping",
+                    n,
+                    self.num_codebooks,
+                )
+                wavs.append(empty)
+                continue
+            frames = n // self.num_codebooks
+            codes_qf = req_ids.reshape(self.num_codebooks, frames)
+            codes_bqf = codes_qf.unsqueeze(0)
+            try:
+                pcm = self.forward_chunk(
+                    codes_bqf,
+                    left_context_size=left_context_size[i],
+                    hop_length=self.hop_length,
+                )
+            except ValueError as exc:
+                logger.warning("HiggsAudioV2Code2Wav: decode skipped (%s)", exc)
+                wavs.append(empty)
+                continue
+            wavs.append(pcm.squeeze(0).squeeze(0).to(torch.float32).cpu())
+
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={
+                "model_outputs": wavs,
+                "sr": [sr_tensor] * len(wavs),
+            },
+        )
+
+    # --------------------------------------------------------------- helpers
+    def _validate_codes(self, audio_codes: torch.Tensor) -> torch.Tensor:
+        """Ensure shape and value range; reject stream specials with ValueError."""
+        if not isinstance(audio_codes, torch.Tensor):
+            raise TypeError(f"audio_codes must be a torch.Tensor, got {type(audio_codes)!r}")
+        if audio_codes.ndim != 3:
+            raise ValueError(
+                f"audio_codes must have shape [B, num_codebooks={self.num_codebooks}, T]; "
+                f"got shape {tuple(audio_codes.shape)}"
+            )
+        if int(audio_codes.shape[1]) != self.num_codebooks:
+            raise ValueError(
+                f"audio_codes second dim must equal num_codebooks={self.num_codebooks}; got {int(audio_codes.shape[1])}"
+            )
+        if audio_codes.numel() > 0:
+            max_val = int(audio_codes.max().item())
+            min_val = int(audio_codes.min().item())
+            if max_val >= self.num_real_codes or min_val < 0:
+                raise ValueError(
+                    "audio_codes contains stream-special or out-of-range IDs: "
+                    f"min={min_val}, max={max_val}; real code range is "
+                    f"[0, {self.num_real_codes - 1}]. Filter audio_stream_bos_id="
+                    f"{self.config.audio_stream_bos_id} and audio_stream_eos_id="
+                    f"{self.config.audio_stream_eos_id} (and anything above) at "
+                    "Stage 0 before sending codes to the codec decoder."
+                )
+        return audio_codes
+
+    @staticmethod
+    def _split_request_ids(ids: torch.Tensor, seq_token_counts: list[int] | None = None) -> list[torch.Tensor]:
+        """Split a concatenated flat-codes tensor into per-request segments."""
+        if seq_token_counts is not None and len(seq_token_counts) > 1:
+            boundaries = [0]
+            for count in seq_token_counts:
+                boundaries.append(boundaries[-1] + int(count))
+            n = int(ids.numel())
+            return [ids[boundaries[i] : min(boundaries[i + 1], n)] for i in range(len(seq_token_counts))]
+        return [ids]
+
+
+# Engine-side architecture identifier alias (mirrors Qwen3TTSCode2Wav usage).
+HiggsAudioV2Code2WavForConditionalGeneration = HiggsAudioV2Code2Wav
diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py
new file mode 100644
index 000000000..5e3883fd2
--- /dev/null
+++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py
@@ -0,0 +1,1684 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Stage-0 talker for higgs-audio v2 (vLLM-native, DualFFN-aware).
+
+Architecture summary (verified against `transformers.models.higgs_audio_v2`
+source and the boson-ai checkpoint config.json):
+
+- Backbone: Llama-3.2-3B (hidden=3072, 28 layers, GQA 24Q/8KV, head_dim=128,
+  vocab=128256, max_position_embeddings=2048, RoPE llama3 scaling factor=32).
+  We reuse vLLM's compiled :class:`vllm.model_executor.models.llama.LlamaModel`
+  for the attention path so PagedAttention scheduling stays intact.
+
+- DualFFN: every transformer block carries a parallel audio expert
+  (``audio_input_layernorm`` + ``audio_post_attention_layernorm`` + ``audio_mlp``)
+  next to the standard text path (``input_layernorm`` + ``post_attention_layernorm``
+  + ``mlp``). A per-position ``audio_token_mask`` of shape ``[B, S]`` selects
+  between the two. The mask is the union of positions where the input token id
+  equals ``audio_token_id=128016`` or ``audio_delay_token_id=128014``.
+
+- Multi-codebook output head: at audio positions, Stage 0 emits one ID per
+  codebook (8 codebooks of vocabulary 1026 each) in parallel via a single
+  fused ``audio_lm_head: Linear(hidden, num_codebooks * codebook_size)``
+  whose output is reshaped to ``[N_audio, num_codebooks, codebook_size]``.
+
+- Delay pattern: codebook k is staggered by k frames using
+  ``audio_delay_token_id=128014`` as a filler at the LM-token positions where
+  codebook k has not yet started emitting real codes. This matches
+  ``HiggsAudioV2DelayPatternLogitsProcessor`` upstream and is the canonical
+  MusicGen pattern ``[0, 1, 2, 3, 4, 5, 6, 7]``.
+
+This file delivers the structural pieces (module classes, weight mapping,
+forward pass, AR sampler dispatch, and the upstream delay-pattern logits
+processor) to reproduce the upstream behavior under vLLM's compiled forward
+path.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+import torch.nn as nn
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.models.llama import LlamaDecoderLayer, LlamaMLP
+
+from vllm_omni.model_executor.models.higgs_audio_v2.configuration_higgs_audio_v2 import (
+    HiggsAudioV2Config,
+)
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+__all__ = [
+    "HiggsAudioV2DecoderLayer",
+    "HiggsAudioV2TalkerForConditionalGeneration",
+]
+
+logger = init_logger(__name__)
+
+
+class HiggsAudioRMSNorm(nn.Module):
+    """Minimal RMSNorm matching the upstream HiggsAudioV2RMSNorm semantics.
+
+    Kept self-contained so the parameter names (``weight``) line up with the
+    upstream state dict for direct load.
+    """
+
+    def __init__(self, hidden_size: int, eps: float = 1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+class HiggsAudioV2DecoderLayer(LlamaDecoderLayer):
+    """One transformer block: vLLM-native attention + DualFFN-routed MLPs.
+
+    Wraps a stock :class:`vllm.model_executor.models.llama.LlamaDecoderLayer`
+    for the self-attention path (so PagedAttention KV caches keep working) and
+    overlays the upstream HiggsAudioV2 DualFFN routing for the layernorm + MLP
+    pairs. The implementation follows the upstream rule from
+    ``transformers.models.higgs_audio_v2.HiggsAudioV2DecoderLayer.forward``:
+
+    1. Pre-attention norm: per-position split into ``audio_input_layernorm``
+       (audio mask True) vs ``input_layernorm`` (audio mask False). The mixed
+       output is fed to a single shared self-attention.
+    2. Post-attention residual + dual MLP: text positions go through
+       ``mlp(post_attention_layernorm(.))`` and audio positions go through
+       ``audio_mlp(audio_post_attention_layernorm(.))``. Both deltas are added
+       to the residual.
+
+    The classical (non-fused-residual) transformer pattern is used so the per-
+    position split is straightforward. vLLM's compiled forward path passes
+    ``audio_token_mask`` through as an extra layer kwarg via
+    :class:`HiggsAudioV2TalkerForConditionalGeneration.forward`.
+    """
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        # Subclass LlamaDecoderLayer so self_attn registers at the
+        # CANONICAL ``layers.{i}.self_attn.attn`` prefix (not nested under
+        # ``.base``). This is the architecturally-correct registration path
+        # for vLLM's compilation_config.static_forward_context.
+        super().__init__(vllm_config=vllm_config, prefix=prefix)
+        config = vllm_config.model_config.hf_config
+        hidden_size = int(config.hidden_size)
+        intermediate_size = int(config.intermediate_size)
+        hidden_act = str(config.hidden_act)
+        mlp_bias = bool(getattr(config, "mlp_bias", False))
+        rms_norm_eps = float(config.rms_norm_eps)
+        # Parallel audio-expert norms + MLP (mirrors upstream parameter names).
+        self.audio_input_layernorm = HiggsAudioRMSNorm(hidden_size, eps=rms_norm_eps)
+        self.audio_post_attention_layernorm = HiggsAudioRMSNorm(hidden_size, eps=rms_norm_eps)
+        self.audio_mlp = LlamaMLP(
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            hidden_act=hidden_act,
+            bias=mlp_bias,
+            prefix=f"{prefix}.audio_mlp",
+        )
+
+    # ------------------------------------------------------------------ helpers
+    def _routed_norm(
+        self,
+        hidden: torch.Tensor,
+        text_norm: nn.Module,
+        audio_norm: HiggsAudioRMSNorm,
+        audio_mask: torch.Tensor | None,
+        residual: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Per-position routed RMSNorm with residual handling.
+
+        Matches the vLLM fused-residual pattern: takes (hidden, residual) and
+        returns (normed_hidden, post_add_residual). Routes per position:
+        audio positions through ``audio_norm`` (HiggsAudioRMSNorm, plain RMS),
+        text positions through ``text_norm`` (vLLM compiled RMSNorm).
+
+        For consistency we explicitly compute the residual sum once, then call
+        the vLLM text_norm on text positions WITHOUT its residual argument
+        (else it'd add the residual again).
+        """
+        # Fold residual into hidden once.
+        if residual is not None:
+            hidden = hidden + residual
+        new_residual = hidden
+        if audio_mask is None:
+            return audio_norm(hidden), new_residual
+        mask_flat = audio_mask.reshape(-1)
+        if hidden.ndim == 3:
+            hidden_flat = hidden.reshape(-1, hidden.shape[-1])
+        else:
+            hidden_flat = hidden
+        out = hidden_flat.clone()
+        if (~mask_flat).any():
+            text_out = text_norm(hidden_flat[~mask_flat])
+            # vLLM compiled RMSNorm.forward(x) returns Tensor; some variants
+            # return (Tensor, residual). Normalize to Tensor.
+            if isinstance(text_out, tuple):
+                text_out = text_out[0]
+            out[~mask_flat] = text_out.to(out.dtype)
+        if mask_flat.any():
+            audio_out = audio_norm(hidden_flat[mask_flat])
+            out[mask_flat] = audio_out.to(out.dtype)
+        return out.reshape_as(hidden), new_residual
+
+    def _routed_mlp(
+        self,
+        hidden: torch.Tensor,
+        audio_mask: torch.Tensor | None,
+    ) -> torch.Tensor:
+        """Apply ``audio_mlp`` to mask positions and ``mlp`` to non-mask positions."""
+        if audio_mask is None:
+            return self.audio_mlp(hidden)
+        mask_flat = audio_mask.reshape(-1)
+        if hidden.ndim == 3:
+            hidden_flat = hidden.reshape(-1, hidden.shape[-1])
+        else:
+            hidden_flat = hidden
+        out = torch.zeros_like(hidden_flat)
+        if (~mask_flat).any():
+            out[~mask_flat] = self.mlp(hidden_flat[~mask_flat]).to(out.dtype)
+        if mask_flat.any():
+            out[mask_flat] = self.audio_mlp(hidden_flat[mask_flat]).to(out.dtype)
+        return out.reshape_as(hidden)
+
+    # ------------------------------------------------------------------ forward
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+        *,
+        audio_token_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # Fused-residual pattern (compatible with vLLM's LlamaModel forward
+        # loop which calls ``hidden, residual = self.norm(hidden, residual)``
+        # at the end).
+        attn_input, residual = self._routed_norm(
+            hidden_states,
+            self.input_layernorm,
+            self.audio_input_layernorm,
+            audio_token_mask,
+            residual,
+        )
+        attn_out = self.self_attn(positions=positions, hidden_states=attn_input)
+
+        mlp_input, residual = self._routed_norm(
+            attn_out,
+            self.post_attention_layernorm,
+            self.audio_post_attention_layernorm,
+            audio_token_mask,
+            residual,
+        )
+        mlp_out = self._routed_mlp(mlp_input, audio_token_mask)
+        return mlp_out, residual
+
+
+class HiggsAudioV2TalkerForConditionalGeneration(nn.Module):
+    """Stage-0 talker class registered under
+    ``HiggsAudioV2ForConditionalGeneration`` (canonical HF arch identifier)
+    AND ``HiggsAudioV2TalkerForConditionalGeneration`` (explicit alias).
+
+    The class wires together:
+      1. vLLM-native ``LlamaModel`` backbone (PagedAttention scheduling).
+      2. A parallel set of DualFFN modules (one per decoder layer) for the
+         audio-expert path.
+      3. Multi-codebook output via a single fused ``audio_lm_head`` of shape
+         ``[num_codebooks * codebook_size, hidden]`` (all codebooks emitted
+         in parallel per audio position).
+      4. ``load_weights`` HF -> vLLM mapping (GQA reshape, llama3 RoPE,
+         text/audio MLP split, audio code heads, audio token embeddings).
+
+    The forward path interleaves DualFFN routing with vLLM's compiled
+    attention path so the talker integrates with the rest of the wiring
+    (pipeline_registry, serving_speech, deploy yaml, stage_input_processor).
+    """
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        hf_config = vllm_config.model_config.hf_config
+        if isinstance(hf_config, HiggsAudioV2Config):
+            self.config: HiggsAudioV2Config = hf_config
+        else:
+            # When loaded via AutoConfig the class is the upstream
+            # transformers config; wrap it in our typed shell preserving
+            # every attribute via PretrainedConfig.
+            self.config = HiggsAudioV2Config(**hf_config.to_dict())
+
+        # ------------------------------------------------------------------ embedding + final norm
+        # directly instantiate the text-side embedding and final norm
+        # so we don't pull in a full LlamaModel.layers list that the engine's
+        # weight loader would require initialized. Each HiggsAudioV2DecoderLayer
+        # below already owns its own LlamaDecoderLayer for the attention path;
+        # duplicating those via a parallel LlamaModel made the loader fail.
+        self.embed_tokens = VocabParallelEmbedding(
+            num_embeddings=self.config.vocab_size,
+            embedding_dim=self.config.hidden_size,
+            org_num_embeddings=self.config.vocab_size,
+            quant_config=None,
+            prefix=f"{prefix}.embed_tokens",
+        )
+        self.norm = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps)
+
+        # ------------------------------------------------------------------ DualFFN-aware layer stack
+        # use vLLM's ``make_layers`` helper to register attention layers
+        # at canonical ``layers.{i}.self_attn.attn`` paths in
+        # ``compilation_config.static_forward_context``. This is what gets
+        # picked up by ``bind_kv_cache`` so prefill writes K/V to the cache
+        # (root cause for empty cache + zero attention output).
+        from vllm.model_executor.models.utils import make_layers
+
+        _, _, self.layers = make_layers(
+            self.config.num_hidden_layers,
+            lambda prefix: HiggsAudioV2DecoderLayer(vllm_config=vllm_config, prefix=prefix),
+            prefix=f"{prefix}.layers",
+        )
+
+        # ------------------------------------------------------------------ audio embedding (shared)
+        # Audio frames are embedded by summing per-codebook lookups; the LM
+        # input_ids stream uses ``audio_token_id`` / ``audio_delay_token_id``
+        # placeholders at the corresponding positions, and the talker
+        # substitutes the audio embedding via the audio_codebook_embeddings
+        # table declared further below.
+        self.register_buffer(
+            "audio_tokens_offsets",
+            torch.arange(self.config.num_codebooks) * self.config.codebook_size,
+            persistent=False,
+        )
+
+        # ------------------------------------------------------------------ heads
+        # Standard text LM head (Llama 128256-wide vocab).
+        self.lm_head = ParallelLMHead(
+            self.config.vocab_size,
+            self.config.hidden_size,
+            org_num_embeddings=self.config.vocab_size,
+            quant_config=None,
+            prefix=f"{prefix}.lm_head",
+        )
+
+        # Audio side: the upstream ``audio_decoder_proj.audio_lm_head`` is a
+        # single fused ``Linear(hidden, num_codebooks * (codebook_size + 2))``.
+        # In our v2 config ``codebook_size = 1026`` ALREADY includes the two
+        # stream specials, so the fused head is shape ``(8 * 1026, 3072) =
+        # (8208, 3072)``. At sample time we run it once and reshape
+        # ``[N, 8208] -> [N, 8, 1026]`` to read each codebook's logits.
+        self.audio_lm_head = nn.Linear(
+            self.config.hidden_size,
+            self.config.num_codebooks * self.config.codebook_size,
+            bias=False,
+        )
+        # Multi-codebook input embedding used both for prompt-side audio frame
+        # substitution and for next-step audio-input-embedding feedback during
+        # decode. Mirrors upstream ``HiggsAudioModel.audio_codebook_embeddings``.
+        self.audio_codebook_embeddings = nn.Embedding(
+            self.config.num_codebooks * self.config.codebook_size,
+            self.config.hidden_size,
+        )
+
+        # vLLM's logits processor wires sampling metadata into the lm_head.
+        self.logits_processor = LogitsProcessor(self.config.vocab_size)
+
+    # ----------------------------------------------------------------- masks
+    @torch.inference_mode()
+    def audio_token_mask(self, input_ids: torch.Tensor) -> torch.Tensor:
+        """Compute the per-position audio_token_mask used by DualFFN routing.
+
+        Matches ``HiggsAudioV2Model.get_placeholder_mask``:
+
+            mask = (input_ids == audio_token_id) | (input_ids == audio_delay_token_id)
+        """
+        return (input_ids == self.config.audio_token_id) | (input_ids == self.config.audio_delay_token_id)
+
+    # ------------------------------------------------------------ ref-audio
+    @torch.inference_mode()
+    def _maybe_apply_ref_audio_substitution(
+        self,
+        hidden_states: torch.Tensor,
+        input_ids: torch.Tensor,
+        info_dicts: Any,
+    ) -> torch.Tensor:
+        """Substitute reference-audio embeddings at prompt-side audio placeholders.
+
+        Voice-clone path: the serving layer threads per-request
+        ``audio_input_ids`` (shape ``[T_frames, num_codebooks]``) +
+        ``audio_input_ids_mask`` (``[T_frames]``) into
+        ``model_intermediate_buffer`` (one dict per request in batch order).
+        For each request that has those tensors, we look up its span in the
+        flat ``input_ids`` via ``query_start_loc`` from the forward context,
+        embed the codes through ``audio_codebook_embeddings`` (with the
+        per-codebook ``audio_tokens_offsets`` shift, mirroring
+        ``HiggsAudioV2Embeddings.forward`` upstream), and overwrite the
+        text-side embedding at every audio_token / audio_delay placeholder
+        position in that span.
+
+        Decode steps (span length 1) are skipped here — the audio-feedback
+        path in ``_maybe_apply_audio_feedback`` takes over from the moment
+        the talker emits its first codebook frame via ``sample()``.
+        """
+        if not info_dicts:
+            return hidden_states
+
+        # Pull the per-request (codes, mask) tensors in batch order. Values
+        # may be on CPU after IPC serialization (msgspec ships tensors as raw
+        # bytes); the loop below moves them to the input_ids device on use.
+        per_req_codes: list[torch.Tensor | None] = []
+        any_codes = False
+        for info in info_dicts:
+            if not isinstance(info, dict):
+                per_req_codes.append(None)
+                continue
+            ids = info.get("audio_input_ids")
+            mask = info.get("audio_input_ids_mask")
+            if isinstance(ids, list):
+                ids = ids[0] if ids else None
+            if isinstance(mask, list):
+                mask = mask[0] if mask else None
+            if not isinstance(ids, torch.Tensor) or not isinstance(mask, torch.Tensor):
+                per_req_codes.append(None)
+                continue
+            if ids.ndim == 3:
+                ids = ids[0]
+            if mask.ndim == 2:
+                mask = mask[0]
+            valid = ids[mask.to(dtype=torch.bool)]  # [T_valid, num_codebooks]
+            if valid.numel() == 0:
+                per_req_codes.append(None)
+                continue
+            per_req_codes.append(valid)
+            any_codes = True
+
+        if not any_codes:
+            return hidden_states
+
+        # Per-request slices in the flat input_ids are recovered from the
+        # forward context's attention metadata (mimo_audio_llm.py:1022-1025
+        # uses the same pattern). Without spans we cannot route codes to the
+        # right placeholder positions, so we no-op instead of risking a
+        # cross-request mix-up.
+        from vllm.forward_context import get_forward_context
+
+        try:
+            attn_metadata = get_forward_context().attn_metadata
+        except Exception:
+            return hidden_states
+        if isinstance(attn_metadata, dict):
+            if not attn_metadata:
+                return hidden_states
+            attn = next(iter(attn_metadata.values()))
+        else:
+            attn = attn_metadata
+        q_start = getattr(attn, "query_start_loc", None)
+        if not isinstance(q_start, torch.Tensor):
+            return hidden_states
+        q_start_cpu = q_start.detach().to("cpu").tolist()
+        if len(q_start_cpu) != len(info_dicts) + 1:
+            return hidden_states
+
+        audio_token_id = int(self.config.audio_token_id)
+        audio_delay_id = int(self.config.audio_delay_token_id)
+        offsets = self.audio_tokens_offsets.to(input_ids.device)
+
+        new_hidden: torch.Tensor | None = None
+        for i, codes in enumerate(per_req_codes):
+            if codes is None:
+                continue
+            s = int(q_start_cpu[i])
+            e = int(q_start_cpu[i + 1])
+            if e - s <= 1:
+                # Decode step — let _maybe_apply_audio_feedback handle this row.
+                continue
+            slice_ids = input_ids[s:e]
+            mask = (slice_ids == audio_token_id) | (slice_ids == audio_delay_id)
+            placeholders = mask.nonzero(as_tuple=True)[0]
+            n_real = int(codes.shape[0])
+            if int(placeholders.numel()) < n_real:
+                # Mismatch (likely a truncated prompt): skip this request rather
+                # than corrupt the batch. The user-visible failure surfaces as
+                # the model not honoring the ref-audio rather than a crash.
+                continue
+            target = placeholders[:n_real] + s
+            embeds = self.audio_codebook_embeddings(codes.to(device=input_ids.device, dtype=torch.long) + offsets).sum(
+                dim=-2
+            )
+            if new_hidden is None:
+                new_hidden = hidden_states.clone()
+            new_hidden[target] = embeds.to(new_hidden.dtype)
+
+        return new_hidden if new_hidden is not None else hidden_states
+
+    # ----------------------------------------------------------------- weights
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Map HF state dict names to the vLLM-native layout.
+
+        The upstream tensor names live under ``model.layers.<L>...`` plus a few
+        top-level entries. The vLLM-native layout uses fused QKV/MLP tensors,
+        so we transform the relevant pairs at load time:
+
+          * Attention: HF ``q_proj/k_proj/v_proj`` -> vLLM ``qkv_proj`` (packed
+            ``[hidden + 2 * kv_dim, hidden]``; for GQA with 24 Q heads, 8 KV
+            heads, head_dim=128 the slabs are
+            ``q_proj[3072, 3072]``, ``k_proj[1024, 3072]``, ``v_proj[1024, 3072]``
+            -> ``qkv_proj[5120, 3072]``).
+          * RoPE parameters are read from ``self.config.rope_parameters`` and
+            consumed by vLLM's LlamaAttention without any weight transform
+            (``rope_type="llama3"``, factor=32, low_freq_factor=0.125,
+            high_freq_factor=0.5, original_max_position_embeddings=1024).
+          * Text MLP: HF ``mlp.gate_proj`` + ``mlp.up_proj`` ->
+            vLLM ``mlp.gate_up_proj`` (concatenated on dim 0).
+          * Audio MLP: HF ``audio_mlp.gate_proj`` + ``audio_mlp.up_proj`` ->
+            vLLM ``dual_ffns.<L>.audio_mlp.gate_up_proj`` (same shape transform).
+          * Layer norms: HF ``input_layernorm`` / ``post_attention_layernorm``
+            -> ``dual_ffns.<L>.{input,post_attention}_layernorm``. Audio norms
+            -> ``dual_ffns.<L>.{audio_input,audio_post_attention}_layernorm``.
+          * Audio token embedding: HF ``model.embed_audio_tokens.weight`` (with
+            optional outer ``embed_audio_tokens.`` prefix from upstream's
+            HiggsAudioV2Embeddings) -> ``self.audio_codebook_embeddings.weight``
+            (see ``_map_simple_name``).
+          * Audio codebook heads: HF ``audio_lm_head.weight`` -> the fused 8x1026
+            head from which Stage 0 reads codebook 0 directly and codebook 1..7
+            via the fast-AR code predictor's residual heads. (Boson-ai's actual
+            checkpoint exposes a single fused ``audio_lm_head`` of shape
+            ``[num_codebooks * codebook_size, hidden]``, not one head per
+            codebook; we split it into the per-codebook heads at load time.)
+
+        this routine covers the simple, fused-QKV, and fused-MLP
+        transcriptions. The remaining gap (full Stage-0 forward integration with
+        vLLM's PagedAttention scheduler so the loaded weights drive a working
+        greedy decode) is tracked as an Open Issue in the goal tracker.
+        """
+        loaded: set[str] = set()
+        own_params = dict(self.named_parameters())
+        # Debug: log the first few HF keys we receive so we can verify the
+        # mapper hits them. Also log our own param-name shape so missing
+        # mappings are obvious.
+        logger.info(
+            "higgs_audio_v2 load_weights: %d own params (sample: %s)",
+            len(own_params),
+            sorted(list(own_params.keys()))[:5],
+        )
+        debug_seen: list[str] = []
+
+        # First pass: collect tensors that need fusing (q/k/v, gate/up for each
+        # MLP slot). We accumulate into per-(layer, slot, kind) buckets and
+        # commit the fused tensor once all parts have arrived.
+        fuse_buckets: dict[tuple[int, str, str], dict[str, torch.Tensor]] = {}
+
+        def _strip_model_prefix(name: str) -> str:
+            """vLLM's DefaultModelLoader strips the leading ``model.`` from HF
+            state-dict keys before invoking ``load_weights``, but we also want
+            to accept raw HF names (for unit tests that pass the state dict
+            directly). Normalize by stripping at most one ``model.`` prefix.
+            """
+            if name.startswith("model."):
+                return name[len("model.") :]
+            return name
+
+        def _try_simple(name: str, tensor: torch.Tensor) -> bool:
+            normalized = _strip_model_prefix(name)
+            # Direct match: many keys (embed_tokens.weight, norm.weight,
+            # layers.<L>.{audio_,}input_layernorm.weight, layers.<L>.{audio_,}post_attention_layernorm.weight,
+            # layers.<L>.self_attn.o_proj.weight, layers.<L>.{,audio_}mlp.down_proj.weight) already match
+            # our named_parameters() keys after the model-prefix strip.
+            if normalized in own_params:
+                target = own_params[normalized]
+                if tuple(target.shape) == tuple(tensor.shape):
+                    with torch.no_grad():
+                        target.copy_(tensor)
+                    loaded.add(normalized)
+                    return True
+                logger.warning(
+                    "higgs_audio_v2 load_weights: shape mismatch %s -> %s: %s vs %s",
+                    name,
+                    normalized,
+                    tuple(tensor.shape),
+                    tuple(target.shape),
+                )
+                return False
+
+            # After making HiggsAudioV2DecoderLayer inherit from
+            # LlamaDecoderLayer, text-side params live directly at
+            # ``layers.<L>.<tail>`` (no ``.base`` indirection). audio side
+            # remains at ``layers.<L>.audio_<...>``.
+            parts = normalized.split(".")
+            if len(parts) >= 3 and parts[0] == "layers":
+                tail = ".".join(parts[2:])
+                # Audio_mlp.down_proj lives at our top-level audio_mlp slot.
+                if tail == "audio_mlp.down_proj.weight":
+                    audio_name = f"layers.{parts[1]}.audio_mlp.down_proj.weight"
+                    if audio_name in own_params:
+                        target = own_params[audio_name]
+                        if tuple(target.shape) == tuple(tensor.shape):
+                            with torch.no_grad():
+                                target.copy_(tensor)
+                            loaded.add(audio_name)
+                            return True
+
+            # Special simple cases (nested audio embedding, text_lm_head alias, etc.)
+            mapped = self._map_simple_name(name)
+            if mapped is None:
+                # Also try after stripping the model. prefix so the _map_simple_name
+                # branches written for un-stripped names continue to work.
+                mapped = self._map_simple_name("model." + normalized)
+            if mapped is None or mapped not in own_params:
+                return False
+            target = own_params[mapped]
+            if tuple(target.shape) != tuple(tensor.shape):
+                logger.warning(
+                    "higgs_audio_v2 load_weights: shape mismatch %s -> %s: %s vs %s",
+                    name,
+                    mapped,
+                    tuple(tensor.shape),
+                    tuple(target.shape),
+                )
+                return False
+            with torch.no_grad():
+                target.copy_(tensor)
+            loaded.add(mapped)
+            return True
+
+        def _stash_fusion(name: str, tensor: torch.Tensor) -> bool:
+            normalized = _strip_model_prefix(name)
+            parts = normalized.split(".")
+            if len(parts) < 4 or parts[0] != "layers":
+                return False
+            try:
+                layer_idx = int(parts[1])
+            except ValueError:
+                return False
+            tail = ".".join(parts[2:])
+            for slot in ("mlp", "audio_mlp"):
+                for kind in ("gate_proj", "up_proj"):
+                    if tail == f"{slot}.{kind}.weight":
+                        fuse_buckets.setdefault((layer_idx, slot, "gate_up_proj"), {})[kind] = tensor
+                        return True
+            for kind in ("q_proj", "k_proj", "v_proj"):
+                if tail == f"self_attn.{kind}.weight":
+                    fuse_buckets.setdefault((layer_idx, "self_attn", "qkv_proj"), {})[kind] = tensor
+                    return True
+            return False
+
+        unhandled: list[str] = []
+        for name, tensor in weights:
+            if len(debug_seen) < 30:
+                debug_seen.append(name)
+            if _try_simple(name, tensor):
+                continue
+            if _stash_fusion(name, tensor):
+                continue
+            # Audio code head: upstream has a fused [num_codebooks*codebook_size, hidden]
+            # head; split into codebook-0 head + residual heads. Accept both
+            # ``audio_lm_head.weight`` and ``model.audio_lm_head.weight``.
+            if name in (
+                "audio_lm_head.weight",
+                "model.audio_lm_head.weight",
+                "audio_decoder_proj.audio_lm_head.weight",
+                "model.audio_decoder_proj.audio_lm_head.weight",
+            ):
+                self._consume_fused_audio_head(tensor, loaded)
+                continue
+            unhandled.append(name)
+
+        # Second pass: emit fused tensors. Targets are the per-layer slots on
+        # the HiggsAudioV2DecoderLayer stack: text-side params at
+        # ``layers.<L>.<tail>`` directly (HiggsAudioV2DecoderLayer inherits
+        # from LlamaDecoderLayer); ``layers.<L>.audio_mlp...`` for the audio
+        # expert.
+        for (layer_idx, slot, fused_kind), parts in fuse_buckets.items():
+            if fused_kind == "gate_up_proj":
+                gate = parts.get("gate_proj")
+                up = parts.get("up_proj")
+                if gate is None or up is None:
+                    continue
+                fused = torch.cat([gate, up], dim=0)
+                if slot == "mlp":
+                    target_name = f"layers.{layer_idx}.mlp.gate_up_proj.weight"
+                else:  # audio_mlp
+                    target_name = f"layers.{layer_idx}.audio_mlp.gate_up_proj.weight"
+            elif fused_kind == "qkv_proj":
+                q = parts.get("q_proj")
+                k = parts.get("k_proj")
+                v = parts.get("v_proj")
+                if q is None or k is None or v is None:
+                    continue
+                fused = torch.cat([q, k, v], dim=0)
+                target_name = f"layers.{layer_idx}.self_attn.qkv_proj.weight"
+            else:
+                continue
+            if target_name in own_params and tuple(own_params[target_name].shape) == tuple(fused.shape):
+                with torch.no_grad():
+                    own_params[target_name].copy_(fused)
+                loaded.add(target_name)
+            else:
+                logger.warning(
+                    "higgs_audio_v2 load_weights: fused %s not found or shape mismatch (%s vs target %s)",
+                    target_name,
+                    tuple(fused.shape),
+                    tuple(own_params[target_name].shape) if target_name in own_params else "<missing>",
+                )
+        logger.info(
+            "higgs_audio_v2 load_weights: %d/%d params initialized "
+            "(sample seen: %s; unhandled count=%d, sample unhandled=%s)",
+            len(loaded),
+            len(own_params),
+            debug_seen[:10],
+            len(unhandled),
+            unhandled[:20],
+        )
+        return loaded
+
+    def _consume_fused_audio_head(self, fused: torch.Tensor, loaded: set[str]) -> None:
+        """Load the fused ``audio_lm_head[num_codebooks*codebook_size, hidden]``
+        directly into ``self.audio_lm_head.weight``.
+
+        Upstream uses a single big head; there is no separate codebook-0
+        head or residual predictor in the actual checkpoint.
+        """
+        own_params = dict(self.named_parameters())
+        num_codebooks = int(self.config.num_codebooks)
+        codebook_size = int(self.config.codebook_size)
+        hidden = int(self.config.hidden_size)
+        target = own_params.get("audio_lm_head.weight")
+        if target is None:
+            logger.warning("higgs_audio_v2: no audio_lm_head.weight slot to load into")
+            return
+        expected = (num_codebooks * codebook_size, hidden)
+        if tuple(fused.shape) != expected:
+            logger.warning(
+                "higgs_audio_v2: unexpected audio_lm_head shape %s; expected %s",
+                tuple(fused.shape),
+                expected,
+            )
+            return
+        with torch.no_grad():
+            target.copy_(fused)
+        loaded.add("audio_lm_head.weight")
+
+    # --------------------------------------------------------------- helpers
+    def _map_simple_name(self, hf_name: str) -> str | None:
+        """Translate the simple (non-fused) HF parameter names to vLLM names."""
+        # Upstream nests the audio embedding as
+        # ``model.embed_audio_tokens.embed_audio_tokens.weight``; the boson-ai
+        # checkpoint also uses this exact key. Accept both nested and flat.
+        # the upstream multi-codebook embedding is ``audio_codebook_embeddings``
+        # (a flat ``Embedding(8 * codebook_size, hidden)``). also had a separate
+        # ``embed_audio_tokens`` (used in prefill-time DualFFN routing helpers);
+        # we keep that slot so the existing prompt-side embedding logic stays
+        # functional. Both slots are loaded from the same upstream tensor.
+        if hf_name in (
+            "embed_audio_tokens.weight",
+            "embed_audio_tokens.embed_audio_tokens.weight",  # stripped form
+            "model.embed_audio_tokens.embed_audio_tokens.weight",  # raw HF form
+            "model.embed_audio_tokens.weight",
+            # Actual key used in the boson-ai checkpoint:
+            "audio_codebook_embeddings.weight",
+            "model.audio_codebook_embeddings.weight",
+        ):
+            return "audio_codebook_embeddings.weight"
+        if hf_name in (
+            "lm_head.weight",
+            "text_lm_head.weight",
+            "model.text_lm_head.weight",
+            # Actual key used in the boson-ai checkpoint:
+            "audio_decoder_proj.text_lm_head.weight",
+            "model.audio_decoder_proj.text_lm_head.weight",
+        ):
+            return "lm_head.weight"
+        # Dropped LlamaModel wrapper; text embedding + final norm now
+        # live directly on the talker (self.embed_tokens / self.norm).
+        if hf_name == "model.embed_tokens.weight":
+            return "embed_tokens.weight"
+        if hf_name == "model.norm.weight":
+            return "norm.weight"
+        # Per-layer audio norms -> our HiggsAudioV2DecoderLayer side.
+        if hf_name.startswith("model.layers.") and hf_name.endswith(".audio_input_layernorm.weight"):
+            layer_idx = hf_name.split(".")[2]
+            return f"layers.{layer_idx}.audio_input_layernorm.weight"
+        if hf_name.startswith("model.layers.") and hf_name.endswith(".audio_post_attention_layernorm.weight"):
+            layer_idx = hf_name.split(".")[2]
+            return f"layers.{layer_idx}.audio_post_attention_layernorm.weight"
+        # Per-layer text-side params now live directly at
+        # ``layers.{i}.<tail>`` (no ``.base`` indirection) since
+        # HiggsAudioV2DecoderLayer inherits from LlamaDecoderLayer.
+        if hf_name.startswith("model.layers.") and hf_name.endswith(".input_layernorm.weight"):
+            layer_idx = hf_name.split(".")[2]
+            return f"layers.{layer_idx}.input_layernorm.weight"
+        if hf_name.startswith("model.layers.") and hf_name.endswith(".post_attention_layernorm.weight"):
+            layer_idx = hf_name.split(".")[2]
+            return f"layers.{layer_idx}.post_attention_layernorm.weight"
+        if hf_name.startswith("model.layers.") and hf_name.endswith(".self_attn.o_proj.weight"):
+            layer_idx = hf_name.split(".")[2]
+            return f"layers.{layer_idx}.self_attn.o_proj.weight"
+        if hf_name.startswith("model.layers.") and hf_name.endswith(".mlp.down_proj.weight"):
+            layer_idx = hf_name.split(".")[2]
+            return f"layers.{layer_idx}.mlp.down_proj.weight"
+        if hf_name.startswith("model.layers.") and hf_name.endswith(".audio_mlp.down_proj.weight"):
+            layer_idx = hf_name.split(".")[2]
+            return f"layers.{layer_idx}.audio_mlp.down_proj.weight"
+        return None
+
+    # ----------------------------------------------------------------- forward
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: Any | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        """DualFFN-routed Stage-0 forward.
+
+        Drives the custom :class:`HiggsAudioV2DecoderLayer` stack with a
+        per-position ``audio_token_mask`` derived from the input ``input_ids``
+        (matching :func:`HiggsAudioV2Model.get_placeholder_mask`). Each layer
+        applies the upstream DualFFN routing internally; this method only
+        composes the embedding and final-norm steps around the layer loop.
+        """
+        if inputs_embeds is None:
+            hidden_states = self.embed_tokens(input_ids)
+        else:
+            hidden_states = inputs_embeds
+
+        # Stash the step's input_ids for :meth:`_apply_audio_mode_bias` to read
+        # in :meth:`sample`. ``sampling_metadata.prompt_token_ids`` is only
+        # populated when penalties or token-id-aware logits processors are
+        # active, so for our use case we have to keep our own copy.
+        if input_ids is not None:
+            self._last_step_input_ids = input_ids
+
+        # Stash query_start_loc here (inside the forward context) so
+        # _apply_audio_mode_bias (which runs LATER in sample(), after the
+        # forward context exits) can recover each request's last token
+        # position from the flat input_ids — see the per-row prev-token
+        # lookup below. Without this, mixed prefill+decode batches would
+        # fall back to ``input_ids[-num_rows:]``, which silently misroutes
+        # the audio_out_bos detection in batch>1 voice clone (was the root
+        # cause of the slot-shifted gibberish at batch=4).
+        try:
+            from vllm.forward_context import get_forward_context
+
+            attn_metadata = get_forward_context().attn_metadata
+            if isinstance(attn_metadata, dict) and attn_metadata:
+                attn = next(iter(attn_metadata.values()))
+            else:
+                attn = attn_metadata
+            qsl = getattr(attn, "query_start_loc", None)
+            if isinstance(qsl, torch.Tensor):
+                self._last_step_query_start_loc = qsl.detach().clone()
+            else:
+                self._last_step_query_start_loc = None
+        except Exception:
+            self._last_step_query_start_loc = None
+
+        # Voice-clone ref-audio substitution: at the prompt-side audio
+        # placeholder span we overwrite the text-embed lookups with the
+        # multi-codebook embedding of the encoded reference clip carried in
+        # ``model_intermediate_buffer``. No-op for plain-text requests and
+        # for decode steps (which fall through to the audio-feedback path).
+        if input_ids is not None and inputs_embeds is None:
+            info_dicts = kwargs.get("model_intermediate_buffer")
+            if info_dicts is None:
+                info_dicts = kwargs.get("runtime_additional_information")
+            hidden_states = self._maybe_apply_ref_audio_substitution(
+                hidden_states,
+                input_ids,
+                info_dicts,
+            )
+
+        # Audio-feedback substitution: vLLM's AR runner calls forward()
+        # with ``inputs_embeds=None`` and lets the model embed its own
+        # input_ids inline. ``embed_input_ids`` is only invoked via the
+        # supports_mm_inputs/prompt_embeds branches, neither of which fires
+        # for the talker. So we apply the substitution directly here: at
+        # decode positions where input_ids == audio_token_id, swap the text
+        # embedding with embed_audio_codes(last audio frame stashed in
+        # ``self._audio_state[req_id]`` by :meth:`sample`).
+        if input_ids is not None and inputs_embeds is None:
+            hidden_states = self._maybe_apply_audio_feedback(hidden_states, input_ids)
+
+        audio_mask = self.audio_token_mask(input_ids) if input_ids is not None else None
+
+        # Fused-residual loop: layers return (hidden, residual)
+        # where ``residual`` carries everything up to the post-MLP add.
+        residual: torch.Tensor | None = None
+        for layer in self.layers:
+            hidden_states, residual = layer(
+                positions,
+                hidden_states,
+                residual,
+                audio_token_mask=audio_mask,
+            )
+
+        # Final fused norm: combine final mlp_out with carried residual.
+        norm_out = self.norm(hidden_states, residual)
+        if isinstance(norm_out, tuple):
+            norm_out = norm_out[0]
+        return norm_out
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+        sampling_metadata: Any = None,
+    ) -> torch.Tensor:
+        """Sampler-compatible logits tensor.
+
+        Returns a single ``[N, vocab_size]`` tensor that the stock
+        vLLM AR runner pipes through ``.contiguous()`` and a generic
+        sampler. The contract is intentionally tensor-only:
+
+        - Text-position logits come from ``lm_head`` (128256-wide Llama vocab).
+        - Audio-position logits at audio_token_id placeholders are also
+          read from ``lm_head``; the per-position codebook routing path
+          lives in :meth:`audio_codebook_logits` and is consumed by a
+          separate codebook sampling adapter.
+        """
+        # Stash for the audio-sampler dispatch in :meth:`sample`. Limited to a
+        # single step at a time; cleared in :meth:`sample` after consumption
+        # so a missed call doesn't leak stale tensors.
+        self._last_logits_hidden = hidden_states
+        return self.logits_processor(self.lm_head, hidden_states, sampling_metadata)
+
+    def audio_codebook_logits(
+        self,
+        hidden_states: torch.Tensor,
+        audio_token_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Per-codebook logits at audio positions, ``[N_audio, num_codebooks, codebook_size]``.
+
+        Runs the single fused ``audio_lm_head: Linear(hidden, num_codebooks *
+        codebook_size)`` and reshapes the output. This mirrors upstream's
+        ``HiggsAudioDecoderProjector.audio_lm_head`` + the
+        ``view(-1, num_codebooks, codebook_size)`` reshape in
+        ``HiggsAudioModel.forward``.
+
+        Returns an empty tensor with the correct trailing shape when the mask
+        selects no positions.
+        """
+        num_codebooks = int(self.config.num_codebooks)
+        codebook_size = int(self.config.codebook_size)
+        mask = audio_token_mask.reshape(-1).to(hidden_states.device)
+        hidden_flat = hidden_states.reshape(-1, hidden_states.shape[-1])
+        if not mask.any():
+            return torch.empty(
+                (0, num_codebooks, codebook_size),
+                dtype=hidden_states.dtype,
+                device=hidden_states.device,
+            )
+        flat = self.audio_lm_head(hidden_flat[mask])
+        return flat.view(-1, num_codebooks, codebook_size)
+
+    def embed_audio_codes(self, audio_ids: torch.Tensor) -> torch.Tensor:
+        """Embed a ``[num_codebooks, T]`` audio-ids slice into a ``[T, hidden]`` tensor.
+
+        Mirrors upstream ``HiggsAudioModel._embed_audio_ids``:
+
+            shift = arange(num_codebooks) * codebook_size  # 1026 per slot
+            audio_embed = sum_over_codebooks(audio_codebook_embeddings(audio_ids + shift))
+
+        The 8 codebooks share a single ``Embedding(8 * codebook_size, hidden)``
+        table; each codebook indexes into its own per-codebook slot via the
+        shift. The output is the per-frame mixture of 8 codebook embeddings
+        and is what the model re-feeds at the next AR step's
+        ``<|AUDIO_OUT|>`` placeholder position.
+        """
+        if audio_ids.ndim != 2:
+            raise ValueError(f"audio_ids must be [num_codebooks, T]; got shape {tuple(audio_ids.shape)}")
+        codebook_size = int(self.config.codebook_size)
+        num_codebooks = int(self.config.num_codebooks)
+        shift = torch.arange(num_codebooks, device=audio_ids.device) * codebook_size
+        shifted = audio_ids + shift.unsqueeze(-1)  # [num_codebooks, T]
+        embed = self.audio_codebook_embeddings(shifted)  # [num_codebooks, T, hidden]
+        return embed.sum(dim=0)  # [T, hidden]
+
+    def _maybe_apply_audio_feedback(self, hidden_states: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor:
+        """Return ``hidden_states`` with audio-token positions replaced by the
+        multi-codebook embedding of the LAST audio frame stashed by
+        :meth:`sample`.
+
+        Called from :meth:`forward` when ``inputs_embeds`` is ``None`` (vLLM's
+        AR runner path) — the equivalent of upstream's
+        ``_embed_audio_ids(audio_out_ids)`` injection at audio positions.
+        We use ``index_copy`` (out-of-place) instead of mutating the original
+        ``hidden_states`` because the embed-tokens output may be a view of a
+        compiled buffer that doesn't tolerate in-place writes.
+        """
+        audio_token_id = int(self.config.audio_token_id)
+        flat_ids = input_ids.reshape(-1)
+        audio_positions = (flat_ids == audio_token_id).nonzero(as_tuple=False).reshape(-1)
+        if audio_positions.numel() == 0:
+            return hidden_states
+        state = getattr(self, "_audio_state", None)
+        if not state:
+            return hidden_states
+        # Build a list of (pos, audio_emb) overrides and apply via index_copy
+        # on a cloned hidden_states.
+        rep_positions: list[int] = []
+        rep_embeds: list[torch.Tensor] = []
+        for pos in audio_positions.tolist():
+            req_state = state.get(int(pos))
+            if req_state is None:
+                continue
+            audio_out_ids = req_state.get("audio_out_ids")
+            if not isinstance(audio_out_ids, torch.Tensor) or audio_out_ids.numel() == 0:
+                continue
+            last_frame = audio_out_ids[:, -1:].to(hidden_states.device)
+            audio_emb = self.embed_audio_codes(last_frame)  # [1, hidden]
+            rep_positions.append(int(pos))
+            rep_embeds.append(audio_emb[0].to(dtype=hidden_states.dtype))
+        if not rep_positions:
+            return hidden_states
+        new_hidden = hidden_states.clone()
+        flat_hidden = new_hidden.reshape(-1, new_hidden.shape[-1])
+        idx = torch.tensor(rep_positions, dtype=torch.long, device=new_hidden.device)
+        rep = torch.stack(rep_embeds, dim=0)
+        flat_hidden.index_copy_(0, idx, rep)
+        return new_hidden
+
+    def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
+        """Embed the LM token stream, swapping in audio embeddings at placeholders.
+
+        Text positions use the standard ``embed_tokens`` lookup. At audio
+        positions (input id equals ``audio_token_id``) we substitute the
+        text embedding with the multi-codebook embedding of the LAST audio
+        frame stashed by :meth:`sample` in ``self._audio_state[req_id]``.
+        This mirrors upstream's ``HiggsAudioModel.forward`` flow where
+        ``_embed_audio_ids(audio_out_ids)`` is plugged into the input at
+        ``<|AUDIO_OUT|>`` placeholder positions.
+
+        Without this feedback the model receives the SAME embedding (just
+        the text ``<|AUDIO_OUT|>`` embedding) at every audio step and has no
+        way to attend to its own audio history — the resulting codes
+        collapse to a fixed loop and the ASR cannot recover the original
+        prompt.
+        """
+        text_embed = self.embed_tokens(input_ids)
+        audio_token_id = int(self.config.audio_token_id)
+        flat_ids = input_ids.reshape(-1)
+        audio_positions = (flat_ids == audio_token_id).nonzero(as_tuple=False).reshape(-1)
+        if audio_positions.numel() == 0:
+            return text_embed
+        state = getattr(self, "_audio_state", None)
+        if not state:
+            return text_embed
+        # We expect one audio position per active request per decode step.
+        # Match rows by walking through the per-request audio_out_ids buffers.
+        # The runner concatenates requests in batch order; ``audio_positions``
+        # lists their indices into the flat ``input_ids``. For each position we
+        # look up the corresponding request's last audio frame.
+        # NOTE: this O(N_audio) loop is small (N_audio == active requests in a
+        # decode step) and runs once per step.
+        flat_embed = text_embed.reshape(-1, text_embed.shape[-1])
+        for pos in audio_positions.tolist():
+            # ``state`` keys are the absolute batch row indices the sampler
+            # saw, which match the positions in flat ``input_ids`` during a
+            # decode step (one token per active request). For a prefill
+            # step's prompt-side audio_token_id positions, no state entry
+            # exists yet — that's fine, we leave the text embed in place.
+            req_state = state.get(int(pos))
+            if req_state is None:
+                continue
+            audio_out_ids = req_state.get("audio_out_ids")
+            if not isinstance(audio_out_ids, torch.Tensor) or audio_out_ids.numel() == 0:
+                continue
+            last_frame = audio_out_ids[:, -1:].to(text_embed.device)  # [Q, 1]
+            audio_emb = self.embed_audio_codes(last_frame)  # [1, hidden]
+            flat_embed[pos] = audio_emb[0].to(dtype=text_embed.dtype)
+        return text_embed
+
+    def make_omni_output(
+        self,
+        model_outputs: torch.Tensor | OmniOutput,
+        **kwargs: Any,
+    ) -> OmniOutput:
+        """Wrap raw decoder outputs into the :class:`OmniOutput` contract.
+
+        Mirrors the canonical Qwen3-TTS / Fish-Speech recovery pattern:
+        the runner threads per-request ``codes.audio`` into
+        ``model_intermediate_buffer`` (a list of dicts in batch order); we
+        concatenate those and trim ``text_hidden_states`` to the emitted
+        audio span. Falls back to the deprecated ``runtime_additional_information``
+        kwarg for older runners, and finally to ``audio_codes`` /
+        ``model_kwargs[audio_codes]`` for direct callers.
+        """
+        if isinstance(model_outputs, OmniOutput):
+            return model_outputs
+        hidden = model_outputs
+
+        # Primary contract: model_intermediate_buffer (Qwen3-TTS / Fish-Speech).
+        info_dicts = kwargs.get("model_intermediate_buffer")
+        if info_dicts is None:
+            info_dicts = kwargs.get("runtime_additional_information")
+        if info_dicts is None:
+            info_dicts = []
+
+        # Build a per-request audio-codes list whose length exactly equals
+        # ``len(info_dicts)`` (== batch size). Requests with no emitted codes
+        # this step get an empty tensor placeholder — that way the runner's
+        # ``to_payload_element`` (which indexes ``element[idx]`` per request)
+        # can never silently fall back to ``element[0]`` for higher slots,
+        # which was the root of the batch>1 voice-clone slot-shift bug
+        # (slot k ended up speaking slot k-1's text).
+        audio_codes_list: list[torch.Tensor] = []
+        any_nonempty = False
+        for info in info_dicts:
+            ac: torch.Tensor | None = None
+            if isinstance(info, dict):
+                codes_field = info.get("codes")
+                if isinstance(codes_field, dict):
+                    ac = codes_field.get("audio")
+                else:
+                    ac = info.get("audio_codes")
+            if isinstance(ac, torch.Tensor) and ac.numel() > 0:
+                audio_codes_list.append(ac)
+                any_nonempty = True
+            else:
+                audio_codes_list.append(torch.empty(0, dtype=torch.long))
+
+        if any_nonempty:
+            return OmniOutput(
+                text_hidden_states=hidden,
+                multimodal_outputs={"codes": {"audio": audio_codes_list}},
+            )
+
+        # Fallbacks: explicit kwarg or wrapped model_kwargs dicts (preserved for
+        # direct-API callers that don't go through the runner buffer).
+        audio_codes = kwargs.get("audio_codes")
+        if audio_codes is None:
+            for source_name in ("model_kwargs", "model_kwargs_extra"):
+                source = kwargs.get(source_name)
+                if isinstance(source, dict) and "audio_codes" in source:
+                    audio_codes = source["audio_codes"]
+                    break
+        if audio_codes is None:
+            audio_codes = torch.empty(0, dtype=torch.long)
+        return OmniOutput(
+            text_hidden_states=hidden,
+            multimodal_outputs={"codes": {"audio": audio_codes}},
+        )
+
+    # ----------------------------------------------------- model-owned sampler
+    # Opt into the AR runner's model-sampler hook (see
+    # vllm_omni/worker/gpu_ar_model_runner.py:_sample). The runner calls
+    # ``self.model.sample(logits, sampling_metadata)`` when this flag is True;
+    # we delegate to the stock vLLM sampler for now and document the
+    # per-position audio-codebook-0 dispatch as the immediate follow-up. The
+    # hook itself is what unblocks an external integrator from plugging in the
+    # final audio-side sampler without touching the runner.
+    prefer_model_sampler: bool = True
+
+    # Tell the runner's :func:`extract_multimodal_outputs` to forward the OmniOutput
+    # produced by :meth:`make_omni_output`. Without this flag the multimodal
+    # payload (audio codes) is dropped and the downstream stage receives only
+    # text hidden states.
+    have_multimodal_outputs: bool = True
+
+    def sample(self, logits: torch.Tensor, sampling_metadata: Any) -> Any:
+        """Model-owned sampler with delay-pattern audio dispatch.
+
+        Each generation step the model emits one LM token. At audio positions
+        (where the previous LM token was ``audio_bos_token_id`` or
+        ``audio_token_id``), the same hidden state also drives the fused
+        ``audio_lm_head`` to produce ``[num_codebooks, codebook_size]``
+        per-codebook logits. We argmax those per-codebook, then APPLY THE
+        UPSTREAM DELAY PATTERN:
+
+          - For the first ``num_codebooks - 1`` audio steps, force codebooks
+            ``[num_delay + 1 :]`` to ``audio_stream_bos_id``; only codebooks
+            up to ``num_delay`` are real.
+          - When ``audio_stream_eos_id`` first appears in a codebook k, force
+            all earlier codebooks to ``eos`` too and start the trailing
+            ``num_codebooks - k - 1`` "ramp-down" steps where successive
+            codebooks also get ``eos``.
+          - When all 8 codebooks have ramped down, switch the LM next-token
+            from ``audio_token_id`` to ``audio_eos_token_id`` so the engine
+            sees a stop signal at the right place.
+
+        Per-request state (``num_delay``, ``num_remaining_delays``,
+        cumulative ``audio_out_ids: [num_codebooks, T_so_far]``) lives in
+        ``self._audio_state[req_id]``; ``embed_input_ids`` reads it on the
+        next step to substitute the audio-token embedding via
+        :meth:`embed_audio_codes`.
+        """
+        sampler = getattr(self, "_stock_sampler", None)
+        if sampler is None:
+            from vllm.v1.sample.sampler import Sampler
+
+            sampler = Sampler()
+            self._stock_sampler = sampler
+
+        # Bias the text-vocab logits so the LM stays on the audio path. This
+        # forces ``audio_token_id`` at audio positions; the delay-pattern
+        # logic below later overrides the LM token to ``audio_eos_token_id``
+        # when the ramp-down completes.
+        self._apply_audio_mode_bias(logits, sampling_metadata)
+        sampler_output = sampler(logits=logits, sampling_metadata=sampling_metadata)
+
+        hidden = getattr(self, "_last_logits_hidden", None)
+        self._last_logits_hidden = None
+        if hidden is None:
+            self._last_audio_codes = None
+            return sampler_output
+
+        sampled = getattr(sampler_output, "sampled_token_ids", None)
+        if sampled is None:
+            self._last_audio_codes = None
+            return sampler_output
+        sampled_flat = sampled.reshape(-1)
+        if int(sampled_flat.numel()) != int(hidden.shape[0]):
+            self._last_audio_codes = None
+            return sampler_output
+
+        audio_token_id = int(self.config.audio_token_id)
+        audio_delay_id = int(self.config.audio_delay_token_id)
+        is_audio = (sampled_flat == audio_token_id) | (sampled_flat == audio_delay_id)
+
+        if not hasattr(self, "_audio_state"):
+            self._audio_state: dict[int, dict[str, Any]] = {}
+        if not hasattr(self, "_slot_output_len"):
+            self._slot_output_len: dict[int, int] = {}
+
+        # Skip the audio-code sampling for rows that just transitioned
+        # OUT of audio_out_bos. The hidden state at this prefill-end step is
+        # at the audio_out_bos position, where upstream's audio_out_mask is
+        # False and no audio_logits are computed. The LM bias above already
+        # forced ``audio_token_id`` at this step, and the next forward will
+        # produce a hidden at the new audio_token_id position — that is what
+        # upstream samples its first audio frame from. Also reset any stale
+        # ``_audio_state`` and ``_slot_output_len`` for the slot, so a fresh
+        # request reusing a finished slot starts clean.
+        first_after_bos = getattr(self, "_last_first_audio_after_bos", None)
+        self._last_first_audio_after_bos = None
+        num_codebooks_local = int(self.config.num_codebooks)
+        bos_local = int(self.config.audio_stream_bos_id)
+        seeded_rows: list[int] = []
+        seeded_frame: torch.Tensor | None = None
+        # HF seeds frame 0 with all-BOS at the audio_bos position and ignores
+        # the model's prediction there; the first real audio frame is sampled
+        # one step later at the audio_token_id position, so we skip the
+        # prediction at the bos row and emit the all-BOS frame instead.
+        if isinstance(first_after_bos, torch.Tensor) and first_after_bos.numel() == is_audio.shape[0]:
+            first_after_bos = first_after_bos.to(is_audio.device)
+            skip_rows = first_after_bos & is_audio
+            if bool(skip_rows.any()):
+                # Upstream's AUDIO_INIT step seeds ``audio_out_ids`` with
+                # a single all-``audio_stream_bos_id`` frame and the NEXT
+                # forward feeds ``embed_audio_codes(all_bos)`` at the new
+                # audio_out_token_idx position. Without this, our audio MLP
+                # receives a text embedding at the audio_token_id position
+                # and produces NaN logits. Seed per-slot state with the
+                # all-BOS frame here, AND emit it as the first output frame
+                # so the downstream delay-pattern revert sees the leading
+                # BOS-only frame upstream emits.
+                seeded_frame = torch.full((num_codebooks_local,), bos_local, dtype=torch.long, device=hidden.device)
+                for bi in torch.nonzero(skip_rows, as_tuple=False).reshape(-1).tolist():
+                    bi = int(bi)
+                    # Reset stale state from a finished prior request.
+                    self._slot_output_len[bi] = 0
+                    self._audio_state[bi] = {
+                        "num_delay": 0,
+                        "num_remaining_delays": None,
+                        "audio_out_ids": seeded_frame.unsqueeze(-1).clone(),  # [Q, 1]
+                    }
+                    seeded_rows.append(bi)
+                is_audio = is_audio & ~first_after_bos
+
+        if not bool(is_audio.any()) and not seeded_rows:
+            self._last_audio_codes = None
+            return sampler_output
+
+        # Map audio-position rows to their absolute batch indices. The
+        # delay-pattern + state-update loop below uses this for per-request
+        # state lookup; the RAS check between the head sample and the
+        # delay-pattern apply also reads it.
+        audio_row_indices = torch.nonzero(is_audio, as_tuple=False).reshape(-1).tolist()
+
+        # Per-codebook logits at audio positions only.
+        cb_logits = self.audio_codebook_logits(hidden, is_audio)  # [N_audio, Q, V]
+
+        # Emulate upstream's ``HiggsAudioV2DelayPatternLogitsProcessor``.
+        # The raw audio_lm_head has very large weights at index 1025
+        # (audio_stream_eos): cb1 |w[1025]|=9.25 vs mean content 1.53 (6x).
+        # Without masking, the model naturally fires 1025 at codebook 1..6
+        # with huge confidence (logit 32 vs runner-up 6) and triggers
+        # premature ramp-down at frame 2. Upstream's processor MASKS
+        # forbidden vocab BEFORE sampling so each codebook k is bound to:
+        #   - emit ``audio_stream_bos_id`` (1024) for k leading frames,
+        #   - emit content tokens [0..1023] thereafter,
+        #   - emit ``audio_stream_eos_id`` (1025) only during ramp-down.
+        bos_pre = int(self.config.audio_stream_bos_id)
+        eos_pre = int(self.config.audio_stream_eos_id)
+        for local_i, batch_i in enumerate(audio_row_indices):
+            state_pre = self._audio_state.get(int(batch_i))
+            num_delay_pre = int(state_pre.get("num_delay", 0)) if state_pre else 0
+            num_rem_pre = state_pre.get("num_remaining_delays") if state_pre else None
+            if num_rem_pre is not None:
+                # Ramp-down: mask all but EOS for the EOS-locked codebooks; the
+                # remaining ones still emit content. Locked range: codebooks
+                # [0 : num_codebooks - num_rem_pre].
+                lock_until = int(self.config.num_codebooks) - int(num_rem_pre)
+                for q in range(int(self.config.num_codebooks)):
+                    row = cb_logits[local_i, q]
+                    if q < lock_until:
+                        mask = torch.full_like(row, float("-inf"))
+                        mask[eos_pre] = row[eos_pre]
+                        cb_logits[local_i, q] = mask
+                    else:
+                        # Disallow stream specials during ramp-down content frames
+                        cb_logits[local_i, q, bos_pre] = float("-inf")
+                        cb_logits[local_i, q, eos_pre] = float("-inf")
+            else:
+                # Leading delay phase. Codebook k starts emitting real content
+                # at step k (delay = k). Steps with k > num_delay_pre are still
+                # in BOS phase — force 1024. Steps with k <= num_delay_pre are
+                # content — disallow 1024 (BOS) and 1025 (EOS).
+                for q in range(int(self.config.num_codebooks)):
+                    row = cb_logits[local_i, q]
+                    if q > num_delay_pre:
+                        mask = torch.full_like(row, float("-inf"))
+                        mask[bos_pre] = row[bos_pre]
+                        cb_logits[local_i, q] = mask
+                    else:
+                        cb_logits[local_i, q, bos_pre] = float("-inf")
+                        # Allow EOS at codebook 0 only (it gates ramp-down).
+                        # Disallow at all other content codebooks.
+                        if q != 0:
+                            cb_logits[local_i, q, eos_pre] = float("-inf")
+
+        # Sample per-codebook, following the upstream
+        # ``HiggsAudioModel._sample_audio_tokens`` pipeline byte-for-byte.
+        cb_logits_2d = cb_logits.reshape(-1, cb_logits.shape[-1])
+        codes_2d = self._sample_audio_codes_upstream(
+            cb_logits_2d,
+            temperature=1.0,
+            top_k=50,
+            top_p=0.95,
+        )
+        codes_flat = codes_2d.view(cb_logits.shape[0], cb_logits.shape[1]).to(torch.long)
+
+        # (PyTorch align): Repetition-Aware Sampling (RAS). Upstream
+        # ``HiggsAudioModel._sample_audio_tokens`` re-samples any codebook
+        # that just emitted the same token >= ras_win_max_num_repeat times
+        # in the last ras_win_len steps. This breaks codebook collapse loops
+        # — a 3B model with greedy/low-temp sampling often locks into one
+        # value otherwise.
+        ras_win_len = 7
+        ras_win_max_num_repeat = 2
+        # Codes_flat is [N_audio, Q]; we re-sample PER (audio_position, codebook)
+        # row that hits the repeat threshold over the last ras_win_len audio
+        # frames of that codebook's history (from state.audio_out_ids).
+        if codes_flat.shape[0] > 0:
+            audio_idx_to_batch = audio_row_indices  # already a list[int]
+            for local_i in range(codes_flat.shape[0]):
+                bi = int(audio_idx_to_batch[local_i])
+                req_state = self._audio_state.get(bi)
+                if req_state is None:
+                    continue
+                aoi = req_state.get("audio_out_ids")
+                if not isinstance(aoi, torch.Tensor) or aoi.shape[-1] < ras_win_len:
+                    continue
+                window = aoi[:, -ras_win_len:].to(codes_flat.device)  # [Q, W]
+                this_codes = codes_flat[local_i]  # [Q]
+                rep_count = (window == this_codes.unsqueeze(-1)).sum(dim=-1)  # [Q]
+                row_indices = torch.nonzero(rep_count >= ras_win_max_num_repeat, as_tuple=False).reshape(-1)
+                if row_indices.numel() == 0:
+                    continue
+                # Re-sample those codebooks from cb_logits directly (no top-k/top-p
+                # in RAS resample per upstream — pure softmax + multinomial).
+                resample_logits = cb_logits[local_i, row_indices].float()
+                resample_probs = resample_logits.softmax(dim=-1)
+                resampled = torch.multinomial(resample_probs, num_samples=1).squeeze(-1)
+                codes_flat[local_i, row_indices] = resampled.to(codes_flat.dtype)
+
+        # Apply the upstream delay pattern + EOS ramp-down per request.
+        num_codebooks = int(self.config.num_codebooks)
+        bos = int(self.config.audio_stream_bos_id)
+        eos_stream = int(self.config.audio_stream_eos_id)
+        audio_eos_vocab = int(getattr(self.config, "audio_eos_token_id", -1))
+
+        # ``audio_row_indices`` and ``self._audio_state`` were initialised
+        # earlier (right after ``is_audio`` was computed) so the RAS check
+        # can reference them.
+
+        # Per-request reset: detect "new request occupied this slot" by
+        # tracking the per-slot output_token_ids length across sample()
+        # calls. Within a request the length monotonically grows; when it
+        # resets to 0 (or goes BELOW the previously-seen high-water mark),
+        # the previous occupant has finished and a new request now owns
+        # the slot. Drop the stale ``_audio_state`` for that slot.
+        if not hasattr(self, "_slot_output_len"):
+            self._slot_output_len: dict[int, int] = {}
+        output_ids = getattr(sampling_metadata, "output_token_ids", None)
+        for batch_i in audio_row_indices:
+            bi = int(batch_i)
+            current_len = len(output_ids[bi]) if output_ids is not None and bi < len(output_ids) else 0
+            prior_len = self._slot_output_len.get(bi, -1)
+            # A fresh request has either never been seen (prior_len == -1
+            # and we're at step 0) or has length less than the prior
+            # high-water mark.
+            is_new_request = (prior_len > current_len) or (prior_len == -1 and current_len == 0)
+            if is_new_request and bi in self._audio_state:
+                self._audio_state.pop(bi, None)
+            self._slot_output_len[bi] = current_len
+
+        # Simplification: we DON'T override ``sampler_output.sampled_token_ids``
+        # back to ``audio_eos_token_id`` at ramp-down completion. The vLLM
+        # ``SamplerOutput.sampled_token_ids`` is plumbed through several
+        # downstream structures that may have already cached the GPU pointer;
+        # an in-place mutation here was crashing the worker with exit code
+        # None (silent CUDA error). Instead, ``max_tokens`` (set in the deploy
+        # yaml) is the stop bound while the audio span ramps down naturally.
+        new_codes_flat: list[torch.Tensor] = []
+        for local_i, batch_i in enumerate(audio_row_indices):
+            state = self._audio_state.setdefault(
+                int(batch_i),
+                {"num_delay": 0, "num_remaining_delays": None, "audio_out_ids": None},
+            )
+            num_delay: int = state["num_delay"]
+            num_remaining_delays: int | None = state["num_remaining_delays"]
+            this_codes = codes_flat[local_i].clone()  # [Q]
+
+            # Leading delay-pattern BOS pad.
+            if num_delay + 1 < num_codebooks:
+                this_codes[num_delay + 1 :] = bos
+                num_delay += 1
+
+            # Trailing eos ramp-down.
+            if num_remaining_delays is not None:
+                this_codes[: num_codebooks - num_remaining_delays] = eos_stream
+                num_remaining_delays -= 1
+            else:
+                eos_positions = (this_codes == eos_stream).nonzero(as_tuple=False).reshape(-1)
+                if eos_positions.numel() > 0:
+                    last_eos_idx = int(eos_positions[-1].item())
+                    this_codes[: last_eos_idx + 1] = eos_stream
+                    num_remaining_delays = num_codebooks - last_eos_idx - 1
+
+            if num_remaining_delays is not None and num_remaining_delays <= 0:
+                # Audio ramp-down finished. Set a per-slot termination flag
+                # so the NEXT step's bias forces audio_eos_token_id (upstream's
+                # ``next_tokens[...] = audio_eos_token_id`` override at the
+                # ramp-completion point — see modeling_higgs_audio.py:1564).
+                _ = audio_eos_vocab
+                num_delay = 0
+                num_remaining_delays = None
+                state["num_delay"] = num_delay
+                state["num_remaining_delays"] = num_remaining_delays
+                state["should_terminate"] = True
+                new_codes_flat.append(torch.full_like(this_codes, -1))
+                continue
+
+            state["num_delay"] = num_delay
+            state["num_remaining_delays"] = num_remaining_delays
+            # Append to the per-request audio_out_ids buffer.
+            if state["audio_out_ids"] is None:
+                state["audio_out_ids"] = this_codes.unsqueeze(-1).clone()
+            else:
+                state["audio_out_ids"] = torch.cat([state["audio_out_ids"], this_codes.unsqueeze(-1)], dim=-1)
+            new_codes_flat.append(this_codes)
+
+        # Reassemble the per-position codes tensor with the post-delay-pattern
+        # values; `-1` rows mark "no audio code at this position".
+        codes_full = torch.full(
+            (int(sampled_flat.numel()), num_codebooks),
+            -1,
+            dtype=torch.long,
+            device=hidden.device,
+        )
+        if new_codes_flat:
+            stacked = torch.stack(new_codes_flat, dim=0).to(hidden.device)
+            codes_full[is_audio] = stacked
+        # Emit the seeded all-BOS frame at the prefill-end (AUDIO_INIT)
+        # rows so the talker output has the leading delay-pattern frame the
+        # downstream revert helper expects.
+        if seeded_rows and seeded_frame is not None:
+            seeded_rows_idx = torch.tensor(seeded_rows, dtype=torch.long, device=hidden.device)
+            codes_full.index_copy_(
+                0,
+                seeded_rows_idx,
+                seeded_frame.unsqueeze(0).expand(len(seeded_rows), -1).to(hidden.device),
+            )
+
+        self._last_audio_codes = codes_full
+        self._postprocess_cursor = 0
+        return sampler_output
+
+    @staticmethod
+    def _sample_audio_codes_upstream(
+        logits: torch.Tensor,
+        temperature: float = 1.0,
+        top_k: int | None = None,
+        top_p: float | None = None,
+    ) -> torch.Tensor:
+        """Replicate the upstream HF ``LogitsProcessorList`` pipeline that
+        ``HiggsAudioModel._sample_audio_tokens`` calls via
+        ``logits_processor(None, next_audio_token_logits)``.
+
+        Order (matching transformers' standard ``TemperatureLogitsWarper``,
+        ``TopKLogitsWarper``, ``TopPLogitsWarper``):
+
+        1. Divide by temperature.
+        2. Top-k: keep the K largest, set others to -inf.
+        3. Top-p (nucleus): keep the smallest prefix whose cumulative
+           softmax probability is <= top_p; mask the rest. Always keep at
+           least the top-1.
+        4. Softmax → multinomial sample.
+
+        ``logits`` shape: ``[N, V]``. Returns ``[N]`` of int sampled indices.
+        """
+        x = logits.float()
+        if temperature is not None and temperature > 0.0 and temperature != 1.0:
+            x = x / temperature
+        if top_k is not None and 0 < top_k < x.shape[-1]:
+            kth = x.topk(top_k, dim=-1).values[..., -1:]
+            x = x.masked_fill(x < kth, float("-inf"))
+        if top_p is not None and 0.0 < top_p < 1.0:
+            sorted_logits, sorted_idx = x.sort(dim=-1, descending=True)
+            cumprobs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
+            sorted_mask = cumprobs > top_p
+            # Shift right by one so we always keep at least the top-1 token.
+            sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
+            sorted_mask[..., 0] = False
+            # Scatter back to original index space and apply.
+            mask = torch.zeros_like(x, dtype=torch.bool)
+            mask.scatter_(-1, sorted_idx, sorted_mask)
+            x = x.masked_fill(mask, float("-inf"))
+        if temperature is None or temperature <= 0.0:
+            return torch.argmax(x, dim=-1)
+        probs = x.softmax(dim=-1)
+        # Guard against any all-masked rows (shouldn't happen with the shift
+        # above) — fall back to argmax of original logits.
+        valid = probs.sum(dim=-1) > 0
+        sampled = torch.empty(x.shape[0], dtype=torch.long, device=x.device)
+        if valid.any():
+            sampled[valid] = torch.multinomial(probs[valid], num_samples=1).squeeze(-1)
+        if (~valid).any():
+            sampled[~valid] = torch.argmax(logits[~valid].float(), dim=-1)
+        return sampled
+
+    def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -> int:
+        """Mask non-audio tokens at audio-mode positions, in-place on ``logits``.
+
+        Heuristic: per-request, find the last token seen so far (last of
+        ``output_token_ids[i]`` if non-empty, else last of
+        ``prompt_token_ids[i]``). If that token is ``audio_bos_token_id`` or
+        ``audio_token_id``, force the next emit to be one of
+        ``{audio_token_id, audio_eos_token_id}``. This unblocks live audio
+        generation since the un-biased argmax over the 128k vocab favours
+        text-token-id ranges (in particular the global ``eos_token_id``).
+
+        We deliberately do NOT include the standard ``eos_token_id`` in the
+        allowed set — letting it through would end the whole sequence at the
+        first step. Stopping inside the audio span is the ``audio_eos`` token's
+        job; once it fires we drop the bias and the next call falls back to
+        the stock sampler.
+        """
+        if logits is None or logits.ndim != 2:
+            return 0
+        audio_bos = int(self.config.audio_bos_token_id)
+        audio_id = int(self.config.audio_token_id)
+        # (PyTorch align, corrected): upstream NEVER lets the model
+        # naturally emit ``audio_eos_token_id`` in the LM stream. The LM
+        # next-token is force-set to ``audio_out_token_idx`` every step;
+        # the audio span only ends when the delay-pattern ramp-down completes
+        # (driven by ``audio_stream_eos_id`` appearing in a codebook), at
+        # which point upstream OVERRIDES the LM token to ``audio_eos``.
+        # Our diag showed the model's ``audio_eos`` logit is very high right
+        # after the first audio step; allowing it via sampling stops the
+        # audio span with 1 frame and the rest of the LM stream goes off-rails.
+        # Force ONLY ``audio_token_id``.
+        allowed_extra: list[int] = []
+        # Walk per-request to decide which rows to mask.
+        prompt_ids = getattr(sampling_metadata, "prompt_token_ids", None)
+        output_ids = getattr(sampling_metadata, "output_token_ids", None)
+        num_rows = int(logits.shape[0])
+        biased = 0
+        # Fallback "previous token" source: read each request's LAST input
+        # position from ``self._last_step_input_ids`` using ``query_start_loc``
+        # from the forward context. The previous version of this code took
+        # ``stash_ids[-num_rows:]`` which only works for pure-decode batches
+        # (where each row contributes exactly one token to the flat input_ids).
+        # In MIXED prefill+decode batches the last N tokens of flat input_ids
+        # all belong to the LAST prefilling request, so the "prev token" for
+        # other rows came back as random tokens from inside that request's
+        # prompt — which caused the audio_out_bos detection to fail for some
+        # rows, the audio_token_id bias to skip them, and the per-row codes to
+        # collapse to -1 silently (see the batch>1 voice-clone gibberish bug).
+        stash_ids = getattr(self, "_last_step_input_ids", None)
+        stash_tail: list[int] | None = None
+        if isinstance(stash_ids, torch.Tensor) and stash_ids.numel() > 0:
+            # forward() stashed query_start_loc for us inside the forward
+            # context — get_forward_context() would fail here because we're
+            # past the with-block by the time sample() runs.
+            q_start = getattr(self, "_last_step_query_start_loc", None)
+            if isinstance(q_start, torch.Tensor) and int(q_start.numel()) == num_rows + 1:
+                # Per-request last position: q_start_loc[i+1] - 1.
+                q_start_cpu = q_start.detach().to("cpu").tolist()
+                tail_idx = [max(0, int(q_start_cpu[i + 1]) - 1) for i in range(num_rows)]
+                flat_ids = stash_ids.detach().to("cpu").tolist()
+                stash_tail = [int(flat_ids[idx]) if idx < len(flat_ids) else -1 for idx in tail_idx]
+            elif int(stash_ids.numel()) >= num_rows:
+                # Pure-decode fallback when query_start_loc is unavailable:
+                # input_ids has exactly num_rows tokens (one per request) and
+                # the slice is order-preserving.
+                stash_tail = stash_ids[-num_rows:].detach().to("cpu").tolist()
+        # Also stash a per-row "is this the first audio-mode step right
+        # after audio_out_bos" flag. Upstream NEVER samples audio codes at
+        # the audio_out_bos position (its audio_out_mask is False there);
+        # the first audio frame is sampled one step later, at the new
+        # audio_out_token_idx position. ``sample()`` reads this flag to
+        # suppress the audio-code sampling at the prefill-end step so the
+        # first audio frame is sampled from the correct hidden state.
+        first_after_bos = torch.zeros(num_rows, dtype=torch.bool, device=logits.device)
+        for i in range(num_rows):
+            prev: int | None = None
+            if output_ids is not None and i < len(output_ids):
+                hist = output_ids[i]
+                if hist:
+                    prev = int(hist[-1])
+            if prev is None and prompt_ids is not None:
+                # Prompt_token_ids may be a tensor or a list-of-lists
+                try:
+                    p_i = prompt_ids[i]
+                    if hasattr(p_i, "tolist"):
+                        p_i = p_i.tolist()
+                    if p_i:
+                        prev = int(p_i[-1])
+                except (IndexError, TypeError):
+                    prev = None
+            if prev is None and stash_tail is not None and i < len(stash_tail):
+                prev = int(stash_tail[i])
+            if prev is None:
+                continue
+            if prev not in (audio_bos, audio_id):
+                continue
+            if prev == audio_bos:
+                first_after_bos[i] = True
+            # If this slot's audio ramp finished last step, force LM token
+            # to audio_eos so the engine stops naturally. Matches upstream's
+            # ``next_tokens[...] = audio_eos_token_id`` override.
+            audio_state = getattr(self, "_audio_state", {}) or {}
+            slot_state = audio_state.get(i)
+            should_terminate = bool(isinstance(slot_state, dict) and slot_state.get("should_terminate"))
+            audio_eos_vocab = int(getattr(self.config, "audio_eos_token_id", -1))
+            if should_terminate and 0 <= audio_eos_vocab < int(logits.shape[-1]):
+                # Force audio_eos at this position; upstream override at
+                # ramp-down completion (modeling_higgs_audio.py:1564).
+                row = logits[i]
+                mask = torch.full_like(row, float("-inf"))
+                mask[audio_eos_vocab] = row[audio_eos_vocab]
+                logits[i].copy_(mask)
+                slot_state["should_terminate"] = False
+                biased += 1
+                continue
+            # Always force ``audio_token_id`` — upstream mirror.
+            allowed: set[int] = {audio_id, *allowed_extra}
+            row = logits[i]
+            mask = torch.full_like(row, float("-inf"))
+            for tok in allowed:
+                if 0 <= tok < row.shape[-1]:
+                    mask[tok] = row[tok]
+            logits[i].copy_(mask)
+            biased += 1
+        self._last_first_audio_after_bos = first_after_bos
+        return biased
+
+    has_postprocess: bool = True
+
+    def postprocess(
+        self,
+        hidden_states_slice: torch.Tensor,
+        multimodal_outputs: Any = None,
+        **req_infos: Any,
+    ) -> dict[str, Any]:
+        """Publish per-request audio codes into model_intermediate_buffer.
+
+        The runner calls this once per request in batch order. We index into
+        ``self._last_audio_codes`` (a per-batch ``[num_requests, num_codebooks]``
+        tensor produced by :meth:`sample`, one row per request slot regardless
+        of prompt span length) by a running cursor — exactly one row per
+        request. Using ``hidden_states_slice.shape[0]`` as the stride is
+        WRONG at prefill: prompt span length is hundreds of tokens, but each
+        request still contributes exactly one sampled row to ``_last_audio_codes``.
+        That bug used to drop the first audio frame at prefill-end (the
+        all-BOS seed) and, in mixed prefill+decode batches, scrambled the
+        per-request → codes mapping so batch=4 output came back shifted /
+        cross-talked across slots.
+        """
+        _ = multimodal_outputs  # consumed by the runner directly
+        codes_full = getattr(self, "_last_audio_codes", None)
+        if codes_full is None:
+            return {}
+
+        # The runner walks requests in batch order; pop from the cursor as we go.
+        # ``_last_audio_codes`` is one row per request slot (== one logits_indices
+        # row), not one row per prompt token, so advance by 1 per call.
+        cursor = int(getattr(self, "_postprocess_cursor", 0))
+        if cursor >= int(codes_full.shape[0]):
+            # Defensive: shape drift between forward and postprocess.
+            self._postprocess_cursor = 0
+            return {}
+        slice_codes = codes_full[cursor : cursor + 1]
+        self._postprocess_cursor = cursor + 1
+        # Drop placeholder rows (-1) — only emit codes for actual audio positions.
+        audio_rows = slice_codes[:, 0] >= 0
+        if not bool(audio_rows.any()):
+            return {}
+        new_codes = slice_codes[audio_rows].to(torch.int32)
+
+        # Return ONLY this step's frames. The engine's
+        # MultimodalOutputProcessor.add_multimodal_tensor accumulates per-step
+        # OmniOutput.multimodal_outputs across all steps and consolidates at
+        # request finish via torch.cat(dim=0). Returning cumulative codes here
+        # caused N(N+1)/2 duplication (prior bug: engine internal accumulation).
+        return {"codes": {"audio": new_codes}}
diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py
new file mode 100644
index 000000000..a51d692b9
--- /dev/null
+++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py
@@ -0,0 +1,332 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""TTS prompt builder + scope validators for higgs-audio v2.
+
+vllm-omni's higgs path supports two request shapes:
+
+* **Plain text -> 24 kHz speech.** :func:`build_plain_text_prompt` runs the
+  upstream HF processor with a bare ``"Generate audio following instruction."``
+  system prompt; the emitted ``input_ids`` are byte-identical to upstream.
+* **Voice clone (shallow).** :func:`build_voice_clone_prompt` runs the
+  upstream HF processor with both a target text turn and a reference
+  ``(ref_audio, ref_text)`` ICL turn; the processor returns ``input_ids``
+  pre-expanded with audio placeholders plus an ``audio_input_ids`` tensor
+  carrying the encoded reference codes (HF tokenizer encodes audio on the
+  spot — no encoder vendored in vllm-omni).
+
+Still out of scope and rejected with explicit 4xx: multi-speaker
+``[SPEAKERn]`` dialogue, ``profile:`` text-only speaker descriptions, the
+``ref_audio_in_system_message`` system-block variant, and chunked long-form
+generation.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+from typing import Any
+
+import numpy as np
+import torch
+
+__all__ = [
+    "UnsupportedInputError",
+    "MULTI_SPEAKER_TAG_PATTERN",
+    "validate_plain_text_input",
+    "build_plain_text_conversation",
+    "build_plain_text_prompt",
+    "build_voice_clone_conversation",
+    "build_voice_clone_prompt",
+    "input_ids_to_python_list",
+]
+
+
+class UnsupportedInputError(ValueError):
+    """Raised when a request asks for an out-of-scope higgs_audio_v2 feature."""
+
+
+# Matches the upstream multi-speaker SPEAKERn tag, e.g. [SPEAKER0], [SPEAKER12].
+MULTI_SPEAKER_TAG_PATTERN = re.compile(r"\[SPEAKER\d+\]", re.IGNORECASE)
+
+
+def validate_plain_text_input(text: str) -> None:
+    """Reject multi-speaker tags inside the user text body.
+
+    Phase-1 explicitly does NOT support multi-speaker dialogue. Catching the
+    pattern here means the rejection happens at the tokenizer boundary and is
+    visible to both offline (`pipeline.py`) and online (`serving_speech.py`)
+    code paths.
+    """
+    if not isinstance(text, str):
+        raise UnsupportedInputError(f"higgs_audio_v2 expects plain text input; got {type(text).__name__}")
+    if MULTI_SPEAKER_TAG_PATTERN.search(text):
+        raise UnsupportedInputError(
+            "higgs_audio_v2 v1 does not support multi-speaker [SPEAKERn] tags; received text contains a speaker tag"
+        )
+
+
+def build_plain_text_conversation(text: str) -> list[dict[str, Any]]:
+    """Build the canonical single-speaker plain-text conversation.
+
+    Uses the bare system prompt ``"Generate audio following instruction."``
+    that matches the upstream HF reference's input formatting; this exact
+    wording is required for input-token parity with the upstream processor.
+    """
+    validate_plain_text_input(text)
+    system_prompt = "Generate audio following instruction."
+    return [
+        {
+            "role": "system",
+            "content": [{"type": "text", "text": system_prompt}],
+        },
+        {
+            "role": "user",
+            "content": [{"type": "text", "text": text}],
+        },
+    ]
+
+
+def build_plain_text_prompt(
+    processor: Any,
+    text: str,
+    *,
+    sampling_rate: int = 24000,
+    return_tensors: str | None = "pt",
+) -> dict[str, Any]:
+    """Run the upstream processor's chat template on a plain-text input.
+
+    Returns the processor output dict (``input_ids`` plus any auxiliary tensors
+    such as ``attention_mask``). The serving layer passes ``input_ids`` to
+    Stage 0 as ``prompt_token_ids`` after a ``.tolist()``.
+
+    Using the upstream processor verbatim (no system-prompt rewriting) is
+    what preserves input-token parity with the HF reference.
+    """
+    conversation = build_plain_text_conversation(text)
+    inputs = processor.apply_chat_template(
+        conversation,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_dict=True,
+        sampling_rate=sampling_rate,
+        return_tensors=return_tensors,
+    )
+    if "input_ids" not in inputs:
+        raise RuntimeError(f"HiggsAudioV2 processor returned no input_ids; got keys {list(inputs.keys())!r}")
+    return inputs
+
+
+def input_ids_to_python_list(inputs: dict[str, Any]) -> list[int]:
+    """Convenience: pull a flat ``list[int]`` of token IDs from a processor output."""
+    ids = inputs["input_ids"]
+    if isinstance(ids, torch.Tensor):
+        if ids.ndim == 2 and int(ids.shape[0]) != 1:
+            raise ValueError(f"expected batch=1 prompt; got input_ids shape {tuple(ids.shape)}")
+        return ids.reshape(-1).tolist()
+    return list(ids)
+
+
+_AUDIO_OUT_TOKEN = "<|AUDIO_OUT|>"
+_AUDIO_OUT_BOS_TOKEN = "<|audio_out_bos|>"
+_AUDIO_EOS_TOKEN = "<|audio_eos|>"
+_AUDIO_DELAY_TOKEN = "<|reserved_special_token_6|>"
+_AUDIO_STREAM_BOS_ID = 1024
+_AUDIO_STREAM_EOS_ID = 1025
+
+
+def _build_delay_pattern(codes: torch.Tensor) -> torch.Tensor:
+    """Apply the upstream delay-pattern wrap to a ``[num_codebooks, T]`` code tensor.
+
+    Mirrors ``HiggsAudioV2Processor.build_delay_pattern``: prepend BOS, append
+    EOS, then arrange in a triangular pattern that stretches the sequence by
+    ``num_codebooks - 1`` frames so codebook k starts emitting real codes at
+    frame k.
+    """
+    num_codebooks, seq_len = codes.shape
+    bos = codes.new_full((num_codebooks, 1), _AUDIO_STREAM_BOS_ID)
+    eos = codes.new_full((num_codebooks, 1), _AUDIO_STREAM_EOS_ID)
+    wrapped = torch.cat([bos, codes, eos], dim=1)
+    wrapped_len = wrapped.shape[1]
+    new_seq_len = wrapped_len + num_codebooks - 1
+
+    output = torch.ones((1, num_codebooks, new_seq_len), dtype=codes.dtype, device=codes.device)
+    bos_mask = torch.tril(output, -1) > 0
+    eos_mask = torch.triu(output, wrapped_len) > 0
+    data_mask = ~(bos_mask | eos_mask)
+    output[bos_mask] = _AUDIO_STREAM_BOS_ID
+    output[data_mask] = wrapped.reshape(-1)
+    output[eos_mask] = _AUDIO_STREAM_EOS_ID
+    return output[0]
+
+
+_ENCODER_CACHE: Any | None = None
+_K2_OMNIVOICE_REPO = "k2-fsa/OmniVoice"
+_K2_OMNIVOICE_SUBDIR = "audio_tokenizer"
+
+
+def _load_audio_tokenizer():
+    """Load the higgs-audio v2 codec via HF ``HiggsAudioV2TokenizerModel``.
+
+    The boson-ai standalone tokenizer repo (``bosonai/higgs-audio-v2-tokenizer``)
+    ships a ``model.safetensors`` that is actually a copy of the 3B talker LM,
+    not the codec — pointing HF's ``from_pretrained`` at it yields an
+    encoder/decoder with randomly initialized structural weights, producing
+    noise codes at encode time. The k2-fsa OmniVoice repo bundles the same
+    codec weights repackaged under
+    ``audio_tokenizer/{config.json,model.safetensors,preprocessor_config.json}``
+    with key naming aligned to ``HiggsAudioV2TokenizerModel``'s class layout
+    (``acoustic_encoder.*`` / ``acoustic_decoder.*`` / ``quantizer.quantizers.*``
+    / ``fc/fc2`` / ``semantic_model.*``), so HF can load it directly.
+
+    Only the ``audio_tokenizer/`` subdirectory (~806 MB) is downloaded.
+    """
+    from huggingface_hub import snapshot_download
+    from transformers import HiggsAudioV2TokenizerModel
+
+    repo_path = snapshot_download(
+        _K2_OMNIVOICE_REPO,
+        allow_patterns=[f"{_K2_OMNIVOICE_SUBDIR}/*"],
+    )
+    audio_tokenizer_dir = os.path.join(repo_path, _K2_OMNIVOICE_SUBDIR)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = HiggsAudioV2TokenizerModel.from_pretrained(audio_tokenizer_dir).to(device)
+    return model.eval()
+
+
+def _encode_ref_audio_codes(
+    wav: np.ndarray,
+    sr: int,
+) -> torch.Tensor:
+    """Encode a single ref clip to codec codes via HF ``HiggsAudioV2TokenizerModel``.
+
+    Returns shape ``[num_codebooks, T_raw]`` (before BOS/EOS + delay-pattern wrap).
+    """
+    import torchaudio
+
+    global _ENCODER_CACHE
+    if _ENCODER_CACHE is None:
+        try:
+            _ENCODER_CACHE = _load_audio_tokenizer()
+        except ImportError as exc:
+            raise RuntimeError(
+                "higgs_audio_v2 voice clone needs `transformers>=5.3.0` "
+                "(which ships the HiggsAudioV2TokenizerModel class). "
+                "Install via `pip install -U 'transformers>=5.3.0'`."
+            ) from exc
+
+    target_sr = int(getattr(_ENCODER_CACHE.config, "sample_rate", 24000))
+    wav_t = torch.as_tensor(wav, dtype=torch.float32).reshape(-1)
+    if sr and sr != target_sr:
+        wav_t = torchaudio.functional.resample(wav_t, sr, target_sr)
+    # HF encode() expects (batch=1, channels=1, num_samples).
+    input_values = wav_t.unsqueeze(0).unsqueeze(0).to(_ENCODER_CACHE.device)
+    with torch.inference_mode():
+        codes = _ENCODER_CACHE.encode(input_values, return_dict=False)
+    # codes shape: (batch=1, num_quantizers, codes_length) -> (num_quantizers, T).
+    if codes.dim() == 3:
+        codes = codes.squeeze(0)
+    return codes.detach().to("cpu").long()
+
+
+def build_voice_clone_conversation(
+    text: str,
+    ref_text: str,
+) -> list[dict[str, Any]]:
+    """ChatML conversation for voice clone (HF jinja-template-compatible).
+
+    The assistant turn is a list with a single ``{"type": "audio"}`` content
+    block — no ``"audio"``/``"url"``/``"path"`` key, so ``apply_chat_template``
+    won't try to extract audio data (that path would collide with our explicit
+    audio encoder). The template renders the assistant block as the literal
+    ``<|audio_out_bos|><|AUDIO_OUT|><|audio_eos|>``; we expand the single
+    ``<|AUDIO_OUT|>`` token to ``N × audio_token + (num_codebooks-1) × delay_token``
+    in :func:`build_voice_clone_prompt` to match the reference clip's frame count.
+    """
+    validate_plain_text_input(text)
+    validate_plain_text_input(ref_text)
+    return [
+        {"role": "system", "content": "Generate audio following instruction."},
+        {"role": "user", "content": ref_text},
+        {"role": "assistant", "content": [{"type": "audio"}]},
+        {"role": "user", "content": text},
+    ]
+
+
+def build_voice_clone_prompt(
+    processor: Any,
+    text: str,
+    ref_audio_wav: np.ndarray | torch.Tensor,
+    ref_audio_sr: int,
+    ref_text: str,
+    *,
+    return_tensors: str | None = "pt",
+) -> dict[str, Any]:
+    """Build a voice-clone prompt using the upstream boson audio tokenizer.
+
+    Returns a dict carrying:
+      - ``prompt_token_ids``: ``list[int]`` — input_ids with the ref-audio
+        ``<|AUDIO_OUT|>`` + ``<|reserved_special_token_6|>`` placeholders
+        already expanded to match the encoded reference clip's frame count.
+      - ``audio_input_ids``: ``Tensor[T_frames, num_codebooks]`` — encoded
+        reference codes with BOS/EOS + delay pattern, transposed to ``[T, Q]``
+        for parity with the HF processor output shape.
+      - ``audio_input_ids_mask``: ``Tensor[T_frames]`` — all-True bool mask.
+    """
+    if isinstance(ref_audio_wav, torch.Tensor):
+        wav = ref_audio_wav.detach().to("cpu").float().reshape(-1).numpy()
+    else:
+        wav = np.asarray(ref_audio_wav, dtype=np.float32).reshape(-1)
+
+    # 1. Encode via the upstream boson tokenizer (real weights from model.pth).
+    codes_qt = _encode_ref_audio_codes(wav, ref_audio_sr)  # [num_codebooks, T_raw]
+    if codes_qt.ndim == 3:
+        codes_qt = codes_qt[0]
+    num_codebooks = int(codes_qt.shape[0])
+
+    # 2. Apply BOS/EOS + delay-pattern wrap.
+    audio_input_ids = _build_delay_pattern(codes_qt)  # [num_codebooks, T_full]
+    T_full = int(audio_input_ids.shape[1])
+
+    # 3. Render the chat template, then expand the single <|AUDIO_OUT|> marker
+    #    in the assistant turn into the full placeholder block. After expansion
+    #    the text contains exactly T_full audio-mask positions, matching the
+    #    audio_input_ids time axis.
+    conversation = build_voice_clone_conversation(text, ref_text)
+    rendered = processor.tokenizer.apply_chat_template(
+        conversation,
+        add_generation_prompt=True,
+        tokenize=False,
+    )
+    n_delay = num_codebooks - 1
+    n_audio = T_full - n_delay
+    if n_audio < 0:
+        raise RuntimeError(f"ref clip too short ({T_full} frames) for delay pattern with num_codebooks={num_codebooks}")
+    placeholders = _AUDIO_OUT_TOKEN * n_audio + _AUDIO_DELAY_TOKEN * n_delay
+    expanded_assistant = f"{_AUDIO_OUT_BOS_TOKEN}{placeholders}{_AUDIO_EOS_TOKEN}"
+    if _AUDIO_OUT_TOKEN not in rendered:
+        raise RuntimeError(
+            f"Voice-clone chat-template render is missing the assistant "
+            f"audio placeholder marker. conversation={conversation!r}"
+        )
+    rendered = rendered.replace(_AUDIO_OUT_TOKEN, expanded_assistant, 1)
+
+    # 4. Tokenize the rendered prompt. ``apply_chat_template`` already emits
+    #    ``<|begin_of_text|>``, so disable add_special_tokens.
+    encoded = processor.tokenizer(
+        rendered,
+        add_special_tokens=False,
+        return_tensors=return_tensors,
+    )
+    prompt_token_ids = encoded["input_ids"]
+    if isinstance(prompt_token_ids, torch.Tensor):
+        prompt_token_ids = prompt_token_ids.reshape(-1).tolist()
+
+    # 5. Transpose codes to [T_full, num_codebooks] for HF parity.
+    audio_input_ids_t = audio_input_ids.transpose(0, 1).contiguous().to(torch.long)
+    audio_input_ids_mask = torch.ones(T_full, dtype=torch.bool)
+
+    return {
+        "prompt_token_ids": prompt_token_ids,
+        "audio_input_ids": audio_input_ids_t,
+        "audio_input_ids_mask": audio_input_ids_mask,
+    }
diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/pipeline.py b/vllm_omni/model_executor/models/higgs_audio_v2/pipeline.py
new file mode 100644
index 000000000..fb6832569
--- /dev/null
+++ b/vllm_omni/model_executor/models/higgs_audio_v2/pipeline.py
@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""higgs-audio v2 pipeline: Talker (text -> 8-codebook codec) -> Code2Wav (codec -> 24 kHz PCM)."""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+_PROC = "vllm_omni.model_executor.stage_input_processors.higgs_audio_v2"
+
+HIGGS_AUDIO_V2_PIPELINE = PipelineConfig(
+    model_type="higgs_audio_v2",
+    model_arch="HiggsAudioV2ForConditionalGeneration",
+    hf_architectures=("HiggsAudioV2ForConditionalGeneration",),
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="higgs_audio_v2",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            owns_tokenizer=True,
+            engine_output_type="latent",
+            async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"),
+            sampling_constraints={
+                "detokenize": False,
+                # Two stop signals for Stage-0:
+                #   * 128009 = standard LM eos_token_id (sequence-level stop)
+                #   * 128012 = audio_eos_token_id (forced by the talker's
+                #     _apply_audio_mode_bias at audio-ramp completion; matches
+                #     upstream HiggsAudioModel._sample_audio_tokens override).
+                "stop_token_ids": [128009, 128012],
+            },
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="code2wav",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="audio",
+            engine_output_type="audio",
+            model_arch="HiggsAudioV2Code2WavForConditionalGeneration",
+            sync_process_input_func=f"{_PROC}.talker2code2wav",
+            sampling_constraints={"detokenize": True},
+        ),
+    ),
+)
diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py
index 3ed3e8f89..03e150409 100644
--- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py
+++ b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py
@@ -84,7 +84,7 @@ from vllm.multimodal.processing import (
     PromptUpdateDetails,
 )
 from vllm.sequence import IntermediateTensors
-from vllm.transformers_utils.tokenizer import get_tokenizer
+from vllm.tokenizers import get_tokenizer
 from vllm.utils.tensor_schema import TensorSchema
 from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py
index fed27c2c5..a640ade3a 100644
--- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py
+++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py
@@ -549,13 +549,17 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
         # do_sample=True here without also routing vLLM's sampled token back into
         # this gate check would cause the two to diverge and corrupt KV-cache state.
         self.global_sampler = MiMoSampler(do_sample=False, temperature=0.6, top_p=0.95)
-        # local_sampler drives audio-code generation inside local_forward.  It is
-        # entirely internal and NOT subject to vLLM's SamplingParams, so stochastic
-        # sampling is safe here and is required to produce natural, varied speech.
-        # Setting do_sample=True also disables the CUDA-graph path (use_cg gate in
-        # local_forward checks `do_sample is False`), preventing MiMoLocalSamplerTensor
-        # from silently forcing argmax even when temperature > 0.
-        self.local_sampler = MiMoSampler(do_sample=True, temperature=0.9, top_p=0.95)
+        # local_sampler drives audio-code generation inside local_forward.  Keep
+        # it greedy (do_sample=False) so the CUDA-graph path (use_cg gate in
+        # local_forward) stays active AND so the audio codes — and therefore the
+        # `new_audio_emb` written into `_cached_new_audio_emb_by_req` — are
+        # deterministic.  That cache is fed back into `inputs_embeds` on the next
+        # decode step (see `_prepare_multimodal_embeddings_with_cache`), so any
+        # stochasticity here propagates into subsequent *text* logits via the
+        # audio-embedding feedback path and destabilises text continuations even
+        # though `global_sampler` is greedy.  Voice diversity must be tackled in
+        # the codec/vocoder path (stage-1), not by randomising local_sampler.
+        self.local_sampler = MiMoSampler(do_sample=False, temperature=0.9, top_p=0.95)
         self.removed_tokens = None
 
         self.speech_vocab_sizes = config.parsed_speech_vocab_sizes()
@@ -818,7 +822,7 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
             device=tokens_device,
         )
         if local_sampler is None:
-            local_sampler = MiMoSampler(do_sample=True, temperature=0.9, top_p=0.95)
+            local_sampler = MiMoSampler(do_sample=False, temperature=0.9, top_p=0.95)
 
         past_key_values = DynamicCache()
         for t in range(delay_iters):
@@ -864,7 +868,7 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
         local_sampler: MiMoSampler | None = None,
     ):
         if local_sampler is None:
-            local_sampler = MiMoSampler(do_sample=True, temperature=0.9, top_p=0.95)
+            local_sampler = MiMoSampler(do_sample=False, temperature=0.9, top_p=0.95)
 
         b = int(local_embeds.shape[0])
         use_cg = (local_sampler.do_sample is None or local_sampler.do_sample is False) and bool(
diff --git a/vllm_omni/model_executor/models/mimo_audio/pipeline.py b/vllm_omni/model_executor/models/mimo_audio/pipeline.py
index 70d14ef78..126c90176 100644
--- a/vllm_omni/model_executor/models/mimo_audio/pipeline.py
+++ b/vllm_omni/model_executor/models/mimo_audio/pipeline.py
@@ -39,6 +39,7 @@ MIMO_AUDIO_PIPELINE = PipelineConfig(
             owns_tokenizer=True,
             engine_output_type="latent",
             async_chunk_process_next_stage_input_func=(f"{_PROC}.llm2code2wav_async_chunk"),
+            custom_process_next_stage_input_func=f"{_PROC}.llm2code2wav_full_payload",
             sampling_constraints={
                 "detokenize": True,
                 # Stop once the speech/text interleaved span ends. Code2Wav
@@ -55,7 +56,8 @@ MIMO_AUDIO_PIPELINE = PipelineConfig(
             final_output=True,
             final_output_type="audio",
             engine_output_type="audio",
-            sync_process_input_func=f"{_PROC}.llm2code2wav",
+            custom_process_input_func=f"{_PROC}.llm2code2wav",
+            sync_process_input_func=f"{_PROC}.llm2code2wav_token_only",
             sampling_constraints={"detokenize": False},
         ),
     ),
diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py
index 0111f8069..4e3da735a 100644
--- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py
+++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py
@@ -126,15 +126,21 @@ class MingFlashOmniForConditionalGeneration(
             self.make_empty_intermediate_tensors = self.thinker.make_empty_intermediate_tensors
 
         elif self.model_stage == "imagegen":
-            # TODO: Implement image generator stage
+            # Image generation is a separate diffusion stage; it does not run
+            # through this AR wrapper. See
+            # `vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py`
+            # and `model_arch: MingImagePipeline` in the deploy YAML.
             raise NotImplementedError(
-                "Image generation stage is not yet implemented. Please use model_stage='thinker' for now."
+                "Image generation stage is not implemented in this AR wrapper. "
+                "Use the DIFFUSION execution_type with model_arch: MingImagePipeline "
+                "for the imagegen stage (see vllm_omni/deploy/ming_flash_omni_image.yaml)."
             )
 
         else:
             raise ValueError(
                 f"Invalid model_stage: {self.model_stage!r}. Must be one of: 'thinker', 'imagegen'. "
-                f"For the talker stage, use MingFlashOmniTalkerForConditionalGeneration directly."
+                f"For the talker stage, use MingFlashOmniTalkerForConditionalGeneration directly. "
+                f"For image generation use stage_type: diffusion with model_arch: MingImagePipeline."
             )
 
     def forward(
@@ -177,9 +183,43 @@ class MingFlashOmniForConditionalGeneration(
         raise NotImplementedError("get_mrope_input_positions not available on current stage")
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stripped = ((name.removeprefix("thinker."), value) for name, value in weights)
-        thinker_loaded = self.thinker.load_weights(stripped)
-        return add_prefix_to_loaded_weights(thinker_loaded, "thinker")
+        loaded_weights: set[str] = set()
+        thinker_weights: list[tuple[str, torch.Tensor]] = []
+
+        for name, value in weights:
+            if name.startswith("thinker."):
+                thinker_weights.append((name, value))
+            elif name.startswith("imagegen."):
+                # Imagegen weights are loaded by MingImagePipeline (diffusion
+                # stage) from its own checkpoint subfolders; silently drop
+                # them here so the AR wrapper does not try to load them.
+                continue
+            elif name.startswith("talker."):
+                # Talker weights are loaded by
+                # MingFlashOmniTalkerForConditionalGeneration directly; drop.
+                continue
+            else:
+                # Weights without prefix go to thinker by default
+                thinker_weights.append((name, value))
+
+        if self.model_stage == "thinker" and thinker_weights:
+            stripped = ((name.removeprefix("thinker."), value) for name, value in thinker_weights)
+            thinker_loaded = self.thinker.load_weights(stripped)
+            thinker_loaded = add_prefix_to_loaded_weights(thinker_loaded, "thinker")
+            loaded_weights.update(thinker_loaded)
+
+            # ``query_tokens_dict.*`` parameters are pre-loaded inside
+            # ``MingFlashOmniThinker.__init__`` from
+            # ``<model>/mlp/model.safetensors`` (not the main shard index),
+            # so vllm's default loader does not see them via the standard
+            # ``load_weights`` path. Report them as loaded here to satisfy the
+            # post-load completeness check.
+            query_tokens_module = getattr(self.thinker, "query_tokens_dict", None)
+            if query_tokens_module is not None:
+                for scale_name, _param in query_tokens_module.items():
+                    loaded_weights.add(f"thinker.query_tokens_dict.{scale_name}")
+
+        return loaded_weights
 
     def get_mm_mapping(self) -> MultiModelKeys:
         return MultiModelKeys.from_string_field(
diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py
index bde7477b9..0a9ce0f5c 100644
--- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py
+++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py
@@ -6,6 +6,7 @@
 
 """Ming-flash-omni-2.0 Thinker stage implementation (multimodal understanding)."""
 
+import os
 from collections.abc import Iterable, Iterator, Mapping, Sequence
 from typing import Annotated, Any
 
@@ -113,7 +114,9 @@ class MingFlashOmniThinkerProcessingInfo(Qwen2VLProcessingInfo):
         return 1
 
     def get_supported_mm_limits(self) -> Mapping[str, int | None]:
-        return {"image": None, "video": None, "audio": None}
+        # ``img2img`` is the serving-layer alias for a reference image in the
+        # image-gen path; treated as a regular image by ``_MingMultiModalDataParser``.
+        return {"image": None, "video": None, "audio": None, "img2img": 1}
 
     def get_mm_max_tokens_per_item(
         self,
@@ -137,7 +140,14 @@ class MingFlashOmniThinkerProcessingInfo(Qwen2VLProcessingInfo):
 
         return mm_max_tokens
 
-    def get_feature_extractor(self, **kwargs: object) -> MingWhisperFeatureExtractor:
+    def get_feature_extractor(self, **kwargs: object):
+        """Return the audio feature extractor from the processor.
+
+        The processor may be loaded via trust_remote_code paths that return
+        a stock transformers ``WhisperFeatureExtractor`` rather than
+        vllm-omni's subclass ``MingWhisperFeatureExtractor``. We accept both
+        as long as the caller-needed attributes (``sampling_rate``) exist.
+        """
         hf_processor = self.get_hf_processor(**kwargs)
         feature_extractor = hf_processor.audio_processor
         assert isinstance(feature_extractor, MingWhisperFeatureExtractor)
@@ -145,13 +155,56 @@ class MingFlashOmniThinkerProcessingInfo(Qwen2VLProcessingInfo):
 
     def get_data_parser(self):
         feature_extractor = self.get_feature_extractor()
-        return MultiModalDataParser(
+        return _MingMultiModalDataParser(
             target_sr=feature_extractor.sampling_rate,
             target_channels=self.get_target_channels(),
             expected_hidden_size=self._get_expected_hidden_size(),
         )
 
 
+class _MingMultiModalDataParser(MultiModalDataParser):
+    """Rewrite ``img2img`` → ``image`` in-place so the thinker embeds the ref.
+
+    Matches Ming's upstream (``processing_bailingmm2.py::__call__``): when
+    ``images`` is provided the image_processor runs → pixel_values + image_grid_thw,
+    and ``_expand_image_tokens`` expands ``<IMAGE>`` placeholders in the text
+    to ``<imagePatch>*N``. The same PIL then ALSO goes to the diffusion stage
+    as ``image_gen_pixel_values_reference``.
+
+    On vllm-omni's serving path serving_chat emits
+    ``multi_modal_data={"img2img": PIL}`` (Bagel convention). We mutate the
+    caller's mm_data dict here — moving the PIL from the ``img2img`` key into
+    the ``image`` key — so that:
+
+      * vLLM's renderer validator sees consistent keys in
+        ``mm_data`` and the parser output (``mm_data_items["image"]``)
+      * The thinker-side image path (``_call_hf_processor`` reads
+        ``mm_data["images"]``, ``_get_prompt_updates`` replaces ``<IMAGE>``
+        placeholders, ``_get_mm_fields_config`` declares image field layout)
+        all run normally with the ref PIL acting as a regular image input
+      * ``thinker2imagegen`` still finds the ref PIL — but now under
+        ``multi_modal_data["image"]`` rather than ``["img2img"]``.
+
+    A companion prepend in ``_apply_hf_processor_main`` ensures the prompt
+    text carries a matching ``<IMAGE>`` placeholder when needed.
+    """
+
+    def parse_mm_data(self, mm_data):
+        if "img2img" in mm_data:
+            img2img = mm_data.pop("img2img")
+            if not isinstance(img2img, list):
+                img2img = [img2img]
+            existing = mm_data.get("image")
+            if existing is None:
+                existing_list: list[Any] = []
+            elif isinstance(existing, list):
+                existing_list = list(existing)
+            else:
+                existing_list = [existing]
+            mm_data["image"] = existing_list + img2img
+        return super().parse_mm_data(mm_data)
+
+
 class MingFlashOmniThinkerDummyInputsBuilder(BaseDummyInputsBuilder[MingFlashOmniThinkerProcessingInfo]):
     def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
         num_images = mm_counts.get("image", 0)
@@ -185,6 +238,8 @@ class MingFlashOmniThinkerDummyInputsBuilder(BaseDummyInputsBuilder[MingFlashOmn
 
         audio_length = int(audio_duration * sample_rate)
 
+        num_img2img = mm_counts.get("img2img", 0)
+
         mm_data: MultiModalDataDict = {
             "image": self._get_dummy_images(
                 width=image_width,
@@ -200,6 +255,13 @@ class MingFlashOmniThinkerDummyInputsBuilder(BaseDummyInputsBuilder[MingFlashOmn
             "audio": [(np.random.randn(audio_length).astype(np.float32), sample_rate) for _ in range(num_audios)],
         }
 
+        if num_img2img > 0:
+            mm_data["img2img"] = self._get_dummy_images(
+                width=image_width,
+                height=image_height,
+                num_images=num_img2img,
+            )
+
         return mm_data
 
 
@@ -396,6 +458,61 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO
     ) -> bool:
         return False
 
+    def _apply_hf_processor_main(
+        self,
+        prompt: str | list[int],
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        tokenization_kwargs: Mapping[str, object],
+        *,
+        enable_hf_prompt_update: bool,
+    ):
+        """Apply Ming-specific prompt rewrites before delegating to vllm's
+        default text / multimodal processing pipeline.
+
+        This is the single point in the base MM-processor call graph where we
+        still have BOTH the full prompt and the request-level
+        ``hf_processor_mm_kwargs`` together. Downstream, vllm's MM caching
+        design splits the work into
+        :meth:`_apply_hf_processor_text_only` (prompt, empty kwargs) and
+        :meth:`_apply_hf_processor_mm_only` (dummy text, real kwargs), so a
+        per-request flag like ``modalities`` cannot reach the real
+        tokenization from inside :meth:`_call_hf_processor`.
+
+        For Ming-flash-omni-2.0, image-generation requests carry
+        ``mm_processor_kwargs["modalities"]`` set to ``["image"]`` (t2i) or
+        ``["img2img"]`` (image edit). We expand the prompt here with the
+        ``<image><imagePatch>*N</image>`` query-token block so the expanded
+        form flows through both the tokenization and (eventual) MM paths.
+        For img2img, we also prepend ``<IMAGE>`` so the thinker-side prompt
+        replacement can locate the reference-image placeholder.
+        """
+        modalities = hf_processor_mm_kwargs.get("modalities") or []
+        is_image_gen = "image" in modalities or "img2img" in modalities
+        if isinstance(prompt, str) and is_image_gen:
+            from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import (
+                DEFAULT_NUM_QUERY_TOKENS,
+                maybe_expand_image_gen_prompt,
+            )
+
+            # img2img: prepend ``<IMAGE>`` so the ref-image placeholder survives
+            # MM caching (when the ref image is cache-warm, the thinker sees 0
+            # missing items and cannot inject the placeholder itself).
+            if "img2img" in modalities:
+                prompt = "<IMAGE>" + prompt
+
+            ig = getattr(self.info.ctx.model_config.hf_config, "image_gen_config", None)
+            num_query_tokens = getattr(ig, "num_query_tokens", DEFAULT_NUM_QUERY_TOKENS)
+            prompt = maybe_expand_image_gen_prompt(prompt, num_query_tokens=int(num_query_tokens))
+
+        return super()._apply_hf_processor_main(
+            prompt=prompt,
+            mm_items=mm_items,
+            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
+            tokenization_kwargs=tokenization_kwargs,
+            enable_hf_prompt_update=enable_hf_prompt_update,
+        )
+
     def _call_hf_processor(
         self,
         prompt: str,
@@ -408,6 +525,11 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO
         We call the image/audio sub-processors directly (instead of going
         through `MingFlashOmniProcessor.__call__`) so that the high-level
         placeholder tokens remain **unexpanded** in the tokenized output.
+
+        NOTE: Image-gen prompt expansion is handled in
+        :meth:`_apply_hf_processor_main` — by the time we get here, vllm's
+        MM caching has already split text/mm processing and ``mm_kwargs``
+        is empty on the text-only branch.
         """
         hf_processor = self.info.get_hf_processor()
         tokenizer = self.info.get_tokenizer()
@@ -425,6 +547,8 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO
 
         videos = mm_data.get("videos", None)
         if videos is not None:
+            # TODO: ``videos=`` on image_processor is deprecated since
+            # transformers v4.57 (removed in v5); migrate to Qwen2VLVideoProcessor.
             video_outputs = hf_processor.image_processor(
                 images=None,
                 videos=videos,
@@ -487,8 +611,13 @@ class MingFlashOmniThinkerForConditionalGeneration(
             return "<VIDEO>"
         elif modality.startswith("audio"):
             return "<AUDIO>"
+        elif modality == "img2img":
+            # Ming's img2img ref image is consumed only by the diffusion stage
+            # (``thinker2imagegen`` reads the raw PIL). No text-side placeholder
+            # is emitted for the thinker.
+            return None
 
-        raise ValueError("Only image, video, or audio modality is supported")
+        raise ValueError("Only image, video, audio, or img2img modality is supported")
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
@@ -548,10 +677,79 @@ class MingFlashOmniThinkerForConditionalGeneration(
             )
         logger.info("Initialized WhisperAudioEncoder and AudioProjector")
 
+        # Image-generation learnable query tokens. Stored as a ParameterDict
+        # keyed by scale name (e.g. "16x16" -> [256, hidden_size]) to mirror
+        # the upstream Ming layout in ``mlp/query_tokens_dict.<scale>``. These
+        # are substituted in at ``<imagePatch>`` token positions during
+        # text-to-image requests; image-comprehension requests continue to
+        # fill those positions with vision features as before.
+        self.query_tokens_dict = nn.ParameterDict()
+        self._load_image_gen_query_tokens(vllm_config.model_config.model)
+
         # Expose interfaces
         self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors
 
-        logger.info("MingFlashOmniThinker initialized with vision and audio towers")
+        logger.info(
+            f"MingFlashOmniThinker initialized with: "
+            f"vision={'yes' if self.vision else 'no'}, "
+            f"audio={'yes' if self.audio else 'no'}, "
+            f"image_gen_scales={list(self.query_tokens_dict.keys())}"
+        )
+
+    def _load_image_gen_query_tokens(self, model_path: str) -> None:
+        """Load learnable query tokens used for image generation.
+
+        Expected layout (observed on inclusionAI/Ming-flash-omni-2.0):
+            <model_path>/mlp/model.safetensors
+                key ``query_tokens_dict.<scale>`` -> [num_tokens, hidden_size]
+                e.g. ``query_tokens_dict.16x16`` -> [256, 4096]
+
+        Missing files are logged as a warning and image-gen is disabled
+        (text-only inference still works).
+        """
+        from pathlib import Path
+
+        p = Path(model_path) / "mlp" / "model.safetensors"
+        if not p.exists() and not os.path.exists(model_path):
+            from vllm_omni.model_executor.model_loader.weight_utils import (
+                download_weights_from_hf_specific,
+            )
+
+            try:
+                local_path = download_weights_from_hf_specific(model_path, None, ["mlp/*"])
+                p = Path(local_path) / "mlp" / "model.safetensors"
+            except Exception:
+                logger.exception("[MingFlashOmniThinker] failed to download mlp/ from %s", model_path)
+        if not p.exists():
+            logger.warning(
+                "[MingFlashOmniThinker] mlp/model.safetensors not found at %s; "
+                "text-to-image will be disabled (text-only inference still works).",
+                p,
+            )
+            return
+        try:
+            from safetensors.torch import load_file  # type: ignore
+
+            state = load_file(str(p))
+        except Exception:
+            logger.exception("[MingFlashOmniThinker] failed to read %s", p)
+            return
+        loaded = 0
+        for key, tensor in state.items():
+            if not key.startswith("query_tokens_dict."):
+                continue
+            scale_name = key.split(".", 1)[1]
+            param = nn.Parameter(tensor.clone(), requires_grad=False)
+            self.query_tokens_dict[scale_name] = param
+            loaded += 1
+            logger.info(
+                "[MingFlashOmniThinker] loaded query_tokens_dict[%s] shape=%s dtype=%s",
+                scale_name,
+                tuple(tensor.shape),
+                tensor.dtype,
+            )
+        if loaded == 0:
+            logger.warning("[MingFlashOmniThinker] no query_tokens_dict.* keys in mlp/model.safetensors")
 
     def extract_image_feature(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor:
         """Extract and project image features.
@@ -732,14 +930,86 @@ class MingFlashOmniThinkerForConditionalGeneration(
     ) -> torch.Tensor:
         inputs_embeds = self.language_model.model.word_embeddings(input_ids)
 
-        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
-            return inputs_embeds
+        # Step 1: if vision/audio features were extracted, merge them into
+        # inputs_embeds at their <patch> positions first. This fills ref-image
+        # and other normal vision-input positions with real features.
+        if multimodal_embeddings is not None and len(multimodal_embeddings) > 0:
+            assert is_multimodal is not None, "`is_multimodal` mask required when `multimodal_embeddings` provided"
+            inputs_embeds = _merge_multimodal_embeddings(
+                inputs_embeds=inputs_embeds,
+                multimodal_embeddings=multimodal_embeddings,
+                is_multimodal=is_multimodal,
+            )
 
-        assert is_multimodal is not None, "`is_multimodal` mask required when `multimodal_embeddings` provided"
-        return _merge_multimodal_embeddings(
-            inputs_embeds=inputs_embeds,
-            multimodal_embeddings=multimodal_embeddings,
-            is_multimodal=is_multimodal,
+        # Step 2: image-gen query tokens are appended at the tail of the prompt
+        # as a contiguous block of ``<imagePatch>`` tokens. Inject the learnable
+        # ``query_tokens_dict`` entries at those positions. Running AFTER step 1
+        # is safe because the query-token block is disjoint from ref-image
+        # patch positions (ref patches are earlier in the prompt, filled by
+        # step 1; query tokens occupy the trailing block and get overwritten
+        # here). Injecting always — including when multimodal embeddings are
+        # present — matches Ming's img2img-edit flow where BOTH a real image
+        # input AND the learnable query tokens need to reach the LLM.
+        if len(self.query_tokens_dict) > 0:
+            self._maybe_inject_image_gen_query_embeds(input_ids, inputs_embeds)
+        return inputs_embeds
+
+    def _maybe_inject_image_gen_query_embeds(
+        self,
+        input_ids: torch.Tensor,
+        inputs_embeds: torch.Tensor,
+    ) -> None:
+        """Overwrite the trailing ``<imagePatch>`` block with learnable query tokens.
+
+        The image-gen query-token block is appended at the tail of the prompt
+        by ``maybe_expand_image_gen_prompt`` as literal text
+        ``<image><imagePatch>*N</image>``. After tokenization:
+
+            ... <image_start_token> <imagePatch>*N <image_end_token>
+
+        is the exact trailing pattern. We verify that signature before
+        injecting — without the check, a pure vision-understanding request
+        whose final ref-image block ends in N or more ``<imagePatch>`` tokens
+        would be misidentified as image-gen and its vision features silently
+        overwritten. Decode-phase calls (``input_ids`` is one next-token)
+        naturally fail the signature check and return early.
+        """
+        if input_ids.numel() == 0:
+            return
+        patch_token_id = self.config.image_patch_token
+        image_end_token_id = getattr(self.config, "image_end_token", None)
+        if image_end_token_id is None:
+            return
+
+        scale_items = list(self.query_tokens_dict.items())
+        if not scale_items:
+            return
+        scale_name, query_param = scale_items[-1]
+        num_query_tokens = int(query_param.shape[0])
+
+        flat = input_ids.view(-1)
+        L = int(flat.numel())
+        # Need at least num_query_tokens patch tokens + 1 <image_end> token.
+        if L < num_query_tokens + 1:
+            return
+
+        # Signature: the prompt must end with <image_end_token>, and the
+        # N positions immediately before it must all be <imagePatch>.
+        if int(flat[-1].item()) != image_end_token_id:
+            return
+        tail_start = L - 1 - num_query_tokens
+        tail_end = L - 1  # exclusive — last index is <image_end_token>
+        if not (flat[tail_start:tail_end] == patch_token_id).all():
+            return
+
+        query_embeds = query_param.to(dtype=inputs_embeds.dtype, device=inputs_embeds.device)
+        inputs_embeds[tail_start:tail_end] = query_embeds
+        logger.debug(
+            "[MingFlashOmniThinker] injected %d image-gen query embeddings (scale=%s) at positions [%d, %d)",
+            num_query_tokens,
+            scale_name,
+            tail_start,
+            tail_end,
         )
 
     def forward(
diff --git a/vllm_omni/model_executor/models/ming_flash_omni/pipeline.py b/vllm_omni/model_executor/models/ming_flash_omni/pipeline.py
index a9d66fbc2..7ebbfd9d8 100644
--- a/vllm_omni/model_executor/models/ming_flash_omni/pipeline.py
+++ b/vllm_omni/model_executor/models/ming_flash_omni/pipeline.py
@@ -58,6 +58,7 @@ MING_FLASH_OMNI_PIPELINE = PipelineConfig(
             engine_output_type="audio",
             tokenizer_subdir="talker/llm",
             custom_process_input_func=f"{_PROC}.thinker2talker",
+            sync_process_input_func=f"{_PROC}.thinker2talker_token_only",
         ),
     ),
 )
@@ -104,3 +105,37 @@ MING_FLASH_OMNI_THINKER_ONLY_PIPELINE = PipelineConfig(
         ),
     ),
 )
+
+
+# Thinker + image-generation (diffusion) variant: text-to-image / img2img.
+MING_FLASH_OMNI_IMAGE_PIPELINE = PipelineConfig(
+    model_type="ming_flash_omni_image",
+    model_arch="MingFlashOmniForConditionalGeneration",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="thinker",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            final_output=False,
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            hf_config_name="thinker_config",
+            engine_output_type="latent",
+            prompt_expand_func=f"{_PROC}.expand_cfg_prompts",
+            sampling_constraints={"detokenize": False},
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="dit",
+            # Resolved in vllm_omni/diffusion/registry.py.
+            model_arch="MingImagePipeline",
+            execution_type=StageExecutionType.DIFFUSION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="image",
+            hf_config_name="image_gen_config",
+            custom_process_input_func=f"{_PROC}.thinker2imagegen",
+        ),
+    ),
+)
diff --git a/vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py b/vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py
index 4271114bc..53c130b90 100644
--- a/vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py
+++ b/vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py
@@ -3,12 +3,86 @@
 # Copyright (c) Ant Group. All rights reserved.
 # Adapted from Ming repo's usage cookbook:
 # https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/cookbook.ipynb
-"""Shared prompt-building helpers for Ming-flash-omni standalone talker."""
+"""Ming-flash-omni-2.0 prompt utilities.
+
+Two unrelated families of helpers live in the same file because they
+both belong to the Ming module and are tightly coupled to Ming-specific
+prompt conventions:
+
+1. **Image-gen query-token expansion** — string-level helpers used by
+   the thinker stage to mark the ``<image><imagePatch>*N</image>`` block
+   that gets substituted with learnable ``query_tokens_dict`` embeddings
+   during forward.
+
+2. **TTS / talker caption builder** — JSON caption template for the
+   Ming-flash-omni standalone talker (TTS) stage.
+"""
+
+from __future__ import annotations
 
 import copy
 import json
 from typing import Any
 
+# ============================================================
+# Image-gen query-token block (thinker stage)
+# ============================================================
+
+# Ming's thinker uses these tokens to mark a learnable image-generation
+# query block inside the text prompt. The thinker substitutes its
+# ``query_tokens_dict`` embeddings at each ``<imagePatch>`` position
+# during forward; see
+# ``MingFlashOmniThinker._maybe_inject_image_gen_query_embeds``.
+_IMAGE_OPEN_TOKEN = "<image>"
+_IMAGE_CLOSE_TOKEN = "</image>"
+IMAGE_PATCH_TOKEN = "<imagePatch>"
+
+# Default query-token count matches ``MingImageGenConfig(img_gen_scales=[16])``
+# (16 * 16 = 256), which is what the released inclusionAI/Ming-flash-omni-2.0
+# checkpoint ships.
+DEFAULT_NUM_QUERY_TOKENS = 256
+
+
+def maybe_expand_image_gen_prompt(
+    prompt: str,
+    num_query_tokens: int = DEFAULT_NUM_QUERY_TOKENS,
+) -> str:
+    """Append the ``<image><imagePatch>*N</image>`` suffix for text-to-image.
+
+    The thinker expects image-generation requests to end with an N-wide
+    block of ``<imagePatch>`` tokens (wrapped in ``<image>``/``</image>``)
+    — those positions get substituted with learnable
+    ``query_tokens_dict`` embeddings during forward.
+
+    This helper is a no-op (returns the input unchanged) when:
+
+      * ``prompt`` is not a non-empty string, or
+      * the prompt already contains an ``<imagePatch>`` block (avoids
+        double expansion for tests / manual calls that pre-format the
+        prompt).
+
+    Args:
+        prompt: Raw user prompt text.
+        num_query_tokens: Total number of query tokens to emit. Defaults
+            to 256 (the released checkpoint's ``img_gen_scales=[16]``).
+
+    Returns:
+        The (possibly expanded) prompt text.
+    """
+    if not isinstance(prompt, str) or not prompt:
+        return prompt
+    if IMAGE_PATCH_TOKEN in prompt:
+        return prompt
+
+    # TODO(multi-scale): single-block emission assumes img_gen_scales=[16].
+    suffix = _IMAGE_OPEN_TOKEN + (IMAGE_PATCH_TOKEN * num_query_tokens) + _IMAGE_CLOSE_TOKEN
+    return prompt + suffix
+
+
+# ============================================================
+# TTS / talker caption builder
+# ============================================================
+
 DEFAULT_PROMPT = "Please generate speech based on the following description.\n"
 
 BASE_CAPTION_TEMPLATE: dict[str, Any] = {
@@ -48,3 +122,13 @@ def create_instruction(user_input: dict[str, Any]) -> str:
         if key in item:
             item[key] = value
     return json.dumps(caption, ensure_ascii=False)
+
+
+__all__ = [
+    "IMAGE_PATCH_TOKEN",
+    "DEFAULT_NUM_QUERY_TOKENS",
+    "maybe_expand_image_gen_prompt",
+    "DEFAULT_PROMPT",
+    "BASE_CAPTION_TEMPLATE",
+    "create_instruction",
+]
diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice_decoder.py b/vllm_omni/model_executor/models/omnivoice/omnivoice_decoder.py
index 59f23d0d0..efb6d9dc7 100644
--- a/vllm_omni/model_executor/models/omnivoice/omnivoice_decoder.py
+++ b/vllm_omni/model_executor/models/omnivoice/omnivoice_decoder.py
@@ -108,7 +108,10 @@ class OmniVoiceDecoder(nn.Module):
         quantized = self.quantizer.decode(codes)
 
         # Project: [B, 1024, T] → fc2 → [B, 256, T]
-        quantized = self.fc2(quantized.transpose(1, 2)).transpose(1, 2)
+        # Cast to fc2 weight dtype (may be fp16 when checkpoint stores weights as fp16),
+        # then upcast back to float32 — acoustic decoder ConvTranspose1d upsampling
+        # produces intermediate values that exceed the fp16 range (~65504), causing NaN.
+        quantized = self.fc2(quantized.transpose(1, 2).to(self.fc2.weight.dtype)).transpose(1, 2).float()
 
         # Acoustic decoder: [B, 256, T] → [B, 1, T*960]
         audio = self.acoustic_decoder(quantized)
@@ -192,6 +195,11 @@ class OmniVoiceDecoder(nn.Module):
                 param.data.copy_(state_dict[higgs_name])
                 loaded += 1
 
+        # Checkpoint weights may be fp16; force float32 so ConvTranspose1d
+        # upsampling doesn't overflow (fp16 max ~65504 is too small for
+        # intermediate activations in the DAC upsampling chain).
+        self.acoustic_decoder.float()
+
         # Apply HiggsAudioV2 output padding adjustment
         self._adjust_output_padding(self.acoustic_decoder)
 
diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py b/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py
index eedb5828a..fc5befeb9 100644
--- a/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py
+++ b/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py
@@ -13,6 +13,9 @@ via FlashAttention/SageAttention/SDPA backends.
 from __future__ import annotations
 
 import math
+import random
+import threading
+from collections import OrderedDict
 
 import torch
 import torch.nn as nn
@@ -24,6 +27,184 @@ from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
 logger = init_logger(__name__)
 
 
+# ---------------------------------------------------------------------------
+# Triton kernels (inference-only; graceful fallback when triton is absent)
+# ---------------------------------------------------------------------------
+
+_TRITON_AVAILABLE = False
+try:
+    import triton
+    import triton.language as tl
+
+    try:
+        from triton.language.extra.libdevice import rsqrt  # noqa: F401
+    except ModuleNotFoundError:
+        from triton.language.extra.cuda.libdevice import rsqrt  # noqa: F401
+
+    def _calculate_settings(n: int) -> tuple[int, int]:
+        MAX_FUSED_SIZE = 65536
+        BLOCK_SIZE = triton.next_power_of_2(n)
+        if BLOCK_SIZE > MAX_FUSED_SIZE:
+            raise RuntimeError(f"n={n} exceeds max Triton block size {MAX_FUSED_SIZE}")
+        num_warps = 4
+        if BLOCK_SIZE >= 32768:
+            num_warps = 32
+        elif BLOCK_SIZE >= 8192:
+            num_warps = 16
+        elif BLOCK_SIZE >= 2048:
+            num_warps = 8
+        return BLOCK_SIZE, num_warps
+
+    @triton.jit
+    def _rms_norm_fwd_kernel(
+        Y_ptr,  # noqa: N803
+        Y_stride,  # noqa: N803
+        X_ptr,  # noqa: N803
+        X_stride,  # noqa: N803
+        W_ptr,  # noqa: N803
+        n_cols,
+        eps,
+        BLOCK_SIZE: tl.constexpr,  # noqa: N803
+    ):
+        row = tl.program_id(0).to(tl.int64)
+        cols = tl.arange(0, BLOCK_SIZE)
+        mask = cols < n_cols
+        x = tl.load(X_ptr + row * X_stride + cols, mask=mask, other=0.0)
+        w = tl.load(W_ptr + cols, mask=mask, other=0.0)
+        x_f32 = x.to(tl.float32)
+        ms = tl.sum(x_f32 * x_f32, axis=0) / n_cols
+        rstd = rsqrt(ms + eps)
+        y = (x_f32 * rstd).to(x.dtype) * w
+        tl.store(Y_ptr + row * Y_stride + cols, y, mask=mask)
+
+    def triton_rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+        shape = x.shape
+        n_cols = shape[-1]
+        x2d = x.contiguous().view(-1, n_cols)
+        w = weight.contiguous()
+        BLOCK_SIZE, num_warps = _calculate_settings(n_cols)
+        y = torch.empty_like(x2d)
+        _rms_norm_fwd_kernel[(x2d.shape[0],)](
+            y,
+            y.stride(0),
+            x2d,
+            x2d.stride(0),
+            w,
+            n_cols,
+            eps,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=num_warps,
+        )
+        return y.view(*shape)
+
+    @triton.jit
+    def _swiglu_fwd_kernel(
+        gate_ptr,
+        up_ptr,
+        out_ptr,
+        stride,
+        n_cols: tl.constexpr,
+        BLOCK_SIZE: tl.constexpr,  # noqa: N803
+    ):
+        pid = tl.program_id(0).to(tl.int64)
+        gate_ptr += pid * stride
+        up_ptr += pid * stride
+        out_ptr += pid * stride
+        cols = tl.arange(0, BLOCK_SIZE)
+        mask = cols < n_cols
+        gate = tl.load(gate_ptr + cols, mask=mask, other=0).to(tl.float32)
+        up = tl.load(up_ptr + cols, mask=mask, other=0)
+        silu_gate = gate * tl.sigmoid(gate)
+        out = silu_gate.cast(up.dtype) * up
+        tl.store(out_ptr + cols, out, mask=mask)
+
+    def triton_swiglu(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
+        gate = gate.contiguous()
+        up = up.contiguous()
+        shape = gate.shape
+        n_cols = shape[-1]
+        g2d = gate.view(-1, n_cols)
+        u2d = up.view(-1, n_cols)
+        out = torch.empty_like(g2d)
+        BLOCK_SIZE, num_warps = _calculate_settings(n_cols)
+        _swiglu_fwd_kernel[(g2d.shape[0],)](
+            g2d,
+            u2d,
+            out,
+            out.stride(0),
+            n_cols=n_cols,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=num_warps,
+        )
+        return out.view(*shape)
+
+    @triton.jit
+    def _fused_add_rms_norm_fwd_kernel(
+        Y_ptr,  # noqa: N803
+        Y_stride,  # noqa: N803
+        S_ptr,  # noqa: N803
+        S_stride,  # noqa: N803
+        X_ptr,  # noqa: N803
+        X_stride,  # noqa: N803
+        R_ptr,  # noqa: N803
+        R_stride,  # noqa: N803
+        W_ptr,  # noqa: N803
+        n_cols,
+        eps,
+        BLOCK_SIZE: tl.constexpr,  # noqa: N803
+    ):
+        row = tl.program_id(0).to(tl.int64)
+        cols = tl.arange(0, BLOCK_SIZE)
+        mask = cols < n_cols
+        x = tl.load(X_ptr + row * X_stride + cols, mask=mask, other=0.0)
+        r = tl.load(R_ptr + row * R_stride + cols, mask=mask, other=0.0)
+        dtype = x.dtype
+        s_f32 = x.to(tl.float32) + r.to(tl.float32)
+        s = s_f32.to(dtype)
+        tl.store(S_ptr + row * S_stride + cols, s, mask=mask)
+        w = tl.load(W_ptr + cols, mask=mask, other=0.0)
+        ms = tl.sum(s_f32 * s_f32, axis=0) / n_cols
+        rstd = rsqrt(ms + eps)
+        y = (s_f32 * rstd).to(dtype) * w
+        tl.store(Y_ptr + row * Y_stride + cols, y, mask=mask)
+
+    def triton_fused_add_rms_norm(
+        x: torch.Tensor,
+        residual: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        shape = x.shape
+        n_cols = shape[-1]
+        x2d = x.contiguous().view(-1, n_cols)
+        r2d = residual.contiguous().view(-1, n_cols)
+        w = weight.contiguous()
+        BLOCK_SIZE, num_warps = _calculate_settings(n_cols)
+        y = torch.empty_like(x2d)
+        s = torch.empty_like(x2d)
+        _fused_add_rms_norm_fwd_kernel[(x2d.shape[0],)](
+            y,
+            y.stride(0),
+            s,
+            s.stride(0),
+            x2d,
+            x2d.stride(0),
+            r2d,
+            r2d.stride(0),
+            w,
+            n_cols,
+            eps,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=num_warps,
+        )
+        return y.view(*shape), s.view(*shape)
+
+    _TRITON_AVAILABLE = True
+    logger.debug("OmniVoice Triton kernels loaded")
+except Exception:
+    logger.debug("Triton not available; using PyTorch fallback for OmniVoice kernels")
+
+
 # ---------------------------------------------------------------------------
 # Unmasking schedule helpers
 # ---------------------------------------------------------------------------
@@ -45,9 +226,13 @@ def _get_time_steps(
     return shifted
 
 
-def _gumbel_sample(logits: torch.Tensor, temperature: float) -> torch.Tensor:
+def _gumbel_sample(logits: torch.Tensor, temperature: float, generator: torch.Generator) -> torch.Tensor:
     """Add Gumbel noise for stochastic position selection."""
-    noise = -torch.log(-torch.log(torch.rand_like(logits).clamp(min=1e-8)))
+    noise = -torch.log(
+        -torch.log(
+            torch.rand(logits.shape, generator=generator, device=logits.device, dtype=logits.dtype).clamp(min=1e-8)
+        )
+    )
     return logits / max(temperature, 1e-8) + noise
 
 
@@ -63,6 +248,8 @@ class OmniVoiceRMSNorm(nn.Module):
         self.eps = eps
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if _TRITON_AVAILABLE:
+            return triton_rms_norm(x, self.weight, self.eps)
         variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True)
         x = x * torch.rsqrt(variance + self.eps)
         return self.weight * x.to(self.weight.dtype)
@@ -158,6 +345,8 @@ class OmniVoiceMLP(nn.Module):
         self.down_proj = nn.Linear(config.llm_intermediate_size, config.llm_hidden_size, bias=False)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if _TRITON_AVAILABLE:
+            return self.down_proj(triton_swiglu(self.gate_proj(x), self.up_proj(x)))
         return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
 
 
@@ -181,13 +370,22 @@ class OmniVoiceTransformerBlock(nn.Module):
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
         hidden_states = self.self_attn(hidden_states, attention_mask=attention_mask, cos=cos, sin=sin)
-        hidden_states = residual + hidden_states
 
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
+        if _TRITON_AVAILABLE:
+            # Fused: (attn_out + residual) + RMSNorm in one kernel
+            hidden_states, residual = triton_fused_add_rms_norm(
+                hidden_states,
+                residual,
+                self.post_attention_layernorm.weight,
+                self.post_attention_layernorm.eps,
+            )
+        else:
+            hidden_states = residual + hidden_states
+            residual = hidden_states
+            hidden_states = self.post_attention_layernorm(hidden_states)
+
         hidden_states = self.mlp(hidden_states)
         hidden_states = residual + hidden_states
-
         return hidden_states
 
 
@@ -222,6 +420,214 @@ def _apply_rotary_pos_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor)
     return x * torch.cat([cos, cos], dim=-1) + rotated * torch.cat([sin, sin], dim=-1)
 
 
+# ---------------------------------------------------------------------------
+# CUDA Graph wrapper
+# ---------------------------------------------------------------------------
+
+
+class _OmniVoiceCUDAGraphForward:
+    """Pre-captures CUDA graphs for predefined sequence-length buckets.
+
+    Memory layout: all graphs share a single per-instance pool handle, which
+    isolates OmniVoice CUDA memory from other vllm modules.  Sequential replay
+    (one step at a time) means pool sharing is safe.
+    """
+
+    # Default bucket count is 10; 16 gives modest headroom for edge cases
+    # (seq_len > max bucket or non-CFG batch) without unbounded GPU growth.
+    _MAX_LAZY_GRAPHS: int = 16
+
+    def __init__(self, generator: OmniVoiceGenerator, capture_sizes: list[int]) -> None:
+        self._gen = generator
+        self._capture_sizes = sorted(capture_sizes)
+        # Pre-warmed graphs keyed by (two_b, bucket); fixed set, never evicted.
+        self._graphs: dict[tuple[int, int], dict] = {}
+        # Lazy-captured graphs for oversized / non-CFG shapes; capped via LRU.
+        self._lazy_graphs: OrderedDict[tuple[int, int], dict] = OrderedDict()
+        self._lock = threading.Lock()
+        # Per-instance pool handle: isolates OmniVoice CUDA memory from other
+        # vllm modules while still allowing safe re-use across sequential replays.
+        self._pool_handle: int | None = None
+
+    def _find_bucket(self, seq_len: int) -> int | None:
+        for bucket in self._capture_sizes:
+            if bucket >= seq_len:
+                return bucket
+        return None
+
+    def _pad_inputs(
+        self,
+        input_ids: torch.Tensor,
+        audio_mask: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        bucket: int,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+        S = input_ids.shape[-1]
+        if S == bucket:
+            return input_ids, audio_mask, attention_mask
+
+        two_b = input_ids.shape[0]
+        num_cb = input_ids.shape[1]
+
+        ids_padded = torch.zeros(two_b, num_cb, bucket, dtype=input_ids.dtype, device=input_ids.device)
+        ids_padded[:, :, :S] = input_ids
+
+        mask_padded = torch.zeros(two_b, bucket, dtype=torch.bool, device=audio_mask.device)
+        mask_padded[:, :S] = audio_mask
+
+        if attention_mask is not None:
+            attn_padded = torch.zeros(
+                two_b,
+                1,
+                bucket,
+                bucket,
+                dtype=attention_mask.dtype,
+                device=attention_mask.device,
+            )
+            attn_padded[:, :, :S, :S] = attention_mask
+        else:
+            attn_padded = None
+
+        return ids_padded, mask_padded, attn_padded
+
+    def _capture_for_key(
+        self,
+        key: tuple[int, int],
+        input_ids: torch.Tensor,
+        audio_mask: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+    ) -> dict:
+        _, bucket = key
+        device = input_ids.device
+
+        self._gen._ensure_rope(bucket, device)
+        model_dtype = self._gen.text_embedding.weight.dtype
+        static_cos = self._gen._rope_cos[:bucket].to(device=device, dtype=model_dtype).contiguous()
+        static_sin = self._gen._rope_sin[:bucket].to(device=device, dtype=model_dtype).contiguous()
+
+        static_input_ids = input_ids.clone()
+        static_audio_mask = audio_mask.clone()
+        static_attn_mask = attention_mask.clone() if attention_mask is not None else None
+
+        with torch.no_grad():
+            _ = self._gen._step_forward(
+                static_input_ids,
+                static_audio_mask,
+                static_attn_mask,
+                static_cos,
+                static_sin,
+            )
+        torch.accelerator.synchronize(device)
+
+        # Lazy-init per-instance pool handle: isolates OmniVoice CUDA Graph
+        # memory from other vllm modules (unlike get_global_graph_pool which
+        # shares a single pool across all captured graphs and can cause memory
+        # aliasing when two graphs replay concurrently).
+        if self._pool_handle is None:
+            self._pool_handle = torch.cuda.graph_pool_handle()
+
+        graph = torch.cuda.CUDAGraph()
+        with torch.no_grad():
+            with torch.cuda.graph(graph, pool=self._pool_handle):
+                static_output = self._gen._step_forward(
+                    static_input_ids,
+                    static_audio_mask,
+                    static_attn_mask,
+                    static_cos,
+                    static_sin,
+                )
+
+        entry = {
+            "graph": graph,
+            "static_input_ids": static_input_ids,
+            "static_audio_mask": static_audio_mask,
+            "static_attn_mask": static_attn_mask,
+            "static_cos": static_cos,
+            "static_sin": static_sin,
+            "static_output": static_output,
+        }
+        logger.info("OmniVoice CUDA Graph captured for key %s", key)
+        return entry
+
+    def warmup(self, device: torch.device) -> None:
+        """Pre-capture graphs for all bucket sizes with B=1 (two_b=2 for CFG)."""
+        if not torch.cuda.is_available():
+            return
+        logger.info(
+            "OmniVoice CUDA Graph warmup: capturing %d bucket sizes %s",
+            len(self._capture_sizes),
+            self._capture_sizes,
+        )
+        two_b = 2
+        num_cb = self._gen.config.num_audio_codebook
+        for bucket in self._capture_sizes:
+            key = (two_b, bucket)
+            dummy_ids = torch.zeros(two_b, num_cb, bucket, dtype=torch.long, device=device)
+            dummy_mask = torch.zeros(two_b, bucket, dtype=torch.bool, device=device)
+            dummy_attn = torch.ones(two_b, 1, bucket, bucket, dtype=torch.bool, device=device)
+            self._graphs[key] = self._capture_for_key(key, dummy_ids, dummy_mask, dummy_attn)
+        logger.info("OmniVoice CUDA Graph warmup complete (%d graphs)", len(self._graphs))
+
+    def __call__(
+        self,
+        input_ids: torch.Tensor,
+        audio_mask: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+    ) -> torch.Tensor:
+        if torch.cuda.is_current_stream_capturing():
+            seq_len = input_ids.shape[-1]
+            self._gen._ensure_rope(seq_len, input_ids.device)
+            dtype = self._gen.text_embedding.weight.dtype
+            cos = self._gen._rope_cos[:seq_len].to(device=input_ids.device, dtype=dtype)
+            sin = self._gen._rope_sin[:seq_len].to(device=input_ids.device, dtype=dtype)
+            return self._gen._step_forward(input_ids, audio_mask, attention_mask, cos, sin)
+
+        seq_len = input_ids.shape[-1]
+        two_b = input_ids.shape[0]
+        bucket = self._find_bucket(seq_len) if two_b == 2 else None
+
+        if bucket is None:
+            # Lazy capture: oversized sequence or non-unit batch (no pre-warmed bucket).
+            # Lock prevents concurrent threads from double-capturing the same key.
+            # _lazy_graphs is capped at _MAX_LAZY_GRAPHS with LRU eviction to
+            # prevent unbounded GPU memory growth when seq_len varies widely.
+            key = (two_b, seq_len)
+            ids_in, mask_in, attn_in = input_ids, audio_mask, attention_mask
+            with self._lock:
+                entry = self._lazy_graphs.get(key)
+                if entry is None:
+                    entry = self._capture_for_key(key, ids_in, mask_in, attn_in)
+                    if len(self._lazy_graphs) >= self._MAX_LAZY_GRAPHS:
+                        evicted_key, _ = self._lazy_graphs.popitem(last=False)
+                        logger.warning("OmniVoice CUDA Graph lazy cache full; evicted key %s", evicted_key)
+                    self._lazy_graphs[key] = entry
+        else:
+            key = (two_b, bucket)
+            ids_in, mask_in, attn_in = self._pad_inputs(input_ids, audio_mask, attention_mask, bucket)
+            with self._lock:
+                entry = self._graphs.get(key)
+                if entry is None:
+                    entry = self._capture_for_key(key, ids_in, mask_in, attn_in)
+                    self._graphs[key] = entry
+
+        entry["static_input_ids"].copy_(ids_in)
+        entry["static_audio_mask"].copy_(mask_in)
+        if attn_in is not None and entry["static_attn_mask"] is not None:
+            entry["static_attn_mask"].copy_(attn_in)
+
+        entry["graph"].replay()
+
+        output = entry["static_output"]
+        if bucket is not None and bucket != seq_len:
+            output = output[:, :, :seq_len, :]
+        return output
+
+    def clear(self) -> None:
+        with self._lock:
+            self._graphs.clear()
+            self._lazy_graphs.clear()
+
+
 # ---------------------------------------------------------------------------
 # Generator model
 # ---------------------------------------------------------------------------
@@ -278,6 +684,11 @@ class OmniVoiceGenerator(nn.Module):
         self._rope_cos = None
         self._rope_sin = None
 
+        # CUDA Graph (bucket-size pre-capture; lazy fallback for oversized shapes)
+        self._cuda_graph_fwd: _OmniVoiceCUDAGraphForward | None = (
+            _OmniVoiceCUDAGraphForward(self, config.cuda_graph_capture_sizes) if config.enable_cuda_graph else None
+        )
+
     def _ensure_rope(self, seq_len: int, device: torch.device) -> None:
         """Lazily compute RoPE cos/sin if needed."""
         if self._rope_cos is None or self._rope_cos.shape[0] < seq_len:
@@ -364,6 +775,20 @@ class OmniVoiceGenerator(nn.Module):
             self.config.audio_vocab_size,
         ).permute(0, 2, 1, 3)  # [B, 8, S, 1025]
 
+    def _step_forward(
+        self,
+        input_ids: torch.Tensor,
+        audio_mask: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        """Single unmasking-step forward using pre-cast RoPE tensors (CUDA graph safe)."""
+        hidden_states = self._prepare_embeddings(input_ids, audio_mask)
+        for layer in self.layers:
+            hidden_states = layer(hidden_states, attention_mask=attention_mask, cos=cos, sin=sin)
+        return self._get_logits(self.norm(hidden_states))
+
     @torch.inference_mode()
     def forward(
         self,
@@ -371,6 +796,7 @@ class OmniVoiceGenerator(nn.Module):
         audio_mask: torch.Tensor,
         attention_mask: torch.Tensor,
         target_lens: list[int],
+        seed: int | None = None,
         num_step: int = 32,
         guidance_scale: float = 2.0,
         t_shift: float = 0.1,
@@ -400,6 +826,9 @@ class OmniVoiceGenerator(nn.Module):
         max_target_len = max(target_lens)
         mask_id = self.config.audio_mask_id
         num_codebooks = self.config.num_audio_codebook
+        if seed is None:
+            seed = random.randint(0, 2**63 - 1)
+        generator = torch.Generator(device=device).manual_seed(seed)
 
         # Initialize all target tokens as [MASK]
         tokens = torch.full(
@@ -440,10 +869,12 @@ class OmniVoiceGenerator(nn.Module):
 
         # Main iterative loop
         for step in range(num_step):
-            # Prepare embeddings and run transformer
-            inputs_embeds = self._prepare_embeddings(input_ids, audio_mask)
-            hidden_states = self._transformer_forward(inputs_embeds, attention_mask)
-            batch_logits = self._get_logits(hidden_states).to(torch.float32)
+            if self._cuda_graph_fwd is not None and input_ids.is_cuda:
+                batch_logits = self._cuda_graph_fwd(input_ids, audio_mask, attention_mask).to(torch.float32)
+            else:
+                inputs_embeds = self._prepare_embeddings(input_ids, audio_mask)
+                hidden_states = self._transformer_forward(inputs_embeds, attention_mask)
+                batch_logits = self._get_logits(hidden_states).to(torch.float32)
             # batch_logits: [2*B, 8, S, 1025]
 
             for i in range(B):
@@ -474,7 +905,7 @@ class OmniVoiceGenerator(nn.Module):
 
                 # Token prediction
                 if class_temperature > 0.0:
-                    pred_tokens = _gumbel_sample(log_probs, class_temperature).argmax(dim=-1)
+                    pred_tokens = _gumbel_sample(log_probs, class_temperature, generator).argmax(dim=-1)
                 else:
                     pred_tokens = log_probs.argmax(dim=-1)  # [1, 8, T]
 
@@ -486,7 +917,7 @@ class OmniVoiceGenerator(nn.Module):
 
                 # Gumbel noise for position selection
                 if position_temperature > 0.0:
-                    scores = _gumbel_sample(scores, position_temperature)
+                    scores = _gumbel_sample(scores, position_temperature, generator)
 
                 # Mask out already unmasked positions
                 sample_tokens = tokens[i : i + 1, :, :t_len]
@@ -586,3 +1017,6 @@ class OmniVoiceGenerator(nn.Module):
             )
         else:
             logger.info("Generator: all %d weights loaded", len(loaded_keys))
+
+        if self._cuda_graph_fwd is not None:
+            self._cuda_graph_fwd.warmup(device)
diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py b/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py
index de0644803..afd0a92a5 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py
+++ b/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py
@@ -30,6 +30,7 @@ QWEN2_5_OMNI_PIPELINE = PipelineConfig(
             requires_multimodal_data=True,
             engine_output_type="latent",
             sampling_constraints={"detokenize": True},
+            custom_process_next_stage_input_func=f"{_PROC}.thinker2talker_full_payload",
         ),
         StagePipelineConfig(
             stage_id=1,
@@ -38,6 +39,8 @@ QWEN2_5_OMNI_PIPELINE = PipelineConfig(
             input_sources=(0,),
             engine_output_type="latent",
             custom_process_input_func=f"{_PROC}.thinker2talker",
+            sync_process_input_func=f"{_PROC}.thinker2talker_token_only",
+            custom_process_next_stage_input_func=f"{_PROC}.talker2code2wav_full_payload",
             sampling_constraints={
                 "detokenize": True,
                 "stop_token_ids": [8294],
@@ -52,6 +55,7 @@ QWEN2_5_OMNI_PIPELINE = PipelineConfig(
             final_output_type="audio",
             engine_output_type="audio",
             custom_process_input_func=f"{_PROC}.talker2code2wav",
+            sync_process_input_func=f"{_PROC}.talker2code2wav_token_only",
             sampling_constraints={"detokenize": True},
         ),
     ),
@@ -74,6 +78,7 @@ QWEN2_5_OMNI_THINKER_ONLY_PIPELINE = PipelineConfig(
             requires_multimodal_data=True,
             engine_output_type="latent",
             sampling_constraints={"detokenize": True},
+            custom_process_next_stage_input_func=f"{_PROC}.thinker2talker_full_payload",
         ),
     ),
 )
diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py
index 34c88bad9..49d3ac114 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py
+++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py
@@ -2,14 +2,12 @@
 #      Start Token2Wav     #
 ############################
 
-import math
 from collections.abc import Iterable
 
 import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.nn import Parameter
 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
 from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
     Qwen2_5OmniBigVGANConfig,
@@ -35,8 +33,9 @@ from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
 from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding
+from vllm_omni.model_executor.models.common.alias_free_activation import AliasFreeActivation1d
+from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta
 from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import cap_and_align_mel_length, resolve_max_mel_frames
-from vllm_omni.platforms import current_omni_platform
 
 
 # Provide a no-op auto_docstring decorator to satisfy annotations if missing
@@ -634,238 +633,9 @@ class DiTDecoderLayer(nn.Module):
         return hidden_states
 
 
-class SnakeBeta(nn.Module):
-    """
-    A modified Snake function which uses separate parameters for the magnitude
-    of the periodic components
-    Shape:
-        - Input: (B, C, T)
-        - Output: (B, C, T), same shape as the input
-    Parameters:
-        - alpha - trainable parameter that controls frequency
-        - beta - trainable parameter that controls magnitude
-    References:
-        - This activation function is a modified version based on this paper
-          by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
-        https://huggingface.co/papers/2006.08195
-    """
-
-    def __init__(self, in_features, alpha=1.0):
-        super().__init__()
-        self.in_features = in_features
-
-        # initialize alpha
-        self.alpha = Parameter(torch.zeros(in_features) * alpha)
-        self.beta = Parameter(torch.zeros(in_features) * alpha)
-
-        self.no_div_by_zero = 0.000000001
-
-    def forward(self, hidden_states):
-        """
-        Forward pass of the function.
-        Applies the function to the input elementwise.
-        SnakeBeta ∶= x + 1/b * sin^2 (xa)
-        """
-        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)  # line up with x to [B, C, T]
-        beta = self.beta.unsqueeze(0).unsqueeze(-1)
-        alpha = torch.exp(alpha)
-        beta = torch.exp(beta)
-        hidden_states = hidden_states + (1.0 / (beta + self.no_div_by_zero)) * torch.pow(
-            torch.sin(hidden_states * alpha), 2
-        )
-
-        return hidden_states
-
-
-def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor:
-    """Generates a 1D Kaiser-windowed sinc filter.
-
-    Args:
-        cutoff (float): Normalized cutoff frequency (0 to 0.5).
-        half_width (float): Transition bandwidth.
-        kernel_size (int): Number of filter taps.
-
-    Returns:
-        torch.Tensor: A tensor of shape (1, 1, kernel_size) representing the filter.
-    """
-    is_even = kernel_size % 2 == 0
-    half_size = kernel_size // 2
-
-    # Compute Kaiser window parameters
-    delta_f = 4 * half_width
-    attenuation = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
-
-    if attenuation > 50.0:
-        beta = 0.1102 * (attenuation - 8.7)
-    elif attenuation >= 21.0:
-        beta = 0.5842 * (attenuation - 21) ** 0.4 + 0.07886 * (attenuation - 21.0)
-    else:
-        beta = 0.0
-
-    # TODO: When torch.kaiser_window supports NPU, remove the device="cpu" argument
-    if current_omni_platform.is_npu():
-        kaiser_window = torch.kaiser_window(
-            kernel_size, beta=beta, periodic=False, dtype=torch.float32, device="cpu"
-        ).to("npu")
-    elif current_omni_platform.is_xpu():
-        kaiser_window = torch.kaiser_window(
-            kernel_size, beta=beta, periodic=False, dtype=torch.float32, device="cpu"
-        ).to("xpu")
-    else:
-        kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32)
-
-    # Compute time indices
-    if is_even:
-        time_indices = torch.arange(-half_size, half_size) + 0.5
-    else:
-        time_indices = torch.arange(kernel_size) - half_size
-
-    # Compute sinc filter
-    if cutoff == 0:
-        return torch.zeros((1, 1, kernel_size), dtype=torch.float32)
-
-    sinc_filter = torch.sinc(2 * cutoff * time_indices)
-    normalized_filter = 2 * cutoff * kaiser_window * sinc_filter
-
-    # Normalize to ensure sum = 1 (avoid leakage of constant component)
-    normalized_filter /= normalized_filter.sum()
-
-    return normalized_filter.view(1, 1, kernel_size)
-
-
-def replication_pad_1d(hidden_states: torch.Tensor, pad_left: int, pad_right: int) -> torch.Tensor:
-    """
-    Manual replicate padding to avoid replication_pad1d kernel limits on NPU.
-    TODO: remove when F.pad supports replicate mode on NPU.
-    """
-    # NOTE: a immature implementation for running in NPU. Need to discuss.
-    if pad_left == 0 and pad_right == 0:
-        return hidden_states
-
-    segments = []
-    if pad_left > 0:
-        left = hidden_states[..., :1].expand(*hidden_states.shape[:-1], pad_left)
-        segments.append(left)
-
-    segments.append(hidden_states)
-
-    if pad_right > 0:
-        right = hidden_states[..., -1:].expand(*hidden_states.shape[:-1], pad_right)
-        segments.append(right)
-
-    return torch.cat(segments, dim=-1)
-
-
-class UpSample1d(nn.Module):
-    def __init__(self, ratio=2, kernel_size=None):
-        super().__init__()
-        self.ratio = ratio
-        self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
-        self.stride = ratio
-        self.pad = self.kernel_size // ratio - 1
-        self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
-        self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
-
-        filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size)
-        self.register_buffer("filter", filter, persistent=False)
-
-    def forward(self, hidden_states):
-        channels = hidden_states.shape[1]
-        if current_omni_platform.is_npu():
-            # TODO: When F.pad supports replicate mode on NPU, remove this branch
-            input_dtype = hidden_states.dtype
-            # F.pad in NPU doesn't support BF16 when mode is replicate.
-            # To ensure the accuracy, manually pad the input tensor.
-            hidden_states = replication_pad_1d(hidden_states.to(self.filter.dtype), self.pad, self.pad)
-            filter_convert_dtype = self.filter.to(hidden_states.dtype)
-            hidden_states = self.ratio * F.conv_transpose1d(
-                hidden_states,
-                filter_convert_dtype.expand(channels, -1, -1),
-                stride=self.stride,
-                groups=channels,
-            ).to(input_dtype)
-        else:
-            hidden_states_dtype = hidden_states.dtype
-            hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate").to(self.filter.dtype)
-            hidden_states = self.ratio * F.conv_transpose1d(
-                hidden_states,
-                self.filter.expand(channels, -1, -1),
-                stride=self.stride,
-                groups=channels,
-            ).to(hidden_states_dtype)
-        hidden_states = hidden_states[..., self.pad_left : -self.pad_right]
-
-        return hidden_states
-
-
-class DownSample1d(nn.Module):
-    def __init__(self, ratio=2, kernel_size=None):
-        super().__init__()
-        cutoff = 0.5 / ratio
-        half_width = 0.6 / ratio
-
-        if cutoff < 0.0:
-            raise ValueError("Minimum cutoff must be larger than zero.")
-        if cutoff > 0.5:
-            raise ValueError("A cutoff above 0.5 does not make sense.")
-
-        self.even = kernel_size % 2 == 0
-        self.pad_left = kernel_size // 2 - int(self.even)
-        self.pad_right = kernel_size // 2
-        self.stride = ratio
-        filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
-        self.register_buffer("filter", filter, persistent=False)
-
-    def forward(self, hidden_states):
-        channels = hidden_states.shape[1]
-        if current_omni_platform.is_npu():
-            input_dtype = hidden_states.dtype
-            # F.pad in NPU doesn't support BF16 when mode is replicate.
-            # To ensure the accuracy, manually pad the input tensor.
-            hidden_states = replication_pad_1d(hidden_states.to(self.filter.dtype), self.pad_left, self.pad_right)
-            filter_on_device = self.filter.to(device=hidden_states.device, dtype=hidden_states.dtype)
-            out = F.conv1d(
-                hidden_states,
-                filter_on_device.expand(channels, -1, -1),
-                stride=self.stride,
-                groups=channels,
-            ).to(input_dtype)
-        else:
-            hidden_states_dtype = hidden_states.dtype
-            hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate").to(
-                self.filter.dtype
-            )
-            out = F.conv1d(
-                hidden_states,
-                self.filter.expand(channels, -1, -1),
-                stride=self.stride,
-                groups=channels,
-            ).to(hidden_states_dtype)
-        return out
-
-
-class TorchActivation1d(nn.Module):
-    def __init__(
-        self,
-        activation,
-        up_ratio: int = 2,
-        down_ratio: int = 2,
-        up_kernel_size: int = 12,
-        down_kernel_size: int = 12,
-    ):
-        super().__init__()
-        if not callable(activation):
-            raise TypeError("Activation function must be callable")
-        self.act = activation
-        self.upsample = UpSample1d(up_ratio, up_kernel_size)
-        self.downsample = DownSample1d(down_ratio, down_kernel_size)
-
-    def forward(self, hidden_states):
-        hidden_states = self.upsample(hidden_states)
-        hidden_states = self.act(hidden_states)
-        hidden_states = self.downsample(hidden_states)
-
-        return hidden_states
+# SnakeBeta, UpSample1d, DownSample1d, AliasFreeActivation1d (formerly
+# TorchActivation1d) are now provided by the shared ``common`` modules.
+# See: vllm_omni/model_executor/models/common/{snake_activation,alias_free_activation}.py
 
 
 class AMPBlock(torch.nn.Module):
@@ -938,7 +708,7 @@ class AMPBlock(torch.nn.Module):
         self.num_layers = len(self.convs1) + len(self.convs2)  # total number of conv layers
 
         self.activations = nn.ModuleList(
-            [TorchActivation1d(activation=SnakeBeta(channels)) for _ in range(self.num_layers)]
+            [AliasFreeActivation1d(activation=SnakeBeta(channels)) for _ in range(self.num_layers)]
         )
 
     def _get_padding(self, kernel_size, dilation=1):
@@ -1002,7 +772,7 @@ class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
             ]
         )
 
-        self.activation_post = TorchActivation1d(
+        self.activation_post = AliasFreeActivation1d(
             activation=SnakeBeta(config.upsample_initial_channel // (2**self.num_upsample_layers))
         )
         self.conv_post = nn.Conv1d(
diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py
index 3c658fc3f..6bd8bf305 100644
--- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py
+++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py
@@ -20,7 +20,6 @@ from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
 from vllm.config import ModelConfig, VllmConfig
 from vllm.inputs import PromptType, TokensPrompt
 from vllm.logger import init_logger
-from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
 from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP, SupportsRealtime
 from vllm.model_executor.models.qwen3_asr_realtime import Qwen3ASRRealtimeBuffer
 from vllm.model_executor.models.qwen3_omni_moe_thinker import (
@@ -180,11 +179,9 @@ class Qwen3OmniMoeForConditionalGeneration(
                 ("hidden_states", "last"),
                 ("hidden_states", "trailing_text"),
                 ("embed", "tts_pad_projected"),
-            }
-            # Keys that need to be accumulated across streaming inputs
-            self.streaming_accumulated_keys: set[tuple[str, str]] = {
-                ("embed", "prefill"),
-                ("hidden_states", "output"),
+                # talker MTP codec codes must stay on GPU to avoid a per-step D2H
+                # sync stall; build_mm_cpu handles the eventual D2H at payload time.
+                ("codes", "audio"),
             }
 
         elif self.model_stage == "code2wav":
@@ -323,7 +320,12 @@ class Qwen3OmniMoeForConditionalGeneration(
                 msg = "Qwen3 Omni thinker get_mrope_input_positions requires mm_features"
                 raise ValueError(msg)
             return self.thinker.get_mrope_input_positions(input_tokens, mm_features)
-        return MRotaryEmbedding.get_input_positions_tensor(input_tokens, **kwargs)
+        # Talker/code2wav stages are text/codec-only and do not need
+        # multimodal M-RoPE position computation. Return a cheap linear
+        # position tensor to avoid unnecessary per-request M-RoPE work.
+        seq_len = len(input_tokens)
+        linear = torch.arange(seq_len, dtype=torch.long).unsqueeze(0).expand(3, seq_len)
+        return linear, 0
 
     def forward(
         self,
@@ -420,13 +422,16 @@ class Qwen3OmniMoeForConditionalGeneration(
                 else:
                     codes = input_ids.reshape(1, 16, -1)
             else:
-                logger.warning(
-                    (
-                        "Input_ids length: %s is not divisible by 16, padding "
-                        "with zeros. This should only happen in warm up."
-                    ),
-                    input_ids.shape[0],
-                )
+                if seq_token_counts is None:
+                    logger.debug(
+                        "Code2Wav warmup input length %s is not divisible by 16; padding with zeros.",
+                        input_ids.shape[0],
+                    )
+                else:
+                    logger.warning_once(
+                        "Code2Wav input length is not divisible by 16; padding with zeros. "
+                        "This is expected only during cudagraph warmup."
+                    )
                 input_ids_flatten = input_ids.reshape(-1)
                 input_ids_flatten = torch.cat(
                     [
@@ -670,8 +675,14 @@ class Qwen3OmniMoeForConditionalGeneration(
 
         span_len = input_ids.shape[0]
         update_dict: OmniPayload = {}
-        if span_len > 1:
-            # prefill
+        # Prefix caching can reduce a new request's remaining prefill span to a
+        # single token. Use the runner-provided phase flag instead of span_len.
+        is_prefill = bool(payload.get("_omni_is_prefill", span_len > 1))
+        if is_prefill:
+            num_computed_tokens = payload.get("_omni_num_computed_tokens")
+            request_resumable = meta.get("resumable", False)
+            if num_computed_tokens is not None and not request_resumable:
+                meta["num_processed_tokens"] = int(num_computed_tokens)
             input_ids, input_embeds, update_dict = self.talker_preprocess_prefill(input_ids, input_embeds, payload)
             code_predictor_codes = torch.zeros(
                 (input_embeds.shape[0], self.talker.num_code_groups),
@@ -1245,7 +1256,7 @@ class Qwen3OmniMoeForConditionalGeneration(
         if (
             getattr(self, "model_stage", None) == "talker"
             and sampling_metadata is not None
-            and (sampling_metadata.temperature is None or (sampling_metadata.temperature <= 0).any())
+            and (sampling_metadata.temperature is None)
         ):
             self._warn_talker_sampling_temperature(sampling_metadata)
 
@@ -1283,6 +1294,22 @@ class Qwen3OmniMoeForConditionalGeneration(
         left_frames = int(extra.get("codec_left_context_frames", 0) or 0)
         return chunk_frames, left_frames
 
+    def _maybe_enable_code2wav_cudagraph(self) -> None:
+        """Enable the inner Code2Wav CUDA graph unless this stage runs in eager mode."""
+        if not self.code2wav or not hasattr(self.code2wav, "enable_cudagraph"):
+            return
+
+        model_cfg = getattr(self.vllm_config, "model_config", None)
+        if getattr(model_cfg, "enforce_eager", False):
+            logger.info("Code2Wav CUDA Graph disabled because enforce_eager is set")
+            return
+
+        chunk_frames, left_frames = self._get_codec_frame_config()
+        self.code2wav.enable_cudagraph(
+            codec_chunk_frames=chunk_frames,
+            codec_left_context_frames=left_frames,
+        )
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load weights for all components of the omni model."""
         loaded_weights = set()
@@ -1319,15 +1346,11 @@ class Qwen3OmniMoeForConditionalGeneration(
             code2wav_loaded = add_prefix_to_loaded_weights(code2wav_loaded, "code2wav")
             loaded_weights.update(code2wav_loaded)
 
-            # Precompute SnakeBeta caches and enable CUDA graph for Code2Wav decoder
+            # Precompute SnakeBeta caches; Code2Wav CUDA graph follows the stage's
+            # enforce_eager setting, the same switch vLLM uses for outer graphs.
             try:
                 self.code2wav.precompute_snake_caches()
-                if hasattr(self.code2wav, "enable_cudagraph"):
-                    chunk_frames, left_frames = self._get_codec_frame_config()
-                    self.code2wav.enable_cudagraph(
-                        codec_chunk_frames=chunk_frames,
-                        codec_left_context_frames=left_frames,
-                    )
+                self._maybe_enable_code2wav_cudagraph()
             except Exception:
                 logger.warning(
                     "Failed to enable CUDA Graph for Code2Wav; falling back to eager.",
diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py
index 8bcba4340..2b1a7ccf2 100644
--- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py
+++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py
@@ -297,9 +297,9 @@ class Qwen3OmniMoeCode2Wav(nn.Module):
                 single tensor with shape ``[1, waveform_len]``.
         """
         if not (left_context_size and seq_token_counts and len(left_context_size) == len(seq_token_counts)):
-            logger.warning(
+            logger.warning_once(
                 "chunked_decode_streaming: missing/invalid left_context_size or seq_token_counts; "
-                "defaulting to left_context_size=zeros(len=codes.shape[0])."
+                "defaulting to left_context_size=zeros(len(codes)). This is expected during cudagraph warmup."
             )
             left_context_size = [0] * codes.shape[0]
         # Decode chunk
diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py
index 8da1ccfec..df26a1faf 100644
--- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py
+++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py
@@ -1205,11 +1205,8 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
     ) -> IntermediateTensors | None:
         if not getattr(self, "deepstack_input_embeds", None):
             return None  # If vision tower is skipped
-        if getattr(self, "deepstack_input_embeds_num_tokens", 0) == 0:
-            return None
-        buf = self.deepstack_input_embeds[0].size(0)
-        if num_tokens > buf:
-            raise ValueError(f"Requested more deepstack tokens than buffer capacity: {num_tokens=} > {buf=}")
+        if num_tokens > self.deepstack_input_embeds[0].size(0):
+            self._resize_deepstack_input_embeds(num_tokens)
         n_valid = self.deepstack_input_embeds_num_tokens
         # embed_input_ids runs on unpadded length; forward uses padded num_tokens
         # (vLLM gpu_model_runner). Tail positions have no vision deepstack and must
@@ -1226,6 +1223,17 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
             }
         )
 
+    def _resize_deepstack_input_embeds(self, num_tokens: int) -> None:
+        self.deepstack_input_embeds = [
+            torch.zeros(
+                num_tokens,
+                self.config.text_config.hidden_size,
+                device=self.deepstack_input_embeds[0].device,
+                dtype=self.deepstack_input_embeds[0].dtype,
+            )
+            for _ in range(self.deepstack_num_level)
+        ]
+
     def _set_deepstack_input_embeds(self, deepstack_input_embeds: torch.Tensor) -> None:
         if not getattr(self, "deepstack_input_embeds", None):
             return
@@ -1233,15 +1241,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
         # set deepstack_input_embeds to buffer
         num_tokens = deepstack_input_embeds.size(1)
         if num_tokens > self.deepstack_input_embeds[0].size(0):
-            self.deepstack_input_embeds = [
-                torch.zeros(
-                    num_tokens,
-                    self.config.text_config.hidden_size,
-                    device=self.deepstack_input_embeds[0].device,
-                    dtype=self.deepstack_input_embeds[0].dtype,
-                )
-                for _ in range(self.deepstack_num_level)
-            ]
+            self._resize_deepstack_input_embeds(num_tokens)
         for idx in range(self.deepstack_num_level):
             self.deepstack_input_embeds[idx][:num_tokens].copy_(deepstack_input_embeds[idx])
         self.deepstack_input_embeds_num_tokens = num_tokens
diff --git a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py
index 5051715ce..7f50931dd 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py
+++ b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py
@@ -26,6 +26,7 @@ QWEN3_TTS_PIPELINE = PipelineConfig(
             owns_tokenizer=True,
             engine_output_type="latent",
             async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"),
+            custom_process_next_stage_input_func=f"{_PROC}.talker2code2wav_full_payload",
             sampling_constraints={
                 "detokenize": False,
                 "stop_token_ids": [2150],
@@ -40,7 +41,14 @@ QWEN3_TTS_PIPELINE = PipelineConfig(
             final_output_type="audio",
             engine_output_type="audio",
             model_arch="Qwen3TTSCode2Wav",
-            sync_process_input_func=f"{_PROC}.talker2code2wav",
+            # ``sync_process_input_func`` is the only input-proc override for
+            # this stage in sync (non-async-chunk) mode: a length-only
+            # ``_token_only`` placeholder.  The bulk codec payload itself
+            # ships via the worker connector from stage 0's
+            # ``talker2code2wav_full_payload`` producer.  Under async_chunk
+            # mode no pre-stage processing is needed -- chunks deliver
+            # directly to the consumer.
+            sync_process_input_func=f"{_PROC}.talker2code2wav_token_only",
             sampling_constraints={"detokenize": True},
             extras={"tts_args": {"max_instructions_length": 500}},
         ),
diff --git a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py
new file mode 100644
index 000000000..c9a0f266d
--- /dev/null
+++ b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py
@@ -0,0 +1,1543 @@
+"""Stand-alone builder for Qwen3-TTS AR talker prompt embeddings.
+
+Factors the prompt-construction logic out of
+:class:`vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker.Qwen3TTSTalkerForConditionalGeneration`
+so alternative AR talker implementations that share the same backbone
+embedding tables and codec / speaker encoders can reuse it without
+duplicating the (substantial) prompt layout logic.
+
+The builder is intentionally *not* coupled to the talker class. The talker
+constructs a builder, supplies its frozen embedding tables, projection
+layers, codec embed lookups, speaker encoder, the precomputed pad-token
+embedding buffer, and a callable that batch-encodes reference audio into
+codec frames (so the builder doesn't have to know whether the talker
+uses :class:`Qwen3TTSTokenizer` or :class:`Qwen3TTSTokenizerV2Encoder`
+under the hood). The single public entry point
+:meth:`Qwen3TTSPromptEmbedsBuilder.build_prompt_embeds` mirrors the
+official ``Qwen3TTSForConditionalGeneration`` prompt layout for the three
+supported task types (``CustomVoice``, ``VoiceDesign``, ``Base``) and
+both streaming / non-streaming modes.
+"""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+import io
+from collections import OrderedDict
+from collections.abc import Callable, Mapping, Sequence
+from functools import lru_cache
+from typing import TYPE_CHECKING, Any
+from urllib.parse import urlparse
+
+import numpy as np
+import soundfile as sf
+import torch
+import torch.nn as nn
+from transformers import AutoTokenizer
+from vllm.logger import init_logger
+from vllm.multimodal.audio import AudioResampler
+
+from vllm_omni.utils.audio import mel_filter_bank
+
+if TYPE_CHECKING:
+    from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSTalkerConfig
+
+logger = init_logger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Constants shared with the talker's preprocess pipeline.
+#
+# These keys mark precomputed artifacts that the batched preprocess pass
+# stashes into ``info_dict`` so the per-request prompt build can skip
+# redundant tokenization / codec encoding work.
+# ---------------------------------------------------------------------------
+
+PRECOMPUTED_REF_CODE_KEY = "precomputed_ref"
+NORMALIZED_REF_AUDIO_KEY = "_qwen3_tts_normalized_ref_audio"
+REF_AUDIO_CACHE_KEY = "_qwen3_tts_ref_audio_cache_key"
+PRECOMPUTED_TEXT_IDS_KEY = "_qwen3_tts_text_ids"
+PRECOMPUTED_REF_IDS_KEY = "_qwen3_tts_ref_ids"
+
+
+# ---------------------------------------------------------------------------
+# Chat-template helpers (must match the official HuggingFace reference).
+# ---------------------------------------------------------------------------
+
+
+def build_assistant_text(text: str) -> str:
+    return f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+
+
+def build_ref_text(text: str) -> str:
+    return f"<|im_start|>assistant\n{text}<|im_end|>\n"
+
+
+def build_instruct_text(instruct: str) -> str:
+    return f"<|im_start|>user\n{instruct}<|im_end|>\n"
+
+
+# ---------------------------------------------------------------------------
+# Type-coercion helpers. All accept the loose shapes that survive
+# msgspec IPC ``additional_information`` round-trips (scalars, single-
+# element lists, ndarrays, tensors).
+# ---------------------------------------------------------------------------
+
+
+def first_value(value: object, default: object = None) -> object:
+    if isinstance(value, list):
+        return value[0] if value else default
+    return value if value is not None else default
+
+
+def coerce_ref_code_tensor(value: object, *, device: torch.device) -> torch.Tensor | None:
+    value = first_value(value)
+    if value is None:
+        return None
+    if isinstance(value, torch.Tensor):
+        ref_code = value
+    elif isinstance(value, np.ndarray):
+        ref_code = torch.from_numpy(value)
+    elif isinstance(value, list) and value:
+        ref_code = torch.as_tensor(value, dtype=torch.long)
+    else:
+        return None
+    if ref_code.ndim == 3:
+        ref_code = ref_code[0]
+    if ref_code.ndim != 2 or ref_code.numel() == 0:
+        return None
+    return ref_code.to(device=device, dtype=torch.long).contiguous()
+
+
+def coerce_token_ids(value: object, *, device: torch.device) -> torch.Tensor | None:
+    value = first_value(value)
+    if value is None:
+        return None
+    if isinstance(value, torch.Tensor):
+        ids = value
+    elif isinstance(value, np.ndarray):
+        ids = torch.from_numpy(value)
+    elif isinstance(value, list) and value and all(isinstance(v, (int, np.integer)) for v in value):
+        ids = torch.tensor(value, dtype=torch.long)
+    else:
+        return None
+    if ids.ndim == 1:
+        ids = ids.unsqueeze(0)
+    if ids.ndim != 2 or ids.numel() == 0:
+        return None
+    return ids.to(device=device, dtype=torch.long).contiguous()
+
+
+# ---------------------------------------------------------------------------
+# Mel-spectrogram utilities used by the speaker encoder front-end.
+# Module-level LRU caches for the (sr, n_fft, n_mels, ...) -> mel_basis and
+# win_size -> hann_window tensors avoid rebuilding them on every request.
+# ---------------------------------------------------------------------------
+
+
+def _dynamic_range_compression(x: torch.Tensor, c: float = 1, clip_val: float = 1e-5) -> torch.Tensor:
+    return torch.log(torch.clamp(x, min=clip_val) * c)
+
+
+@lru_cache(maxsize=8)
+def _cached_mel_filter_bank(sampling_rate: int, n_fft: int, n_mels: int, fmin: int, fmax: int | None) -> torch.Tensor:
+    return mel_filter_bank(sr=sampling_rate, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax)
+
+
+@lru_cache(maxsize=8)
+def _cached_hann_window(win_size: int) -> torch.Tensor:
+    return torch.hann_window(win_size)
+
+
+def mel_spectrogram(
+    y: torch.Tensor,
+    n_fft: int,
+    num_mels: int,
+    sampling_rate: int,
+    hop_size: int,
+    win_size: int,
+    fmin: int,
+    fmax: int | None = None,
+    center: bool = False,
+    mel_basis: torch.Tensor | None = None,
+    hann_window: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Mel spectrogram via torch STFT and a (cached) mel filterbank."""
+    if torch.min(y) < -1.0:
+        logger.warning("Min value of input waveform signal is %s", torch.min(y))
+    if torch.max(y) > 1.0:
+        logger.warning("Max value of input waveform signal is %s", torch.max(y))
+    device = y.device
+    if mel_basis is None:
+        mel_basis = _cached_mel_filter_bank(sampling_rate, n_fft, num_mels, fmin, fmax).to(device)
+    elif mel_basis.device != device:
+        mel_basis = mel_basis.to(device)
+    if hann_window is None:
+        hann_window = _cached_hann_window(win_size).to(device)
+    elif hann_window.device != device:
+        hann_window = hann_window.to(device)
+    padding = (n_fft - hop_size) // 2
+    y = torch.nn.functional.pad(y.unsqueeze(1), (padding, padding), mode="reflect").squeeze(1)
+    spec = torch.stft(
+        y,
+        n_fft,
+        hop_length=hop_size,
+        win_length=win_size,
+        window=hann_window,
+        center=center,
+        pad_mode="reflect",
+        normalized=False,
+        onesided=True,
+        return_complex=True,
+    )
+    spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9)
+    mel_spec = torch.matmul(mel_basis, spec)
+    return _dynamic_range_compression(mel_spec)
+
+
+# ---------------------------------------------------------------------------
+# Local audio loading helpers (URL / base64 / local path).
+# ---------------------------------------------------------------------------
+
+
+def _is_probably_base64(s: str) -> bool:
+    if s.startswith("data:audio"):
+        return True
+    if ("/" not in s and "\\" not in s) and len(s) > 256:
+        return True
+    return False
+
+
+def _is_url(s: str) -> bool:
+    try:
+        u = urlparse(s)
+        if u.scheme in ("http", "https"):
+            return bool(u.netloc)
+        return u.scheme == "file"
+    except Exception:
+        return False
+
+
+def _decode_base64_to_wav_bytes(b64: str) -> bytes:
+    if "," in b64 and b64.strip().startswith("data:"):
+        b64 = b64.split(",", 1)[1]
+    return base64.b64decode(b64)
+
+
+def load_audio_to_np(x: str) -> tuple[np.ndarray, int]:
+    """Load mono float32 audio from a local path, URL, or base64 data URI.
+
+    Uses upstream vLLM's MediaConnector for http(s) URLs and ``file:`` URIs
+    with unrestricted local access (offline inference is trusted).
+    """
+    from vllm.multimodal.media.audio import load_audio
+
+    if _is_url(x):
+        from vllm.multimodal.media import MediaConnector
+
+        connector = MediaConnector(allowed_local_media_path="/")
+        audio, sr = connector.fetch_audio(x)
+    elif _is_probably_base64(x):
+        wav_bytes = _decode_base64_to_wav_bytes(x)
+        with io.BytesIO(wav_bytes) as f:
+            audio, sr = sf.read(f, dtype="float32", always_2d=False)
+    else:
+        audio, sr = load_audio(x, sr=None, mono=True)
+
+    if isinstance(audio, np.ndarray) and audio.ndim > 1:
+        audio = np.mean(audio, axis=-1)
+
+    return np.asarray(audio, dtype=np.float32), int(sr)
+
+
+# ---------------------------------------------------------------------------
+# Prompt-embeds builder.
+# ---------------------------------------------------------------------------
+
+
+class Qwen3TTSPromptEmbedsBuilder:
+    """Build prefill prompt embeddings for the Qwen3-TTS AR talker.
+
+    Stand-alone wrapper around the embedding layers, encoders, and speaker
+    cache required to assemble a talker prompt. Mirrors the layout of the
+    official HuggingFace ``Qwen3TTSForConditionalGeneration`` pipeline for
+    ``task_type in {CustomVoice, VoiceDesign, Base}`` and both streaming /
+    non-streaming modes.
+
+    The builder does **not** own the speech / codec tokenizer used to turn
+    raw reference audio into codec frames. The caller passes in a
+    ``encode_ref_audio_batch`` callable so the same builder works against
+    either the legacy :class:`Qwen3TTSTokenizer` HF wrapper or the native
+    :class:`Qwen3TTSTokenizerV2Encoder` exposed on the talker.
+
+    Args:
+        config: Top-level Qwen3-TTS config; supplies ``tts_*_token_id`` and
+            ``speaker_encoder_config``.
+        talker_config: Talker sub-config; supplies the codec control-token
+            ids, ``num_code_groups``, ``spk_id``, ``spk_is_dialect`` and
+            ``codec_language_id`` maps.
+        model_path: HuggingFace repo or local path used to lazily load the
+            text tokenizer.
+        text_embedding: Text-side embedding table.
+        text_projection: Callable projecting text-side embeddings into the
+            talker hidden size.
+        codec_embed: Callable performing codec-side embedding lookup
+            (normally the talker backbone's ``embed_input_ids``).
+        residual_code_embeddings: Callable returning the per-codebook
+            embedding modules used by the code predictor; called lazily so
+            the predictor doesn't need to be constructed before the builder.
+        speaker_encoder: ECAPA-TDNN speaker encoder used by the Base task.
+        tts_pad_embed: Pre-computed ``tts_pad`` embedding (shape
+            ``[1, hidden]``) reused at every prefill to avoid recomputing.
+        encode_ref_audio_batch: Callable ``(wavs, sr, *, device) -> list[Tensor]``
+            that returns one ``[T, num_quantizers]`` long tensor per input
+            waveform. Single-waveform encoding is implemented in terms of
+            this callable.
+        speaker_cache: Optional :class:`SpeakerEmbeddingCache`. When
+            provided, Base ``CustomVoice``-style results are cached by
+            ``(speaker_name, mode, created_at)``.
+        ref_audio_artifact_cache_max_entries: Upper bound on the SHA1
+            keyed LRU cache of per-ref-audio ``(ref_code, ref_spk_embedding)``
+            artifacts. ``0`` disables the cache.
+    """
+
+    def __init__(
+        self,
+        *,
+        config: Qwen3TTSConfig,
+        talker_config: Qwen3TTSTalkerConfig,
+        model_path: str,
+        text_embedding: nn.Module,
+        text_projection: Callable[[torch.Tensor], torch.Tensor],
+        codec_embed: Callable[[torch.Tensor], torch.Tensor],
+        residual_code_embeddings: Callable[[], Sequence[Callable[[torch.Tensor], torch.Tensor]]],
+        speaker_encoder: nn.Module,
+        tts_pad_embed: torch.Tensor,
+        encode_ref_audio_batch: Callable[..., list[torch.Tensor]],
+        speaker_cache: Any | None = None,
+        ref_audio_artifact_cache_max_entries: int = 256,
+    ):
+        self._config = config
+        self._talker_config = talker_config
+        self._model_path = model_path
+        self._text_embedding = text_embedding
+        self._text_projection = text_projection
+        self._codec_embed = codec_embed
+        self._residual_code_embeddings = residual_code_embeddings
+        self._speaker_encoder = speaker_encoder
+        self._tts_pad_embed_buffer = tts_pad_embed
+        self._encode_ref_audio_batch_fn = encode_ref_audio_batch
+        self._speaker_cache = speaker_cache
+
+        self._text_tokenizer: Any | None = None
+
+        self._ref_audio_artifact_cache_max_entries = int(ref_audio_artifact_cache_max_entries)
+        self._ref_audio_artifact_cache: OrderedDict[str, dict[str, torch.Tensor | bool]] = OrderedDict()
+
+        # Bounded LRU; caller-supplied orig_sr can otherwise grow this without limit.
+        self._resampler_cache: OrderedDict[tuple[int, int], AudioResampler] = OrderedDict()
+        self._resampler_cache_max = 16
+
+    # -------------------- device / buffer helpers --------------------
+
+    def _device(self) -> torch.device:
+        # text_embedding is always present and lives on the talker device.
+        return next(self._text_embedding.parameters()).device
+
+    def _pad_embed(self, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+        """Return the pad-token embedding on the requested device/dtype, shape ``[1, 1, H]``."""
+        return self._tts_pad_embed_buffer.to(device=device, dtype=dtype).reshape(1, 1, -1)
+
+    def _get_resampler(self, orig_sr: int, target_sr: int) -> AudioResampler:
+        key = (int(orig_sr), int(target_sr))
+        cache = self._resampler_cache
+        if key in cache:
+            cache.move_to_end(key)
+            return cache[key]
+        cache[key] = AudioResampler(target_sr=int(target_sr))
+        if len(cache) > self._resampler_cache_max:
+            cache.popitem(last=False)
+        return cache[key]
+
+    # -------------------- text tokenizer (lazy) --------------------
+
+    def get_text_tokenizer(self):
+        if self._text_tokenizer is None:
+            import transformers
+
+            kwargs: dict[str, Any] = dict(trust_remote_code=True, use_fast=True)
+            if transformers.__version__ < "5":
+                kwargs["fix_mistral_regex"] = True
+            tok = AutoTokenizer.from_pretrained(self._model_path, **kwargs)
+            tok.padding_side = "left"
+            self._text_tokenizer = tok
+        return self._text_tokenizer
+
+    # -------------------- ref-audio handling --------------------
+
+    def load_audio_to_np(self, x: str) -> tuple[np.ndarray, int]:
+        return load_audio_to_np(x)
+
+    def normalize_ref_audio(self, ref_audio: object) -> tuple[np.ndarray, int]:
+        """Coerce a serialized ref_audio payload into ``(wav_np, sample_rate)``.
+
+        ``additional_information`` may serialize ``(wav, sr)`` into (nested)
+        lists across processes, so we accept and unwrap a variety of shapes.
+        """
+        if isinstance(ref_audio, str):
+            return self.load_audio_to_np(ref_audio)
+
+        def _is_sr(x: object) -> bool:
+            try:
+                v = int(x)  # type: ignore[arg-type]
+            except Exception:
+                return False
+            return 1_000 <= v <= 200_000
+
+        def _is_number_sequence(xs: list[object]) -> bool:
+            if not xs:
+                return False
+            for v in xs[:8]:
+                if not isinstance(v, (int, float, np.number)):
+                    return False
+            return True
+
+        wav_candidates: list[object] = []
+        sr_candidates: list[int] = []
+
+        def _summarize(obj: object, depth: int = 0) -> str:
+            if depth > 2:
+                if isinstance(obj, (int, np.integer)):
+                    return f"int({int(obj)})"
+                return type(obj).__name__
+            if obj is None:
+                return "None"
+            if isinstance(obj, str):
+                if len(obj) <= 16:
+                    return f"str({obj!r})"
+                return f"str(len={len(obj)})"
+            if isinstance(obj, (int, float, np.number)):
+                return f"{type(obj).__name__}({obj})"
+            if isinstance(obj, np.ndarray):
+                return f"ndarray(shape={obj.shape}, dtype={obj.dtype})"
+            if isinstance(obj, torch.Tensor):
+                return f"Tensor(shape={tuple(obj.shape)}, dtype={obj.dtype}, device={obj.device})"
+            if isinstance(obj, dict):
+                keys = list(obj.keys())
+                return f"dict(keys={keys[:8]})"
+            if isinstance(obj, (tuple, list)):
+                items = list(obj)
+                head = ", ".join(_summarize(x, depth + 1) for x in items[:3])
+                return f"{type(obj).__name__}(len={len(items)}; head=[{head}])"
+            return f"{type(obj).__name__}"
+
+        def _scan(obj: object, depth: int = 0) -> None:
+            if depth > 4:
+                return
+            if obj is None:
+                return
+            if _is_sr(obj):
+                sr_candidates.append(int(obj))  # type: ignore[arg-type]
+                return
+            if isinstance(obj, np.ndarray) and obj.size > 0:
+                wav_candidates.append(obj)
+                return
+            if isinstance(obj, torch.Tensor) and obj.numel() > 0:
+                wav_candidates.append(obj)
+                return
+            if isinstance(obj, dict):
+                # Inlined ndarray/tensor payloads from the input processor.
+                if obj.get("__ndarray__") and "data" in obj and "dtype" in obj and "shape" in obj:
+                    try:
+                        data = obj["data"]
+                        dtype = obj["dtype"]
+                        shape = obj["shape"]
+                        if isinstance(data, (bytes, bytearray, memoryview)):
+                            arr = np.frombuffer(data, dtype=dtype).reshape(shape)
+                            if arr.size > 0:
+                                wav_candidates.append(arr)
+                                return
+                    except Exception:
+                        pass
+                if obj.get("__tensor__") and "data" in obj and "dtype" in obj and "shape" in obj:
+                    try:
+                        data = obj["data"]
+                        dtype = obj["dtype"]
+                        shape = obj["shape"]
+                        if isinstance(data, (bytes, bytearray, memoryview)):
+                            # Stored as raw CPU bytes; interpret as numpy for audio.
+                            np_dtype = np.dtype(dtype)
+                            arr = np.frombuffer(data, dtype=np_dtype).reshape(shape)
+                            if arr.size > 0:
+                                wav_candidates.append(arr)
+                                return
+                    except Exception:
+                        pass
+                wav_obj = obj.get("array") or obj.get("wav") or obj.get("audio")
+                sr_obj = obj.get("sampling_rate") or obj.get("sr") or obj.get("sample_rate")
+                if wav_obj is not None:
+                    _scan(wav_obj, depth + 1)
+                if sr_obj is not None:
+                    _scan(sr_obj, depth + 1)
+                return
+            if isinstance(obj, (tuple, list)):
+                obj_list = list(obj)
+                # Unwrap singleton nesting ([[wav, sr]]).
+                while isinstance(obj_list, list) and len(obj_list) == 1:
+                    inner = obj_list[0]
+                    if isinstance(inner, np.ndarray) and inner.size > 0:
+                        wav_candidates.append(inner)
+                        return
+                    if isinstance(inner, torch.Tensor) and inner.numel() > 0:
+                        wav_candidates.append(inner)
+                        return
+                    if isinstance(inner, dict):
+                        _scan(inner, depth + 1)
+                        return
+                    if isinstance(inner, (tuple, list)):
+                        obj_list = list(inner)  # type: ignore[list-item]
+                        continue
+                    break
+
+                # If the *unwrapped* list is a long list of numbers, treat it as waveform.
+                if len(obj_list) >= 512 and _is_number_sequence(obj_list):
+                    wav_candidates.append(obj_list)
+                    return
+
+                # Otherwise, recurse into elements (but avoid descending into huge numeric lists).
+                for item in obj_list:
+                    if isinstance(item, list) and len(item) >= 512 and _is_number_sequence(item):  # type: ignore[arg-type]
+                        wav_candidates.append(item)
+                        continue
+                    _scan(item, depth + 1)
+                return
+
+        _scan(ref_audio)
+        if not sr_candidates:
+            raise TypeError(f"ref_audio missing sample_rate: {_summarize(ref_audio)}")
+        sr = int(sr_candidates[0])
+
+        def _wav_len(x: object) -> int:
+            try:
+                if isinstance(x, np.ndarray):
+                    return int(x.size)
+                if isinstance(x, torch.Tensor):
+                    return int(x.numel())
+                if isinstance(x, list):
+                    return int(len(x))
+            except Exception:
+                pass
+            return 0
+
+        if not wav_candidates:
+            raise TypeError(f"ref_audio missing waveform: {_summarize(ref_audio)}")
+        wav_obj = max(wav_candidates, key=_wav_len)
+
+        def _to_np(x: object) -> np.ndarray:
+            if isinstance(x, np.ndarray):
+                return x.astype(np.float32).reshape(-1)
+            if isinstance(x, torch.Tensor):
+                return x.detach().to("cpu").float().contiguous().numpy().reshape(-1)
+            if isinstance(x, dict) and x.get("__ndarray__") and "data" in x and "dtype" in x and "shape" in x:
+                data = x["data"]
+                dtype = x["dtype"]
+                shape = x["shape"]
+                if isinstance(data, (bytes, bytearray, memoryview)):
+                    return np.frombuffer(data, dtype=dtype).reshape(shape).astype(np.float32).reshape(-1)
+            if isinstance(x, list):
+                # list of numbers
+                if len(x) >= 2 and _is_number_sequence(x):  # type: ignore[arg-type]
+                    return np.asarray(x, dtype=np.float32).reshape(-1)
+                # list of chunks
+                parts: list[np.ndarray] = []
+                for part in x:
+                    if isinstance(part, (np.ndarray, torch.Tensor, list)):
+                        parts.append(_to_np(part))
+                if parts:
+                    return np.concatenate(parts, axis=0)
+            raise TypeError(f"Unsupported waveform type: {type(x)}")
+
+        wav_np = _to_np(wav_obj)
+        if wav_np.size < 1024:
+            raise ValueError(f"ref_audio waveform too short: {wav_np.size} samples")
+        return wav_np, sr
+
+    # -------------------- ref-audio artifact cache --------------------
+
+    @staticmethod
+    def make_ref_audio_cache_key(wav: np.ndarray, sr: int) -> str:
+        wav_f32 = wav.astype(np.float32, copy=False).reshape(-1)
+        h = hashlib.sha1()
+        h.update(int(sr).to_bytes(4, byteorder="little", signed=False))
+        h.update(int(wav_f32.size).to_bytes(8, byteorder="little", signed=False))
+        h.update(wav_f32.tobytes(order="C"))
+        return h.hexdigest()
+
+    def get_ref_audio_artifacts(self, cache_key: str) -> dict[str, torch.Tensor | bool] | None:
+        if self._ref_audio_artifact_cache_max_entries <= 0:
+            return None
+        entry = self._ref_audio_artifact_cache.get(cache_key)
+        if entry is None:
+            return None
+        self._ref_audio_artifact_cache.move_to_end(cache_key, last=True)
+        return entry
+
+    def put_ref_audio_artifacts(
+        self,
+        cache_key: str,
+        *,
+        ref_code: torch.Tensor | None = None,
+        ref_spk_embedding: torch.Tensor | None = None,
+    ) -> None:
+        if self._ref_audio_artifact_cache_max_entries <= 0 or not cache_key:
+            return
+        entry = self._ref_audio_artifact_cache.get(cache_key)
+        if entry is None:
+            entry = {}
+        if isinstance(ref_code, torch.Tensor):
+            entry["ref_code"] = ref_code.detach().to("cpu", dtype=torch.long).contiguous()
+        if isinstance(ref_spk_embedding, torch.Tensor):
+            entry["ref_spk_embedding"] = ref_spk_embedding.detach().to("cpu", dtype=torch.bfloat16).reshape(-1)
+        if not entry:
+            return
+        self._ref_audio_artifact_cache[cache_key] = entry
+        self._ref_audio_artifact_cache.move_to_end(cache_key, last=True)
+        while len(self._ref_audio_artifact_cache) > self._ref_audio_artifact_cache_max_entries:
+            self._ref_audio_artifact_cache.popitem(last=False)
+
+    # -------------------- speaker encoder / codec encoder --------------------
+
+    def extract_speaker_embedding(self, wav: np.ndarray, sr: int) -> torch.Tensor:
+        # vLLM workers do not automatically move arbitrary torch.nn.Modules to
+        # CUDA. Ensure the speaker encoder is on the same device/dtype as the
+        # main model before running it.
+        dev = self._device()
+        try:
+            spk_param = next(self._speaker_encoder.parameters())
+            if spk_param.device != dev or spk_param.dtype != torch.bfloat16:
+                self._speaker_encoder.to(device=dev, dtype=torch.bfloat16)
+        except StopIteration:
+            pass
+
+        # Resample to 24kHz for speaker encoder.
+        target_sr = int(getattr(self._config.speaker_encoder_config, "sample_rate", 24000))
+        if sr != target_sr:
+            resampler = self._get_resampler(int(sr), target_sr)
+            wav = resampler.resample(wav.astype(np.float32), orig_sr=int(sr))
+            sr = target_sr
+
+        # Follow official implementation: mel_spectrogram expects 24kHz. Move
+        # the waveform first so STFT/mel computation stays on the model device
+        # instead of materializing a CPU mel tensor and copying it per request.
+        wav_tensor = torch.from_numpy(wav).to(device=dev, dtype=torch.float32).unsqueeze(0)
+        mels = mel_spectrogram(
+            wav_tensor,
+            n_fft=1024,
+            num_mels=128,
+            sampling_rate=24000,
+            hop_size=256,
+            win_size=1024,
+            fmin=0,
+            fmax=12000,
+        ).transpose(1, 2)
+        spk = self._speaker_encoder(mels.to(dtype=torch.bfloat16))[0]
+        return spk.to(dtype=torch.bfloat16)
+
+    def encode_ref_audio_batch(
+        self,
+        wavs: list[np.ndarray],
+        sr: int,
+        *,
+        device: torch.device,
+    ) -> list[torch.Tensor]:
+        """Delegate to the talker-provided batched ref-audio encoder."""
+        return self._encode_ref_audio_batch_fn(wavs, int(sr), device=device)
+
+    def encode_ref_audio_to_code(self, wav: np.ndarray, sr: int) -> torch.Tensor:
+        """Single-waveform convenience wrapper around :meth:`encode_ref_audio_batch`."""
+        codes = self.encode_ref_audio_batch([wav], int(sr), device=self._device())
+        if not codes:
+            raise ValueError("encode_ref_audio_batch returned no codes")
+        return codes[0]
+
+    # -------------------- batched preprocess (cross-request) --------------------
+
+    @staticmethod
+    def _voice_clone_prompt_dict(raw: object) -> dict[str, object] | None:
+        raw = first_value(raw)
+        if isinstance(raw, dict):
+            return raw
+        if isinstance(raw, list) and raw and isinstance(raw[0], dict):
+            return raw[0]
+        return None
+
+    @staticmethod
+    def _has_ref_code_like(value: object) -> bool:
+        value = first_value(value)
+        if isinstance(value, torch.Tensor):
+            return value.numel() > 0 and value.ndim in (2, 3)
+        if isinstance(value, np.ndarray):
+            return value.size > 0 and value.ndim in (2, 3)
+        if isinstance(value, list):
+            return bool(value)
+        return False
+
+    @staticmethod
+    def _needs_initial_prompt_preprocess(info_dict: dict[str, Any]) -> bool:
+        embed = info_dict.get("embed")
+        return not (isinstance(embed, dict) and "prefill" in embed)
+
+    def _needs_batched_ref_code(self, info_dict: dict[str, Any]) -> bool:
+        if not self._needs_initial_prompt_preprocess(info_dict):
+            return False
+        if first_value(info_dict.get("task_type"), "CustomVoice") != "Base":
+            return False
+        codes = info_dict.get("codes")
+        if isinstance(codes, dict) and self._has_ref_code_like(codes.get(PRECOMPUTED_REF_CODE_KEY)):
+            return False
+
+        voice_clone_prompt = self._voice_clone_prompt_dict(info_dict.get("voice_clone_prompt"))
+        if voice_clone_prompt is not None:
+            if self._has_ref_code_like(voice_clone_prompt.get("ref_code")):
+                return False
+            icl_flag = first_value(voice_clone_prompt.get("icl_mode"))
+            if isinstance(icl_flag, bool) and not icl_flag:
+                return False
+
+        xvec_only = first_value(info_dict.get("x_vector_only_mode"), False)
+        if isinstance(xvec_only, bool) and xvec_only:
+            return False
+
+        ref_audio_list = info_dict.get("ref_audio")
+        return isinstance(ref_audio_list, list) and bool(ref_audio_list)
+
+    @torch.inference_mode()
+    def preprocess_batch(
+        self,
+        *,
+        req_ids: list[str],
+        model_intermediate_buffer: dict[str, dict[str, Any]],
+        device: torch.device,
+    ) -> None:
+        """Batch Base voice-clone ref-audio codec extraction for current prefill requests.
+
+        Tokenizes assistant / ref text and runs the talker-supplied
+        ``encode_ref_audio_batch`` callable once per sample-rate group,
+        stashing the precomputed artifacts under
+        :data:`PRECOMPUTED_TEXT_IDS_KEY`, :data:`PRECOMPUTED_REF_IDS_KEY`,
+        :data:`PRECOMPUTED_REF_CODE_KEY`, and :data:`NORMALIZED_REF_AUDIO_KEY`
+        so the per-request :meth:`build_prompt_embeds` call can skip the
+        equivalent serial work.
+        """
+        pending_text: list[tuple[dict[str, Any], str]] = []
+        pending_ref_text: list[tuple[dict[str, Any], str]] = []
+        groups: dict[int, list[tuple[dict[str, Any], np.ndarray, int]]] = {}
+        for req_id in req_ids:
+            info_dict = model_intermediate_buffer.get(req_id)
+            if not isinstance(info_dict, dict):
+                continue
+
+            if (
+                self._needs_initial_prompt_preprocess(info_dict)
+                and first_value(info_dict.get("task_type"), "CustomVoice") == "Base"
+            ):
+                if PRECOMPUTED_TEXT_IDS_KEY not in info_dict:
+                    text = first_value(info_dict.get("text"), "")
+                    if isinstance(text, str):
+                        pending_text.append((info_dict, build_assistant_text(text)))
+
+                if PRECOMPUTED_REF_IDS_KEY not in info_dict:
+                    ref_text = first_value(info_dict.get("ref_text"))
+                    if isinstance(ref_text, str) and ref_text.strip():
+                        pending_ref_text.append((info_dict, build_ref_text(ref_text)))
+
+            if not self._needs_batched_ref_code(info_dict):
+                continue
+            ref_audio_list = info_dict.get("ref_audio")
+            if not isinstance(ref_audio_list, list) or not ref_audio_list:
+                continue
+            try:
+                wav, sr = self.normalize_ref_audio(ref_audio_list[0])
+            except Exception:
+                # Keep the original per-request path responsible for surfacing
+                # invalid ref_audio errors with its existing messages.
+                continue
+            info_dict[NORMALIZED_REF_AUDIO_KEY] = (wav, sr)
+            cache_key = self.make_ref_audio_cache_key(wav, sr)
+            info_dict[REF_AUDIO_CACHE_KEY] = cache_key
+            cached = self.get_ref_audio_artifacts(cache_key)
+            if cached is not None:
+                cached_ref_code = coerce_ref_code_tensor(cached.get("ref_code"), device=device)
+                if isinstance(cached_ref_code, torch.Tensor):
+                    info_dict.setdefault("codes", {})[PRECOMPUTED_REF_CODE_KEY] = cached_ref_code
+                    continue
+            groups.setdefault(int(sr), []).append((info_dict, wav, int(sr)))
+
+        if pending_text or pending_ref_text:
+            tok_text = self.get_text_tokenizer()
+            for items, key in (
+                (pending_text, PRECOMPUTED_TEXT_IDS_KEY),
+                (pending_ref_text, PRECOMPUTED_REF_IDS_KEY),
+            ):
+                if not items:
+                    continue
+                try:
+                    tokenized = tok_text([text for _, text in items], padding=False)
+                except Exception as exc:
+                    logger.debug("Qwen3-TTS batched text tokenization failed; falling back to serial path: %s", exc)
+                    continue
+                input_ids = tokenized.get("input_ids") if isinstance(tokenized, dict) else None
+                if not isinstance(input_ids, list) or len(input_ids) != len(items):
+                    continue
+                for (info_dict, _), ids in zip(items, input_ids, strict=True):
+                    if isinstance(ids, list) and ids:
+                        info_dict[key] = torch.tensor(ids, dtype=torch.long)
+
+        groups = {sr: items for sr, items in groups.items() if len(items) >= 2}
+        if not groups:
+            return
+
+        for sr, items in groups.items():
+            wavs = [wav for _, wav, _ in items]
+            try:
+                ref_codes = self.encode_ref_audio_batch(wavs, int(sr), device=device)
+            except Exception as exc:
+                logger.debug("Qwen3-TTS batched ref_code encode failed; falling back to serial path: %s", exc)
+                continue
+            for (info_dict, wav, item_sr), ref_code in zip(items, ref_codes, strict=True):
+                ref_code_t = coerce_ref_code_tensor(ref_code, device=device)
+                if ref_code_t is None:
+                    continue
+                info_dict.setdefault("codes", {})[PRECOMPUTED_REF_CODE_KEY] = ref_code_t
+                info_dict[NORMALIZED_REF_AUDIO_KEY] = (wav, item_sr)
+                cache_key = info_dict.get(REF_AUDIO_CACHE_KEY)
+                if isinstance(cache_key, str) and cache_key:
+                    self.put_ref_audio_artifacts(cache_key, ref_code=ref_code_t)
+
+    # -------------------- ICL prompt assembly --------------------
+
+    def _generate_icl_prompt(
+        self,
+        *,
+        text_id: torch.Tensor,
+        ref_id: torch.Tensor,
+        ref_code: torch.Tensor,
+        tts_pad_embed: torch.Tensor,
+        tts_eos_embed: torch.Tensor,
+        non_streaming_mode: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Port of the official ``Qwen3TTSForConditionalGeneration.generate_icl_prompt``."""
+        text_embed = self._text_projection(self._text_embedding(torch.cat([ref_id, text_id], dim=-1)))
+        text_embed = torch.cat([text_embed, tts_eos_embed], dim=1)
+
+        # codec embed (codec bos + codec) 1 T2 D
+        residual_embeds = self._residual_code_embeddings()
+        codec_embed: list[torch.Tensor] = []
+        for i in range(int(self._talker_config.num_code_groups)):
+            if i == 0:
+                codec_embed.append(self._codec_embed(ref_code[:, :1]))
+            else:
+                codec_embed.append(residual_embeds[i - 1](ref_code[:, i : i + 1]))
+        codec_embed_sum = torch.cat(codec_embed, dim=1).sum(1).unsqueeze(0)  # [1,T,H]
+        codec_embed_sum = torch.cat(
+            [
+                self._codec_embed(
+                    torch.tensor([[self._talker_config.codec_bos_id]], device=codec_embed_sum.device, dtype=torch.long)
+                ),
+                codec_embed_sum,
+            ],
+            dim=1,
+        )
+
+        text_lens = int(text_embed.shape[1])
+        codec_lens = int(codec_embed_sum.shape[1])
+        if non_streaming_mode:
+            # Official non-streaming mode: append the full text conditioning in
+            # prefill, and use PAD in decode steps.
+            icl_input_embed = text_embed + self._codec_embed(
+                torch.tensor(
+                    [[self._talker_config.codec_pad_id] * text_lens],
+                    device=codec_embed_sum.device,
+                    dtype=torch.long,
+                )
+            )
+            icl_input_embed = torch.cat([icl_input_embed, codec_embed_sum + tts_pad_embed], dim=1)
+            return icl_input_embed, tts_pad_embed
+        if text_lens > codec_lens:
+            return text_embed[:, :codec_lens] + codec_embed_sum, text_embed[:, codec_lens:]
+        text_embed = torch.cat([text_embed] + [tts_pad_embed] * (codec_lens - text_lens), dim=1)
+        return text_embed + codec_embed_sum, tts_pad_embed
+
+    # -------------------- main entry point --------------------
+
+    def build_prompt_embeds(
+        self,
+        *,
+        task_type: str,
+        info_dict: dict[str, Any],
+    ) -> tuple[torch.Tensor, torch.Tensor, int | None, torch.Tensor | None]:
+        """Build the prefill prompt embeddings for a single TTS request.
+
+        Args:
+            task_type: One of ``"CustomVoice"``, ``"VoiceDesign"``, ``"Base"``.
+            info_dict: Per-request ``additional_information`` plus any
+                precomputed artifacts stashed by the batched preprocess pass
+                (``PRECOMPUTED_TEXT_IDS_KEY``, ``PRECOMPUTED_REF_IDS_KEY``,
+                ``PRECOMPUTED_REF_CODE_KEY``, ``NORMALIZED_REF_AUDIO_KEY``).
+                Consumed keys are popped so the caller doesn't re-feed them
+                on subsequent prefill chunks.
+
+        Returns:
+            ``(talker_prompt, trailing_text_hidden, ref_code_len, ref_code)``:
+
+            * ``talker_prompt``: ``[prompt_len, hidden]`` prefill embedding.
+            * ``trailing_text_hidden``: ``[T, hidden]`` queue of text-side
+                embeddings consumed one-per-decode-step (streaming mode) or a
+                single pad row (non-streaming).
+            * ``ref_code_len``: Number of ref codec frames consumed when the
+                Base task ran in in-context mode (``None`` otherwise).
+            * ``ref_code``: The ``[T, Q]`` ref codec tensor used for ICL
+                (``None`` outside Base / ICL).
+        """
+        config = self._config
+        talker_config = self._talker_config
+        text_embedding = self._text_embedding
+        text_projection = self._text_projection
+        codec_embed = self._codec_embed
+
+        text = (info_dict.get("text") or [""])[0]
+        language = (info_dict.get("language") or ["Auto"])[0]
+        non_streaming_mode_val = info_dict.get("non_streaming_mode")
+        if isinstance(non_streaming_mode_val, list):
+            non_streaming_mode_raw = non_streaming_mode_val[0] if non_streaming_mode_val else None
+        else:
+            non_streaming_mode_raw = non_streaming_mode_val
+        if isinstance(non_streaming_mode_raw, bool):
+            non_streaming_mode = non_streaming_mode_raw
+        else:
+            # Match official inference defaults:
+            # - CustomVoice/VoiceDesign: non_streaming_mode=True
+            # - Base: non_streaming_mode=False
+            non_streaming_mode = task_type in ("CustomVoice", "VoiceDesign")
+
+        # Text ids for assistant template (always).
+        tok = self.get_text_tokenizer()
+        dev = self._device()
+        input_ids = coerce_token_ids(info_dict.pop(PRECOMPUTED_TEXT_IDS_KEY, None), device=dev)
+        if input_ids is None:
+            input_ids = tok(build_assistant_text(text), return_tensors="pt", padding=False)["input_ids"].to(device=dev)
+
+        # Optional instruct prefix.
+        instruct = (info_dict.get("instruct") or [""])[0]
+        instruct_embed = None
+        if isinstance(instruct, str) and instruct.strip():
+            instruct_ids = tok(build_instruct_text(instruct), return_tensors="pt", padding=False)["input_ids"].to(
+                device=input_ids.device
+            )
+            instruct_embed = text_projection(text_embedding(instruct_ids))
+
+        # tts special token embeds (projected into talker hidden).
+        # ``tts_pad_embed`` is precomputed (request-independent), so we only
+        # need bos/eos here.
+        tts_tokens = torch.tensor(
+            [[config.tts_bos_token_id, config.tts_eos_token_id]],
+            device=input_ids.device,
+            dtype=input_ids.dtype,
+        )
+        tts_bos_embed, tts_eos_embed = text_projection(text_embedding(tts_tokens)).chunk(2, dim=1)
+        tts_pad_embed = self._pad_embed(input_ids.device, tts_bos_embed.dtype)
+
+        # Codec prefill tags.
+        language_id = None
+        if isinstance(language, str) and language.lower() != "auto":
+            language_id = talker_config.codec_language_id.get(language.lower())
+        # Match official dialect override:
+        # If language is Chinese/Auto and the selected speaker is a dialect voice,
+        # set language_id to that dialect to improve code generation stability.
+        if language_id is None and isinstance(language, str) and language.lower() in ("chinese", "auto"):
+            speaker_for_dialect = None
+            if task_type == "CustomVoice":
+                speaker_for_dialect = (info_dict.get("speaker") or [""])[0]
+            if isinstance(speaker_for_dialect, str) and speaker_for_dialect.strip():
+                spk_is_dialect = getattr(talker_config, "spk_is_dialect", None) or {}
+                dialect = spk_is_dialect.get(speaker_for_dialect.lower())
+                if isinstance(dialect, str) and dialect:
+                    language_id = talker_config.codec_language_id.get(dialect)
+        if language_id is None:
+            codec_prefill_list = [
+                [
+                    talker_config.codec_nothink_id,
+                    talker_config.codec_think_bos_id,
+                    talker_config.codec_think_eos_id,
+                ]
+            ]
+        else:
+            codec_prefill_list = [
+                [
+                    talker_config.codec_think_id,
+                    talker_config.codec_think_bos_id,
+                    int(language_id),
+                    talker_config.codec_think_eos_id,
+                ]
+            ]
+
+        codec_input_0 = codec_embed(torch.tensor(codec_prefill_list, device=input_ids.device, dtype=torch.long))
+        codec_input_1 = codec_embed(
+            torch.tensor([[talker_config.codec_pad_id, talker_config.codec_bos_id]], device=input_ids.device)
+        )
+
+        # Speaker embedding/token (task-dependent)
+        speaker_embed: torch.Tensor | None = None
+        ref_code_len: int | None = None
+        ref_code_prompt: torch.Tensor | None = None
+
+        def _as_singleton(x: object) -> object:
+            if isinstance(x, list):
+                return x[0] if x else None
+            return x
+
+        def _to_long_tensor(x: object, *, device: torch.device) -> torch.Tensor | None:
+            x = _as_singleton(x)
+            if x is None:
+                return None
+            if isinstance(x, torch.Tensor):
+                t = x
+            elif isinstance(x, np.ndarray):
+                t = torch.from_numpy(x)
+            elif isinstance(x, list) and x and all(isinstance(v, (int, np.integer)) for v in x):
+                t = torch.tensor(x, dtype=torch.long)
+            else:
+                return None
+            if t.ndim == 1:
+                t = t.unsqueeze(0)
+            return t.to(device=device, dtype=torch.long)
+
+        def _normalize_voice_clone_prompt(raw: object) -> dict[str, object] | None:
+            raw = _as_singleton(raw)
+            if raw is None:
+                return None
+            if isinstance(raw, dict):
+                return raw
+            # Some callers may pass list[dict] directly.
+            if isinstance(raw, list) and raw and isinstance(raw[0], dict):
+                return raw[0]
+            return None
+
+        if task_type == "Base":
+            # Base supports voice clone prompt with in-context mode.
+            xvec_only = bool((info_dict.get("x_vector_only_mode") or [False])[0])
+            in_context_mode = not xvec_only
+            voice_clone_prompt = _normalize_voice_clone_prompt(info_dict.get("voice_clone_prompt"))
+            ref_audio_wav: np.ndarray | None = None
+            ref_audio_sr: int | None = None
+
+            def _get_ref_audio() -> tuple[np.ndarray, int]:
+                nonlocal ref_audio_wav, ref_audio_sr
+                if ref_audio_wav is None or ref_audio_sr is None:
+                    normalized_ref_audio = info_dict.pop(NORMALIZED_REF_AUDIO_KEY, None)
+                    if (
+                        isinstance(normalized_ref_audio, tuple)
+                        and len(normalized_ref_audio) == 2
+                        and isinstance(normalized_ref_audio[0], np.ndarray)
+                    ):
+                        ref_audio_wav = normalized_ref_audio[0]
+                        ref_audio_sr = int(normalized_ref_audio[1])
+                        return ref_audio_wav, ref_audio_sr
+                    ref_audio_list = info_dict.get("ref_audio")
+                    if not isinstance(ref_audio_list, list) or not ref_audio_list:
+                        raise ValueError("Base requires `ref_audio`.")
+                    ref_audio_wav, ref_audio_sr = self.normalize_ref_audio(ref_audio_list[0])
+                return ref_audio_wav, ref_audio_sr
+
+            # Per-ref-audio artifact cache: usually pre-populated by
+            # preprocess_batch. Fall back to on-demand key creation for
+            # paths it skips but that still extract an ECAPA spk_embed
+            # from ref_audio: x_vector_only_mode, precomputed ref_code
+            # with no ref_spk_embedding, voice_clone_prompt.icl_mode=False.
+            ref_audio_cache_key = info_dict.pop(REF_AUDIO_CACHE_KEY, None)
+            if not (isinstance(ref_audio_cache_key, str) and ref_audio_cache_key):
+                normalized = info_dict.get(NORMALIZED_REF_AUDIO_KEY)
+                wav_peek: np.ndarray | None = None
+                sr_peek: int | None = None
+                if isinstance(normalized, tuple) and len(normalized) == 2 and isinstance(normalized[0], np.ndarray):
+                    wav_peek, sr_peek = normalized[0], int(normalized[1])
+                else:
+                    ref_audio_list = info_dict.get("ref_audio")
+                    if isinstance(ref_audio_list, list) and ref_audio_list:
+                        try:
+                            wav_peek, sr_peek = self.normalize_ref_audio(ref_audio_list[0])
+                            info_dict[NORMALIZED_REF_AUDIO_KEY] = (wav_peek, int(sr_peek))
+                        except Exception:
+                            wav_peek = None
+                if isinstance(wav_peek, np.ndarray) and sr_peek is not None:
+                    ref_audio_cache_key = self.make_ref_audio_cache_key(wav_peek, sr_peek)
+            cached_artifacts = None
+            if isinstance(ref_audio_cache_key, str) and ref_audio_cache_key:
+                cached_artifacts = self.get_ref_audio_artifacts(ref_audio_cache_key)
+
+            # Speaker cache: only for uploaded (named) speakers
+            _speaker_cache_key = None
+            if voice_clone_prompt is None and self._speaker_cache is not None:
+                _speaker_list = info_dict.get("speaker")
+                if isinstance(_speaker_list, list) and _speaker_list:
+                    _voice_name = str(_speaker_list[0]).lower()
+                    # Per-mode namespace — xvec and icl produce different artifacts
+                    # for the same voice, so they must not share a cache slot.
+                    _mode = "xvec" if xvec_only else "icl"
+                    _voice_created_at = int((info_dict.get("voice_created_at") or [0])[0])
+                    _speaker_cache_key = self._speaker_cache.make_cache_key(
+                        _voice_name,
+                        model_type=f"qwen3_tts_{_mode}",
+                        created_at=_voice_created_at,
+                    )
+                    _cached = self._speaker_cache.get(_speaker_cache_key)
+                    if _cached is not None:
+                        ref_code_cached = _cached.get("ref_code")
+                        ref_spk_embed_cached = _cached.get("ref_spk_embedding")
+                        if isinstance(ref_code_cached, torch.Tensor):
+                            ref_code_cached = ref_code_cached.to(device=input_ids.device)
+                        if isinstance(ref_spk_embed_cached, torch.Tensor):
+                            ref_spk_embed_cached = ref_spk_embed_cached.to(device=input_ids.device)
+                        voice_clone_prompt = {
+                            "ref_code": ref_code_cached,
+                            "ref_spk_embedding": ref_spk_embed_cached,
+                            "icl_mode": _cached.get("icl_mode"),
+                        }
+                        _speaker_cache_key = None  # hit → don't store again
+
+            # Official implementation may pass `voice_clone_prompt.icl_mode`.
+            if voice_clone_prompt is not None and "icl_mode" in voice_clone_prompt:
+                icl_flag = _as_singleton(voice_clone_prompt.get("icl_mode"))
+                if isinstance(icl_flag, bool):
+                    in_context_mode = icl_flag
+                    xvec_only = not in_context_mode
+            ref_code = None
+            if voice_clone_prompt is not None:
+                ref_code = _as_singleton(voice_clone_prompt.get("ref_code"))
+            ref_code_t = None
+            if isinstance(ref_code, torch.Tensor):
+                ref_code_t = ref_code
+            elif isinstance(ref_code, np.ndarray):
+                ref_code_t = torch.from_numpy(ref_code)
+            if isinstance(ref_code_t, torch.Tensor):
+                if ref_code_t.ndim == 3:
+                    ref_code_t = ref_code_t[0]
+                ref_code_t = ref_code_t.to(device=input_ids.device, dtype=torch.long)
+                ref_code_len = int(ref_code_t.shape[0])
+            else:
+                codes = info_dict.get("codes")
+                precomputed_ref_code = codes.get(PRECOMPUTED_REF_CODE_KEY) if isinstance(codes, dict) else None
+                ref_code_t = coerce_ref_code_tensor(precomputed_ref_code, device=input_ids.device)
+                if isinstance(ref_code_t, torch.Tensor):
+                    ref_code_len = int(ref_code_t.shape[0])
+            if ref_code_t is None and in_context_mode and cached_artifacts is not None:
+                cached_ref_code = coerce_ref_code_tensor(cached_artifacts.get("ref_code"), device=input_ids.device)
+                if isinstance(cached_ref_code, torch.Tensor):
+                    ref_code_t = cached_ref_code
+                    ref_code_len = int(ref_code_t.shape[0])
+            if ref_code_t is None and in_context_mode:
+                # Compute ref_code from ref_audio if not provided.
+                wav_np, sr = _get_ref_audio()
+                ref_code_t = self.encode_ref_audio_to_code(wav_np, sr).to(device=input_ids.device)
+                ref_code_len = int(ref_code_t.shape[0])
+            if isinstance(ref_code_t, torch.Tensor):
+                ref_code_prompt = ref_code_t
+
+            # Speaker embedding: use prompt embed if provided; otherwise extract from audio.
+            # NOTE: Do NOT use _as_singleton here — the embedding may be a plain
+            # float list (from API via msgspec IPC) that _as_singleton would
+            # destructively unwrap to a single scalar.
+            spk = None
+            if voice_clone_prompt is not None:
+                spk = voice_clone_prompt.get("ref_spk_embedding")
+            if isinstance(spk, torch.Tensor):
+                speaker_embed = spk.to(device=input_ids.device, dtype=torch.bfloat16).view(1, 1, -1)
+            elif isinstance(spk, (list, np.ndarray)):
+                speaker_embed = torch.tensor(spk, dtype=torch.bfloat16, device=input_ids.device).view(1, 1, -1)
+            elif cached_artifacts is not None and isinstance(cached_artifacts.get("ref_spk_embedding"), torch.Tensor):
+                speaker_embed = (
+                    cached_artifacts["ref_spk_embedding"]
+                    .to(device=input_ids.device, dtype=torch.bfloat16)
+                    .view(1, 1, -1)
+                )
+            else:
+                wav_np, sr = _get_ref_audio()
+                speaker_embed = self.extract_speaker_embedding(wav_np, sr).view(1, 1, -1)
+
+            # Cache miss: store extraction result in the speaker cache.
+            if _speaker_cache_key is not None and speaker_embed is not None and self._speaker_cache is not None:
+                self._speaker_cache.put(
+                    _speaker_cache_key,
+                    {
+                        "ref_code": ref_code_prompt.detach().cpu()
+                        if isinstance(ref_code_prompt, torch.Tensor)
+                        else None,
+                        "ref_spk_embedding": speaker_embed.detach().cpu().reshape(-1),
+                        "icl_mode": in_context_mode,
+                    },
+                )
+
+            # Cache miss: store extraction result in the per-ref-audio artifact cache.
+            if isinstance(ref_audio_cache_key, str) and ref_audio_cache_key:
+                self.put_ref_audio_artifacts(
+                    ref_audio_cache_key,
+                    ref_code=ref_code_prompt if isinstance(ref_code_prompt, torch.Tensor) else None,
+                    ref_spk_embedding=speaker_embed.reshape(-1) if isinstance(speaker_embed, torch.Tensor) else None,
+                )
+
+            codec_input = torch.cat([codec_input_0, speaker_embed, codec_input_1], dim=1)
+
+            # Role header (<|im_start|>assistant\n) -> projected text embeds.
+            role_embed = text_projection(text_embedding(input_ids[:, :3]))
+
+            codec_prefix = torch.cat((tts_pad_embed.expand(-1, codec_input.shape[1] - 2, -1), tts_bos_embed), dim=1)
+            codec_prefix = codec_prefix + codec_input[:, :-1]
+            talker_prompt = torch.cat((role_embed, codec_prefix), dim=1)
+
+            if in_context_mode:
+                # Prefer explicit tokenized `ref_ids` if provided (matches official signature).
+                ref_ids = coerce_token_ids(info_dict.pop(PRECOMPUTED_REF_IDS_KEY, None), device=input_ids.device)
+                if ref_ids is None:
+                    ref_ids = _to_long_tensor(info_dict.get("ref_ids"), device=input_ids.device)
+                if ref_ids is None and voice_clone_prompt is not None:
+                    ref_ids = _to_long_tensor(
+                        voice_clone_prompt.get("ref_ids") or voice_clone_prompt.get("ref_id"),
+                        device=input_ids.device,
+                    )
+                if ref_ids is None:
+                    ref_text = _as_singleton(info_dict.get("ref_text"))
+                    if isinstance(ref_text, str) and ref_text.strip():
+                        ref_ids = tok(
+                            build_ref_text(ref_text),
+                            return_tensors="pt",
+                            padding=False,
+                        )["input_ids"].to(device=input_ids.device)
+                    else:
+                        logger.warning("Base ICL: ref_text/ref_ids missing, falling back to x-vector-only mode.")
+                        in_context_mode = False
+            if in_context_mode:
+                icl_input_embed, trailing_text_hidden = self._generate_icl_prompt(
+                    text_id=input_ids[:, 3:-5],
+                    ref_id=ref_ids[:, 3:-2],
+                    ref_code=ref_code_t,  # type: ignore[arg-type]
+                    tts_pad_embed=tts_pad_embed,
+                    tts_eos_embed=tts_eos_embed,
+                    non_streaming_mode=non_streaming_mode,
+                )
+                talker_prompt = torch.cat([talker_prompt, icl_input_embed], dim=1)
+            else:
+                # First text token (+ codec_bos).
+                if non_streaming_mode:
+                    # Official non-streaming mode: put the full text into the
+                    # prefill prompt and use PAD for decode steps.
+                    text_all = text_projection(text_embedding(input_ids[:, 3:-5]))
+                    text_all = torch.cat([text_all, tts_eos_embed], dim=1)
+                    pad_ids = torch.full(
+                        (1, int(text_all.shape[1])),
+                        int(talker_config.codec_pad_id),
+                        device=input_ids.device,
+                        dtype=torch.long,
+                    )
+                    talker_prompt = torch.cat(
+                        [
+                            talker_prompt,
+                            text_all + codec_embed(pad_ids),
+                            tts_pad_embed
+                            + codec_embed(torch.tensor([[talker_config.codec_bos_id]], device=input_ids.device)),
+                        ],
+                        dim=1,
+                    )
+                    trailing_text_hidden = tts_pad_embed
+                else:
+                    first_text = text_projection(text_embedding(input_ids[:, 3:4])) + codec_input[:, -1:]
+                    talker_prompt = torch.cat([talker_prompt, first_text], dim=1)
+                    trailing_text_hidden = torch.cat(
+                        (
+                            text_projection(text_embedding(input_ids[:, 4:-5])),
+                            tts_eos_embed,
+                        ),
+                        dim=1,
+                    )
+
+        elif task_type == "CustomVoice":
+            _speaker_raw = info_dict.get("speaker") or [""]
+            speaker = (
+                ((_speaker_raw[0] if isinstance(_speaker_raw, (list, tuple)) else _speaker_raw) or "").lower().strip()
+            )
+            if not speaker:
+                raise ValueError("CustomVoice requires additional_information.speaker.")
+            spk_id_map = {k.lower(): v for k, v in (getattr(talker_config, "spk_id", None) or {}).items()}
+            if speaker not in spk_id_map:
+                raise ValueError(f"Unsupported speaker: {speaker}")
+            spk_id = spk_id_map[speaker]
+            # Keep it at least 1D; embedding on a 0-d tensor can return 1D.
+            spk_tensor = torch.tensor([spk_id], device=input_ids.device, dtype=torch.long)
+            spk_embed = codec_embed(spk_tensor)
+            if spk_embed.ndim == 1:
+                spk_embed = spk_embed.view(1, 1, -1)
+            elif spk_embed.ndim == 2:
+                spk_embed = spk_embed.view(1, 1, -1)
+            speaker_embed = spk_embed
+            codec_input = torch.cat([codec_input_0, speaker_embed, codec_input_1], dim=1)
+
+            role_embed = text_projection(text_embedding(input_ids[:, :3]))
+            codec_prefix = torch.cat((tts_pad_embed.expand(-1, codec_input.shape[1] - 2, -1), tts_bos_embed), dim=1)
+            codec_prefix = codec_prefix + codec_input[:, :-1]
+            talker_prompt = torch.cat((role_embed, codec_prefix), dim=1)
+
+            if non_streaming_mode:
+                text_all = text_projection(text_embedding(input_ids[:, 3:-5]))
+                text_all = torch.cat([text_all, tts_eos_embed], dim=1)
+                pad_ids = torch.full(
+                    (1, int(text_all.shape[1])),
+                    int(talker_config.codec_pad_id),
+                    device=input_ids.device,
+                    dtype=torch.long,
+                )
+                talker_prompt = torch.cat(
+                    [
+                        talker_prompt,
+                        text_all + codec_embed(pad_ids),
+                        tts_pad_embed
+                        + codec_embed(torch.tensor([[talker_config.codec_bos_id]], device=input_ids.device)),
+                    ],
+                    dim=1,
+                )
+                trailing_text_hidden = tts_pad_embed
+            else:
+                first_text = text_projection(text_embedding(input_ids[:, 3:4])) + codec_input[:, -1:]
+                talker_prompt = torch.cat([talker_prompt, first_text], dim=1)
+                trailing_text_hidden = torch.cat(
+                    (
+                        text_projection(text_embedding(input_ids[:, 4:-5])),
+                        tts_eos_embed,
+                    ),
+                    dim=1,
+                )
+
+        elif task_type == "VoiceDesign":
+            # No known speaker identity; only codec tags + text.
+            codec_input = torch.cat([codec_input_0, codec_input_1], dim=1)
+
+            role_embed = text_projection(text_embedding(input_ids[:, :3]))
+            codec_prefix = torch.cat((tts_pad_embed.expand(-1, codec_input.shape[1] - 2, -1), tts_bos_embed), dim=1)
+            codec_prefix = codec_prefix + codec_input[:, :-1]
+            talker_prompt = torch.cat((role_embed, codec_prefix), dim=1)
+
+            if non_streaming_mode:
+                text_all = text_projection(text_embedding(input_ids[:, 3:-5]))
+                text_all = torch.cat([text_all, tts_eos_embed], dim=1)
+                pad_ids = torch.full(
+                    (1, int(text_all.shape[1])),
+                    int(talker_config.codec_pad_id),
+                    device=input_ids.device,
+                    dtype=torch.long,
+                )
+                talker_prompt = torch.cat(
+                    [
+                        talker_prompt,
+                        text_all + codec_embed(pad_ids),
+                        tts_pad_embed
+                        + codec_embed(torch.tensor([[talker_config.codec_bos_id]], device=input_ids.device)),
+                    ],
+                    dim=1,
+                )
+                trailing_text_hidden = tts_pad_embed
+            else:
+                first_text = text_projection(text_embedding(input_ids[:, 3:4])) + codec_input[:, -1:]
+                talker_prompt = torch.cat([talker_prompt, first_text], dim=1)
+                trailing_text_hidden = torch.cat(
+                    (
+                        text_projection(text_embedding(input_ids[:, 4:-5])),
+                        tts_eos_embed,
+                    ),
+                    dim=1,
+                )
+        else:
+            raise ValueError(f"Unsupported task_type={task_type}")
+
+        if instruct_embed is not None:
+            talker_prompt = torch.cat([instruct_embed, talker_prompt], dim=1)
+
+        return (
+            talker_prompt.squeeze(0),  # [prompt_len, H]
+            trailing_text_hidden.squeeze(0),  # [T, H]
+            ref_code_len,
+            ref_code_prompt.contiguous() if isinstance(ref_code_prompt, torch.Tensor) else None,
+        )
+
+    # -------------------- prompt-length estimator --------------------
+
+    @staticmethod
+    def estimate_prompt_len_from_additional_information(
+        additional_information: dict[str, Any] | None,
+        *,
+        task_type: str,
+        tokenize_prompt: Callable[[str], list[int]],
+        codec_language_id: Mapping[str, int] | None,
+        spk_is_dialect: Mapping[str, object] | None,
+        estimate_ref_code_len: Callable[[object], int | None] | None = None,
+    ) -> int:
+        """Compute Stage-0 placeholder prompt length (length-only mirror of
+        :meth:`build_prompt_embeds`).
+
+        It must match the model-side ``inputs_embeds`` length to avoid extra
+        padding and quality drop.
+        """
+
+        def _first(x: object, default: object) -> object:
+            if isinstance(x, list):
+                return x[0] if x else default
+            return x if x is not None else default
+
+        info: dict[str, Any] = additional_information or {}
+        text = _first(info.get("text"), "")
+        language = _first(info.get("language"), "Auto")
+        speaker = _first(info.get("speaker"), "").lower().strip()
+        instruct = _first(info.get("instruct"), "")
+        non_streaming_mode_raw = _first(info.get("non_streaming_mode"), None)
+
+        if isinstance(non_streaming_mode_raw, bool):
+            non_streaming_mode = non_streaming_mode_raw
+        else:
+            # Official defaults: CustomVoice/VoiceDesign -> non_streaming_mode=True; Base -> False.
+            non_streaming_mode = task_type in ("CustomVoice", "VoiceDesign")
+
+        if not isinstance(text, str):
+            text = ""
+        if not isinstance(instruct, str):
+            instruct = ""
+        if not isinstance(language, str):
+            language = "Auto"
+
+        instruct_len = 0
+        if instruct.strip():
+            instruct_len = len(tokenize_prompt(build_instruct_text(instruct)))
+
+        # ---- codec prefix portion (matches build_prompt_embeds) ----
+        language_id = None
+        if language.lower() != "auto" and codec_language_id:
+            language_id = codec_language_id.get(language.lower())
+        if (
+            language_id is None
+            and codec_language_id
+            and spk_is_dialect
+            and isinstance(language, str)
+            and language.lower() in ("chinese", "auto")
+            and isinstance(speaker, str)
+            and speaker.strip()
+        ):
+            dialect = spk_is_dialect.get(speaker.lower())
+            if isinstance(dialect, str) and dialect:
+                language_id = codec_language_id.get(dialect)
+        prefill_len = 3 if language_id is None else 4
+
+        speaker_len = 1 if task_type in ("CustomVoice", "Base") else 0
+        codec_input_len = prefill_len + speaker_len + 2  # + [codec_pad, codec_bos]
+        codec_prefix_len = codec_input_len - 1  # codec_input[:-1] + tts_bos
+
+        # Role header: input_ids[:, :3] in model.
+        role_len = 3
+        prompt_len = instruct_len + role_len + codec_prefix_len
+
+        # ---- text conditioning portion (matches build_prompt_embeds) ----
+        assistant_len = len(tokenize_prompt(build_assistant_text(text)))
+        if assistant_len < 8:
+            raise ValueError(f"Unexpected assistant prompt length: {assistant_len}")
+
+        if task_type in ("CustomVoice", "VoiceDesign"):
+            if non_streaming_mode:
+                # model: full text ids (input_ids[:, 3:-5]) + eos + codec_bos step
+                prompt_len += assistant_len - 6
+            else:
+                # model: only first text token in prefill
+                prompt_len += 1
+
+        if task_type == "Base":
+            xvec_only = bool(_first(info.get("x_vector_only_mode"), False))
+            in_context_mode = not xvec_only
+
+            voice_clone_prompt = _first(info.get("voice_clone_prompt"), None)
+            if isinstance(voice_clone_prompt, dict):
+                icl_flag = _first(voice_clone_prompt.get("icl_mode"), None)
+                if isinstance(icl_flag, bool):
+                    in_context_mode = icl_flag
+
+            if in_context_mode:
+                ref_code = None
+                if isinstance(voice_clone_prompt, dict):
+                    # Do NOT apply `_first(...)` here. `_first` unwraps the singleton-list
+                    # batching convention used at the top level of `additional_information`;
+                    # values inside `voice_clone_prompt` are per-item payloads. `ref_code`
+                    # is a `[num_frames, num_codebooks]` 2D list, so `_first` would strip
+                    # its outer dim and report `num_codebooks` (~8) instead of `num_frames`
+                    # (~hundreds), silently truncating the prefill embeddings downstream.
+                    ref_code = voice_clone_prompt.get("ref_code")
+
+                ref_code_len: int | None = None
+                if isinstance(ref_code, list):
+                    if ref_code and isinstance(ref_code[0], list):
+                        ref_code_len = len(ref_code)
+                    elif ref_code:
+                        ref_code_len = len(ref_code)
+                elif hasattr(ref_code, "shape"):
+                    try:
+                        shape = getattr(ref_code, "shape")
+                        if shape and len(shape) >= 1:
+                            ref_code_len = int(shape[0])
+                    except Exception:
+                        ref_code_len = None
+
+                if ref_code_len is None and estimate_ref_code_len is not None:
+                    ref_code_len = estimate_ref_code_len(info.get("ref_audio"))
+                if ref_code_len is None:
+                    raise ValueError(
+                        "Base in-context voice cloning requires either `voice_clone_prompt.ref_code` "
+                        "or a readable `ref_audio` that can be mapped to a codec frame length."
+                    )
+
+                codec_lens = 1 + int(ref_code_len)  # codec_bos + ref_code
+                if non_streaming_mode:
+                    # Mirrors _generate_icl_prompt(non_streaming_mode=True):
+                    # text_embed = ref_ids + text_ids + eos.
+                    ref_ids = _first(info.get("ref_ids"), None)
+                    if isinstance(voice_clone_prompt, dict) and ref_ids is None:
+                        # Same reasoning as ref_code above: values inside `voice_clone_prompt`
+                        # are per-item payloads, not singleton-wrapped batches. Applying
+                        # `_first` to a list-of-int `ref_ids` collapses it to an int and the
+                        # branches below fall through to `ref_ids_len = 0`.
+                        ref_ids = voice_clone_prompt.get("ref_ids") or voice_clone_prompt.get("ref_id")
+
+                    if ref_ids is None:
+                        ref_text = _first(info.get("ref_text"), "")
+                        if not isinstance(ref_text, str) or not ref_text.strip():
+                            raise ValueError(
+                                "Base in-context non-streaming requires `ref_text` or tokenized `ref_ids`."
+                            )
+                        ref_ids_len = len(tokenize_prompt(build_ref_text(ref_text)))
+                    elif hasattr(ref_ids, "shape"):
+                        shape = getattr(ref_ids, "shape", None)
+                        ref_ids_len = int(shape[-1]) if shape else 0
+                    elif isinstance(ref_ids, list):
+                        ref_ids_len = len(ref_ids)
+                    else:
+                        ref_ids_len = 0
+
+                    # model uses ref_ids[:, 3:-2] (strip 5 tokens) and text_id=input_ids[:, 3:-5] (strip 8).
+                    ref_id_len = max(0, int(ref_ids_len) - 5)
+                    text_id_len = max(0, int(assistant_len) - 8)
+                    text_embed_len = ref_id_len + text_id_len + 1  # + eos
+                    prompt_len += text_embed_len + codec_lens
+                else:
+                    # Mirrors _generate_icl_prompt(non_streaming_mode=False):
+                    # aligned to codec_lens.
+                    prompt_len += codec_lens
+            else:
+                # Base without ICL behaves like CustomVoice.
+                if non_streaming_mode:
+                    prompt_len += assistant_len - 6
+                else:
+                    prompt_len += 1
+
+        return max(2, int(prompt_len))
diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py
index 8c200141a..b7c3c9d6e 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py
+++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py
@@ -25,6 +25,27 @@ from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
 logger = init_logger(__name__)
 
 
+def _codec_ids_from_payload_or_input(
+    input_ids: torch.Tensor,
+    runtime_info: dict[str, Any] | None,
+) -> torch.Tensor:
+    """Prefer connector-delivered codec ids over token placeholders.
+
+    In non-async full-payload mode, the scheduler only needs placeholder
+    token ids for allocation.  The real codec sequence is delivered through
+    model_intermediate_buffer as ``codes.audio``.
+    """
+    if isinstance(runtime_info, dict):
+        codes = runtime_info.get("codes")
+        if isinstance(codes, dict):
+            audio = codes.get("audio")
+            if isinstance(audio, torch.Tensor) and audio.numel() > 0:
+                return audio.reshape(-1).to(device=input_ids.device, dtype=torch.long)
+            if isinstance(audio, (list, tuple)) and audio:
+                return torch.as_tensor(audio, device=input_ids.device, dtype=torch.long).reshape(-1)
+    return input_ids.reshape(-1).to(dtype=torch.long)
+
+
 class Qwen3TTSCode2Wav(nn.Module):
     """Stage-1 code2wav model for Qwen3-TTS (GenerationModelRunner).
     Consumes frame-aligned codec tokens from input_ids and decodes waveform
@@ -111,6 +132,56 @@ class Qwen3TTSCode2Wav(nn.Module):
                 return [ids[boundaries[i] : boundaries[i + 1]] for i in range(len(boundaries) - 1)]
         return [ids]
 
+    def _maybe_enable_decoder_cudagraph(
+        self,
+        *,
+        device: torch.device,
+        codec_chunk_frames: int,
+        codec_left_context_frames: int,
+        decode_cudagraph_capture_sizes: list[int] | None,
+        decode_cudagraph_batch_sizes: list[int] | None,
+        decode_cudagraph_extra_capture_shapes: list[tuple[int, int]] | None,
+        decode_compile_shapes: list[tuple[int, int]] | None,
+    ) -> None:
+        """Enable inner Code2Wav CUDA graph unless stage is enforce_eager."""
+        if not hasattr(self.decoder, "enable_cudagraph") or device.type != "cuda":
+            return
+
+        model_cfg = getattr(self.vllm_config, "model_config", None)
+        if getattr(model_cfg, "enforce_eager", False):
+            logger.info("Qwen3-TTS Code2Wav CUDA Graph disabled because enforce_eager is set")
+            return
+
+        if (
+            codec_chunk_frames > 0
+            and codec_left_context_frames > 0
+            and self._decoder_sliding_window
+            and codec_left_context_frames < self._decoder_sliding_window
+        ):
+            logger.warning(
+                "Qwen3-TTS streaming codec_left_context_frames=%d "
+                "is smaller than decoder sliding_window=%d; "
+                "chunk-boundary distortion may occur. "
+                "Increase codec_left_context_frames to at least "
+                "%d for streaming.",
+                codec_left_context_frames,
+                self._decoder_sliding_window,
+                self._decoder_sliding_window,
+            )
+
+        self.decoder.enable_cudagraph(
+            capture_sizes=decode_cudagraph_capture_sizes,
+            capture_batch_sizes=decode_cudagraph_batch_sizes,
+            extra_capture_shapes=decode_cudagraph_extra_capture_shapes,
+            compile_shapes=decode_compile_shapes,
+            device=device,
+            codec_chunk_frames=codec_chunk_frames,
+            codec_left_context_frames=codec_left_context_frames,
+            decode_chunk_size=self._decode_chunk_frames,
+            decode_left_context=self._decode_left_context_frames,
+        )
+        logger.info("Code2Wav decoder CUDA Graph enabled")
+
     def _get_decode_batch_bucket_frames(self, actual_frames: int) -> int:
         for bucket_frames in self._decode_batch_bucket_frames:
             if actual_frames <= bucket_frames:
@@ -189,6 +260,7 @@ class Qwen3TTSCode2Wav(nn.Module):
                 multimodal_outputs={"model_outputs": [empty], "sr": [sr_tensor]},
             )
 
+        runtime_infos = runtime_additional_information or []
         ids = input_ids.reshape(-1).to(dtype=torch.long)
         request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts"))
 
@@ -196,14 +268,18 @@ class Qwen3TTSCode2Wav(nn.Module):
         valid_codes_qf: list[torch.Tensor] = []
         valid_indices: list[int] = []
         left_context_size = [0] * len(request_ids_list)
-        if runtime_additional_information is not None:
-            for i, info in enumerate(runtime_additional_information):
+        if runtime_infos:
+            for i, info in enumerate(runtime_infos):
                 if i >= len(left_context_size):
                     break
+                if not isinstance(info, dict):
+                    continue
                 meta = info.get("meta", {})
                 if "left_context_size" in meta:
                     left_context_size[i] = meta["left_context_size"]
         for i, req_ids in enumerate(request_ids_list):
+            runtime_info = runtime_infos[i] if i < len(runtime_infos) else None
+            req_ids = _codec_ids_from_payload_or_input(req_ids, runtime_info)
             if req_ids.numel() < 1:
                 parsed.append((0, 0))
                 continue
@@ -581,35 +657,15 @@ class Qwen3TTSCode2Wav(nn.Module):
 
         if hasattr(self.decoder, "enable_cudagraph") and device.type == "cuda":
             try:
-                if (
-                    codec_chunk_frames > 0
-                    and codec_left_context_frames > 0
-                    and self._decoder_sliding_window
-                    and codec_left_context_frames < self._decoder_sliding_window
-                ):
-                    logger.warning(
-                        "Qwen3-TTS streaming codec_left_context_frames=%d "
-                        "is smaller than decoder sliding_window=%d; "
-                        "chunk-boundary distortion may occur. "
-                        "Increase codec_left_context_frames to at least "
-                        "%d for streaming.",
-                        codec_left_context_frames,
-                        self._decoder_sliding_window,
-                        self._decoder_sliding_window,
-                    )
-
-                self.decoder.enable_cudagraph(
-                    capture_sizes=decode_cudagraph_capture_sizes,
-                    capture_batch_sizes=decode_cudagraph_batch_sizes,
-                    extra_capture_shapes=decode_cudagraph_extra_capture_shapes,
-                    compile_shapes=decode_compile_shapes,
+                self._maybe_enable_decoder_cudagraph(
                     device=device,
                     codec_chunk_frames=codec_chunk_frames,
                     codec_left_context_frames=codec_left_context_frames,
-                    decode_chunk_size=self._decode_chunk_frames,
-                    decode_left_context=self._decode_left_context_frames,
+                    decode_cudagraph_capture_sizes=decode_cudagraph_capture_sizes,
+                    decode_cudagraph_batch_sizes=decode_cudagraph_batch_sizes,
+                    decode_cudagraph_extra_capture_shapes=decode_cudagraph_extra_capture_shapes,
+                    decode_compile_shapes=decode_compile_shapes,
                 )
-                logger.info("Code2Wav decoder CUDA Graph enabled")
             except Exception:
                 logger.warning(
                     "Failed to enable CUDA Graph for Code2Wav decoder",
diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py
index a705a1cd4..ca1c4b395 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py
+++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py
@@ -1,26 +1,21 @@
 from __future__ import annotations
 
-import base64
 import copy
-import io
-import os
-from collections.abc import Callable, Iterable, Mapping
+from collections.abc import Iterable, Mapping
 from typing import Any
-from urllib.parse import urlparse
 
 import numpy as np
-import soundfile as sf
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers import AutoTokenizer
+from transformers import AutoFeatureExtractor
 from transformers.activations import ACT2FN
-from transformers.utils.hub import cached_file
 from vllm.config import VllmConfig
 from vllm.distributed import get_pp_group
 from vllm.logger import init_logger
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
+from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
 from vllm.model_executor.models.qwen3 import Qwen3Model
 from vllm.model_executor.models.utils import AutoWeightsLoader, PPMissingLayer, WeightsMapper, maybe_prefix
 from vllm.multimodal.audio import AudioResampler
@@ -28,20 +23,17 @@ from vllm.sequence import IntermediateTensors
 
 from vllm_omni.data_entry_keys import OmniPayload
 from vllm_omni.model_executor.models.output_templates import OmniOutput
-from vllm_omni.utils.audio import mel_filter_bank
 from vllm_omni.utils.speaker_cache import get_speaker_cache
 
 from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSSpeakerEncoderConfig, Qwen3TTSTalkerConfig
+from .prompt_embeds_builder import Qwen3TTSPromptEmbedsBuilder
 from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM
-from .qwen3_tts_tokenizer import Qwen3TTSTokenizer
+from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config
+from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Encoder
 
 logger = init_logger(__name__)
 
 _TRAILING_TEXT_COMPACT_MIN_FRAMES = 64
-_PRECOMPUTED_REF_CODE_KEY = "precomputed_ref"
-_NORMALIZED_REF_AUDIO_KEY = "_qwen3_tts_normalized_ref_audio"
-_PRECOMPUTED_TEXT_IDS_KEY = "_qwen3_tts_text_ids"
-_PRECOMPUTED_REF_IDS_KEY = "_qwen3_tts_ref_ids"
 
 
 # ---------------------------------------------------------------------------
@@ -248,65 +240,6 @@ class Qwen3TTSSpeakerEncoder(torch.nn.Module):
         return hidden_states.squeeze(-1)
 
 
-# ---- Audio utilities ----
-
-
-def _dynamic_range_compression(x, c=1, clip_val=1e-5):
-    return torch.log(torch.clamp(x, min=clip_val) * c)
-
-
-def mel_spectrogram(
-    y: torch.Tensor,
-    n_fft: int,
-    num_mels: int,
-    sampling_rate: int,
-    hop_size: int,
-    win_size: int,
-    fmin: int,
-    fmax: int | None = None,
-    center: bool = False,
-    mel_basis: torch.Tensor | None = None,
-    hann_window: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """Calculate mel spectrogram of an input signal using torchaudio mel filterbank and torch STFT."""
-    if torch.min(y) < -1.0:
-        logger.warning("Min value of input waveform signal is %s", torch.min(y))
-    if torch.max(y) > 1.0:
-        logger.warning("Max value of input waveform signal is %s", torch.max(y))
-    device = y.device
-    if mel_basis is None:
-        mel_basis = mel_filter_bank(
-            sr=sampling_rate,
-            n_fft=n_fft,
-            n_mels=num_mels,
-            fmin=fmin,
-            fmax=fmax,
-        ).to(device)
-    elif mel_basis.device != device:
-        mel_basis = mel_basis.to(device)
-    if hann_window is None:
-        hann_window = torch.hann_window(win_size, device=device)
-    elif hann_window.device != device:
-        hann_window = hann_window.to(device)
-    padding = (n_fft - hop_size) // 2
-    y = torch.nn.functional.pad(y.unsqueeze(1), (padding, padding), mode="reflect").squeeze(1)
-    spec = torch.stft(
-        y,
-        n_fft,
-        hop_length=hop_size,
-        win_length=win_size,
-        window=hann_window,
-        center=center,
-        pad_mode="reflect",
-        normalized=False,
-        onesided=True,
-        return_complex=True,
-    )
-    spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9)
-    mel_spec = torch.matmul(mel_basis, spec)
-    return _dynamic_range_compression(mel_spec)
-
-
 # ---------------------------------------------------------------------------
 # Main AR Talker model
 # ---------------------------------------------------------------------------
@@ -354,6 +287,17 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         self.have_multimodal_outputs = True
         self.has_preprocess = True
         self.has_postprocess = True
+        # Qwen3-TTS postprocess() only reads hidden_states[-1, :]. On a prefix-
+        # cache hit, the last hidden state is in the newly computed tail, so
+        # reconstructing the full cached_prefix + new_tail span is wasted work.
+        # Opt out of the per-step GPU->CPU hidden-state cache write and merged-
+        # tensor read; postprocess receives the tail-only slice instead, which
+        # avoids ~18 ms merge + ~6 ms write per step (Sy0307 profile, #3665).
+        self.requires_full_prefix_cached_hidden_states = False
+        # ``codes.audio`` is only needed for future prefix-hit reconstruction
+        # after a request has produced codec rows. Keep per-step rows on GPU and
+        # materialize the CPU OmniTensorPrefixCache entry once at completion.
+        self.deferred_prefix_cache_mm_keys = {"codes.audio"}
         # Used by OmniGPUModelRunner for the GPU-side MTP fast-path.
         self.mtp_hidden_size = int(self.talker_config.hidden_size)
         # OmniGPUModelRunner will store talker_mtp output under this key in
@@ -434,13 +378,40 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         self.gpu_resident_buffer_keys: set[tuple[str, str]] = {
             ("codes", "audio"),
             ("hidden_states", "last"),
-            ("embed", "tts_pad"),
             ("hidden_states", "trailing_text"),
         }
 
-        # Tokenizer for prompt building.
-        self._tokenizer = None
-        self._speech_tokenizer: Qwen3TTSTokenizer | None = None
+        # ``text_proj(text_emb(tts_pad_token_id))`` is request-independent —
+        # it depends only on frozen ``text_embedding`` / ``text_projection``
+        # weights, so we precompute it once in :meth:`_init_runtime_buffers`
+        # (called from :meth:`load_weights`) and reuse the same buffer at
+        # every prefill and decode step instead of round-tripping it through
+        # the per-request ``info_dict``. Declared here as zeros so the
+        # attribute exists under ``load_format: dummy`` (which bypasses
+        # ``load_weights`` entirely and leaves the value uninitialized).
+        model_dtype = getattr(vllm_config.model_config, "dtype", torch.bfloat16)
+        self.register_buffer(
+            "_tts_pad_embed",
+            torch.zeros(1, int(self.talker_config.hidden_size), dtype=model_dtype),
+            persistent=False,
+        )
+
+        tokenizer_config = Qwen3TTSTokenizerV2Config.from_pretrained(
+            self.model_path,
+            subfolder="speech_tokenizer",
+        )
+        self.encoder = Qwen3TTSTokenizerV2Encoder._from_config(
+            tokenizer_config.encoder_config,
+        )
+        self.encoder.eval()
+        self.encoder.to(dtype=torch.bfloat16)
+        self._encoder_valid_num_quantizers = int(tokenizer_config.encoder_valid_num_quantizers)
+        self._encoder_downsample_rate = int(tokenizer_config.encode_downsample_rate)
+
+        self._encoder_feature_extractor = AutoFeatureExtractor.from_pretrained(
+            self.model_path,
+            subfolder="speech_tokenizer",
+        )
 
         self._speaker_cache = get_speaker_cache()
         raw_subtalker_sampling = getattr(vllm_config.model_config, "subtalker_sampling_params", None)
@@ -448,6 +419,29 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             dict(raw_subtalker_sampling) if isinstance(raw_subtalker_sampling, Mapping) else {}
         )
 
+        self._stacked_codec_embed: torch.Tensor | None = None
+
+        # Stand-alone wrapper around the embedding layers, encoders and
+        # tokenizers required to assemble a talker prefill prompt. Owns the
+        # text + speech tokenizers (loaded lazily on first use) and the
+        # per-ref-audio / resampler caches, so they don't leak onto the
+        # talker class. Other talker variants can construct their own builder
+        # with the same set of dependencies and reuse ``build_prompt_embeds``
+        # verbatim.
+        self._prompt_builder = Qwen3TTSPromptEmbedsBuilder(
+            config=self.config,
+            talker_config=self.talker_config,
+            model_path=self.model_path,
+            text_embedding=self.text_embedding,
+            text_projection=self.text_projection,
+            codec_embed=self.embed_input_ids,
+            residual_code_embeddings=lambda: self.code_predictor.get_input_embeddings(),
+            speaker_encoder=self.speaker_encoder,
+            tts_pad_embed=self._tts_pad_embed,
+            encode_ref_audio_batch=self._encode_ref_audio_batch,
+            speaker_cache=self._speaker_cache,
+        )
+
     # -------------------- vLLM required hooks --------------------
 
     def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
@@ -519,21 +513,22 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             ref_len = meta.get("ref_code_len")
             if ref_len is None:
                 continue
+            if not isinstance(ac, torch.Tensor):
+                continue
+            span_len = int(ac.shape[0])
             if isinstance(ref_len, torch.Tensor):
                 if ref_len.numel() == 0:
                     raise ValueError("ref_code_len is an empty tensor")
-                ref_len_val = int(ref_len.reshape(-1)[-1].item())
-            elif isinstance(ref_len, list):
+                ref_len_tail = ref_len.reshape(-1)[-1:].to(dtype=torch.int32, device=ac.device)
+                ref_code_len_list.append(ref_len_tail.expand(span_len).contiguous())
+                continue
+            if isinstance(ref_len, list):
                 if len(ref_len) != 1:
                     raise ValueError(f"ref_code_len must be scalar or 1-element list, got len={len(ref_len)}")
                 ref_len_val = int(ref_len[0])
             else:
                 ref_len_val = int(ref_len)
-            if isinstance(ac, torch.Tensor):
-                # Emit ref_code_len per-token span for runner slicing (consumer takes the last value).
-                ref_code_len_list.append(
-                    torch.full((int(ac.shape[0]),), ref_len_val, dtype=torch.int32, device=ac.device)
-                )
+            ref_code_len_list.append(torch.full((span_len,), ref_len_val, dtype=torch.int32, device=ac.device))
 
         if not audio_codes_list:
             return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
@@ -598,32 +593,30 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         else:
             codec_streaming = task_type == "Base"
 
+        # ``tts_pad_embed`` is a request-independent constant — see
+        # :meth:`_init_runtime_buffers`. Materialize once on the right
+        # device/dtype and reuse for both the prefill placeholder padding
+        # and the decode text-step fallback below.
+        tts_pad_embed = self._tts_pad_embed.to(device=input_ids.device, dtype=torch.bfloat16).reshape(1, -1)
+
         if is_prefill:
             # Prefill (prompt embeddings)
             prompt_embeds_cpu = embed.get("prefill")
-            tts_pad_embed_cpu = embed.get("tts_pad")
-            tts_pad_embed = None
-            if isinstance(tts_pad_embed_cpu, torch.Tensor) and tts_pad_embed_cpu.numel() > 0:
-                tts_pad_embed = tts_pad_embed_cpu.to(device=input_ids.device, dtype=torch.bfloat16).reshape(1, -1)
 
             # First prefill round: prompt_embeds_cpu is not yet populated.
             # Subsequent prefill rounds (multi-chunk): prompt_embeds_cpu is a Tensor stored by the first round.
             is_first_prefill = not isinstance(prompt_embeds_cpu, torch.Tensor) or prompt_embeds_cpu.ndim != 2
             if is_first_prefill:
-                full_prompt_embeds, tailing_text_hidden, tts_pad_embed, ref_code_len, ref_code = (
-                    self._build_prompt_embeds(task_type=task_type, info_dict=info_dict)
+                full_prompt_embeds, tailing_text_hidden, ref_code_len, ref_code = (
+                    self._prompt_builder.build_prompt_embeds(task_type=task_type, info_dict=info_dict)
                 )
                 # Store full prompt embeddings on CPU (large, prefill-only).
-                # tailing_text_hidden and tts_pad_embed stay on GPU (gpu_resident_buffer_keys).
+                # tailing_text_hidden stays on GPU (gpu_resident_buffer_keys).
                 prompt_embeds_cpu = full_prompt_embeds.detach().to("cpu").contiguous()
                 info_update: OmniPayload = {
-                    "embed": {
-                        "prefill": prompt_embeds_cpu,
-                        "tts_pad": tts_pad_embed.detach(),
-                    },
+                    "embed": {"prefill": prompt_embeds_cpu},
                     "hidden_states": {"trailing_text": tailing_text_hidden.detach()},
                     "meta": {
-                        "talker_prefill_offset": 0,
                         "talker_text_offset": 0,
                         "codec_streaming": codec_streaming,
                     },
@@ -632,39 +625,27 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                     info_update.setdefault("codes", {})["ref"] = ref_code.detach().to("cpu").contiguous()
                 if ref_code_len is not None:
                     info_update["meta"]["ref_code_len"] = int(ref_code_len)
-                # Always return a span_len slice; if the scheduled placeholder is longer, pad with tts_pad_embed.
-                # This preserves placeholder/embedding alignment.
-                offset = 0
-                s = 0
-                e = span_len
-                take = prompt_embeds_cpu[s:e]
-                if int(take.shape[0]) < span_len:
-                    pad_n = int(span_len - int(take.shape[0]))
-                    pad_rows = tts_pad_embed.reshape(1, -1).to("cpu").expand(pad_n, -1)
-                    take = torch.cat([take, pad_rows], dim=0)
-                prompt_embeds = take.to(device=input_ids.device, dtype=torch.bfloat16)
-                info_update["meta"]["talker_prefill_offset"] = int(offset + span_len)
+                # First prefill: source the slice offset from `_omni_num_computed_tokens`
+                # so cache-recovery (prefill replay at a later offset) lands on the right
+                # slice of the stored embeddings. Subsequent chunks below advance from
+                # `talker_prefill_offset` written here.
+                offset = max(0, int(info_dict.get("_omni_num_computed_tokens", 0) or 0))
             else:
                 # Subsequent prefill chunk: slice from stored embeddings at running offset.
-                if tts_pad_embed is None:
-                    raise RuntimeError("Missing `tts_pad_embed` in additional_information; prefill must initialize it.")
-                offset = int(meta.get("talker_prefill_offset", 0) or 0)
-                if offset < 0:
-                    offset = 0
-                s = max(0, min(offset, int(prompt_embeds_cpu.shape[0])))
-                e = max(0, min(offset + span_len, int(prompt_embeds_cpu.shape[0])))
-                take = prompt_embeds_cpu[s:e]
-                if int(take.shape[0]) < span_len:
-                    pad_n = int(span_len - int(take.shape[0]))
-                    pad_rows = tts_pad_embed.reshape(1, -1).to("cpu").expand(pad_n, -1)
-                    take = torch.cat([take, pad_rows], dim=0)
-                prompt_embeds = take.to(device=input_ids.device, dtype=torch.bfloat16)
-                info_update = {
-                    "meta": {
-                        "talker_prefill_offset": int(offset + span_len),
-                        "codec_streaming": codec_streaming,
-                    }
-                }
+                offset = max(0, int(meta.get("talker_prefill_offset", 0) or 0))
+                info_update = {"meta": {"codec_streaming": codec_streaming}}
+
+            # Always return a span_len slice; if the scheduled placeholder is longer than what
+            # the prompt actually fills, pad with tts_pad_embed (preserves placeholder/embedding alignment).
+            s = max(0, min(offset, int(prompt_embeds_cpu.shape[0])))
+            e = max(0, min(offset + span_len, int(prompt_embeds_cpu.shape[0])))
+            take = prompt_embeds_cpu[s:e]
+            if int(take.shape[0]) < span_len:
+                pad_n = int(span_len - int(take.shape[0]))
+                pad_rows = tts_pad_embed.reshape(1, -1).to("cpu").expand(pad_n, -1)
+                take = torch.cat([take, pad_rows], dim=0)
+            prompt_embeds = take.to(device=input_ids.device, dtype=torch.bfloat16)
+            info_update["meta"]["talker_prefill_offset"] = int(offset + span_len)
 
             # When inputs_embeds is set, token ids are ignored by the model but must stay in-vocab for vLLM bookkeeping.
             input_ids_out = input_ids.clone()
@@ -680,11 +661,8 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
 
         # Decode: span_len == 1
         # Pop one text-step vector from tailing_text_hidden queue.
-        # These tensors stay on GPU via gpu_resident_buffer_keys - .to() is a no-op.
-        tts_pad_embed_buf = embed.get("tts_pad")
-        if not isinstance(tts_pad_embed_buf, torch.Tensor):
-            raise RuntimeError("Missing `tts_pad_embed` in additional_information; prefill must run first.")
-        tts_pad_embed = tts_pad_embed_buf.to(device=input_ids.device, dtype=torch.bfloat16).reshape(1, -1)
+        # ``tts_pad_embed`` was materialized above from :attr:`_tts_pad_embed`
+        # (request-independent buffer) — no per-request fetch needed.
 
         tail = hs.get("trailing_text")
         text_offset = max(0, int(meta.get("talker_text_offset", 0) or 0))
@@ -720,9 +698,11 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             next_text_offset = text_offset
 
         last_hidden = hs.get("last")
-        if not isinstance(last_hidden, torch.Tensor):
-            raise RuntimeError("Missing hidden_states['last'] in additional_information; postprocess must run.")
-        past_hidden = last_hidden.to(device=input_ids.device, dtype=torch.bfloat16).reshape(1, -1)
+        if isinstance(last_hidden, torch.Tensor):
+            past_hidden = last_hidden.to(device=input_ids.device, dtype=torch.bfloat16).reshape(1, -1)
+        else:
+            # Defensive: EOS step row is zeroed by the invalid-layer-0 mask and filtered downstream.
+            past_hidden = torch.zeros_like(text_step)
 
         # Use OmniGPUModelRunner talker_mtp fast-path for residual codebooks and per-step inputs_embeds update.
         last_id_hidden = self.embed_input_ids(input_ids.reshape(1, 1).to(torch.long)).to(
@@ -760,6 +740,10 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
 
         device = input_ids_flat.device
         dtype = torch.bfloat16
+        # Request-independent constant (see :meth:`_init_runtime_buffers`) —
+        # compute once for the batch instead of fetching it per request from
+        # ``info_dict["embed"]["tts_pad"]``.
+        tts_pad_embed = self._tts_pad_embed.to(device=device, dtype=dtype).reshape(1, -1)
         past_hidden_list: list[torch.Tensor] = []
         text_step_list: list[torch.Tensor] = []
         updates: list[dict[str, Any]] = []
@@ -773,7 +757,6 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                 info_dict = merged
 
             payload: OmniPayload = info_dict
-            embed = payload.get("embed", {})
             hs = payload.get("hidden_states", {})
             meta = payload.get("meta", {})
 
@@ -792,11 +775,6 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             else:
                 codec_streaming = task_type == "Base"
 
-            tts_pad_embed_buf = embed.get("tts_pad")
-            if not isinstance(tts_pad_embed_buf, torch.Tensor):
-                raise RuntimeError("Missing `tts_pad_embed` in additional_information; prefill must run first.")
-            tts_pad_embed = tts_pad_embed_buf.to(device=device, dtype=dtype).reshape(1, -1)
-
             tail = hs.get("trailing_text")
             text_offset = max(0, int(meta.get("talker_text_offset", 0) or 0))
             trailing_text_update = None
@@ -860,647 +838,6 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         last = hidden_states[-1, :].detach()
         return {"hidden_states": {"last": last}}
 
-    # -------------------- prompt construction helpers --------------------
-
-    def _get_tokenizer(self):
-        if self._tokenizer is None:
-            import transformers
-
-            kwargs = dict(trust_remote_code=True, use_fast=True)
-            if transformers.__version__ < "5":
-                kwargs["fix_mistral_regex"] = True
-            self._tokenizer = AutoTokenizer.from_pretrained(self.model_path, **kwargs)
-            self._tokenizer.padding_side = "left"
-        return self._tokenizer
-
-    @staticmethod
-    def _build_assistant_text(text: str) -> str:
-        return f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
-
-    @staticmethod
-    def _build_ref_text(text: str) -> str:
-        return f"<|im_start|>assistant\n{text}<|im_end|>\n"
-
-    @staticmethod
-    def _build_instruct_text(instruct: str) -> str:
-        return f"<|im_start|>user\n{instruct}<|im_end|>\n"
-
-    @staticmethod
-    def estimate_prompt_len_from_additional_information(
-        additional_information: dict[str, Any] | None,
-        *,
-        task_type: str,
-        tokenize_prompt: Callable[[str], list[int]],
-        codec_language_id: Mapping[str, int] | None,
-        spk_is_dialect: Mapping[str, object] | None,
-        estimate_ref_code_len: Callable[[object], int | None] | None = None,
-    ) -> int:
-        """Compute Stage-0 placeholder prompt length (length-only mirror of `_build_prompt_embeds()`).
-        It must match the model-side `inputs_embeds` length to avoid extra padding and quality drop."""
-
-        def _first(x: object, default: object) -> object:
-            if isinstance(x, list):
-                return x[0] if x else default
-            return x if x is not None else default
-
-        info: dict[str, Any] = additional_information or {}
-        text = _first(info.get("text"), "")
-        language = _first(info.get("language"), "Auto")
-        speaker = _first(info.get("speaker"), "").lower().strip()
-        instruct = _first(info.get("instruct"), "")
-        non_streaming_mode_raw = _first(info.get("non_streaming_mode"), None)
-
-        if isinstance(non_streaming_mode_raw, bool):
-            non_streaming_mode = non_streaming_mode_raw
-        else:
-            # Official defaults: CustomVoice/VoiceDesign -> non_streaming_mode=True; Base -> False.
-            non_streaming_mode = task_type in ("CustomVoice", "VoiceDesign")
-
-        if not isinstance(text, str):
-            text = ""
-        if not isinstance(instruct, str):
-            instruct = ""
-        if not isinstance(language, str):
-            language = "Auto"
-
-        instruct_len = 0
-        if instruct.strip():
-            instruct_text = Qwen3TTSTalkerForConditionalGeneration._build_instruct_text(instruct)
-            instruct_len = len(tokenize_prompt(instruct_text))
-
-        # ---- codec prefix portion (matches _build_prompt_embeds) ----
-        language_id = None
-        if language.lower() != "auto" and codec_language_id:
-            language_id = codec_language_id.get(language.lower())
-        if (
-            language_id is None
-            and codec_language_id
-            and spk_is_dialect
-            and isinstance(language, str)
-            and language.lower() in ("chinese", "auto")
-            and isinstance(speaker, str)
-            and speaker.strip()
-        ):
-            dialect = spk_is_dialect.get(speaker.lower())
-            if isinstance(dialect, str) and dialect:
-                language_id = codec_language_id.get(dialect)
-        prefill_len = 3 if language_id is None else 4
-
-        speaker_len = 1 if task_type in ("CustomVoice", "Base") else 0
-        codec_input_len = prefill_len + speaker_len + 2  # + [codec_pad, codec_bos]
-        codec_prefix_len = codec_input_len - 1  # codec_input[:-1] + tts_bos
-
-        # Role header: input_ids[:, :3] in model.
-        role_len = 3
-        prompt_len = instruct_len + role_len + codec_prefix_len
-
-        # ---- text conditioning portion (matches _build_prompt_embeds) ----
-        assistant_text = Qwen3TTSTalkerForConditionalGeneration._build_assistant_text(text)
-        assistant_len = len(tokenize_prompt(assistant_text))
-        if assistant_len < 8:
-            raise ValueError(f"Unexpected assistant prompt length: {assistant_len}")
-
-        if task_type in ("CustomVoice", "VoiceDesign"):
-            if non_streaming_mode:
-                # model: full text ids (input_ids[:, 3:-5]) + eos + codec_bos step
-                prompt_len += assistant_len - 6
-            else:
-                # model: only first text token in prefill
-                prompt_len += 1
-
-        if task_type == "Base":
-            xvec_only = bool(_first(info.get("x_vector_only_mode"), False))
-            in_context_mode = not xvec_only
-
-            voice_clone_prompt = _first(info.get("voice_clone_prompt"), None)
-            if isinstance(voice_clone_prompt, dict):
-                icl_flag = _first(voice_clone_prompt.get("icl_mode"), None)
-                if isinstance(icl_flag, bool):
-                    in_context_mode = icl_flag
-
-            if in_context_mode:
-                ref_code = None
-                if isinstance(voice_clone_prompt, dict):
-                    ref_code = _first(voice_clone_prompt.get("ref_code"), None)
-
-                ref_code_len: int | None = None
-                if isinstance(ref_code, list):
-                    if ref_code and isinstance(ref_code[0], list):
-                        ref_code_len = len(ref_code)
-                    elif ref_code:
-                        ref_code_len = len(ref_code)
-                elif hasattr(ref_code, "shape"):
-                    try:
-                        shape = getattr(ref_code, "shape")
-                        if shape and len(shape) >= 1:
-                            ref_code_len = int(shape[0])
-                    except Exception:
-                        ref_code_len = None
-
-                if ref_code_len is None and estimate_ref_code_len is not None:
-                    ref_code_len = estimate_ref_code_len(info.get("ref_audio"))
-                if ref_code_len is None:
-                    raise ValueError(
-                        "Base in-context voice cloning requires either `voice_clone_prompt.ref_code` "
-                        "or a readable `ref_audio` that can be mapped to a codec frame length."
-                    )
-
-                codec_lens = 1 + int(ref_code_len)  # codec_bos + ref_code
-                if non_streaming_mode:
-                    # _generate_icl_prompt(non_streaming_mode=True):
-                    # text_embed = ref_ids + text_ids + eos.
-                    ref_ids = _first(info.get("ref_ids"), None)
-                    if isinstance(voice_clone_prompt, dict) and ref_ids is None:
-                        ref_ids = _first(voice_clone_prompt.get("ref_ids") or voice_clone_prompt.get("ref_id"), None)
-
-                    if ref_ids is None:
-                        ref_text = _first(info.get("ref_text"), "")
-                        if not isinstance(ref_text, str) or not ref_text.strip():
-                            raise ValueError(
-                                "Base in-context non-streaming requires `ref_text` or tokenized `ref_ids`."
-                            )
-                        ref_text_ids = tokenize_prompt(Qwen3TTSTalkerForConditionalGeneration._build_ref_text(ref_text))
-                        ref_ids_len = len(ref_text_ids)
-                    elif hasattr(ref_ids, "shape"):
-                        shape = getattr(ref_ids, "shape", None)
-                        ref_ids_len = int(shape[-1]) if shape else 0
-                    elif isinstance(ref_ids, list):
-                        ref_ids_len = len(ref_ids)
-                    else:
-                        ref_ids_len = 0
-
-                    # model uses ref_ids[:, 3:-2] (strip 5 tokens) and text_id=input_ids[:, 3:-5] (strip 8).
-                    ref_id_len = max(0, int(ref_ids_len) - 5)
-                    text_id_len = max(0, int(assistant_len) - 8)
-                    text_embed_len = ref_id_len + text_id_len + 1  # + eos
-                    prompt_len += text_embed_len + codec_lens
-                else:
-                    # _generate_icl_prompt(non_streaming_mode=False): aligned to codec_lens.
-                    prompt_len += codec_lens
-            else:
-                # Base without ICL behaves like CustomVoice.
-                if non_streaming_mode:
-                    prompt_len += assistant_len - 6
-                else:
-                    prompt_len += 1
-
-        return max(2, int(prompt_len))
-
-    def _is_probably_base64(self, s: str) -> bool:
-        if s.startswith("data:audio"):
-            return True
-        if ("/" not in s and "\\" not in s) and len(s) > 256:
-            return True
-        return False
-
-    def _is_url(self, s: str) -> bool:
-        try:
-            u = urlparse(s)
-            if u.scheme in ("http", "https"):
-                return bool(u.netloc)
-            return u.scheme == "file"
-        except Exception:
-            return False
-
-    def _decode_base64_to_wav_bytes(self, b64: str) -> bytes:
-        if "," in b64 and b64.strip().startswith("data:"):
-            b64 = b64.split(",", 1)[1]
-        return base64.b64decode(b64)
-
-    def _load_audio_to_np(self, x: str) -> tuple[np.ndarray, int]:
-        """Load audio from local path, URL, or base64 data URI.
-
-        Uses upstream vLLM's MediaConnector for http(s) URLs and ``file:``
-        URIs, with unrestricted local access (offline inference is trusted).
-        """
-        from vllm.multimodal.media.audio import load_audio
-
-        if self._is_url(x):
-            from vllm.multimodal.media import MediaConnector
-
-            connector = MediaConnector(allowed_local_media_path="/")
-            audio, sr = connector.fetch_audio(x)
-        elif self._is_probably_base64(x):
-            wav_bytes = self._decode_base64_to_wav_bytes(x)
-            with io.BytesIO(wav_bytes) as f:
-                audio, sr = sf.read(f, dtype="float32", always_2d=False)
-        else:
-            audio, sr = load_audio(x, sr=None, mono=True)
-
-        if isinstance(audio, np.ndarray) and audio.ndim > 1:
-            audio = np.mean(audio, axis=-1)
-
-        return np.asarray(audio, dtype=np.float32), int(sr)
-
-    def _normalize_ref_audio(self, ref_audio: object) -> tuple[np.ndarray, int]:
-        # NOTE: additional_information may serialize (wav, sr) into (nested) lists across processes; be tolerant.
-        if isinstance(ref_audio, str):
-            return self._load_audio_to_np(ref_audio)
-
-        def _is_sr(x: object) -> bool:
-            try:
-                v = int(x)  # type: ignore[arg-type]
-            except Exception:
-                return False
-            return 1_000 <= v <= 200_000
-
-        def _is_number_sequence(xs: list[object]) -> bool:
-            if not xs:
-                return False
-            for v in xs[:8]:
-                if not isinstance(v, (int, float, np.number)):
-                    return False
-            return True
-
-        wav_candidates: list[object] = []
-        sr_candidates: list[int] = []
-
-        def _summarize(obj: object, depth: int = 0) -> str:
-            if depth > 2:
-                if isinstance(obj, (int, np.integer)):
-                    return f"int({int(obj)})"
-                return type(obj).__name__
-            if obj is None:
-                return "None"
-            if isinstance(obj, str):
-                if len(obj) <= 16:
-                    return f"str({obj!r})"
-                return f"str(len={len(obj)})"
-            if isinstance(obj, (int, float, np.number)):
-                return f"{type(obj).__name__}({obj})"
-            if isinstance(obj, np.ndarray):
-                return f"ndarray(shape={obj.shape}, dtype={obj.dtype})"
-            if isinstance(obj, torch.Tensor):
-                return f"Tensor(shape={tuple(obj.shape)}, dtype={obj.dtype}, device={obj.device})"
-            if isinstance(obj, dict):
-                keys = list(obj.keys())
-                return f"dict(keys={keys[:8]})"
-            if isinstance(obj, (tuple, list)):
-                items = list(obj)
-                head = ", ".join(_summarize(x, depth + 1) for x in items[:3])
-                return f"{type(obj).__name__}(len={len(items)}; head=[{head}])"
-            return f"{type(obj).__name__}"
-
-        def _scan(obj: object, depth: int = 0) -> None:
-            if depth > 4:
-                return
-            if obj is None:
-                return
-            if _is_sr(obj):
-                sr_candidates.append(int(obj))  # type: ignore[arg-type]
-                return
-            if isinstance(obj, np.ndarray) and obj.size > 0:
-                wav_candidates.append(obj)
-                return
-            if isinstance(obj, torch.Tensor) and obj.numel() > 0:
-                wav_candidates.append(obj)
-                return
-            if isinstance(obj, dict):
-                # Inlined ndarray/tensor payloads from the input processor.
-                if obj.get("__ndarray__") and "data" in obj and "dtype" in obj and "shape" in obj:
-                    try:
-                        data = obj["data"]
-                        dtype = obj["dtype"]
-                        shape = obj["shape"]
-                        if isinstance(data, (bytes, bytearray, memoryview)):
-                            arr = np.frombuffer(data, dtype=dtype).reshape(shape)
-                            if arr.size > 0:
-                                wav_candidates.append(arr)
-                                return
-                    except Exception:
-                        pass
-                if obj.get("__tensor__") and "data" in obj and "dtype" in obj and "shape" in obj:
-                    try:
-                        data = obj["data"]
-                        dtype = obj["dtype"]
-                        shape = obj["shape"]
-                        if isinstance(data, (bytes, bytearray, memoryview)):
-                            # Stored as raw CPU bytes; interpret as numpy for audio.
-                            np_dtype = np.dtype(dtype)
-                            arr = np.frombuffer(data, dtype=np_dtype).reshape(shape)
-                            if arr.size > 0:
-                                wav_candidates.append(arr)
-                                return
-                    except Exception:
-                        pass
-                wav_obj = obj.get("array") or obj.get("wav") or obj.get("audio")
-                sr_obj = obj.get("sampling_rate") or obj.get("sr") or obj.get("sample_rate")
-                if wav_obj is not None:
-                    _scan(wav_obj, depth + 1)
-                if sr_obj is not None:
-                    _scan(sr_obj, depth + 1)
-                return
-            if isinstance(obj, (tuple, list)):
-                obj_list = list(obj)
-                # Unwrap singleton nesting ([[wav, sr]]).
-                while isinstance(obj_list, list) and len(obj_list) == 1:
-                    inner = obj_list[0]
-                    if isinstance(inner, np.ndarray) and inner.size > 0:
-                        wav_candidates.append(inner)
-                        return
-                    if isinstance(inner, torch.Tensor) and inner.numel() > 0:
-                        wav_candidates.append(inner)
-                        return
-                    if isinstance(inner, dict):
-                        _scan(inner, depth + 1)
-                        return
-                    if isinstance(inner, (tuple, list)):
-                        obj_list = list(inner)  # type: ignore[list-item]
-                        continue
-                    break
-
-                # If the *unwrapped* list is a long list of numbers, treat it as waveform.
-                if len(obj_list) >= 512 and _is_number_sequence(obj_list):
-                    wav_candidates.append(obj_list)
-                    return
-
-                # Otherwise, recurse into elements (but avoid descending into huge numeric lists).
-                for item in obj_list:
-                    if isinstance(item, list) and len(item) >= 512 and _is_number_sequence(item):  # type: ignore[arg-type]
-                        wav_candidates.append(item)
-                        continue
-                    _scan(item, depth + 1)
-                return
-
-        _scan(ref_audio)
-        if not sr_candidates:
-            raise TypeError(f"ref_audio missing sample_rate: {_summarize(ref_audio)}")
-        sr = int(sr_candidates[0])
-
-        def _wav_len(x: object) -> int:
-            try:
-                if isinstance(x, np.ndarray):
-                    return int(x.size)
-                if isinstance(x, torch.Tensor):
-                    return int(x.numel())
-                if isinstance(x, list):
-                    return int(len(x))
-            except Exception:
-                pass
-            return 0
-
-        if not wav_candidates:
-            raise TypeError(f"ref_audio missing waveform: {_summarize(ref_audio)}")
-        wav_obj = max(wav_candidates, key=_wav_len)
-
-        def _to_np(x: object) -> np.ndarray:
-            if isinstance(x, np.ndarray):
-                return x.astype(np.float32).reshape(-1)
-            if isinstance(x, torch.Tensor):
-                return x.detach().to("cpu").float().contiguous().numpy().reshape(-1)
-            if isinstance(x, dict) and x.get("__ndarray__") and "data" in x and "dtype" in x and "shape" in x:
-                data = x["data"]
-                dtype = x["dtype"]
-                shape = x["shape"]
-                if isinstance(data, (bytes, bytearray, memoryview)):
-                    return np.frombuffer(data, dtype=dtype).reshape(shape).astype(np.float32).reshape(-1)
-            if isinstance(x, list):
-                # list of numbers
-                if len(x) >= 2 and _is_number_sequence(x):  # type: ignore[arg-type]
-                    return np.asarray(x, dtype=np.float32).reshape(-1)
-                # list of chunks
-                parts: list[np.ndarray] = []
-                for part in x:
-                    if isinstance(part, (np.ndarray, torch.Tensor, list)):
-                        parts.append(_to_np(part))
-                if parts:
-                    return np.concatenate(parts, axis=0)
-            raise TypeError(f"Unsupported waveform type: {type(x)}")
-
-        wav_np = _to_np(wav_obj)
-        if wav_np.size < 1024:
-            raise ValueError(f"ref_audio waveform too short: {wav_np.size} samples")
-        return wav_np, sr
-
-    def _extract_speaker_embedding(self, wav: np.ndarray, sr: int) -> torch.Tensor:
-        # vLLM workers do not automatically move arbitrary torch.nn.Modules to
-        # CUDA. Ensure the speaker encoder is on the same device/dtype as the
-        # main model before running it.
-        dev = next(self.parameters()).device
-        try:
-            spk_param = next(self.speaker_encoder.parameters())
-            if spk_param.device != dev or spk_param.dtype != torch.bfloat16:
-                self.speaker_encoder.to(device=dev, dtype=torch.bfloat16)
-        except StopIteration:
-            pass
-        # Resample to 24kHz for speaker encoder.
-        target_sr = int(getattr(self.config.speaker_encoder_config, "sample_rate", 24000))
-        if sr != target_sr:
-            resampler = AudioResampler(target_sr=target_sr)
-            wav = resampler.resample(wav.astype(np.float32), orig_sr=int(sr))
-            sr = target_sr
-
-        # Follow official implementation: mel_spectrogram expects 24kHz.  Move
-        # the waveform first so STFT/mel computation stays on the model device
-        # instead of materializing a CPU mel tensor and copying it per request.
-        wav_tensor = torch.from_numpy(wav).to(device=dev, dtype=torch.float32).unsqueeze(0)
-        mel_basis, hann_window = self._get_speaker_mel_buffers(
-            device=dev,
-            sampling_rate=24000,
-            n_fft=1024,
-            num_mels=128,
-            win_size=1024,
-            fmin=0,
-            fmax=12000,
-        )
-        mels = mel_spectrogram(
-            wav_tensor,
-            n_fft=1024,
-            num_mels=128,
-            sampling_rate=24000,
-            hop_size=256,
-            win_size=1024,
-            fmin=0,
-            fmax=12000,
-            mel_basis=mel_basis,
-            hann_window=hann_window,
-        ).transpose(1, 2)
-        spk = self.speaker_encoder(mels.to(dtype=torch.bfloat16))[0]
-        return spk.to(dtype=torch.bfloat16)
-
-    def _get_speaker_mel_buffers(
-        self,
-        *,
-        device: torch.device,
-        sampling_rate: int,
-        n_fft: int,
-        num_mels: int,
-        win_size: int,
-        fmin: int,
-        fmax: int | None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        cache = getattr(self, "_speaker_mel_buffer_cache", None)
-        if cache is None:
-            cache = {}
-            self._speaker_mel_buffer_cache = cache
-        key = (str(device), int(sampling_rate), int(n_fft), int(num_mels), int(win_size), int(fmin), fmax)
-        cached = cache.get(key)
-        if cached is not None:
-            return cached
-        mel_basis = mel_filter_bank(
-            sr=int(sampling_rate),
-            n_fft=int(n_fft),
-            n_mels=int(num_mels),
-            fmin=int(fmin),
-            fmax=fmax,
-        ).to(device)
-        hann_window = torch.hann_window(int(win_size), device=device)
-        cached = (mel_basis, hann_window)
-        cache[key] = cached
-        return cached
-
-    def _ensure_speech_tokenizer_loaded(self) -> Qwen3TTSTokenizer:
-        if self._speech_tokenizer is not None:
-            return self._speech_tokenizer
-        speech_tokenizer_path = cached_file(self.model_path, "speech_tokenizer/config.json")
-        if speech_tokenizer_path is None:
-            raise ValueError(f"{self.model_path}/speech_tokenizer/config.json not found")
-        # Ensure the HF feature extractor config is present. Transformers'
-        # AutoFeatureExtractor does not proactively fetch this file.
-        preprocessor_config_path = cached_file(self.model_path, "speech_tokenizer/preprocessor_config.json")
-        if preprocessor_config_path is None:
-            raise ValueError(f"{self.model_path}/speech_tokenizer/preprocessor_config.json not found")
-        speech_tokenizer_dir = os.path.dirname(speech_tokenizer_path)
-        tok = Qwen3TTSTokenizer.from_pretrained(
-            speech_tokenizer_dir,
-            torch_dtype=torch.bfloat16,
-        )
-        # Only move encoder to GPU; the decoder is unused by Talker (which
-        # only calls tok.encode()) and would otherwise waste bf16 VRAM.
-        # NOTE: after this point the tokenizer instance is encode-only;
-        # calling tok.decode() will fail because tok.model.decoder is None.
-        dev = next(self.parameters()).device
-        if dev.type != "cpu":
-            try:
-                del tok.model.decoder
-                tok.model.decoder = None
-                tok.model.encoder.to(dev)
-                tok.device = dev
-            except Exception as e:
-                raise RuntimeError(f"Failed to move speech tokenizer encoder to {dev}: {e}") from e
-        else:
-            tok.device = dev
-        self._speech_tokenizer = tok
-        return tok
-
-    def _encode_ref_audio_to_code(self, wav: np.ndarray, sr: int) -> torch.Tensor:
-        tok = self._ensure_speech_tokenizer_loaded()
-        enc = tok.encode(wav, sr=int(sr), return_dict=True)
-        ref_code = getattr(enc, "audio_codes", None)
-        if isinstance(ref_code, list):
-            ref_code = ref_code[0] if ref_code else None
-        if isinstance(ref_code, torch.Tensor):
-            # 12Hz: likely [T, Q] or [B, T, Q]
-            if ref_code.ndim == 3:
-                ref_code = ref_code[0]
-            return ref_code.to(device=next(self.parameters()).device, dtype=torch.long)
-        raise ValueError("SpeechTokenizer.encode did not return audio_codes tensor")
-
-    @staticmethod
-    def _first_value(value: object, default: object = None) -> object:
-        if isinstance(value, list):
-            return value[0] if value else default
-        return value if value is not None else default
-
-    @staticmethod
-    def _coerce_ref_code_tensor(value: object, *, device: torch.device) -> torch.Tensor | None:
-        value = Qwen3TTSTalkerForConditionalGeneration._first_value(value)
-        if value is None:
-            return None
-        if isinstance(value, torch.Tensor):
-            ref_code = value
-        elif isinstance(value, np.ndarray):
-            ref_code = torch.from_numpy(value)
-        elif isinstance(value, list) and value:
-            ref_code = torch.as_tensor(value, dtype=torch.long)
-        else:
-            return None
-        if ref_code.ndim == 3:
-            ref_code = ref_code[0]
-        if ref_code.ndim != 2 or ref_code.numel() == 0:
-            return None
-        return ref_code.to(device=device, dtype=torch.long).contiguous()
-
-    @staticmethod
-    def _coerce_token_ids(value: object, *, device: torch.device) -> torch.Tensor | None:
-        value = Qwen3TTSTalkerForConditionalGeneration._first_value(value)
-        if value is None:
-            return None
-        if isinstance(value, torch.Tensor):
-            ids = value
-        elif isinstance(value, np.ndarray):
-            ids = torch.from_numpy(value)
-        elif isinstance(value, list) and value and all(isinstance(v, (int, np.integer)) for v in value):
-            ids = torch.tensor(value, dtype=torch.long)
-        else:
-            return None
-        if ids.ndim == 1:
-            ids = ids.unsqueeze(0)
-        if ids.ndim != 2 or ids.numel() == 0:
-            return None
-        return ids.to(device=device, dtype=torch.long).contiguous()
-
-    @staticmethod
-    def _split_ref_code_batch(ref_codes: object, expected: int) -> list[object] | None:
-        if isinstance(ref_codes, torch.Tensor):
-            if ref_codes.ndim == 3 and int(ref_codes.shape[0]) == expected:
-                return [ref_codes[i] for i in range(expected)]
-            if expected == 1:
-                return [ref_codes]
-            return None
-        if isinstance(ref_codes, (list, tuple)) and len(ref_codes) == expected:
-            return list(ref_codes)
-        return None
-
-    @staticmethod
-    def _voice_clone_prompt_dict(raw: object) -> dict[str, object] | None:
-        raw = Qwen3TTSTalkerForConditionalGeneration._first_value(raw)
-        if isinstance(raw, dict):
-            return raw
-        if isinstance(raw, list) and raw and isinstance(raw[0], dict):
-            return raw[0]
-        return None
-
-    @staticmethod
-    def _has_ref_code_like(value: object) -> bool:
-        value = Qwen3TTSTalkerForConditionalGeneration._first_value(value)
-        if isinstance(value, torch.Tensor):
-            return value.numel() > 0 and value.ndim in (2, 3)
-        if isinstance(value, np.ndarray):
-            return value.size > 0 and value.ndim in (2, 3)
-        if isinstance(value, list):
-            return bool(value)
-        return False
-
-    @staticmethod
-    def _needs_initial_prompt_preprocess(info_dict: dict[str, Any]) -> bool:
-        embed = info_dict.get("embed")
-        return not (isinstance(embed, dict) and "prefill" in embed)
-
-    def _needs_batched_ref_code(self, info_dict: dict[str, Any]) -> bool:
-        if not self._needs_initial_prompt_preprocess(info_dict):
-            return False
-        if self._first_value(info_dict.get("task_type"), "CustomVoice") != "Base":
-            return False
-        codes = info_dict.get("codes")
-        if isinstance(codes, dict) and self._has_ref_code_like(codes.get(_PRECOMPUTED_REF_CODE_KEY)):
-            return False
-
-        voice_clone_prompt = self._voice_clone_prompt_dict(info_dict.get("voice_clone_prompt"))
-        if voice_clone_prompt is not None:
-            if self._has_ref_code_like(voice_clone_prompt.get("ref_code")):
-                return False
-            icl_flag = self._first_value(voice_clone_prompt.get("icl_mode"))
-            if isinstance(icl_flag, bool) and not icl_flag:
-                return False
-
-        xvec_only = self._first_value(info_dict.get("x_vector_only_mode"), False)
-        if isinstance(xvec_only, bool) and xvec_only:
-            return False
-
-        ref_audio_list = info_dict.get("ref_audio")
-        return isinstance(ref_audio_list, list) and bool(ref_audio_list)
-
     @torch.inference_mode()
     def preprocess_batch(
         self,
@@ -1509,568 +846,42 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         model_intermediate_buffer: dict[str, dict[str, Any]],
         device: torch.device,
     ) -> None:
-        """Batch Base voice-clone ref-audio codec extraction for current prefill requests."""
-        pending_text: list[tuple[dict[str, Any], str]] = []
-        pending_ref_text: list[tuple[dict[str, Any], str]] = []
-        groups: dict[int, list[tuple[dict[str, Any], np.ndarray, int]]] = {}
-        for req_id in req_ids:
-            info_dict = model_intermediate_buffer.get(req_id)
-            if not isinstance(info_dict, dict):
-                continue
-
-            if (
-                self._needs_initial_prompt_preprocess(info_dict)
-                and self._first_value(info_dict.get("task_type"), "CustomVoice") == "Base"
-            ):
-                if _PRECOMPUTED_TEXT_IDS_KEY not in info_dict:
-                    text = self._first_value(info_dict.get("text"), "")
-                    if isinstance(text, str):
-                        pending_text.append((info_dict, self._build_assistant_text(text)))
-
-                if _PRECOMPUTED_REF_IDS_KEY not in info_dict:
-                    ref_text = self._first_value(info_dict.get("ref_text"))
-                    if isinstance(ref_text, str) and ref_text.strip():
-                        pending_ref_text.append((info_dict, self._build_ref_text(ref_text)))
-
-            if not self._needs_batched_ref_code(info_dict):
-                continue
-            ref_audio_list = info_dict.get("ref_audio")
-            if not isinstance(ref_audio_list, list) or not ref_audio_list:
-                continue
-            try:
-                wav, sr = self._normalize_ref_audio(ref_audio_list[0])
-            except Exception:
-                # Keep the original per-request path responsible for surfacing
-                # invalid ref_audio errors with its existing messages.
-                continue
-            info_dict[_NORMALIZED_REF_AUDIO_KEY] = (wav, sr)
-            groups.setdefault(int(sr), []).append((info_dict, wav, int(sr)))
-
-        if pending_text or pending_ref_text:
-            tok_text = self._get_tokenizer()
-            for items, key in (
-                (pending_text, _PRECOMPUTED_TEXT_IDS_KEY),
-                (pending_ref_text, _PRECOMPUTED_REF_IDS_KEY),
-            ):
-                if not items:
-                    continue
-                try:
-                    tokenized = tok_text([text for _, text in items], padding=False)
-                except Exception as exc:
-                    logger.debug("Qwen3-TTS batched text tokenization failed; falling back to serial path: %s", exc)
-                    continue
-                input_ids = tokenized.get("input_ids") if isinstance(tokenized, dict) else None
-                if not isinstance(input_ids, list) or len(input_ids) != len(items):
-                    continue
-                for (info_dict, _), ids in zip(items, input_ids, strict=True):
-                    if isinstance(ids, list) and ids:
-                        info_dict[key] = torch.tensor(ids, dtype=torch.long)
-
-        groups = {sr: items for sr, items in groups.items() if len(items) >= 2}
-        if not groups:
-            return
-
-        tok = self._ensure_speech_tokenizer_loaded()
-        for sr, items in groups.items():
-            wavs = [wav for _, wav, _ in items]
-            try:
-                enc = tok.encode(wavs, sr=int(sr), return_dict=True)
-                ref_codes = self._split_ref_code_batch(getattr(enc, "audio_codes", None), len(items))
-            except Exception as exc:
-                logger.debug("Qwen3-TTS batched ref_code encode failed; falling back to serial path: %s", exc)
-                continue
-            if ref_codes is None:
-                continue
-            for (info_dict, wav, item_sr), ref_code in zip(items, ref_codes, strict=True):
-                ref_code_t = self._coerce_ref_code_tensor(ref_code, device=device)
-                if ref_code_t is None:
-                    continue
-                info_dict.setdefault("codes", {})[_PRECOMPUTED_REF_CODE_KEY] = ref_code_t
-                info_dict[_NORMALIZED_REF_AUDIO_KEY] = (wav, item_sr)
-
-    def _generate_icl_prompt(
-        self,
-        *,
-        text_id: torch.Tensor,
-        ref_id: torch.Tensor,
-        ref_code: torch.Tensor,
-        tts_pad_embed: torch.Tensor,
-        tts_eos_embed: torch.Tensor,
-        non_streaming_mode: bool,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        # Ported from official Qwen3TTSForConditionalGeneration.generate_icl_prompt
-        text_embed = self.text_projection(self.text_embedding(torch.cat([ref_id, text_id], dim=-1)))
-        text_embed = torch.cat([text_embed, tts_eos_embed], dim=1)
-
-        # codec embed (codec bos + codec) 1 T2 D
-        codec_embed: list[torch.Tensor] = []
-        for i in range(int(self.talker_config.num_code_groups)):
-            if i == 0:
-                codec_embed.append(self.embed_input_ids(ref_code[:, :1]))
-            else:
-                codec_embed.append(self.code_predictor.get_input_embeddings()[i - 1](ref_code[:, i : i + 1]))
-        codec_embed_sum = torch.cat(codec_embed, dim=1).sum(1).unsqueeze(0)  # [1,T,H]
-        codec_embed_sum = torch.cat(
-            [
-                self.embed_input_ids(
-                    torch.tensor([[self.talker_config.codec_bos_id]], device=codec_embed_sum.device, dtype=torch.long)
-                ),
-                codec_embed_sum,
-            ],
-            dim=1,
+        """Delegate batched preprocess to :class:`Qwen3TTSPromptEmbedsBuilder`."""
+        self._prompt_builder.preprocess_batch(
+            req_ids=req_ids,
+            model_intermediate_buffer=model_intermediate_buffer,
+            device=device,
         )
 
-        text_lens = int(text_embed.shape[1])
-        codec_lens = int(codec_embed_sum.shape[1])
-        if non_streaming_mode:
-            # Official non-streaming mode: append the full text conditioning in
-            # prefill, and use PAD in decode steps.
-            icl_input_embed = text_embed + self.embed_input_ids(
-                torch.tensor(
-                    [[self.talker_config.codec_pad_id] * text_lens],
-                    device=codec_embed_sum.device,
-                    dtype=torch.long,
-                )
-            )
-            icl_input_embed = torch.cat([icl_input_embed, codec_embed_sum + tts_pad_embed], dim=1)
-            return icl_input_embed, tts_pad_embed
-        if text_lens > codec_lens:
-            return text_embed[:, :codec_lens] + codec_embed_sum, text_embed[:, codec_lens:]
-        text_embed = torch.cat([text_embed] + [tts_pad_embed] * (codec_lens - text_lens), dim=1)
-        return text_embed + codec_embed_sum, tts_pad_embed
+    def _encode_ref_audio_batch(self, wavs: list[np.ndarray], sr: int, *, device: torch.device) -> list[torch.Tensor]:
+        fe = self._encoder_feature_extractor
+        target_sr = int(fe.sampling_rate)
+        if int(sr) != target_sr:
+            resampler = AudioResampler(target_sr=target_sr)
+            wavs = [resampler.resample(w.astype(np.float32), orig_sr=int(sr)) for w in wavs]
 
-    def _build_prompt_embeds(
-        self,
-        *,
-        task_type: str,
-        info_dict: dict[str, Any],
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int | None, torch.Tensor | None]:
-        text = (info_dict.get("text") or [""])[0]
-        language = (info_dict.get("language") or ["Auto"])[0]
-        non_streaming_mode_val = info_dict.get("non_streaming_mode")
-        if isinstance(non_streaming_mode_val, list):
-            non_streaming_mode_raw = non_streaming_mode_val[0] if non_streaming_mode_val else None
-        else:
-            non_streaming_mode_raw = non_streaming_mode_val
-        if isinstance(non_streaming_mode_raw, bool):
-            non_streaming_mode = non_streaming_mode_raw
-        else:
-            # Match official inference defaults:
-            # - CustomVoice/VoiceDesign: non_streaming_mode=True
-            # - Base: non_streaming_mode=False
-            non_streaming_mode = task_type in ("CustomVoice", "VoiceDesign")
+        inputs = fe(
+            raw_audio=wavs,
+            sampling_rate=target_sr,
+            return_tensors="pt",
+        )
+        inputs = inputs.to(device).to(torch.bfloat16)
 
-        # Text ids for assistant template (always).
-        tok = self._get_tokenizer()
-        dev = next(self.parameters()).device
-        input_ids = self._coerce_token_ids(info_dict.pop(_PRECOMPUTED_TEXT_IDS_KEY, None), device=dev)
-        if input_ids is None:
-            input_ids = tok(self._build_assistant_text(text), return_tensors="pt", padding=False)["input_ids"].to(
-                device=dev
+        input_values = inputs["input_values"].squeeze(1)
+        padding_mask = inputs["padding_mask"].squeeze(1)
+
+        with torch.inference_mode():
+            encoded = self.encoder.encode(
+                input_values=input_values.unsqueeze(1),
+                return_dict=True,
             )
 
-        # Optional instruct prefix.
-        instruct = (info_dict.get("instruct") or [""])[0]
-        instruct_embed = None
-        if isinstance(instruct, str) and instruct.strip():
-            instruct_ids = tok(self._build_instruct_text(instruct), return_tensors="pt", padding=False)["input_ids"].to(
-                device=input_ids.device
-            )
-            instruct_embed = self.text_projection(self.text_embedding(instruct_ids))
-
-        # tts special token embeds (projected into talker hidden).
-        tts_tokens = torch.tensor(
-            [[self.config.tts_bos_token_id, self.config.tts_eos_token_id, self.config.tts_pad_token_id]],
-            device=input_ids.device,
-            dtype=input_ids.dtype,
-        )
-        tts_bos_embed, tts_eos_embed, tts_pad_embed = self.text_projection(self.text_embedding(tts_tokens)).chunk(
-            3, dim=1
-        )
-
-        # Codec prefill tags.
-        language_id = None
-        if isinstance(language, str) and language.lower() != "auto":
-            language_id = self.talker_config.codec_language_id.get(language.lower())
-        # Match official dialect override:
-        # If language is Chinese/Auto and the selected speaker is a dialect voice,
-        # set language_id to that dialect to improve code generation stability.
-        if language_id is None and isinstance(language, str) and language.lower() in ("chinese", "auto"):
-            speaker_for_dialect = None
-            if task_type == "CustomVoice":
-                speaker_for_dialect = (info_dict.get("speaker") or [""])[0]
-            if isinstance(speaker_for_dialect, str) and speaker_for_dialect.strip():
-                spk_is_dialect = getattr(self.talker_config, "spk_is_dialect", None) or {}
-                dialect = spk_is_dialect.get(speaker_for_dialect.lower())
-                if isinstance(dialect, str) and dialect:
-                    language_id = self.talker_config.codec_language_id.get(dialect)
-        if language_id is None:
-            codec_prefill_list = [
-                [
-                    self.talker_config.codec_nothink_id,
-                    self.talker_config.codec_think_bos_id,
-                    self.talker_config.codec_think_eos_id,
-                ]
-            ]
-        else:
-            codec_prefill_list = [
-                [
-                    self.talker_config.codec_think_id,
-                    self.talker_config.codec_think_bos_id,
-                    int(language_id),
-                    self.talker_config.codec_think_eos_id,
-                ]
-            ]
-
-        codec_input_0 = self.embed_input_ids(
-            torch.tensor(codec_prefill_list, device=input_ids.device, dtype=torch.long)
-        )
-        codec_input_1 = self.embed_input_ids(
-            torch.tensor([[self.talker_config.codec_pad_id, self.talker_config.codec_bos_id]], device=input_ids.device)
-        )
-
-        # Speaker embedding/token (task-dependent)
-        speaker_embed = None
-        ref_code_len: int | None = None
-        ref_code_prompt: torch.Tensor | None = None
-
-        def _as_singleton(x: object) -> object:
-            if isinstance(x, list):
-                return x[0] if x else None
-            return x
-
-        def _to_long_tensor(x: object, *, device: torch.device) -> torch.Tensor | None:
-            x = _as_singleton(x)
-            if x is None:
-                return None
-            if isinstance(x, torch.Tensor):
-                t = x
-            elif isinstance(x, np.ndarray):
-                t = torch.from_numpy(x)
-            elif isinstance(x, list) and x and all(isinstance(v, (int, np.integer)) for v in x):
-                t = torch.tensor(x, dtype=torch.long)
-            else:
-                return None
-            if t.ndim == 1:
-                t = t.unsqueeze(0)
-            return t.to(device=device, dtype=torch.long)
-
-        def _normalize_voice_clone_prompt(raw: object) -> dict[str, object] | None:
-            raw = _as_singleton(raw)
-            if raw is None:
-                return None
-            if isinstance(raw, dict):
-                return raw
-            # Some callers may pass list[dict] directly.
-            if isinstance(raw, list) and raw and isinstance(raw[0], dict):
-                return raw[0]
-            return None
-
-        if task_type == "Base":
-            # Base supports voice clone prompt with in-context mode.
-            xvec_only = bool((info_dict.get("x_vector_only_mode") or [False])[0])
-            in_context_mode = not xvec_only
-            voice_clone_prompt = _normalize_voice_clone_prompt(info_dict.get("voice_clone_prompt"))
-            ref_audio_wav: np.ndarray | None = None
-            ref_audio_sr: int | None = None
-
-            def _get_ref_audio() -> tuple[np.ndarray, int]:
-                nonlocal ref_audio_wav, ref_audio_sr
-                if ref_audio_wav is None or ref_audio_sr is None:
-                    normalized_ref_audio = info_dict.pop(_NORMALIZED_REF_AUDIO_KEY, None)
-                    if (
-                        isinstance(normalized_ref_audio, tuple)
-                        and len(normalized_ref_audio) == 2
-                        and isinstance(normalized_ref_audio[0], np.ndarray)
-                    ):
-                        ref_audio_wav = normalized_ref_audio[0]
-                        ref_audio_sr = int(normalized_ref_audio[1])
-                        return ref_audio_wav, ref_audio_sr
-                    ref_audio_list = info_dict.get("ref_audio")
-                    if not isinstance(ref_audio_list, list) or not ref_audio_list:
-                        raise ValueError("Base requires `ref_audio`.")
-                    ref_audio_wav, ref_audio_sr = self._normalize_ref_audio(ref_audio_list[0])
-                return ref_audio_wav, ref_audio_sr
-
-            # Speaker cache: only for uploaded (named) speakers
-            _speaker_cache_key = None
-            if voice_clone_prompt is None:
-                _speaker_list = info_dict.get("speaker")
-                if isinstance(_speaker_list, list) and _speaker_list:
-                    _voice_name = str(_speaker_list[0]).lower()
-                    # Per-mode namespace — xvec and icl produce different artifacts
-                    # for the same voice, so they must not share a cache slot.
-                    _mode = "xvec" if xvec_only else "icl"
-                    _voice_created_at = int((info_dict.get("voice_created_at") or [0])[0])
-                    _speaker_cache_key = self._speaker_cache.make_cache_key(
-                        _voice_name,
-                        model_type=f"qwen3_tts_{_mode}",
-                        created_at=_voice_created_at,
-                    )
-                    _cached = self._speaker_cache.get(_speaker_cache_key)
-                    if _cached is not None:
-                        # Transfer cached tensors to current device
-                        ref_code_cached = _cached.get("ref_code")
-                        ref_spk_embed_cached = _cached.get("ref_spk_embedding")
-                        if isinstance(ref_code_cached, torch.Tensor):
-                            ref_code_cached = ref_code_cached.to(device=input_ids.device)
-                        if isinstance(ref_spk_embed_cached, torch.Tensor):
-                            ref_spk_embed_cached = ref_spk_embed_cached.to(device=input_ids.device)
-                        voice_clone_prompt = {
-                            "ref_code": ref_code_cached,
-                            "ref_spk_embedding": ref_spk_embed_cached,
-                            "icl_mode": _cached.get("icl_mode"),
-                        }
-                        _speaker_cache_key = None  # hit → don't store again
-
-            # Official implementation may pass `voice_clone_prompt.icl_mode`.
-            if voice_clone_prompt is not None and "icl_mode" in voice_clone_prompt:
-                icl_flag = _as_singleton(voice_clone_prompt.get("icl_mode"))
-                if isinstance(icl_flag, bool):
-                    in_context_mode = icl_flag
-                    xvec_only = not in_context_mode
-            ref_code = None
-            if voice_clone_prompt is not None:
-                ref_code = _as_singleton(voice_clone_prompt.get("ref_code"))
-            ref_code_t = None
-            if isinstance(ref_code, torch.Tensor):
-                ref_code_t = ref_code
-            elif isinstance(ref_code, np.ndarray):
-                ref_code_t = torch.from_numpy(ref_code)
-            if isinstance(ref_code_t, torch.Tensor):
-                if ref_code_t.ndim == 3:
-                    ref_code_t = ref_code_t[0]
-                ref_code_t = ref_code_t.to(device=input_ids.device, dtype=torch.long)
-                ref_code_len = int(ref_code_t.shape[0])
-            else:
-                codes = info_dict.get("codes")
-                precomputed_ref_code = codes.get(_PRECOMPUTED_REF_CODE_KEY) if isinstance(codes, dict) else None
-                ref_code_t = self._coerce_ref_code_tensor(precomputed_ref_code, device=input_ids.device)
-                if isinstance(ref_code_t, torch.Tensor):
-                    ref_code_len = int(ref_code_t.shape[0])
-            if ref_code_t is None and in_context_mode:
-                # Compute ref_code from ref_audio if not provided.
-                wav_np, sr = _get_ref_audio()
-                ref_code_t = self._encode_ref_audio_to_code(wav_np, sr).to(device=input_ids.device)
-                ref_code_len = int(ref_code_t.shape[0])
-            if isinstance(ref_code_t, torch.Tensor):
-                ref_code_prompt = ref_code_t
-
-            # Speaker embedding: use prompt embed if provided; otherwise extract from audio.
-            # NOTE: Do NOT use _as_singleton here — the embedding may be a plain
-            # float list (from API via msgspec IPC) that _as_singleton would
-            # destructively unwrap to a single scalar.
-            spk = None
-            if voice_clone_prompt is not None:
-                spk = voice_clone_prompt.get("ref_spk_embedding")
-            if isinstance(spk, torch.Tensor):
-                speaker_embed = spk.to(device=input_ids.device, dtype=torch.bfloat16).view(1, 1, -1)
-            elif isinstance(spk, (list, np.ndarray)):
-                # Plain list/array from API (survived msgspec IPC serialization).
-                speaker_embed = torch.tensor(spk, dtype=torch.bfloat16, device=input_ids.device).view(1, 1, -1)
-            else:
-                wav_np, sr = _get_ref_audio()
-                speaker_embed = self._extract_speaker_embedding(wav_np, sr).view(1, 1, -1)
-
-            # Cache miss: store extraction result
-            if _speaker_cache_key is not None and speaker_embed is not None:
-                self._speaker_cache.put(
-                    _speaker_cache_key,
-                    {
-                        "ref_code": ref_code_prompt.detach().cpu()
-                        if isinstance(ref_code_prompt, torch.Tensor)
-                        else None,
-                        "ref_spk_embedding": speaker_embed.detach().cpu().reshape(-1),
-                        "icl_mode": in_context_mode,
-                    },
-                )
-
-            codec_input = torch.cat([codec_input_0, speaker_embed, codec_input_1], dim=1)
-
-            # Role header (<|im_start|>assistant\n) -> projected text embeds.
-            role_embed = self.text_projection(self.text_embedding(input_ids[:, :3]))
-
-            codec_prefix = torch.cat((tts_pad_embed.expand(-1, codec_input.shape[1] - 2, -1), tts_bos_embed), dim=1)
-            codec_prefix = codec_prefix + codec_input[:, :-1]
-            talker_prompt = torch.cat((role_embed, codec_prefix), dim=1)
-
-            if in_context_mode:
-                # Prefer explicit tokenized `ref_ids` if provided (matches official signature).
-                ref_ids = self._coerce_token_ids(info_dict.pop(_PRECOMPUTED_REF_IDS_KEY, None), device=input_ids.device)
-                if ref_ids is None:
-                    ref_ids = _to_long_tensor(info_dict.get("ref_ids"), device=input_ids.device)
-                if ref_ids is None and voice_clone_prompt is not None:
-                    ref_ids = _to_long_tensor(
-                        voice_clone_prompt.get("ref_ids") or voice_clone_prompt.get("ref_id"), device=input_ids.device
-                    )
-                if ref_ids is None:
-                    ref_text = _as_singleton(info_dict.get("ref_text"))
-                    if isinstance(ref_text, str) and ref_text.strip():
-                        ref_ids = tok(
-                            self._build_ref_text(ref_text),
-                            return_tensors="pt",
-                            padding=False,
-                        )["input_ids"].to(device=input_ids.device)
-                    else:
-                        logger.warning("Base ICL: ref_text/ref_ids missing, falling back to x-vector-only mode.")
-                        in_context_mode = False
-            if in_context_mode:
-                icl_input_embed, trailing_text_hidden = self._generate_icl_prompt(
-                    text_id=input_ids[:, 3:-5],
-                    ref_id=ref_ids[:, 3:-2],
-                    ref_code=ref_code_t,  # type: ignore[arg-type]
-                    tts_pad_embed=tts_pad_embed,
-                    tts_eos_embed=tts_eos_embed,
-                    non_streaming_mode=non_streaming_mode,
-                )
-                talker_prompt = torch.cat([talker_prompt, icl_input_embed], dim=1)
-            else:
-                # First text token (+ codec_bos).
-                if non_streaming_mode:
-                    # Official non-streaming mode: put the full text into the
-                    # prefill prompt and use PAD for decode steps.
-                    text_all = self.text_projection(self.text_embedding(input_ids[:, 3:-5]))
-                    text_all = torch.cat([text_all, tts_eos_embed], dim=1)
-                    pad_ids = torch.full(
-                        (1, int(text_all.shape[1])),
-                        int(self.talker_config.codec_pad_id),
-                        device=input_ids.device,
-                        dtype=torch.long,
-                    )
-                    talker_prompt = torch.cat(
-                        [
-                            talker_prompt,
-                            text_all + self.embed_input_ids(pad_ids),
-                            tts_pad_embed
-                            + self.embed_input_ids(
-                                torch.tensor([[self.talker_config.codec_bos_id]], device=input_ids.device)
-                            ),
-                        ],
-                        dim=1,
-                    )
-                    trailing_text_hidden = tts_pad_embed
-                else:
-                    first_text = self.text_projection(self.text_embedding(input_ids[:, 3:4])) + codec_input[:, -1:]
-                    talker_prompt = torch.cat([talker_prompt, first_text], dim=1)
-                    trailing_text_hidden = torch.cat(
-                        (
-                            self.text_projection(self.text_embedding(input_ids[:, 4:-5])),
-                            tts_eos_embed,
-                        ),
-                        dim=1,
-                    )
-
-        elif task_type == "CustomVoice":
-            _speaker_raw = info_dict.get("speaker") or [""]
-            speaker = (
-                ((_speaker_raw[0] if isinstance(_speaker_raw, (list, tuple)) else _speaker_raw) or "").lower().strip()
-            )
-            if not speaker:
-                raise ValueError("CustomVoice requires additional_information.speaker.")
-            spk_id_map = {k.lower(): v for k, v in (getattr(self.talker_config, "spk_id", None) or {}).items()}
-            if speaker not in spk_id_map:
-                raise ValueError(f"Unsupported speaker: {speaker}")
-            spk_id = spk_id_map[speaker]
-            # Keep it at least 1D; embedding on a 0-d tensor can return 1D.
-            spk_tensor = torch.tensor([spk_id], device=input_ids.device, dtype=torch.long)
-            spk_embed = self.embed_input_ids(spk_tensor)
-            if spk_embed.ndim == 1:
-                spk_embed = spk_embed.view(1, 1, -1)
-            elif spk_embed.ndim == 2:
-                spk_embed = spk_embed.view(1, 1, -1)
-            speaker_embed = spk_embed
-            codec_input = torch.cat([codec_input_0, speaker_embed, codec_input_1], dim=1)
-
-            role_embed = self.text_projection(self.text_embedding(input_ids[:, :3]))
-            codec_prefix = torch.cat((tts_pad_embed.expand(-1, codec_input.shape[1] - 2, -1), tts_bos_embed), dim=1)
-            codec_prefix = codec_prefix + codec_input[:, :-1]
-            talker_prompt = torch.cat((role_embed, codec_prefix), dim=1)
-
-            if non_streaming_mode:
-                text_all = self.text_projection(self.text_embedding(input_ids[:, 3:-5]))
-                text_all = torch.cat([text_all, tts_eos_embed], dim=1)
-                pad_ids = torch.full(
-                    (1, int(text_all.shape[1])),
-                    int(self.talker_config.codec_pad_id),
-                    device=input_ids.device,
-                    dtype=torch.long,
-                )
-                talker_prompt = torch.cat(
-                    [
-                        talker_prompt,
-                        text_all + self.embed_input_ids(pad_ids),
-                        tts_pad_embed
-                        + self.embed_input_ids(
-                            torch.tensor([[self.talker_config.codec_bos_id]], device=input_ids.device)
-                        ),
-                    ],
-                    dim=1,
-                )
-                trailing_text_hidden = tts_pad_embed
-            else:
-                first_text = self.text_projection(self.text_embedding(input_ids[:, 3:4])) + codec_input[:, -1:]
-                talker_prompt = torch.cat([talker_prompt, first_text], dim=1)
-                trailing_text_hidden = torch.cat(
-                    (
-                        self.text_projection(self.text_embedding(input_ids[:, 4:-5])),
-                        tts_eos_embed,
-                    ),
-                    dim=1,
-                )
-
-        elif task_type == "VoiceDesign":
-            # No known speaker identity; only codec tags + text.
-            codec_input = torch.cat([codec_input_0, codec_input_1], dim=1)
-
-            role_embed = self.text_projection(self.text_embedding(input_ids[:, :3]))
-            codec_prefix = torch.cat((tts_pad_embed.expand(-1, codec_input.shape[1] - 2, -1), tts_bos_embed), dim=1)
-            codec_prefix = codec_prefix + codec_input[:, :-1]
-            talker_prompt = torch.cat((role_embed, codec_prefix), dim=1)
-
-            if non_streaming_mode:
-                text_all = self.text_projection(self.text_embedding(input_ids[:, 3:-5]))
-                text_all = torch.cat([text_all, tts_eos_embed], dim=1)
-                pad_ids = torch.full(
-                    (1, int(text_all.shape[1])),
-                    int(self.talker_config.codec_pad_id),
-                    device=input_ids.device,
-                    dtype=torch.long,
-                )
-                talker_prompt = torch.cat(
-                    [
-                        talker_prompt,
-                        text_all + self.embed_input_ids(pad_ids),
-                        tts_pad_embed
-                        + self.embed_input_ids(
-                            torch.tensor([[self.talker_config.codec_bos_id]], device=input_ids.device)
-                        ),
-                    ],
-                    dim=1,
-                )
-                trailing_text_hidden = tts_pad_embed
-            else:
-                first_text = self.text_projection(self.text_embedding(input_ids[:, 3:4])) + codec_input[:, -1:]
-                talker_prompt = torch.cat([talker_prompt, first_text], dim=1)
-                trailing_text_hidden = torch.cat(
-                    (
-                        self.text_projection(self.text_embedding(input_ids[:, 4:-5])),
-                        tts_eos_embed,
-                    ),
-                    dim=1,
-                )
-        else:
-            raise ValueError(f"Unsupported task_type={task_type}")
-
-        if instruct_embed is not None:
-            talker_prompt = torch.cat([instruct_embed, talker_prompt], dim=1)
-
-        return (
-            talker_prompt.squeeze(0),  # [prompt_len, H]
-            trailing_text_hidden.squeeze(0),  # [T, H]
-            tts_pad_embed.squeeze(0),  # [1, H]
-            ref_code_len,
-            ref_code_prompt.contiguous() if isinstance(ref_code_prompt, torch.Tensor) else None,
-        )
+        audio_codes = encoded.audio_codes[:, : self._encoder_valid_num_quantizers]
+        downsample = self._encoder_downsample_rate
+        return [
+            code[..., : -(-mask.sum() // downsample)].transpose(0, 1).to(device=device, dtype=torch.long)
+            for code, mask in zip(audio_codes, padding_mask)
+        ]
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         # Consume talker weights, and conditionally consume speaker encoder
@@ -2096,9 +907,76 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             # Some checkpoints do not include speaker_encoder weights; keep the
             # eagerly initialized module and satisfy the strict loader check.
             loaded |= {name for name, _ in self.named_parameters() if name.startswith("speaker_encoder.")}
+        # Load speech tokenizer encoder weights from speech_tokenizer/
+        # subfolder.  Skip decoder weights — the Talker only uses the
+        # encoder for ref_audio encoding.
+        model_loader = DefaultModelLoader(self.vllm_config.load_config)
+        source = DefaultModelLoader.Source(
+            model_or_path=self.model_path,
+            revision=self.vllm_config.model_config.revision,
+            subfolder="speech_tokenizer",
+        )
+        subfolder_weights = model_loader._get_weights_iterator(source)
+        enc_loaded = AutoWeightsLoader(
+            self,
+            skip_prefixes=["decoder."],
+        ).load_weights(subfolder_weights)
+        loaded |= enc_loaded
+
+        # AutoWeightsLoader only loads parameters; the encoder's VQ
+        # codebook state (embed, embed_sum, cluster_usage, initialized)
+        # are registered as buffers. Load them from the checkpoint
+        # directly so the quantizer produces correct codes.
+        encoder_buffers = dict(self.encoder.named_buffers())
+        source2 = DefaultModelLoader.Source(
+            model_or_path=self.model_path,
+            revision=self.vllm_config.model_config.revision,
+            subfolder="speech_tokenizer",
+        )
+        for name, tensor in model_loader._get_weights_iterator(source2):
+            if not name.startswith("encoder."):
+                continue
+            buf_name = name[len("encoder.") :]
+            if buf_name in encoder_buffers:
+                encoder_buffers[buf_name].copy_(tensor)
+                loaded.add(name)
+
+        device = self.vllm_config.device_config.device
+        self.encoder.to(device=device, dtype=torch.bfloat16)
+
+        self._init_runtime_buffers()
+
         logger.info("Loaded %d weights for Qwen3TTSTalkerForConditionalGeneration", len(loaded))
+        self._build_stacked_codec_embed()
         return loaded
 
+    def _build_stacked_codec_embed(self) -> None:
+        embeds = self.code_predictor.get_input_embeddings()
+        if not embeds:
+            return
+        w = embeds[0].weight
+        self._stacked_codec_embed = torch.stack([e.weight.detach() for e in embeds], dim=0).to(
+            device=w.device, dtype=w.dtype
+        )
+
+    @torch.no_grad()
+    def _init_runtime_buffers(self) -> None:
+        """Populate request-independent runtime buffers from frozen weights.
+
+        Currently this only computes :attr:`_tts_pad_embed`
+        (``text_proj(text_emb(tts_pad_token_id))``), which the prefill
+        prompt builder and every decode step mix into the input embedding
+        in place of an actual text token. The value depends only on the
+        ``text_embedding`` / ``text_projection`` weights and is therefore
+        the same for every request — computing it once here avoids
+        recomputing (and round-tripping through ``info_dict``) on each
+        forward call.
+        """
+        device = next(self.parameters()).device
+        pad_ids = torch.tensor([[int(self.config.tts_pad_token_id)]], device=device, dtype=torch.long)
+        pad_proj = self.text_projection(self.text_embedding(pad_ids)).reshape(1, -1)
+        self._tts_pad_embed.copy_(pad_proj.to(device=self._tts_pad_embed.device, dtype=self._tts_pad_embed.dtype))
+
     # -------------------- GPU-side MTP fast-path --------------------
 
     @torch.inference_mode()
@@ -2158,11 +1036,13 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         invalid0 = (layer0 < 0) | (layer0 >= int(self._codebook_vocab_size))
         audio_codes = torch.where(invalid0.expand_as(audio_codes), torch.zeros_like(audio_codes), audio_codes)
 
-        # Sum embeddings of all code groups, then add the current text step.
+        # Single gather over stacked [Q-1, V, H] replaces Q-1 serial embedding kernels.
         residual_ids_t = audio_codes[:, 1:]
-        embeds: list[torch.Tensor] = [last_id_hidden]
-        for i in range(max_steps):
-            embeds.append(self.code_predictor.get_input_embeddings()[i](residual_ids_t[:, i : i + 1]))
-        summed = torch.cat(embeds, dim=1).sum(1, keepdim=True)  # [B,1,H]
+        if self._stacked_codec_embed is None:
+            self._build_stacked_codec_embed()
+        embed_weight = self._stacked_codec_embed.to(device=dev)
+        row_idx = torch.arange(max_steps, device=dev).unsqueeze(0).expand(bsz, -1)
+        gathered = embed_weight[row_idx, residual_ids_t]
+        summed = (last_id_hidden.squeeze(1) + gathered.sum(dim=1)).unsqueeze(1)
         inputs_embeds_out = (summed + text_step).reshape(bsz, -1)
         return inputs_embeds_out, audio_codes.to(dtype=torch.long)
diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py
index 87acb6002..99b26ea1e 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py
+++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 """PyTorch Qwen3TTSTokenizerV2 model."""
 
+import inspect
 import math
 from collections.abc import Callable
 from dataclasses import dataclass
@@ -565,12 +566,17 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr
             # Prepare mask arguments
             mask_kwargs = {
                 "config": self.config,
-                "input_embeds": inputs_embeds,
                 "attention_mask": attention_mask,
-                "cache_position": cache_position,
                 "past_key_values": past_key_values,
                 "position_ids": position_ids,
             }
+            # Handle API changes across transformers versions
+            sig = inspect.signature(create_causal_mask)
+            if "input_embeds" in sig.parameters:
+                mask_kwargs["input_embeds"] = inputs_embeds
+                mask_kwargs["cache_position"] = cache_position
+            else:
+                mask_kwargs["inputs_embeds"] = inputs_embeds
             # Create the masks
             causal_mask_mapping = {
                 "full_attention": create_causal_mask(**mask_kwargs),
diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py
index 7bc0030bc..6ddbedd9b 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py
+++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py
@@ -13,13 +13,11 @@
 # limitations under the License.
 """PyTorch Qwen3TTSTokenizerV1 model."""
 
-import math
 from dataclasses import dataclass
 
 import numpy as np
 import torch
 from torch import nn
-from torch.nn import Parameter
 from torch.nn import functional as F
 from torch.nn.utils.rnn import pad_sequence
 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
@@ -27,6 +25,8 @@ from transformers.utils import ModelOutput, auto_docstring, logging
 from transformers.utils.hub import cached_file
 
 from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding
+from vllm_omni.model_executor.models.common.alias_free_activation import AliasFreeActivation1d
+from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta
 
 from .configuration_qwen3_tts_tokenizer_v1 import (
     Qwen3TTSTokenizerV1Config,
@@ -663,166 +663,9 @@ class DiTDecoderLayer(nn.Module):
         return hidden_states
 
 
-class SnakeBeta(nn.Module):
-    """
-    A modified Snake function which uses separate parameters for the magnitude of the periodic components
-    Shape:
-        - Input: (B, C, T)
-        - Output: (B, C, T), same shape as the input
-    Parameters:
-        - alpha - trainable parameter that controls frequency
-        - beta - trainable parameter that controls magnitude
-    References:
-        - This activation function is a modified version based on this paper
-          by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
-          https://huggingface.co/papers/2006.08195
-    """
-
-    def __init__(self, in_features, alpha=1.0):
-        super().__init__()
-        self.in_features = in_features
-
-        # initialize alpha
-        self.alpha = Parameter(torch.zeros(in_features) * alpha)
-        self.beta = Parameter(torch.zeros(in_features) * alpha)
-
-        self.no_div_by_zero = 0.000000001
-
-    def forward(self, hidden_states):
-        """
-        Forward pass of the function.
-        Applies the function to the input elementwise.
-        SnakeBeta ∶= x + 1/b * sin^2 (xa)
-        """
-        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)  # line up with x to [B, C, T]
-        beta = self.beta.unsqueeze(0).unsqueeze(-1)
-        alpha = torch.exp(alpha)
-        beta = torch.exp(beta)
-        hidden_states = hidden_states + (1.0 / (beta + self.no_div_by_zero)) * torch.pow(
-            torch.sin(hidden_states * alpha), 2
-        )
-
-        return hidden_states
-
-
-def kaiser_sinc_filter1d(cutoff, half_width, kernel_size) -> torch.Tensor:
-    """Generates a 1D Kaiser-windowed sinc filter.
-
-    Args:
-        cutoff (float): Normalized cutoff frequency (0 to 0.5).
-        half_width (float): Transition bandwidth.
-        kernel_size (int): Number of filter taps.
-
-    Returns:
-        torch.Tensor: A tensor of shape (1, 1, kernel_size) representing the filter.
-    """
-    is_even = kernel_size % 2 == 0
-    half_size = kernel_size // 2
-
-    # Compute Kaiser window parameters
-    delta_f = 4 * half_width
-    attenuation = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
-
-    if attenuation > 50.0:
-        beta = 0.1102 * (attenuation - 8.7)
-    elif attenuation >= 21.0:
-        beta = 0.5842 * (attenuation - 21) ** 0.4 + 0.07886 * (attenuation - 21.0)
-    else:
-        beta = 0.0
-
-    kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32)
-
-    # Compute time indices
-    if is_even:
-        time_indices = torch.arange(-half_size, half_size) + 0.5
-    else:
-        time_indices = torch.arange(kernel_size) - half_size
-
-    # Compute sinc filter
-    if cutoff == 0:
-        return torch.zeros((1, 1, kernel_size), dtype=torch.float32)  # Ensures correct shape
-
-    sinc_filter = torch.sinc(2 * cutoff * time_indices)
-    normalized_filter = 2 * cutoff * kaiser_window * sinc_filter
-
-    # Normalize to ensure sum = 1 (avoid leakage of constant component)
-    normalized_filter /= normalized_filter.sum()
-
-    return normalized_filter.view(1, 1, kernel_size)
-
-
-class UpSample1d(nn.Module):
-    def __init__(self, ratio=2, kernel_size=None):
-        super().__init__()
-        self.ratio = ratio
-        self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
-        self.stride = ratio
-        self.pad = self.kernel_size // ratio - 1
-        self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
-        self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
-
-        filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size)
-        self.register_buffer("filter", filter, persistent=False)
-
-    def forward(self, hidden_states):
-        channels = hidden_states.shape[1]
-
-        hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate")
-        hidden_states = self.ratio * F.conv_transpose1d(
-            hidden_states, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels
-        )
-        hidden_states = hidden_states[..., self.pad_left : -self.pad_right]
-
-        return hidden_states
-
-
-class DownSample1d(nn.Module):
-    def __init__(self, ratio=2, kernel_size=None):
-        super().__init__()
-        cutoff = 0.5 / ratio
-        half_width = 0.6 / ratio
-
-        if cutoff < 0.0:
-            raise ValueError("Minimum cutoff must be larger than zero.")
-        if cutoff > 0.5:
-            raise ValueError("A cutoff above 0.5 does not make sense.")
-
-        self.even = kernel_size % 2 == 0
-        self.pad_left = kernel_size // 2 - int(self.even)
-        self.pad_right = kernel_size // 2
-        self.stride = ratio
-        filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
-        self.register_buffer("filter", filter, persistent=False)
-
-    def forward(self, hidden_states):
-        channels = hidden_states.shape[1]
-        hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate")
-        out = F.conv1d(hidden_states, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels)
-        return out
-
-
-class TorchActivation1d(nn.Module):
-    def __init__(
-        self,
-        activation,
-        up_ratio: int = 2,
-        down_ratio: int = 2,
-        up_kernel_size: int = 12,
-        down_kernel_size: int = 12,
-    ):
-        super().__init__()
-        if not callable(activation):
-            raise TypeError("Activation function must be callable")
-        self.act = activation
-        self.upsample = UpSample1d(up_ratio, up_kernel_size)
-        self.downsample = DownSample1d(down_ratio, down_kernel_size)
-
-    def forward(self, hidden_states):
-        hidden_states = self.upsample(hidden_states)
-        hidden_states = self.act(hidden_states)
-        hidden_states = self.downsample(hidden_states)
-
-        return hidden_states
+# SnakeBeta, UpSample1d, DownSample1d, AliasFreeActivation1d (formerly
+# TorchActivation1d) are now provided by the shared ``common`` modules.
+# See: vllm_omni/model_executor/models/common/{snake_activation,alias_free_activation}.py
 
 
 class CausalConv1d(nn.Conv1d):
@@ -929,7 +772,7 @@ class AMPBlock(torch.nn.Module):
         self.num_layers = len(self.convs1) + len(self.convs2)  # total number of conv layers
 
         self.activations = nn.ModuleList(
-            [TorchActivation1d(activation=SnakeBeta(channels)) for _ in range(self.num_layers)]
+            [AliasFreeActivation1d(activation=SnakeBeta(channels)) for _ in range(self.num_layers)]
         )
 
         if causal_type == "2":
@@ -940,7 +783,7 @@ class AMPBlock(torch.nn.Module):
                 stride=1,
                 padding=self._get_padding(kernel_size, 1),
             )
-            self.pre_act = TorchActivation1d(activation=SnakeBeta(channels))
+            self.pre_act = AliasFreeActivation1d(activation=SnakeBeta(channels))
         else:
             self.pre_conv = nn.Identity()
             self.pre_act = nn.Identity()
@@ -1002,7 +845,7 @@ class Qwen3TTSTokenizerV1DecoderBigVGANModel(Qwen3TTSTokenizerV1DecoderPreTraine
             ]
         )
 
-        self.activation_post = TorchActivation1d(
+        self.activation_post = AliasFreeActivation1d(
             activation=SnakeBeta(config.upsample_initial_channel // (2**self.num_upsample_layers))
         )
         self.conv_post = nn.Conv1d(
diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py
index 4f29c0ae3..ff61d6e6c 100644
--- a/vllm_omni/model_executor/models/registry.py
+++ b/vllm_omni/model_executor/models/registry.py
@@ -102,6 +102,22 @@ _OMNI_MODELS = {
         "qwen3_tts_code2wav",
         "Qwen3TTSCode2Wav",
     ),
+    ## higgs-audio v2
+    "HiggsAudioV2ForConditionalGeneration": (
+        "higgs_audio_v2",
+        "higgs_audio_v2_talker",
+        "HiggsAudioV2TalkerForConditionalGeneration",
+    ),
+    "HiggsAudioV2TalkerForConditionalGeneration": (
+        "higgs_audio_v2",
+        "higgs_audio_v2_talker",
+        "HiggsAudioV2TalkerForConditionalGeneration",
+    ),
+    "HiggsAudioV2Code2WavForConditionalGeneration": (
+        "higgs_audio_v2",
+        "higgs_audio_v2_code2wav",
+        "HiggsAudioV2Code2WavForConditionalGeneration",
+    ),
     ## mimo_audio
     "MiMoAudioModel": (
         "mimo_audio",
@@ -129,6 +145,12 @@ _OMNI_MODELS = {
         "glm_image_ar",
         "GlmImageForConditionalGeneration",
     ),
+    ## glm_tts
+    "GLMTTSForConditionalGeneration": (
+        "glm_tts",
+        "glm_tts",
+        "GLMTTSForConditionalGeneration",
+    ),
     "OmniBagelForConditionalGeneration": (
         "bagel",
         "bagel",
diff --git a/vllm_omni/model_executor/stage_configs/dynin_omni.yaml b/vllm_omni/model_executor/stage_configs/dynin_omni.yaml
index 131a0d1cd..024443e8d 100644
--- a/vllm_omni/model_executor/stage_configs/dynin_omni.yaml
+++ b/vllm_omni/model_executor/stage_configs/dynin_omni.yaml
@@ -6,6 +6,7 @@ stage_args:
       max_batch_size: 1
     engine_args:
       model_stage: token2text
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image_full_payload
       model_arch: DyninOmniForConditionalGeneration
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
diff --git a/vllm_omni/model_executor/stage_configs/dynin_omni_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/dynin_omni_multiconnector.yaml
index 4a54f8188..1ab65f0fa 100644
--- a/vllm_omni/model_executor/stage_configs/dynin_omni_multiconnector.yaml
+++ b/vllm_omni/model_executor/stage_configs/dynin_omni_multiconnector.yaml
@@ -6,6 +6,7 @@ stage_args:
       max_batch_size: 1
     engine_args:
       model_stage: token2text
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image_full_payload
       model_arch: DyninOmniForConditionalGeneration
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
diff --git a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py
index cf1ca39ee..4c7245e77 100644
--- a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py
+++ b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py
@@ -7,6 +7,7 @@ from typing import Any
 import numpy as np
 import torch
 from vllm.inputs import TextPrompt
+from vllm.logger import init_logger
 
 from vllm_omni.data_entry_keys import (
     CodesStruct,
@@ -16,6 +17,10 @@ from vllm_omni.data_entry_keys import (
 )
 from vllm_omni.inputs.data import OmniTokensPrompt
 
+logger = init_logger(__name__)
+
+_COSYVOICE3_SPEECH_TOKEN_SIZE = 6561
+
 
 def _build_prompt_embed_struct(prompt_payload: dict[str, Any]) -> EmbeddingsStruct | None:
     """Wrap prompt_payload's flat speech_token/speech_feat/embedding tensors into EmbeddingsStruct."""
@@ -46,6 +51,52 @@ def _ensure_list(x: Any) -> list[Any]:
         return [x]
 
 
+def _to_token_id_list(value: Any) -> list[int]:
+    if value is None:
+        return []
+    if isinstance(value, torch.Tensor):
+        value = value.detach().to("cpu").reshape(-1).tolist()
+    token_ids: list[int] = []
+    for item in _ensure_list(value):
+        if isinstance(item, torch.Tensor):
+            token_ids.extend(_to_token_id_list(item))
+            continue
+        if isinstance(item, (list, tuple)):
+            token_ids.extend(_to_token_id_list(item))
+            continue
+        token_ids.append(int(item))
+    return token_ids
+
+
+def _strip_prompt_prefix(output_ids: list[Any], prefix_ids: list[Any]) -> list[Any]:
+    if prefix_ids and len(output_ids) >= len(prefix_ids) and output_ids[: len(prefix_ids)] == prefix_ids:
+        return output_ids[len(prefix_ids) :]
+    return output_ids
+
+
+def _prompt_speech_token_ids(multi_modal_data: dict[str, Any]) -> list[int]:
+    speech_token = multi_modal_data.get("speech_token")
+    if speech_token is None:
+        embed = multi_modal_data.get("embed")
+        if isinstance(embed, dict):
+            speech_token = embed.get("speech_token")
+    return _to_token_id_list(speech_token)
+
+
+def _has_speech_stop_token(output_ids: list[Any]) -> bool:
+    return any(token_id >= _COSYVOICE3_SPEECH_TOKEN_SIZE for token_id in _to_token_id_list(output_ids))
+
+
+def _set_non_stream_prompt_trim(additional_info: dict[str, Any], prompt_speech_len: int) -> None:
+    if prompt_speech_len <= 0:
+        return
+    meta = additional_info.get("meta")
+    if not isinstance(meta, dict):
+        meta = {}
+        additional_info["meta"] = meta
+    meta["talker_prefill_offset"] = prompt_speech_len
+
+
 def _to_cpu_tensor(x: Any) -> torch.Tensor | None:
     if isinstance(x, list):
         if not x:
@@ -95,9 +146,14 @@ def text2flow(
         if multi_modal_data is None:
             raise RuntimeError(f"Missing multimodal_output for request {source_output.request_id}")
 
-        output_ids = _ensure_list(output.cumulative_token_ids)
         prefix_ids = _ensure_list(source_output.prompt_token_ids)
+        raw_output_ids = _ensure_list(output.cumulative_token_ids)
+        prompt_speech_ids = _prompt_speech_token_ids(multi_modal_data)
+        output_ids = _strip_prompt_prefix(raw_output_ids, prefix_ids)
+        output_ids = _strip_prompt_prefix(output_ids, prompt_speech_ids)
         additional_info = dict(multi_modal_data)
+        if _has_speech_stop_token(raw_output_ids):
+            _set_non_stream_prompt_trim(additional_info, len(prompt_speech_ids))
         additional_info.setdefault("ids", {})["prompt"] = prefix_ids
         engine_inputs.append(OmniTokensPrompt(prompt_token_ids=output_ids, additional_information=additional_info))
     return engine_inputs
@@ -271,3 +327,107 @@ def talker2code2wav_async_chunk(
 
         state["emitted_chunks"] = int(state.get("emitted_chunks", 0)) + 1
         return payload
+
+
+# ============================================================================
+# Worker-connector data plane (non-async-chunk path).
+# cosyvoice3 talker emits `multimodal_outputs={"embed": {"speech_token": t,
+# "speech_feat": t, "embedding": t}}` ONLY at prefill (decode steps emit
+# `{}`).  After flatten_payload these become flat top-level keys
+# `embed.speech_token` etc., persisted across decode steps by the
+# full-payload accumulator (decode doesn't re-emit them).  Shipping via the connector
+# keeps the orchestrator off the heavy-tensor path.
+# ============================================================================
+
+# All three embed tensors are emitted once at prefill and must REPLACE-not-
+# CONCAT across the (already trivial) per-request accumulator history so a
+# regression where decode unexpectedly re-emits them does not silently
+# duplicate the prefill tensor.  See mixin._FULL_PAYLOAD_REPLACE_KEYS.
+_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset({"embed.speech_token", "embed.speech_feat", "embed.embedding"})
+
+
+def text2flow_token_only(
+    source_outputs: list,
+    prompt: OmniTokensPrompt | TextPrompt = None,
+    _requires_multimodal_data: bool = True,
+):
+    """Sync-side builder for the non-async-chunk text→flow path.
+
+    CosyVoice3 sync keeps codec ids on the legacy token path.  Some vLLM v1
+    histories include the source prompt prefix, so strip it only when it is an
+    exact leading match.
+    """
+    del prompt
+    engine_inputs: list[OmniTokensPrompt] = []
+    for source_output in source_outputs:
+        if not source_output.finished:
+            continue
+        output = source_output.outputs[0]
+        prefix_ids = _ensure_list(source_output.prompt_token_ids)
+        raw_output_ids = _ensure_list(output.cumulative_token_ids)
+        output_ids = _strip_prompt_prefix(raw_output_ids, prefix_ids)
+        multi_modal_data = output.multimodal_output
+        if multi_modal_data is None:
+            raise RuntimeError(f"Missing multimodal_output for request {source_output.request_id}")
+        prompt_speech_ids = _prompt_speech_token_ids(multi_modal_data)
+        output_ids = _strip_prompt_prefix(output_ids, prompt_speech_ids)
+        additional_info: dict[str, Any] = dict(multi_modal_data)
+        if _has_speech_stop_token(raw_output_ids):
+            _set_non_stream_prompt_trim(additional_info, len(prompt_speech_ids))
+        additional_info.setdefault("ids", {})["prompt"] = prefix_ids
+        engine_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=output_ids,
+                additional_information=additional_info,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+            )
+        )
+    return engine_inputs
+
+
+def text2flow_full_payload(
+    transfer_manager,
+    pooling_output,
+    request,
+):
+    """Producer-side payload builder.
+
+    Reads prefill-emitted `embed.{speech_token, speech_feat, embedding}` from
+    the accumulator and ships prompt conditioning as a connector payload.
+    The downstream flow stage reads these from `model_intermediate_buffer`
+    (see cosyvoice3.py:671 in the code2wav forward — runtime_info pickup).
+    """
+    del transfer_manager
+    rid = getattr(request, "external_req_id", None) or getattr(request, "request_id", "?")
+    if not isinstance(pooling_output, dict):
+        logger.warning(
+            "cosyvoice3.text2flow_full_payload: pooling_output not a dict "
+            "(type=%s) for req=%s; consumer wait gate may hang.",
+            type(pooling_output).__name__,
+            rid,
+        )
+        return None
+    embed_out: dict[str, Any] = {}
+    for key in ("speech_token", "speech_feat", "embedding"):
+        v = pooling_output.get(f"embed.{key}")
+        if v is None:
+            nested = pooling_output.get("embed")
+            if isinstance(nested, dict):
+                v = nested.get(key)
+        if isinstance(v, torch.Tensor) and v.numel() > 0:
+            embed_out[key] = v
+    if not embed_out:
+        logger.warning(
+            "cosyvoice3.text2flow_full_payload: no embed.{speech_token,speech_feat,embedding} "
+            "found in pooling_output (keys=%s) for req=%s; consumer wait gate may hang.",
+            list(pooling_output.keys()),
+            rid,
+        )
+        return None
+    return {
+        "meta": {
+            "finished": torch.tensor(True, dtype=torch.bool),
+        },
+        "embed": embed_out,
+    }
diff --git a/vllm_omni/model_executor/stage_input_processors/covo_audio.py b/vllm_omni/model_executor/stage_input_processors/covo_audio.py
index a0a964bdd..7b5ed8c26 100644
--- a/vllm_omni/model_executor/stage_input_processors/covo_audio.py
+++ b/vllm_omni/model_executor/stage_input_processors/covo_audio.py
@@ -1,27 +1,44 @@
 # Copyright 2026 Tencent.
 from typing import Any
 
+import torch
+from vllm.logger import init_logger
+
 from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.models.covo_audio.config_covo_audio import COVO_AUDIO_TOKEN_INDEX
 
+logger = init_logger(__name__)
+
+# Per-model REPLACE-keys for the full-payload accumulator (none for covo_audio:
+# the producer side does not emit per-step hidden_states / model_outputs;
+# llm2code2wav_full_payload reads token_ids directly from `request`).
+_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset()
+
+
+def _filter_audio_codes(token_ids: list[int]) -> list[int]:
+    """Filter codec-range token ids and rebase by COVO_AUDIO_TOKEN_INDEX."""
+    audio_codes = [t - COVO_AUDIO_TOKEN_INDEX for t in token_ids if t >= COVO_AUDIO_TOKEN_INDEX]
+    if not audio_codes:
+        audio_codes = [-1]
+    return audio_codes
+
 
 def llm2code2wav(
     source_outputs: list[Any],
     prompt: Any = None,
     requires_multimodal_data: bool = False,
 ) -> list[OmniTokensPrompt]:
+    """Legacy orchestrator-path builder (retained for async_chunk + back-compat).
+
+    The non-async-chunk path now goes through ``llm2code2wav_token_only`` +
+    worker connector + ``llm2code2wav_full_payload``.
+    """
     talker_outputs = source_outputs
     code2wav_inputs = []
 
-    for i, talker_output in enumerate(talker_outputs):
+    for talker_output in talker_outputs:
         output = talker_output.outputs[0]
-        token_ids = output.token_ids
-
-        audio_codes = [t - COVO_AUDIO_TOKEN_INDEX for t in token_ids if t >= COVO_AUDIO_TOKEN_INDEX]
-
-        if not audio_codes:
-            audio_codes = [-1]
-
+        audio_codes = _filter_audio_codes(list(output.token_ids))
         code2wav_inputs.append(
             OmniTokensPrompt(
                 prompt_token_ids=audio_codes,
@@ -31,3 +48,56 @@ def llm2code2wav(
         )
 
     return code2wav_inputs
+
+
+def llm2code2wav_token_only(
+    source_outputs: list[Any],
+    prompt: Any = None,
+    requires_multimodal_data: bool = False,
+) -> list[OmniTokensPrompt]:
+    """Sync-side placeholder for the non-async-chunk Stage-1 input.
+
+    Returns an OmniTokensPrompt sized to the code2wav stage's expected
+    prefill length (one slot per audio code).  The actual codec ids are
+    delivered via the worker connector payload built by
+    ``llm2code2wav_full_payload``.
+    """
+    code2wav_inputs: list[OmniTokensPrompt] = []
+    for output_wrapper in source_outputs:
+        output = output_wrapper.outputs[0]
+        audio_codes = _filter_audio_codes(list(output.token_ids))
+        code2wav_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[0] * len(audio_codes),
+                additional_information=None,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+            )
+        )
+    return code2wav_inputs
+
+
+def llm2code2wav_full_payload(
+    transfer_manager: Any,
+    pooling_output: dict[str, Any],
+    request: Any,
+) -> dict[str, Any] | None:
+    """Producer-side payload builder for the worker connector data plane.
+
+    covo_audio's fused_thinker_talker stage emits codec ids via
+    ``request.output_token_ids`` (token-ids only -- no
+    hidden_states or embed tensors), so the connector payload is
+    just the filtered audio codes plus a finished marker.
+    """
+    output_token_ids = list(getattr(request, "output_token_ids", None) or [])
+    if not output_token_ids:
+        logger.warning(
+            "covo_audio.llm2code2wav_full_payload: empty output_token_ids for req=%s; consumer wait gate may hang.",
+            getattr(request, "request_id", "?"),
+        )
+        return None
+    audio_codes = _filter_audio_codes(output_token_ids)
+    return {
+        "codes": {"audio": audio_codes},
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+    }
diff --git a/vllm_omni/model_executor/stage_input_processors/dynin_omni.py b/vllm_omni/model_executor/stage_input_processors/dynin_omni.py
index 87cecc103..5a5804de3 100644
--- a/vllm_omni/model_executor/stage_input_processors/dynin_omni.py
+++ b/vllm_omni/model_executor/stage_input_processors/dynin_omni.py
@@ -5,9 +5,12 @@ from typing import Any
 
 import torch
 from vllm.inputs import TextPrompt
+from vllm.logger import init_logger
 
 from vllm_omni.inputs.data import OmniTokensPrompt
 
+logger = init_logger(__name__)
+
 
 def _to_prompt_dict(prompt_item: OmniTokensPrompt | TextPrompt | str | None) -> dict[str, Any]:
     if isinstance(prompt_item, dict):
@@ -146,3 +149,128 @@ def token2image_to_token2audio(
     requires_multimodal_data: bool = False,
 ):
     return _bridge_tokens(source_outputs, prompt, requires_multimodal_data)
+
+
+# ============================================================================
+# Worker-connector data plane (non-async-chunk path).
+# ============================================================================
+
+# Per-model REPLACE-keys for the full-payload accumulator.  dynin_omni's
+# producer model emits new chunks per step (token_ids / runtime_info_json),
+# all of which use the default CONCAT/replace semantics — no model_outputs
+# entry needs explicit REPLACE.
+_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset()
+
+
+def _build_full_payload(pooling_output: dict[str, Any] | None, request: Any) -> dict[str, Any] | None:
+    """Producer-side payload builder: assemble dynin_omni connector payload.
+
+    Reads token_ids from ``pooling_output["token_ids"]`` (preferred) or
+    ``request.output_token_ids`` (fallback).  Reads structured non-tensor
+    metadata from ``pooling_output["runtime_info_json"]`` (JSON-in-uint8)
+    if present, falling back to ``pooling_output["runtime_info"]`` dict.
+    Carries forward ``request.additional_information`` so prompt-side
+    metadata (speaker / language / detok_id) survives the IPC boundary.
+    """
+    if not isinstance(pooling_output, dict):
+        pooling_output = {}
+
+    token_ids = _to_token_id_list(pooling_output.get("token_ids"))
+    if not token_ids:
+        token_ids = _to_token_id_list(pooling_output.get("text_tokens"))
+    if not token_ids and request is not None:
+        token_ids = _to_token_id_list(getattr(request, "output_token_ids", None))
+    if not token_ids:
+        logger.warning(
+            "dynin_omni._build_full_payload: no token_ids found in pooling_output "
+            "(keys=%s) or request.output_token_ids for req=%s; consumer wait gate may hang.",
+            list(pooling_output.keys()),
+            getattr(request, "request_id", "?"),
+        )
+        return None
+
+    src_additional_info = getattr(request, "additional_information", {}) if request is not None else {}
+    if not isinstance(src_additional_info, dict):
+        src_additional_info = {}
+
+    runtime_bridge_info = _decode_runtime_bridge_info(pooling_output.get("runtime_info_json"))
+    if not runtime_bridge_info:
+        runtime_bridge_info = pooling_output.get("runtime_info", {}) or {}
+
+    payload = _normalize_additional_info(src_additional_info)
+    payload.update(_normalize_additional_info(runtime_bridge_info))
+    payload["detok_id"] = [_to_int(pooling_output.get("detok_id"), default=_to_int(payload.get("detok_id"), default=0))]
+    # Use nested OmniPayload shape so the scheduling-metadata extractor in
+    # OmniConnectorModelRunnerMixin reads codes.audio and meta.finished
+    # (flat keys at the top level are silently dropped with a warning).
+    payload["codes"] = {"audio": token_ids}
+    payload["meta"] = {"finished": torch.tensor(True, dtype=torch.bool)}
+    return payload
+
+
+def token2text_to_token2image_full_payload(
+    transfer_manager: Any,
+    pooling_output: dict[str, Any],
+    request: Any,
+) -> dict[str, Any] | None:
+    """Producer-side payload builder for the Stage-0 → Stage-1 (text → image) transition."""
+    del transfer_manager
+    return _build_full_payload(pooling_output, request)
+
+
+def token2image_to_token2audio_full_payload(
+    transfer_manager: Any,
+    pooling_output: dict[str, Any],
+    request: Any,
+) -> dict[str, Any] | None:
+    """Producer-side payload builder for the Stage-1 → Stage-2 (image → audio) transition."""
+    del transfer_manager
+    return _build_full_payload(pooling_output, request)
+
+
+def _token_only_from_source(source_outputs: list[Any]) -> list[OmniTokensPrompt]:
+    """Length-only placeholder list mirroring ``_bridge_tokens`` token counts."""
+    inputs: list[OmniTokensPrompt] = []
+    for source_output in source_outputs:
+        output = source_output.outputs[0]
+        mm_out = getattr(output, "multimodal_output", None) or {}
+        token_ids = _to_token_id_list(mm_out.get("token_ids"))
+        if not token_ids:
+            token_ids = _to_token_id_list(mm_out.get("text_tokens"))
+        if not token_ids:
+            token_ids = list(getattr(output, "token_ids", []) or [])
+        if not token_ids:
+            token_ids = [0]
+        inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[0] * len(token_ids),
+                additional_information=None,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+            )
+        )
+    return inputs
+
+
+def token2text_to_token2image_token_only(
+    stage_list,
+    engine_input_source,
+    prompt: OmniTokensPrompt | TextPrompt = None,
+    requires_multimodal_data: bool = False,
+) -> list[OmniTokensPrompt]:
+    """Sync-side placeholder for Stage-1 input (token2image)."""
+    source_stage_id = engine_input_source[0] if engine_input_source else 0
+    source_outputs = stage_list[source_stage_id].engine_outputs
+    return _token_only_from_source(source_outputs)
+
+
+def token2image_to_token2audio_token_only(
+    stage_list,
+    engine_input_source,
+    prompt: OmniTokensPrompt | TextPrompt = None,
+    requires_multimodal_data: bool = False,
+) -> list[OmniTokensPrompt]:
+    """Sync-side placeholder for Stage-2 input (token2audio)."""
+    source_stage_id = engine_input_source[0] if engine_input_source else 0
+    source_outputs = stage_list[source_stage_id].engine_outputs
+    return _token_only_from_source(source_outputs)
diff --git a/vllm_omni/model_executor/stage_input_processors/fish_speech.py b/vllm_omni/model_executor/stage_input_processors/fish_speech.py
index 6eac326a3..ac12e6035 100644
--- a/vllm_omni/model_executor/stage_input_processors/fish_speech.py
+++ b/vllm_omni/model_executor/stage_input_processors/fish_speech.py
@@ -3,7 +3,6 @@
 from typing import Any
 
 import torch
-from vllm.logger import init_logger
 
 from vllm_omni.data_entry_keys import (
     CodesStruct,
@@ -11,7 +10,85 @@ from vllm_omni.data_entry_keys import (
     OmniPayloadStruct,
 )
 
-logger = init_logger(__name__)
+
+def _get_connector_extra(transfer_manager: Any) -> dict[str, Any]:
+    connector = getattr(transfer_manager, "connector", None)
+    raw_cfg = getattr(connector, "config", {}) or {}
+    return raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
+
+
+def _use_tensor_code_payload(transfer_manager: Any) -> bool:
+    cfg = _get_connector_extra(transfer_manager)
+    return bool(cfg.get("fish_speech_tensor_codes", False))
+
+
+def _cfg_bool(cfg: dict[str, Any], key: str, default: bool = False) -> bool:
+    return bool(cfg.get(key, default))
+
+
+def _cfg_int(cfg: dict[str, Any], key: str, default: int = 0) -> int:
+    value = cfg.get(key, default)
+    try:
+        return int(value)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"Invalid Fish Speech integer config {key}={value!r}") from exc
+
+
+def _cfg_float(cfg: dict[str, Any], key: str, default: float) -> float:
+    value = cfg.get(key, default)
+    try:
+        return float(value)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"Invalid Fish Speech float config {key}={value!r}") from exc
+
+
+def _active_codec_requests(transfer_manager: Any) -> int:
+    return sum(1 for value in transfer_manager.code_prompt_token_ids.values() if len(value) > 0)
+
+
+def _ready_codec_requests(transfer_manager: Any, min_frames: int) -> int:
+    return sum(1 for value in transfer_manager.code_prompt_token_ids.values() if len(value) >= min_frames)
+
+
+def _select_backlog_chunk_size(
+    transfer_manager: Any,
+    cfg: dict[str, Any],
+    base_chunk_size: int,
+    current_length: int,
+) -> int:
+    backlog_chunk_size = _cfg_int(
+        cfg,
+        "fish_speech_backlog_codec_chunk_frames",
+        0,
+    )
+    if backlog_chunk_size <= base_chunk_size:
+        return base_chunk_size
+
+    capacity = max(1, int(getattr(transfer_manager, "scheduler_max_num_seqs", 1)))
+    active = _active_codec_requests(transfer_manager)
+    load_factor = min(active / capacity, 1.0)
+    threshold = _cfg_float(
+        cfg,
+        "fish_speech_backlog_load_threshold",
+        0.75,
+    )
+    if load_factor < threshold:
+        return base_chunk_size
+
+    ready_min = _cfg_int(
+        cfg,
+        "fish_speech_backlog_min_ready",
+        max(2, min(capacity, int(capacity * threshold))),
+    )
+    # Include the current request even when the caller has not yet committed
+    # the latest frame into the transfer-manager map.
+    ready = _ready_codec_requests(transfer_manager, base_chunk_size)
+    if current_length >= base_chunk_size:
+        ready = max(ready, 1)
+    if ready < ready_min:
+        return base_chunk_size
+
+    return backlog_chunk_size
 
 
 def _extract_last_frame(pooling_output: dict[str, Any]) -> torch.Tensor | None:
@@ -21,11 +98,18 @@ def _extract_last_frame(pooling_output: dict[str, Any]) -> torch.Tensor | None:
         return None
     if audio_codes.ndim == 2:
         frame = audio_codes[-1]
-        if frame.numel() == 0 or not bool(frame.any().item()):
+        valid = pooling_output.get("audio_code_valid")
+        if isinstance(valid, torch.Tensor) and valid.numel() > 0:
+            is_valid = bool(valid.reshape(-1)[-1].item())
+        elif valid is not None:
+            is_valid = bool(valid)
+        else:
+            is_valid = bool(frame.any().item())
+        if frame.numel() == 0 or not is_valid:
             return None
-        return frame.to(torch.long).reshape(-1)
+        return frame.to(device="cpu", dtype=torch.long).reshape(-1)
     if audio_codes.ndim == 1:
-        return audio_codes.to(torch.long).reshape(-1)
+        return audio_codes.to(device="cpu", dtype=torch.long).reshape(-1)
     raise ValueError(f"Invalid audio_codes shape for Fish Speech async_chunk: {tuple(audio_codes.shape)}")
 
 
@@ -74,20 +158,20 @@ def slow_ar_to_dac_decoder_async_chunk(
     """
     request_id = request.external_req_id
     finished = bool(is_finished or request.is_finished())
+    cfg = _get_connector_extra(transfer_manager)
 
     if isinstance(pooling_output, dict):
         frame = _extract_last_frame(pooling_output)
         if frame is not None:
-            transfer_manager.code_prompt_token_ids[request_id].append(frame.detach().to(device="cpu", dtype=torch.long))
+            transfer_manager.code_prompt_token_ids[request_id].append(frame.detach())
     elif not finished:
         return None
 
-    connector = getattr(transfer_manager, "connector", None)
-    raw_cfg = getattr(connector, "config", {}) or {}
-    cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
     chunk_size = int(cfg.get("codec_chunk_frames", 25))
     left_context_size_config = int(cfg.get("codec_left_context_frames", 25))
-    initial_chunk_size = int(cfg.get("initial_codec_chunk_frames", 0))
+    configured_initial_chunk_size = int(cfg.get("initial_codec_chunk_frames", 0))
+
+    initial_chunk_size = configured_initial_chunk_size
 
     # Per-request override.
     additional_information = getattr(request, "additional_information", None)
@@ -100,7 +184,7 @@ def slow_ar_to_dac_decoder_async_chunk(
         if entry.list_data is not None and len(entry.list_data) == 1:
             initial_chunk_size = int(entry.list_data[0])
 
-    if chunk_size <= 0 or left_context_size_config < 0 or initial_chunk_size < 0:
+    if chunk_size <= 0 or left_context_size_config < 0 or configured_initial_chunk_size < 0 or initial_chunk_size < 0:
         raise ValueError(
             f"Invalid codec chunk config: codec_chunk_frames={chunk_size}, "
             f"codec_left_context_frames={left_context_size_config}, "
@@ -110,6 +194,7 @@ def slow_ar_to_dac_decoder_async_chunk(
         initial_chunk_size = chunk_size
 
     length = len(transfer_manager.code_prompt_token_ids[request_id])
+    steady_chunk_size = _select_backlog_chunk_size(transfer_manager, cfg, chunk_size, length)
 
     if length <= 0:
         if finished:
@@ -119,9 +204,30 @@ def slow_ar_to_dac_decoder_async_chunk(
             )
         return None
 
-    in_initial_phase = initial_chunk_size > 0 and length <= chunk_size
+    single_initial_chunk = _cfg_bool(
+        cfg,
+        "fish_speech_single_initial_chunk",
+        False,
+    )
+    use_first_chunk = initial_chunk_size > 0 and initial_chunk_size < steady_chunk_size
 
-    if in_initial_phase:
+    if single_initial_chunk and use_first_chunk:
+        if length <= initial_chunk_size:
+            if not finished and length < initial_chunk_size:
+                return None
+            context_length = length if finished and length < initial_chunk_size else initial_chunk_size
+        else:
+            adjusted = length - initial_chunk_size
+            if adjusted <= 0:
+                return None
+            if not finished and adjusted % steady_chunk_size != 0:
+                return None
+            chunk_length = adjusted % steady_chunk_size
+            context_length = chunk_length if chunk_length != 0 else steady_chunk_size
+        end_index = min(length, left_context_size_config + context_length)
+        left_context_size = max(0, int(end_index - context_length))
+        window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
+    elif initial_chunk_size > 0 and length <= chunk_size:
         already_sent = transfer_manager.put_req_chunk[request_id] * initial_chunk_size
         pending = length - already_sent
         if pending <= 0:
@@ -134,17 +240,23 @@ def slow_ar_to_dac_decoder_async_chunk(
     else:
         initial_coverage = (chunk_size // initial_chunk_size) * initial_chunk_size if initial_chunk_size > 0 else 0
         adjusted = length - initial_coverage
-        chunk_length = adjusted % chunk_size
+        chunk_length = adjusted % steady_chunk_size
         if chunk_length != 0 and not finished:
             return None
-        context_length = chunk_length if chunk_length != 0 else chunk_size
+        context_length = chunk_length if chunk_length != 0 else steady_chunk_size
         end_index = min(length, left_context_size_config + context_length)
         left_context_size = max(0, int(end_index - context_length))
         window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
 
-    # Pack into codebook-major flat codes.
+    # Pack into codebook-major codes. The tensor path avoids expanding codec
+    # indices into Python ints across the connector boundary; Stage1 schedules
+    # a single placeholder token and consumes the real codes from runtime info.
     stacked_frames = torch.stack(window_frames, dim=0)
-    code_predictor_codes = stacked_frames.transpose(0, 1).reshape(-1)
+    codes_qf = stacked_frames.transpose(0, 1).contiguous()
+    if _use_tensor_code_payload(transfer_manager):
+        code_predictor_codes = codes_qf
+    else:
+        code_predictor_codes = codes_qf.reshape(-1)
 
     return OmniPayloadStruct(
         codes=CodesStruct(audio=code_predictor_codes),
diff --git a/vllm_omni/model_executor/stage_input_processors/glm_tts.py b/vllm_omni/model_executor/stage_input_processors/glm_tts.py
new file mode 100644
index 000000000..caafbf86c
--- /dev/null
+++ b/vllm_omni/model_executor/stage_input_processors/glm_tts.py
@@ -0,0 +1,362 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Stage input processor for GLM-TTS: AR → DiT Pipeline.
+
+Supports both sync (non-streaming) and async_chunk (streaming) modes.
+Adapted for LLM_GENERATION execution type on stage 1 (DiT).
+"""
+
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.data_entry_keys import (
+    CodesStruct,
+    EmbeddingsStruct,
+    MetaStruct,
+    OmniPayloadStruct,
+)
+from vllm_omni.engine.serialization import deserialize_additional_information
+
+logger = init_logger(__name__)
+
+_GLM_TTS_PRIVATE_KEY = "glm_tts"
+
+
+def _copy_voice_clone_payload(
+    src: dict[str, Any],
+    dst: dict[str, Any],
+    *,
+    to_cpu: bool = False,
+    skip_existing: bool = False,
+) -> None:
+    if not skip_existing or "prompt_speech_token" not in dst:
+        prompt = src.get("prompt_speech_token")
+        if prompt is None:
+            prompt = src.get("prompt_token")
+        if prompt is not None:
+            if to_cpu:
+                prompt = _to_cpu_tensor(prompt)
+            if prompt is not None:
+                dst["prompt_speech_token"] = prompt
+
+    for key in ("prompt_feat", "embedding"):
+        if skip_existing and key in dst:
+            continue
+        val = src.get(key)
+        if val is not None:
+            if to_cpu:
+                val = _to_cpu_tensor(val)
+            if val is not None:
+                dst[key] = val
+
+
+def _build_voice_clone_embed_struct(payload: dict[str, Any]) -> EmbeddingsStruct | None:
+    prompt_token = payload.get("prompt_speech_token")
+    if prompt_token is None:
+        prompt_token = payload.get("prompt_token")
+    prompt_feat = payload.get("prompt_feat")
+    embedding = payload.get("embedding")
+    if prompt_token is None and prompt_feat is None and embedding is None:
+        return None
+    return EmbeddingsStruct(
+        speech_token=prompt_token,
+        speech_feat=prompt_feat,
+        embedding=embedding,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Sync processor: collect all speech tokens from AR, pass to DiT
+# ---------------------------------------------------------------------------
+
+
+def ar_to_dit(
+    source_outputs: list[Any],
+    prompt: Any = None,
+    _requires_multimodal_data: bool = False,
+) -> list[Any]:
+    """Non-streaming: collect speech tokens from AR, pass to DiT.
+
+    Propagates voice cloning data (prompt_token, prompt_feat, embedding)
+    from the AR model's multimodal_output to the DiT stage using the same
+    nested ``embed`` dict format as the async_chunk path.  This ensures
+    ``nested_get(info, "embed", ...)`` in the DiT wrapper reads all three
+    fields without relying on backward-compat flat-key fallbacks.
+    """
+    from vllm_omni.inputs.data import OmniTokensPrompt
+
+    dit_inputs: list[OmniTokensPrompt] = []
+
+    for output in source_outputs:
+        mm = output.outputs[0].multimodal_output
+        speech_tokens = mm.get("speech_tokens")
+
+        # Extract valid token list (filter -1 placeholder markers from prefill)
+        if isinstance(speech_tokens, torch.Tensor):
+            valid = speech_tokens.to(torch.long).reshape(-1)
+            token_list = valid[valid >= 0].cpu().tolist()
+        elif speech_tokens is not None:
+            token_list = [t for t in speech_tokens if t >= 0]
+        else:
+            token_list = []
+
+        if not token_list:
+            logger.warning("No valid speech tokens for request %s", output.request_id)
+
+        additional_info: dict[str, Any] = {"speech_tokens": token_list}
+
+        # Build nested embed dict (unified with async_chunk path).
+        # _build_voice_clone_embed_struct reads prompt_speech_token/prompt_feat/
+        # embedding from *mm* and returns an EmbeddingsStruct; we convert to a
+        # plain dict so flatten_payload/unflatten_payload round-trips correctly
+        # (key "embed" is in _NESTED_KEYS).
+        embed = _build_voice_clone_embed_struct(mm)
+        if embed is not None:
+            additional_info["embed"] = {
+                "speech_token": embed.speech_token,
+                "speech_feat": embed.speech_feat,
+                "embedding": embed.embedding,
+            }
+
+        # Populate meta dict so DiT wrapper uses block-causal attention.
+        # Without meta, forward() sets uses_streaming=False → bidirectional
+        # attention → garbage audio. The full token sequence is treated as
+        # a single "chunk" with stream_finished=True.
+        n_tokens = len(token_list) if token_list else 1
+        additional_info["meta"] = {
+            "req_id": [output.request_id],
+            "left_context_size": 0,
+            "stream_finished": True,
+        }
+        additional_info["kv_metadata"] = {
+            _GLM_TTS_PRIVATE_KEY: {
+                "chunk_sizes_history": [n_tokens],
+                "block_pattern": [n_tokens],
+            }
+        }
+
+        dit_inputs.append(
+            OmniTokensPrompt(
+                # LLM_GENERATION scheduler needs at least 1 token
+                prompt_token_ids=token_list or [0],
+                additional_information=additional_info,
+            )
+        )
+
+    return dit_inputs
+
+
+# ---------------------------------------------------------------------------
+# Helper: extract last speech token from AR model pooling output
+# ---------------------------------------------------------------------------
+
+
+def _extract_last_speech_token(pooling_output: dict[str, Any]) -> int | None:
+    """Extract the last valid speech token from AR model output.
+
+    GLM-TTS AR produces one speech token per decode step.
+    Returns the token ID (relative to ATS, i.e. 0-based), or None.
+    """
+    speech_tokens = pooling_output.get("speech_tokens")
+    if not isinstance(speech_tokens, torch.Tensor) or speech_tokens.numel() == 0:
+        return None
+    token_val = int(speech_tokens.reshape(-1).to(torch.long)[-1].item())
+    # -1 = invalid/EOA marker
+    if token_val < 0:
+        return None
+    return token_val
+
+
+def _to_cpu_tensor(value: Any) -> torch.Tensor | None:
+    """Convert value to CPU tensor if possible."""
+    if isinstance(value, torch.Tensor):
+        return value.detach().cpu()
+    if isinstance(value, list):
+        if not value:
+            return None
+        if isinstance(value[0], torch.Tensor):
+            return value[0].detach().cpu()
+    return None
+
+
+# ---------------------------------------------------------------------------
+# Async streaming processor: emit speech token chunks as AR produces them
+# ---------------------------------------------------------------------------
+
+
+def ar_to_dit_async_chunk(
+    transfer_manager: Any,
+    pooling_output: dict[str, Any] | None,
+    request: Any,
+    is_finished: bool = False,
+) -> OmniPayloadStruct | None:
+    """Async streaming processor: emit speech token chunks as AR produces them.
+
+    Accumulates per-step speech tokens and emits fixed-size cumulative prefixes
+    for GLM-TTS flow-cache streaming.
+
+    Follows the CosyVoice3 talker2code2wav_async_chunk transfer pattern:
+    - Per-request state tracking (seen tokens, prompt sent flag)
+    - First chunk carries voice clone conditioning payload
+    - Each chunk includes left_context_size, stream_finished, req_id
+    - code_predictor_codes for chunk_transfer_adapter consumption
+
+    Unlike codec left-context decoders, official GLM-TTS sends the cumulative
+    AR token prefix to the flow stage on every chunk and reuses diffusion
+    latents internally.  Here left_context_size means the stable token prefix
+    that has already been emitted, so DiT can crop regenerated audio after
+    sampling.
+
+    GLM-TTS produces single-token-per-step (no multi-codebook), so each entry
+    in code_prompt_token_ids is a plain int, not a list of codebook values.
+    """
+    request_id = getattr(request, "external_req_id", None) or getattr(request, "request_id", None)
+    if request_id is None:
+        raise ValueError("GLM-TTS async chunk request is missing request id")
+    finished = bool(is_finished or request.is_finished())
+
+    # Read connector chunk config (supports progressive list or single int)
+    connector = getattr(transfer_manager, "connector", None)
+    raw_cfg = getattr(connector, "config", {}) or {}
+    cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
+    chunk_frames_cfg = cfg.get("codec_chunk_frames", 25)
+    if isinstance(chunk_frames_cfg, list):
+        progressive_chunk_sizes = [int(c) for c in chunk_frames_cfg]
+    else:
+        progressive_chunk_sizes = [int(chunk_frames_cfg)]
+    left_context_size_config = int(cfg.get("codec_left_context_frames", 25))
+    crossfade_sec = float(cfg.get("crossfade_sec", 0.1))
+
+    if not progressive_chunk_sizes or any(c <= 0 for c in progressive_chunk_sizes) or left_context_size_config < 0:
+        raise ValueError(
+            f"Invalid codec chunk config: codec_chunk_frames={chunk_frames_cfg}, "
+            f"codec_left_context_frames={left_context_size_config}"
+        )
+
+    # Initialize per-request state (like CosyVoice3)
+    request_payload = getattr(transfer_manager, "request_payload", None)
+    if request_payload is None:
+        request_payload = {}
+        transfer_manager.request_payload = request_payload
+    code_prompt_token_ids = getattr(transfer_manager, "code_prompt_token_ids", None)
+    if code_prompt_token_ids is None:
+        code_prompt_token_ids = {}
+        transfer_manager.code_prompt_token_ids = code_prompt_token_ids
+    code_prompt_token_ids.setdefault(request_id, [])
+    request_state = request_payload.get(request_id)
+    if not isinstance(request_state, dict) or "_glm_tts_async_state" not in request_state:
+        # Extract voice clone conditioning from request additional_information
+        info = deserialize_additional_information(getattr(request, "additional_information", None))
+        prompt_payload: dict[str, Any] = {}
+        _copy_voice_clone_payload(info, prompt_payload, to_cpu=True)
+
+        # Also try to extract from pooling_output (first call)
+        if isinstance(pooling_output, dict):
+            _copy_voice_clone_payload(pooling_output, prompt_payload, to_cpu=True, skip_existing=True)
+
+        request_state = {
+            "_glm_tts_async_state": {
+                "seen_len": 0,
+                "sent_prompt": False,
+                "emitted_chunks": 0,
+                "emitted_token_len": 0,
+                "terminal_sent": False,
+                "prompt_payload": prompt_payload,
+                "chunk_sizes_history": [],
+                "block_pattern": progressive_chunk_sizes,
+            }
+        }
+        request_payload[request_id] = request_state
+
+    state = request_state["_glm_tts_async_state"]
+    if state.get("terminal_sent", False):
+        return None
+
+    # Helper: get voice clone embed (once per request)
+    def _pop_embed() -> EmbeddingsStruct | None:
+        if state.get("sent_prompt", False):
+            return None
+        state["sent_prompt"] = True
+        return _build_voice_clone_embed_struct(state.get("prompt_payload", {}))
+
+    # Helper: build terminal or chunk payload
+    def _make_payload(
+        codes: list[int],
+        *,
+        tok_offset: int,
+        is_terminal: bool,
+    ) -> OmniPayloadStruct:
+        return OmniPayloadStruct(
+            codes=CodesStruct(
+                audio=(torch.tensor(codes, dtype=torch.long) if codes else torch.empty(0, dtype=torch.long))
+            ),
+            meta=MetaStruct(
+                finished=torch.tensor(is_terminal, dtype=torch.bool),
+                stream_finished=torch.tensor(is_terminal, dtype=torch.bool),
+                req_id=[request_id],
+                left_context_size=tok_offset,
+            ),
+            kv_metadata={
+                _GLM_TTS_PRIVATE_KEY: {
+                    "chunk_sizes_history": list(state.get("chunk_sizes_history", [])),
+                    "block_pattern": list(state.get("block_pattern", progressive_chunk_sizes)),
+                    "crossfade_sec": crossfade_sec,
+                }
+            },
+            embed=_pop_embed(),
+        )
+
+    # Accumulate new speech token from this step
+    if isinstance(pooling_output, dict):
+        token = _extract_last_speech_token(pooling_output)
+        if token is not None:
+            code_prompt_token_ids[request_id].append(token)
+    elif not finished:
+        return None
+
+    token_frames = code_prompt_token_ids[request_id]
+    length = len(token_frames)
+    emitted_token_len = int(state.get("emitted_token_len", 0))
+
+    # Terminal: no tokens, or finished on/before chunk boundary
+    if finished and length <= emitted_token_len:
+        state["terminal_sent"] = True
+        # Boundary finish is metadata-only: the cumulative prefix has already
+        # been emitted, so resending it would make the DiT/vocoder redo work
+        # and can duplicate audio.
+        return _make_payload([], tok_offset=emitted_token_len, is_terminal=True)
+
+    if length <= 0:
+        return None
+
+    # Progressive chunk size: 25 → 50 → 200 (official GLM-TTS pattern)
+    chunk_count = int(state.get("emitted_chunks", 0))
+    idx = min(chunk_count, len(progressive_chunk_sizes) - 1)
+    current_chunk_size = progressive_chunk_sizes[idx]
+
+    # Check if enough tokens accumulated for next chunk
+    if not finished and (length - emitted_token_len) < current_chunk_size:
+        return None
+
+    # Determine end index for cumulative prefix
+    if emitted_token_len == 0:
+        end_index = min(length, current_chunk_size)
+        token_offset = 0
+    else:
+        end_index = length if finished else min(length, emitted_token_len + current_chunk_size)
+        token_offset = emitted_token_len
+
+    # Track chunk sizes for cache slicing
+    chunk_sizes_history: list[int] = list(state.get("chunk_sizes_history", []))
+    chunk_sizes_history.append(end_index - emitted_token_len)
+    state["chunk_sizes_history"] = chunk_sizes_history
+    state["emitted_chunks"] = chunk_count + 1
+
+    if finished:
+        state["terminal_sent"] = True
+    else:
+        state["emitted_token_len"] = max(emitted_token_len, end_index)
+
+    return _make_payload(list(token_frames[:end_index]), tok_offset=token_offset, is_terminal=finished)
diff --git a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py
new file mode 100644
index 000000000..0d6781586
--- /dev/null
+++ b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py
@@ -0,0 +1,252 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Stage-input processor for higgs-audio v2: Talker -> Code2Wav.
+
+Mirrors the qwen3_tts contract:
+
+- ``talker2code2wav`` is the sync (non-async-chunk) adapter that runs once,
+  after Stage 0 has finished, collecting every audio frame into a single
+  ``OmniTokensPrompt`` for Stage 1.
+- ``talker2code2wav_async_chunk`` is the streaming adapter invoked once per
+  emitted codec frame; it accumulates frames in a per-request buffer and
+  flushes an ``OmniPayloadStruct`` once a chunk is ready (or at end-of-stream).
+
+Stage 1 strictly rejects any code id >= ``audio_stream_bos_id`` (1024). This
+processor filters frames that still carry stream specials so the codec sees
+only real codes in [0, 1023], regardless of where Stage 0 emitted them.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.data_entry_keys import (
+    CodesStruct,
+    MetaStruct,
+    OmniPayload,
+    OmniPayloadStruct,
+)
+
+__all__ = ["talker2code2wav", "talker2code2wav_async_chunk"]
+
+logger = init_logger(__name__)
+
+# Default codec chunk-shape knobs for async_chunk mode.  These match the
+# canonical Qwen3-TTS defaults (25 frames = 1 s at 25 fps).  Real values
+# come from the connector config; these are only fallbacks.
+_DEFAULT_CODEC_CHUNK_FRAMES = 25
+_DEFAULT_CODEC_LEFT_CONTEXT_FRAMES = 25
+
+# Boson-ai constants. Kept local (rather than imported from the model package)
+# so this module remains usable without circular imports during
+# stage-input-processor discovery.
+_NUM_CODEBOOKS = 8
+_AUDIO_STREAM_BOS_ID = 1024
+_AUDIO_STREAM_EOS_ID = 1025
+_NUM_REAL_CODES = _AUDIO_STREAM_BOS_ID  # codes in [0, 1023] are real
+
+
+def _filter_real_code_frames(audio_codes: torch.Tensor) -> torch.Tensor:
+    """Keep only frames whose codes are entirely in [0, 1023].
+
+    ``audio_codes`` has shape ``[num_frames, num_codebooks]`` after the
+    standard Stage-0 transpose. We drop frames containing stream specials
+    or anything padded outside the real-code range.
+    """
+    if audio_codes.numel() == 0:
+        return audio_codes
+    if audio_codes.ndim != 2:
+        raise ValueError(f"expected [num_frames, num_codebooks] audio_codes; got shape {tuple(audio_codes.shape)}")
+    valid = (audio_codes >= 0).all(dim=1) & (audio_codes < _NUM_REAL_CODES).all(dim=1)
+    return audio_codes[valid]
+
+
+def _revert_delay_pattern(audio_codes_qt: torch.Tensor) -> torch.Tensor:
+    """Reverse the delay-pattern shift, mirroring upstream's
+    ``boson_multimodal.model.higgs_audio.utils.revert_delay_pattern``.
+
+    Input shape: ``[num_codebooks, seq_len + num_codebooks - 1]``.
+    Output shape: ``[num_codebooks, seq_len]``.
+
+    For each codebook ``i``, slice rows ``[i:i+1]`` from columns
+    ``[i : seq_len + i]`` to remove ``i`` leading BOS pads and ``Q-1-i``
+    trailing PAD entries.
+    """
+    if audio_codes_qt.ndim != 2:
+        raise ValueError(f"_revert_delay_pattern expects [Q, T] input; got {tuple(audio_codes_qt.shape)}")
+    q, t = audio_codes_qt.shape
+    if t < q:
+        # Not enough frames to revert delay pattern; return as-is.
+        return audio_codes_qt
+    seq_len = t - q + 1
+    out_l = []
+    for i in range(q):
+        out_l.append(audio_codes_qt[i : i + 1, i : seq_len + i])
+    return torch.cat(out_l, dim=0)
+
+
+def talker2code2wav(
+    source_outputs: list[Any],
+    prompt: Any = None,
+    _requires_multimodal_data: bool = False,
+) -> list[Any]:
+    """Sync: collect all talker codes, then pass to code2wav at once."""
+    from vllm_omni.inputs.data import OmniTokensPrompt
+
+    code2wav_inputs: list[OmniTokensPrompt] = []
+    for talker_output in source_outputs:
+        if not talker_output.finished:
+            continue
+        output = talker_output.outputs[0]
+        mm = output.multimodal_output
+        mm_codes = mm.get("codes", {})
+
+        audio_codes = mm_codes.get("audio")
+        if audio_codes is None or not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
+            # Nothing to decode for this request; emit an empty payload so the
+            # downstream stage can still close the response.
+            code2wav_inputs.append(
+                OmniTokensPrompt(
+                    prompt_token_ids=[],
+                    multi_modal_data=None,
+                    mm_processor_kwargs=None,
+                    additional_information=None,
+                )
+            )
+            continue
+
+        audio_codes = audio_codes.to(torch.long)
+        if audio_codes.ndim == 1:
+            # Stage 0 emitted a flat sequence; reshape to [num_frames, num_codebooks].
+            if audio_codes.numel() % _NUM_CODEBOOKS != 0:
+                raise ValueError(
+                    f"flat audio_codes length {audio_codes.numel()} not divisible by num_codebooks={_NUM_CODEBOOKS}"
+                )
+            audio_codes = audio_codes.reshape(-1, _NUM_CODEBOOKS)
+
+        if audio_codes.ndim != 2:
+            raise ValueError(f"audio_codes must be 1D or 2D; got shape {tuple(audio_codes.shape)}")
+
+        # Audio_codes from Stage-0 are emitted in DELAY-PATTERN layout.
+        # Revert the shift to recover the canonical ``[Q, seq_len]`` form,
+        # then clip to the real-code range and trim the first/last frame
+        # (the codec doesn't consume the BOS/EOS edge frames). This mirrors
+        # the upstream ``HiggsAudioServeEngine.generate`` decode line:
+        #     vq_code = revert_delay_pattern(out).clip(0, real-1)[:, 1:-1]
+        codes_qt = audio_codes.transpose(0, 1).contiguous().cpu()
+        codes_qt = _revert_delay_pattern(codes_qt)
+        codes_qt = codes_qt.clamp_(min=0, max=_NUM_REAL_CODES - 1)
+        if codes_qt.shape[-1] >= 3:
+            codes_qt = codes_qt[:, 1:-1]
+        # Code2Wav expects codebook-major flat: [Q * num_frames].
+        codec_codes = codes_qt.reshape(-1).tolist()
+
+        code2wav_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=codec_codes,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+                additional_information=None,
+            )
+        )
+    return code2wav_inputs
+
+
+def _extract_last_frame(pooling_output: OmniPayload) -> torch.Tensor | None:
+    audio_codes = pooling_output.get("codes", {}).get("audio")
+    if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
+        return None
+    if audio_codes.ndim == 2:
+        frame = audio_codes[-1]
+    elif audio_codes.ndim == 1:
+        frame = audio_codes
+    else:
+        raise ValueError(f"unexpected audio_codes shape for higgs_audio_v2 async_chunk: {tuple(audio_codes.shape)}")
+    if frame.numel() == 0:
+        return None
+    frame = frame.to(torch.long).reshape(-1)
+    # Filter frames that still carry stream specials.
+    if int(frame.min().item()) < 0 or int(frame.max().item()) >= _NUM_REAL_CODES:
+        return None
+    return frame
+
+
+def talker2code2wav_async_chunk(
+    transfer_manager: Any,
+    pooling_output: OmniPayload | None,
+    request: Any,
+    is_finished: bool = False,
+) -> OmniPayloadStruct | None:
+    """Streaming adapter: buffer frames per request and flush a chunk at a time.
+
+    Emits codes in codebook-major flat layout (``[Q * num_frames]``) plus a
+    ``MetaStruct(left_context_size, finished)`` so Stage 1 can stitch chunks
+    together with the conventional left-context overlap.
+    """
+    request_id = request.external_req_id
+    finished = bool(is_finished or request.is_finished())
+
+    if isinstance(pooling_output, dict):
+        frame = _extract_last_frame(pooling_output)
+        if frame is not None:
+            transfer_manager.code_prompt_token_ids[request_id].append(frame.cpu().tolist())
+    elif not finished:
+        return None
+
+    connector = getattr(transfer_manager, "connector", None)
+    raw_cfg = getattr(connector, "config", {}) or {}
+    cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
+    chunk_size = int(cfg.get("codec_chunk_frames", _DEFAULT_CODEC_CHUNK_FRAMES))
+    left_context_size_config = int(cfg.get("codec_left_context_frames", _DEFAULT_CODEC_LEFT_CONTEXT_FRAMES))
+
+    if chunk_size <= 0 or left_context_size_config < 0:
+        raise ValueError(
+            f"Invalid codec chunk config: codec_chunk_frames={chunk_size}, "
+            f"codec_left_context_frames={left_context_size_config}"
+        )
+
+    length = len(transfer_manager.code_prompt_token_ids[request_id])
+    if length <= 0:
+        if finished:
+            return OmniPayloadStruct(
+                codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
+                meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)),
+            )
+        return None
+
+    chunk_remainder = length % chunk_size
+    if not finished and chunk_remainder != 0:
+        # Wait until we have a clean chunk boundary.
+        return None
+
+    context_length = chunk_remainder if (finished and chunk_remainder > 0) else chunk_size
+    end_index = min(length, left_context_size_config + context_length)
+    left_context_size = max(0, end_index - context_length)
+    window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
+
+    if not window_frames:
+        # Defensive: should not happen because length > 0, but keeps the contract.
+        return OmniPayloadStruct(
+            codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
+            meta=MetaStruct(finished=torch.tensor(finished, dtype=torch.bool)),
+        )
+
+    num_codebooks = len(window_frames[0])
+    if num_codebooks != _NUM_CODEBOOKS:
+        raise ValueError(f"expected {_NUM_CODEBOOKS} codebooks per frame; got {num_codebooks}")
+    num_frames = len(window_frames)
+    flat = torch.tensor(
+        [window_frames[f][q] for q in range(num_codebooks) for f in range(num_frames)],
+        dtype=torch.long,
+    )
+
+    return OmniPayloadStruct(
+        codes=CodesStruct(audio=flat),
+        meta=MetaStruct(
+            left_context_size=left_context_size,
+            finished=torch.tensor(finished, dtype=torch.bool),
+        ),
+    )
diff --git a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py b/vllm_omni/model_executor/stage_input_processors/mimo_audio.py
index 55aa003f6..2443ad2c4 100644
--- a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py
+++ b/vllm_omni/model_executor/stage_input_processors/mimo_audio.py
@@ -147,6 +147,18 @@ def llm2code2wav_async_chunk(
     Accumulates codes in connector per request_id,
     returns payload only when chunk_size is full or request is finished; returns None when waiting.
     """
+    # Null guard: chunk_transfer_adapter calls this every emit step
+    # including no-output steps where pooling_output is None.
+    if pooling_output is None or not isinstance(pooling_output, dict):
+        if is_finished:
+            connector = getattr(transfer_manager, "connector", None)
+            raw_cfg = getattr(connector, "config", {}) or {}
+            cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
+            chunk_size = int(cfg.get("codec_chunk_frames", 3))
+            left_context_size = int(cfg.get("codec_left_context_frames", 3))
+            request_id = getattr(request, "external_req_id", None)
+            return _flush_remaining_codes(transfer_manager, request_id, chunk_size, left_context_size)
+        return None
     connector = getattr(transfer_manager, "connector", None)
     raw_cfg = getattr(connector, "config", {}) or {}
     cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
@@ -173,7 +185,12 @@ def llm2code2wav_async_chunk(
 
     request_id = getattr(request, "external_req_id", None)
 
-    po_codes = pooling_output.get("codes", {})
+    # Text-only paths (e.g. modalities=["text"]) yield no codec pooling output;
+    # stage-0 still drives the chunk transfer adapter, so treat None as "no codes
+    # this step" rather than letting `.get()` raise AttributeError — an unhandled
+    # error here drops the chunk, starves stage-1 of the finished payload, and
+    # the stage subprocesses die before the final token is emitted.
+    po_codes = pooling_output.get("codes", {}) if pooling_output is not None else {}
     if "audio" not in po_codes:
         if is_finished:
             return _flush_remaining_codes(transfer_manager, request_id, chunk_size, left_context_size)
@@ -295,3 +312,137 @@ def llm2code2wav(
         )
 
     return code2wav_inputs
+
+
+# ============================================================================
+# Worker-connector data plane (non-async-chunk path).
+# AR runner's `flatten_payload` converts the model emit
+# `multimodal_outputs={"codes": {"audio": ...}}` to flat
+# `pooling_output["codes.audio"]` before the full-payload accumulator runs, so default
+# CONCAT semantics build the full codec tensor across all decode steps.
+# ============================================================================
+
+# Per-model REPLACE-keys for the full-payload accumulator.  mimo_audio's
+# producer side emits per-step codec frames that should be CONCAT'd across
+# steps (not REPLACE'd), so this stays empty.
+_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset()
+
+
+def _filter_zero_codec_rows(codec_codes: torch.Tensor) -> torch.Tensor:
+    """Drop zero-padded codec rows from a 4-D `[N, 1, 8, 4]` tensor.
+
+    Mirrors the zero-row filter in the orchestrator-path `llm2code2wav`
+    body (see this file's ``llm2code2wav`` around line 224).
+    """
+    if codec_codes.ndim != 4 or codec_codes.numel() == 0:
+        return codec_codes
+    is_all_zero = (codec_codes == 0).all(dim=(1, 2, 3))
+    nonzero_idx = (~is_all_zero).nonzero(as_tuple=True)[0]
+    if len(nonzero_idx) == 0:
+        # All rows are zero-padded; return an empty tensor so the caller
+        # can detect this via numel()==0 and skip the request.
+        return codec_codes[:0]
+    if len(nonzero_idx) < codec_codes.shape[0]:
+        return codec_codes[nonzero_idx]
+    return codec_codes
+
+
+def llm2code2wav_token_only(
+    source_outputs: list,
+    _prompt=None,
+    _requires_multimodal_data: bool = False,
+) -> list:
+    """Sync-side placeholder for the non-async-chunk Stage-1 (code2wav) input.
+
+    Returns an ``OmniTokensPrompt`` sized to the orchestrator-shape codec
+    length so the consumer runtime allocates the right number of slots.
+    The actual codec ids are delivered via the worker connector payload
+    built by ``llm2code2wav_full_payload``.
+    """
+    from vllm_omni.inputs.data import OmniTokensPrompt
+
+    code2wav_inputs: list = []
+    for output_wrapper in source_outputs:
+        out = output_wrapper.outputs[0]
+        mm = out.multimodal_output if hasattr(out, "multimodal_output") else None
+        mm = mm if isinstance(mm, dict) else {}
+        mm_codes = mm.get("codes", {}) if isinstance(mm, dict) else {}
+        prompt_len = 0
+        if isinstance(mm_codes, dict) and "audio" in mm_codes:
+            audio = mm_codes["audio"]
+            if isinstance(audio, torch.Tensor) and audio.numel() > 0:
+                audio = audio.to(torch.long)
+                audio = _filter_zero_codec_rows(audio)
+                # +B*4 per batch row for the prepended pad_vec (see prepend_and_flatten_colmajor)
+                batch_size = int(audio.shape[0]) if audio.ndim >= 1 else 1
+                prompt_len = int(audio.numel()) + batch_size * 4
+        if prompt_len > MAX_CODE2WAV_TOKENS:
+            prompt_len = MAX_CODE2WAV_TOKENS
+        code2wav_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[0] * prompt_len,
+                additional_information=None,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+            )
+        )
+    return code2wav_inputs
+
+
+def llm2code2wav_full_payload(
+    transfer_manager,
+    pooling_output: dict,
+    request,
+) -> dict | None:
+    """Producer-side payload builder for the worker connector data plane.
+
+    AR runner's ``flatten_payload`` converts the per-step model emit
+    ``{"codes": {"audio": ...}}`` to ``pooling_output["codes.audio"]``.
+    The accumulator CONCATs per-step tensors along dim 0, so by flush
+    time this holds the full ``[total_steps, 1, 8, 4]`` codec tensor.
+
+    A back-compat fallback to nested ``pooling_output["codes"]["audio"]``
+    is kept in case a future runtime path bypasses `flatten_payload`.
+    """
+    del transfer_manager
+    rid = getattr(request, "request_id", "?")
+    if not isinstance(pooling_output, dict):
+        logger.warning(
+            "mimo_audio.llm2code2wav_full_payload: pooling_output not a dict "
+            "(type=%s) for req=%s; consumer wait gate may hang.",
+            type(pooling_output).__name__,
+            rid,
+        )
+        return None
+    codec_codes = pooling_output.get("codes.audio")
+    if codec_codes is None:
+        # Back-compat fallback for un-flattened pooler emits.
+        codes = pooling_output.get("codes")
+        if isinstance(codes, dict):
+            codec_codes = codes.get("audio")
+    if not isinstance(codec_codes, torch.Tensor) or codec_codes.numel() == 0:
+        logger.warning(
+            "mimo_audio.llm2code2wav_full_payload: missing/empty codes.audio "
+            "(keys=%s) for req=%s; consumer wait gate may hang.",
+            list(pooling_output.keys()),
+            rid,
+        )
+        return None
+    codec_codes = codec_codes.to(torch.long)
+    codec_codes = _filter_zero_codec_rows(codec_codes)
+    if codec_codes.numel() == 0:
+        logger.warning(
+            "mimo_audio.llm2code2wav_full_payload: codec_codes empty after _filter_zero_codec_rows for req=%s.",
+            rid,
+        )
+        return None
+
+    pad_vec = torch.tensor([TALKER_CODEC_PAD_TOKEN_ID] * 4)
+    code_final = prepend_and_flatten_colmajor(codec_codes, pad_vec).tolist()
+    if len(code_final) > MAX_CODE2WAV_TOKENS:
+        code_final = code_final[:MAX_CODE2WAV_TOKENS]
+
+    return {
+        "codes": {"audio": code_final},
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+    }
diff --git a/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py b/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py
index dddca3a9e..938018856 100644
--- a/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py
+++ b/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py
@@ -4,27 +4,458 @@
 
 from __future__ import annotations
 
+import logging
+import re
+from dataclasses import dataclass
 from typing import Any
 
+import torch
 from vllm.inputs import TextPrompt
 
 from vllm_omni.inputs.data import OmniTokensPrompt
 
+logger = logging.getLogger(__name__)
+
+
+CFG_TEXT_SUFFIX = "__cfg_text"
+
+
+# Fallback when stage config introspection fails; matches
+# llm_config.image_patch_token on the released Ming-flash-omni-2.0 checkpoint.
+_DEFAULT_IMAGE_PATCH_TOKEN_ID = 157157
+
+
+# Ming's byte5 glyph text is auto-extracted from the user prompt's quoted
+# spans (ASCII double quotes / Chinese curly quotes). Patterns and regex
+# taken verbatim from Ming ``processing_bailingmm2.py::get_text_from_prompt``.
+_GLYPH_QUOTE_PATTERNS = [r"\"(.*?)\"", r"‘(.*?)’", r"“(.*?)”"]
+_GLYPH_REMOVE_KEYWORDS = ("remove", "delete", "erase")
+
+
+def _check_single_quotes(s: str) -> bool:
+    """Decide whether ASCII single quotes in ``s`` act as glyph delimiters.
+
+    Mirrors Ming's heuristic: count Chinese chars as weight 3, others as 1;
+    if any paired-single-quote span crosses weight 20 we assume the quotes are
+    apostrophes (e.g. "don't"), not glyph markers.
+    """
+    if s.count("'") % 2 != 0:
+        return False
+    positions = [i for i, ch in enumerate(s) if ch == "'"]
+    for i in range(0, len(positions), 2):
+        start, end = positions[i], positions[i + 1]
+        inner = s[start + 1 : end]
+        chinese = sum(1 for c in inner if "\u4e00" <= c <= "\u9fff")
+        total = 3 * chinese + (len(inner) - chinese)
+        if total >= 20:
+            return False
+    return True
+
+
+def _extract_byte5_glyph_text(prompt: str) -> str:
+    """Return ``Ming``-style ``'Text "<glyph>". '`` or ``""`` when no glyph."""
+    if not isinstance(prompt, str) or not prompt:
+        return ""
+    if "'" in prompt and _check_single_quotes(prompt):
+        prompt = prompt.replace("'", '"')
+
+    texts: list[str] = []
+    for pattern in _GLYPH_QUOTE_PATTERNS:
+        texts.extend(re.findall(pattern, prompt))
+
+    if len(texts) == 1:
+        # Treat "remove/delete/erase ..." as a glyph-removal intent, not generation.
+        text_start = min(
+            (prompt.find(q) for q in ('"', "‘", "“") if prompt.find(q) >= 0),
+            default=-1,
+        )
+        lower = prompt.lower()
+        for kw in _GLYPH_REMOVE_KEYWORDS:
+            idx = lower.find(kw)
+            if 0 <= idx < text_start:
+                return ""
+
+    if not texts:
+        return ""
+    # Only the last quoted span is used (Ming's choice; keeps the most recent edit target).
+    return f'Text "{texts[-1]}". '
+
+
+# ---------------------------------------------------------------------------
+# CFG prompt expansion (stage 0: prompt_expand_func)
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class _CfgExpandedPrompt:
+    """Minimal structural object consumed by ``AsyncOmniEngine._enqueue_cfg_companions``."""
+
+    prompt: dict[str, Any]
+    role: str
+    request_id_suffix: str
+
+    def apply_overrides(self, base_params: Any, base_spl: list[Any]) -> tuple[Any, list[Any]]:
+        return base_params, base_spl
+
+
+def expand_cfg_prompts(
+    prompt: dict[str, Any] | str,
+    sampling_params: Any,
+) -> list[_CfgExpandedPrompt]:
+    """Expand a text-to-image request into one CFG-text companion (opt-in).
+
+    Triggers only when a non-empty
+    ``sampling_params.extra_args["image_gen"]["negative_prompt"]`` is set on
+    the stage-0 params; otherwise returns ``[]`` and the pipeline falls back
+    to zero negative (Ming's default behavior).
+    """
+    if not isinstance(prompt, dict):
+        return []
+    if prompt.get("modalities") != ["image"]:
+        return []
+
+    extra_args = getattr(sampling_params, "extra_args", None) or {}
+    image_gen_args = extra_args.get("image_gen") or {}
+    negative = image_gen_args.get("negative_prompt")
+    if not isinstance(negative, str) or not negative.strip():
+        return []
+
+    neg_prompt_dict: dict[str, Any] = {
+        "prompt": negative,
+        "modalities": prompt.get("modalities"),
+    }
+    mm_kwargs = prompt.get("mm_processor_kwargs")
+    if mm_kwargs:
+        neg_prompt_dict["mm_processor_kwargs"] = dict(mm_kwargs)
+
+    return [_CfgExpandedPrompt(prompt=neg_prompt_dict, role="cfg_text", request_id_suffix=CFG_TEXT_SUFFIX)]
+
+
+# ---------------------------------------------------------------------------
+# Thinker → imagegen bridge (stage 1: custom_process_input_func)
+# ---------------------------------------------------------------------------
+
+
+def _resolve_num_query_tokens(stage: Any) -> int | None:
+    """Return the image-gen ``num_query_tokens`` from the source stage's config.
+
+    Falls back to 256 (``img_gen_scales=[16]``) when the stage config lacks
+    a ``MingImageGenConfig``. Cached on the stage object for O(1) re-reads.
+    """
+    cached = getattr(stage, "_ming_num_query_tokens", None)
+    if isinstance(cached, int):
+        return cached
+    n = 256  # Ming-flash-omni-2.0 default (img_gen_scales=[16])
+    try:
+        hf_config = stage.vllm_config.model_config.hf_config
+        ig = getattr(hf_config, "image_gen_config", None)
+        resolved = getattr(ig, "num_query_tokens", None)
+        if isinstance(resolved, int) and resolved > 0:
+            n = resolved
+    except AttributeError:
+        pass
+    try:
+        stage._ming_num_query_tokens = n
+    except AttributeError:
+        pass
+    return n
+
+
+def _resolve_image_end_token_id(stage: Any) -> int | None:
+    """Return the ``<image_end>`` token id from *stage*'s HF config, cached on first call."""
+    cached = getattr(stage, "_image_end_token_id", None)
+    if isinstance(cached, int):
+        return cached
+    token_id: int | None = None
+    try:
+        hf_config = stage.vllm_config.model_config.hf_config
+        llm_config = getattr(hf_config, "llm_config", None)
+        resolved = getattr(llm_config, "image_end_token", None)
+        if isinstance(resolved, int):
+            token_id = resolved
+    except AttributeError:
+        pass
+    try:
+        stage._image_end_token_id = token_id
+    except AttributeError:
+        pass
+    return token_id
+
+
+def _resolve_image_patch_token_id(stage: Any) -> int:
+    """Return the ``<imagePatch>`` token id from *stage*'s HF config, cached on first call."""
+    cached = getattr(stage, "_image_patch_token_id", None)
+    if isinstance(cached, int):
+        return cached
+
+    token_id = _DEFAULT_IMAGE_PATCH_TOKEN_ID
+    try:
+        hf_config = stage.vllm_config.model_config.hf_config
+        llm_config = getattr(hf_config, "llm_config", None)
+        resolved = getattr(llm_config, "image_patch_token", None)
+        if isinstance(resolved, int):
+            token_id = resolved
+    except AttributeError:
+        pass
+
+    try:
+        stage._image_patch_token_id = token_id
+    except AttributeError:
+        pass
+    return token_id
+
 
 def _validate_stage_inputs(stage_list, engine_input_source):
-    """Validate stage inputs and return the source engine outputs."""
     if not engine_input_source:
         raise ValueError("engine_input_source cannot be empty")
-
     stage_id = engine_input_source[0]
     if stage_id >= len(stage_list):
         raise IndexError(f"Invalid stage_id: {stage_id}")
-
     stage = stage_list[stage_id]
     if stage.engine_outputs is None:
         raise RuntimeError(f"Stage {stage_id} has no outputs yet")
+    return stage, stage.engine_outputs
 
-    return stage.engine_outputs
+
+def _ensure_list(x) -> list[int]:
+    """Convert ConstantList / tensor-like to plain list."""
+    if hasattr(x, "_x"):
+        return list(x._x)
+    if isinstance(x, list):
+        return x
+    if hasattr(x, "tolist"):
+        return x.tolist()
+    return list(x)
+
+
+def _slice_patch_hidden(
+    thinker_output: Any,
+    image_patch_token_id: int,
+    tag: str,
+    num_query_tokens: int | None = None,
+    image_end_token_id: int | None = None,
+) -> torch.Tensor | None:
+    """Return ``[N, H]`` hidden at the image-gen query-token block.
+
+    The image-gen block is always appended at the prompt tail as
+    ``<image_start><imagePatch>*N<image_end>`` by
+    ``maybe_expand_image_gen_prompt``. We find the exact trailing window by
+    anchoring on the final ``image_end_token`` and walking back N positions,
+    verifying they're all ``image_patch_token``. Without that signature check
+    we'd risk slicing ref-image patches (img2img case) or comprehension-only
+    patch blocks.
+
+    Falls back to "all patch positions" if ``image_end_token_id`` or
+    ``num_query_tokens`` is not provided — matches the pre-img2img t2i-only
+    behavior used by older callers.
+    """
+    output = thinker_output.outputs[0]
+    mm_out = getattr(output, "multimodal_output", None) or {}
+    full_hidden = mm_out.get("final_hidden_states")
+    if full_hidden is None:
+        logger.warning("[thinker2imagegen] %s: missing final_hidden_states (keys=%s)", tag, list(mm_out.keys()))
+        return None
+
+    prompt_ids = _ensure_list(thinker_output.prompt_token_ids)
+    prompt_ids_t = torch.tensor(prompt_ids, dtype=torch.long, device=full_hidden.device)
+    patch_indices = (prompt_ids_t == image_patch_token_id).nonzero(as_tuple=False).squeeze(-1)
+    total_patches = int(patch_indices.numel())
+    if total_patches == 0:
+        logger.warning("[thinker2imagegen] %s: no <imagePatch> tokens in prompt (len=%d)", tag, len(prompt_ids))
+        return None
+
+    if full_hidden.dim() == 3:
+        assert full_hidden.shape[0] == 1, f"expected batch=1, got {full_hidden.shape}"
+        full_hidden = full_hidden[0]
+    if full_hidden.dim() != 2 or full_hidden.shape[0] != prompt_ids_t.shape[0]:
+        logger.warning(
+            "[thinker2imagegen] %s: hidden shape %s inconsistent with prompt len %d",
+            tag,
+            tuple(full_hidden.shape),
+            prompt_ids_t.shape[0],
+        )
+        return None
+
+    selected_indices = patch_indices
+    if num_query_tokens is not None and image_end_token_id is not None:
+        L = int(prompt_ids_t.numel())
+        if L >= num_query_tokens + 1 and int(prompt_ids_t[-1].item()) == image_end_token_id:
+            tail_start = L - 1 - num_query_tokens
+            tail_end = L - 1  # exclusive
+            tail_slice = prompt_ids_t[tail_start:tail_end]
+            if (tail_slice == image_patch_token_id).all():
+                selected_indices = torch.arange(tail_start, tail_end, dtype=torch.long, device=full_hidden.device)
+            else:
+                logger.warning(
+                    "[thinker2imagegen] %s: tail signature mismatch (expected N patches "
+                    "before <image_end>); falling back to all patch positions",
+                    tag,
+                )
+        else:
+            logger.warning(
+                "[thinker2imagegen] %s: image-gen block signature not found at prompt tail "
+                "(len=%d, last_tok=%s); falling back to all patch positions",
+                tag,
+                L,
+                int(prompt_ids_t[-1].item()) if L else None,
+            )
+
+    hidden = full_hidden[selected_indices].detach().contiguous()
+    if logger.isEnabledFor(logging.DEBUG):
+        f = hidden.float()
+        logger.debug(
+            "[thinker2imagegen] %s sliced=%s (%d of %d patches) mean=%+.4f std=%.4f |x|/tok=%.3f",
+            tag,
+            tuple(hidden.shape),
+            int(selected_indices.numel()),
+            total_patches,
+            f.mean().item(),
+            f.std().item(),
+            f.norm(dim=-1).mean().item(),
+        )
+    return hidden
+
+
+def _resolve_token_ids_from_stage_or_defaults(
+    stage: Any | None,
+) -> tuple[int, int | None, int | None]:
+    """Return (image_patch_token_id, image_end_token_id, num_query_tokens).
+
+    Tries to read from the stage's HF config when available (old API).
+    Falls back to Ming-flash-omni-2.0 defaults when stage is None (new API).
+    """
+    if stage is not None:
+        return (
+            _resolve_image_patch_token_id(stage),
+            _resolve_image_end_token_id(stage),
+            _resolve_num_query_tokens(stage),
+        )
+    # Defaults from Ming-flash-omni-2.0:
+    #   llm_config.image_patch_token = 157157
+    #   llm_config.image_end_token   = 157159
+    #   img_gen_scales=[16] -> 16*16 = 256 query tokens
+    return (_DEFAULT_IMAGE_PATCH_TOKEN_ID, 157159, 256)
+
+
+def _extract_byte5_from_sampling_params(sampling_params: Any) -> list[str] | None:
+    """Read ``byte5_text`` from the diffusion-stage sampling_params.
+
+    Looks up ``sampling_params.extra_args["image_gen"]["byte5_text"]`` (the
+    explicit API surface for ByT5 glyph text). Returns ``None`` if absent or
+    malformed, so callers can fall back to other sources.
+    """
+    if sampling_params is None:
+        return None
+    extra = getattr(sampling_params, "extra_args", None)
+    if not isinstance(extra, dict):
+        return None
+    image_gen = extra.get("image_gen")
+    if not isinstance(image_gen, dict):
+        return None
+    texts = image_gen.get("byte5_text")
+    if isinstance(texts, list) and texts:
+        return [t for t in texts if isinstance(t, str)]
+    return None
+
+
+def thinker2imagegen(
+    source_outputs: list[Any],
+    prompt: Any | None = None,
+    requires_multimodal_data: bool = False,  # noqa: ARG001
+    sampling_params: Any | None = None,
+) -> list[dict[str, Any]]:
+    """Bridge thinker AR outputs into image-generation DiT inputs.
+
+    The orchestrator passes ``source_outputs`` as
+    ``[parent_output, *companion_outputs]``. Parent outputs feed
+    ``extra[thinker_hidden_states]``; the cfg_text companion feeds
+    ``extra[negative_thinker_hidden_states]`` used by MingImagePipeline as real
+    CFG negative conditioning. Unknown-suffix outputs are skipped.
+
+    ``sampling_params`` is the diffusion stage's own SamplingParams, supplied
+    by the orchestrator. ByT5 explicit ``byte5_text`` is read from
+    ``sampling_params.extra_args.image_gen.byte5_text`` (preferred); falls
+    back to ``prompt.image_gen_extra_args.byte5_text``, then to auto-extraction
+    from quoted prompt text.
+    """
+    thinker_outputs = source_outputs
+    image_patch_token_id, image_end_token_id, num_query_tokens = _resolve_token_ids_from_stage_or_defaults(stage=None)
+
+    parent_output = None
+    negative_output = None
+    for o in thinker_outputs:
+        rid = getattr(o, "request_id", "")
+        if rid.endswith(CFG_TEXT_SUFFIX):
+            negative_output = o
+        elif parent_output is None:
+            parent_output = o
+
+    if parent_output is None:
+        logger.warning("[thinker2imagegen] no parent output in engine_outputs; skipping")
+        return []
+
+    parent_hidden = _slice_patch_hidden(
+        parent_output,
+        image_patch_token_id,
+        tag="parent",
+        num_query_tokens=num_query_tokens,
+        image_end_token_id=image_end_token_id,
+    )
+    if parent_hidden is None:
+        return []
+
+    extra: dict[str, Any] = {"thinker_hidden_states": parent_hidden}
+    if negative_output is not None:
+        neg_hidden = _slice_patch_hidden(
+            negative_output,
+            image_patch_token_id,
+            tag="cfg_text",
+            num_query_tokens=num_query_tokens,
+            image_end_token_id=image_end_token_id,
+        )
+        if neg_hidden is not None:
+            extra["negative_thinker_hidden_states"] = neg_hidden
+
+    # img2img: forward the reference image PIL/tensor to the diffusion stage.
+    if isinstance(prompt, dict):
+        mm_data = prompt.get("multi_modal_data") or {}
+        ref_image = mm_data.get("image")
+        if isinstance(ref_image, list) and ref_image:
+            ref_image = ref_image[0]
+        if ref_image is None:
+            ref_image = mm_data.get("img2img")
+            if isinstance(ref_image, list) and ref_image:
+                ref_image = ref_image[0]
+        if ref_image is not None:
+            extra["reference_image"] = ref_image
+
+        # ByT5 glyph text: prefer sampling_params (stage-1 explicit API), then
+        # prompt.image_gen_extra_args (serving_chat path), then auto-extract
+        # from quoted prompt text.
+        prompt_text = prompt.get("prompt", "")
+        sp_byte5 = _extract_byte5_from_sampling_params(sampling_params)
+        if sp_byte5:
+            byte5_texts: list[str] | None = sp_byte5
+        else:
+            ig_extra = prompt.get("image_gen_extra_args") or {}
+            cand = ig_extra.get("byte5_text")
+            byte5_texts = cand if isinstance(cand, list) and cand else None
+
+        if byte5_texts:
+            extra["byte5_text"] = [
+                t if t.startswith("Text ") else f'Text "{t}". ' for t in byte5_texts if isinstance(t, str)
+            ]
+        else:
+            glyph = _extract_byte5_glyph_text(prompt_text)
+            if glyph:
+                extra["byte5_text"] = [glyph]
+    else:
+        # prompt is not a dict — still honor explicit byte5_text from sampling_params.
+        sp_byte5 = _extract_byte5_from_sampling_params(sampling_params)
+        if sp_byte5:
+            extra["byte5_text"] = [t if t.startswith("Text ") else f'Text "{t}". ' for t in sp_byte5]
+
+    return [{"prompt": "", "extra": extra}]
 
 
 def thinker2talker(
@@ -39,7 +470,7 @@ def thinker2talker(
     a talker input prompt with the text and any speaker/instruction info
     from the original request.
     """
-    source_outputs = _validate_stage_inputs(stage_list, engine_input_source)
+    _, source_outputs = _validate_stage_inputs(stage_list, engine_input_source)
 
     if not isinstance(prompt, list):
         prompt = [prompt]
@@ -102,3 +533,99 @@ def thinker2talker(
         )
 
     return talker_inputs
+
+
+# ming_flash_omni is not in ``_OMNI_CONNECTOR_INIT_ARCHS`` or
+# ``_FULL_PAYLOAD_INPUT_STAGES``, so the worker connector is not
+# initialised for this arch and the consumer never waits on a connector
+# payload.  Data flows through ``additional_information`` written by
+# ``thinker2talker_token_only`` (wired as ``sync_process_input_func``
+# in the pipeline) or the legacy ``thinker2talker`` (wired as
+# ``custom_process_input_func``).
+
+
+def thinker2talker_token_only(
+    source_outputs: list[Any],
+    prompt: OmniTokensPrompt | TextPrompt | None = None,
+    _requires_multimodal_data: bool = False,
+) -> list[OmniTokensPrompt]:
+    """Sync-side builder for the non-async-chunk thinker→talker path.
+
+    Ports the legacy ``thinker2talker`` body to the new stage-input-processor signature
+    (``source_outputs`` instead of ``stage_list, engine_input_source``).
+    Body is otherwise identical: extracts the
+    generated text from each thinker output and packages it with the
+    request's voice/speaker additional_information for the talker.
+    """
+    if not isinstance(prompt, list):
+        prompt = [prompt]
+
+    talker_inputs: list[OmniTokensPrompt] = []
+    for i, source_output in enumerate(source_outputs):
+        output = source_output.outputs[0]
+
+        generated_text = output.text if hasattr(output, "text") and output.text else ""
+
+        original_prompt = prompt[i] if i < len(prompt) else None
+        additional_info: dict[str, Any] = {}
+        if original_prompt is not None and hasattr(original_prompt, "additional_information"):
+            additional_info = original_prompt.additional_information or {}
+
+        spk_emb = additional_info.get("spk_emb", None)
+        if isinstance(spk_emb, list) and spk_emb and not hasattr(spk_emb[0], "device"):
+            import torch
+
+            spk_emb = torch.tensor(spk_emb, dtype=torch.float32).unsqueeze(0)
+
+        talker_info = {
+            "ming_task": "omni",
+            "text": generated_text,
+            "spk_emb": spk_emb,
+            "voice_name": additional_info.get("voice_name", "DB30"),
+            "prompt_text": additional_info.get("prompt_text", None),
+            "prompt_wav_lat": additional_info.get("prompt_wav_lat", None),
+            "prompt_wav_emb": additional_info.get("prompt_wav_emb", None),
+            "max_text_length": additional_info.get("max_text_length", 50),
+        }
+
+        talker_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[0],
+                additional_information=talker_info,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+            )
+        )
+
+    return talker_inputs
+
+
+thinker2talker_token_only._is_sync_input = True
+
+
+def thinker2talker_full_payload(
+    transfer_manager,
+    pooling_output,
+    request,
+):
+    """Producer-side payload builder — no-op.
+
+    ming_flash_omni's thinker emits no heavy tensor to ship via the
+    worker connector (the bridge passes text only, and speaker metadata
+    arrives through the USER request's additional_information).
+    ming_flash_omni is not in ``_OMNI_CONNECTOR_INIT_ARCHS`` so this
+    function is never invoked at runtime; it is retained for forward
+    compatibility with the connector path.
+    """
+    del transfer_manager, pooling_output, request
+    return None
+
+
+__all__ = [
+    "CFG_TEXT_SUFFIX",
+    "expand_cfg_prompts",
+    "thinker2imagegen",
+    "thinker2talker",
+    "thinker2talker_full_payload",
+    "thinker2talker_token_only",
+]
diff --git a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py
index 472dcb933..225674360 100644
--- a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py
+++ b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py
@@ -1,3 +1,5 @@
+import logging
+
 import torch
 from vllm.inputs import TextPrompt
 
@@ -11,6 +13,8 @@ from vllm_omni.data_entry_keys import (
 )
 from vllm_omni.inputs.data import OmniTokensPrompt
 
+logger = logging.getLogger(__name__)
+
 TALKER_CODEC_PAD_TOKEN_ID = 8292
 TALKER_CODEC_START_TOKEN_ID = 8293
 TALKER_CODEC_END_TOKEN_ID = 8294
@@ -86,3 +90,347 @@ def talker2code2wav(
             )
         )
     return code2wav_inputs
+
+
+# ============================================================================
+# Worker-connector data plane (non-async-chunk path).
+# Both transitions ship payloads via the worker connector
+# (registered in ``_FULL_PAYLOAD_INPUT_STAGES`` in
+# omni_scheduling_coordinator):
+# - thinker->talker reads accumulated ``pooling_output["hidden"]`` and
+#   packs an OmniPayload-shaped dict (embed.prefill /
+#   hidden_states.output / ids.prompt / ids.output) for the talker, which
+#   the talker's ``talker_preprocess`` reads from
+#   ``model_intermediate_buffer``.  The shape matches what legacy
+#   ``thinker2talker`` writes into ``additional_information`` as a debug
+#   fallback; ``thinker2talker_token_only`` only allocates prompt slots.
+# - talker->code2wav strips TALKER_CODEC_{START,END} boundary tokens
+#   and ships the codec token ids.
+# ============================================================================
+
+# Per-model REPLACE-keys for the full-payload accumulator.  qwen2_5_omni's
+# producer side does not emit model_outputs through pooler_output (it ships
+# token_ids on the request directly), so the empty set preserves correctness.
+_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset()
+
+
+def _strip_codec_boundaries(token_ids: list[int]) -> list[int]:
+    """Keep only real codec ids for the code2wav stage.
+
+    The talker stream can contain prompt/control ids (START/PAD/END/MASK) in
+    addition to sampled codec ids.  Code2wav expects codec ids only; carrying
+    the prompt PAD span forward can inflate the sequence enough to OOM on L4.
+    Async scheduling may also leave trailing ``-1`` placeholders, so preserve
+    their length by repeating the last valid codec id.
+    """
+    tids = list(token_ids)
+    trailing_placeholder_count = 0
+    while trailing_placeholder_count < len(tids) and tids[-1 - trailing_placeholder_count] == -1:
+        trailing_placeholder_count += 1
+
+    if tids and tids[-1] == TALKER_CODEC_END_TOKEN_ID:
+        tids = tids[:-1]
+        trailing_placeholder_count = 0
+
+    codec_ids = [tid for tid in tids if 0 <= tid < TALKER_CODEC_PAD_TOKEN_ID]
+    if trailing_placeholder_count > 0 and codec_ids:
+        codec_ids.extend([codec_ids[-1]] * trailing_placeholder_count)
+    return codec_ids
+
+
+def talker2code2wav_token_only(
+    source_outputs,
+    _prompt: OmniTokensPrompt | TextPrompt = None,
+    _requires_multimodal_data: bool = False,
+):
+    """Sync-side placeholder for Stage-2 input (code2wav).
+
+    Returns OmniTokensPrompt sized to the stripped codec token count.
+    Actual codec ids are delivered via the worker connector payload built
+    by ``talker2code2wav_full_payload``.
+    """
+    code2wav_inputs = []
+    for talker_output in source_outputs:
+        output = talker_output.outputs[0]
+        token_ids = _strip_codec_boundaries(list(output.cumulative_token_ids))
+        if not token_ids:
+            continue
+        code2wav_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[0] * len(token_ids),
+                additional_information=None,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+            )
+        )
+    return code2wav_inputs
+
+
+def talker2code2wav_full_payload(
+    transfer_manager,
+    pooling_output: dict,
+    request,
+) -> dict | None:
+    """Producer-side payload builder: ship the stripped codec ids via connector.
+
+    Token-ids-only shape.  The talker stage's output already
+    carries the codec ids on ``request.output_token_ids``; we strip the
+    boundary tokens and pack a minimal payload.
+    """
+    del transfer_manager
+    rid = getattr(request, "request_id", "?")
+    token_ids = list(getattr(request, "output_token_ids", None) or [])
+    if not token_ids:
+        logger.warning(
+            "qwen2_5_omni.talker2code2wav_full_payload: empty output_token_ids "
+            "for req=%s; consumer wait gate may hang.",
+            rid,
+        )
+        return None
+    token_ids = _strip_codec_boundaries(token_ids)
+    if not token_ids:
+        logger.warning(
+            "qwen2_5_omni.talker2code2wav_full_payload: codec ids empty after "
+            "stripping boundary tokens for req=%s; consumer wait gate may hang.",
+            rid,
+        )
+        return None
+    return {
+        "codes": {"audio": token_ids},
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+    }
+
+
+# ============================================================================
+# Worker-connector data plane (non-async-chunk path) -- thinker->talker.
+#
+# qwen2_5_omni's talker consumes the thinker's last-layer hidden state
+# via Linear(3584, 896).  The AR runner publishes those hidden states
+# per decode step on ``pooling_output["hidden"]`` (unpacked from
+# ``OmniOutput.text_hidden_states``); the full-payload accumulator
+# concatenates them so ``thinker2talker_full_payload`` sees the full
+# prefill+decode trajectory and packs an OmniPayload-shaped dict that
+# the talker's ``talker_preprocess`` reads from
+# ``model_intermediate_buffer``.  ``thinker2talker_token_only`` only
+# allocates the talker's codec prompt slots; legacy
+# ``thinker2talker`` above remains as a debug fallback that bundles the
+# same shape into ``additional_information``.
+# ============================================================================
+
+
+def thinker2talker_token_only(
+    source_outputs,
+    prompt: OmniTokensPrompt | TextPrompt = None,
+    requires_multimodal_data: bool = False,
+):
+    """Placeholder builder for the connector-driven thinker->talker path.
+
+    Allocates the TALKER_CODEC_{START,PAD,END} prompt slots sized to the
+    thinker prompt length and forwards ``multi_modal_data``.  The bulk
+    payload (hidden_states / embed / ids) ships exclusively through
+    ``thinker2talker_full_payload`` via the worker connector and lands
+    in ``model_intermediate_buffer`` before the talker's forward() runs.
+
+    Consumer-wait gating is whitelist-driven via
+    ``_FULL_PAYLOAD_INPUT_STAGES`` (see the mixin
+    ``should_accumulate_full_payload_output`` docstring).
+    """
+    thinker_outputs = source_outputs
+    talker_inputs = []
+    if not isinstance(prompt, list):
+        prompt = [prompt]
+    multi_modal_data = {
+        thinker_output.request_id: p.get("multi_modal_data", None) for thinker_output, p in zip(thinker_outputs, prompt)
+    }
+
+    for thinker_output in thinker_outputs:
+        prompt_token_ids = thinker_output.prompt_token_ids
+        talker_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[TALKER_CODEC_START_TOKEN_ID]
+                + [TALKER_CODEC_PAD_TOKEN_ID] * (len(prompt_token_ids))
+                + [TALKER_CODEC_END_TOKEN_ID],
+                additional_information=None,
+                multi_modal_data=(
+                    multi_modal_data[thinker_output.request_id]
+                    if requires_multimodal_data and multi_modal_data is not None
+                    else None
+                ),
+                mm_processor_kwargs=None,
+            )
+        )
+
+    return talker_inputs
+
+
+def thinker2talker_full_payload(
+    transfer_manager,
+    pooling_output,
+    request,
+):
+    """Producer-side payload builder for the worker-connector data plane.
+
+    The AR runner emits per-step ``pooling_output["hidden"]`` (the
+    thinker's last-layer hidden states for the request span, unpacked
+    from ``OmniOutput.text_hidden_states``).  The full-payload
+    accumulator concatenates those per-step rows across decode steps, so
+    by the time this builder fires the materialized
+    ``pooling_output["hidden"]`` contains the full prefill+decode
+    hidden-state trajectory of size
+    ``len(prompt_token_ids) + len(output_token_ids)``.
+
+    We split it at ``len(prompt_token_ids)`` into prefill embeddings and
+    decode hidden states, then pack the ``OmniPayload``-shaped dict that
+    the talker's ``thinker_to_talker_process`` reads from
+    ``model_intermediate_buffer`` (keys ``embed.prefill`` /
+    ``hidden_states.output`` / ``ids.prompt`` / ``ids.output``).  Shape
+    matches what legacy ``thinker2talker`` writes into
+    ``additional_information`` as a debug fallback, so the talker
+    consumes the same payload layout from either path.
+
+    Like ``qwen3_omni.thinker2talker_full_payload``, we apply a
+    finish-reason-aware stop-row trim: vLLM v1 appends the sampled
+    token to ``output_token_ids`` before ``check_stop``, so a request
+    that finished via ``FINISHED_STOPPED`` has one extra accumulated
+    hidden-state row that the talker must not consume.  Max-token
+    finishes need no drop.  Status is read from the request when
+    available; otherwise we fall back to a last-token-in-stop-set
+    heuristic.
+    """
+    del transfer_manager
+    rid = getattr(request, "request_id", "?")
+    if not isinstance(pooling_output, dict):
+        logger.warning(
+            "qwen2_5_omni.thinker2talker_full_payload: pooling_output not a dict "
+            "(type=%s) for req=%s; consumer wait gate may hang.",
+            type(pooling_output).__name__,
+            rid,
+        )
+        return None
+
+    hidden = pooling_output.get("hidden")
+    if not isinstance(hidden, torch.Tensor):
+        logger.warning(
+            "qwen2_5_omni.thinker2talker_full_payload: missing 'hidden' tensor "
+            "(keys=%s) for req=%s; consumer wait gate may hang.",
+            list(pooling_output.keys()),
+            rid,
+        )
+        return None
+
+    def _ensure_list(x):
+        if x is None:
+            return []
+        if hasattr(x, "_x"):
+            # vLLM wraps cached token-id lists in ConstantList-like objects.
+            return list(x._x)
+        if isinstance(x, list):
+            return list(x)
+        return list(x)
+
+    prompt_token_ids = _ensure_list(getattr(request, "prompt_token_ids", None))
+    output_token_ids = _ensure_list(getattr(request, "output_token_ids", None))
+    all_token_ids = _ensure_list(getattr(request, "all_token_ids", None) or [])
+    if not all_token_ids:
+        all_token_ids = list(prompt_token_ids) + list(output_token_ids)
+
+    # Length-aware trim of accumulated thinker output, finish-reason-aware.
+    # Mirror qwen3_omni.thinker2talker_full_payload's logic so a stop-finish
+    # does not leak an extra hidden-state row to the talker.
+    status = getattr(request, "status", None)
+    status_name = getattr(status, "name", None) or ""
+    if not status_name and status is not None:
+        status_name = str(status).rsplit(".", 1)[-1]
+    stop_emission_drop = 1 if status_name == "FINISHED_STOPPED" else 0
+    if stop_emission_drop == 0 and not status_name and output_token_ids:
+        # Worker-side CachedRequestState has no `.status` field in vLLM v1;
+        # fall back to a last-token-in-stop-set heuristic.
+        sampling_params = getattr(request, "sampling_params", None)
+        if sampling_params is not None:
+            stop_ids: set[int] = set()
+            ignore_eos = bool(getattr(sampling_params, "ignore_eos", False))
+            for sid in getattr(sampling_params, "stop_token_ids", None) or ():
+                if isinstance(sid, int):
+                    stop_ids.add(sid)
+            if not ignore_eos:
+                for eos in (
+                    getattr(sampling_params, "eos_token_id", None),
+                    getattr(sampling_params, "_eos_token_id", None),
+                ):
+                    if isinstance(eos, int):
+                        stop_ids.add(eos)
+                for sid in (
+                    getattr(sampling_params, "all_stop_token_ids", None)
+                    or getattr(sampling_params, "_all_stop_token_ids", None)
+                    or ()
+                ):
+                    if isinstance(sid, int):
+                        stop_ids.add(sid)
+            if stop_ids and output_token_ids[-1] in stop_ids:
+                stop_emission_drop = 1
+
+    # Trim accumulated thinker output based on stop_emission_drop computed
+    # above.  Mirror qwen3_omni.thinker2talker_full_payload's contract:
+    #   target_rows = len(all_token_ids) - stop_emission_drop
+    # which excludes the stop-emission row for FINISHED_STOPPED but keeps
+    # all rows for FINISHED_LENGTH_CAPPED (max_tokens) finishes.
+    if stop_emission_drop > 0 and len(output_token_ids) >= stop_emission_drop:
+        output_token_ids = output_token_ids[:-stop_emission_drop]
+    h = hidden.detach().cpu().to(torch.float32)
+    target_rows = max(0, len(all_token_ids) - stop_emission_drop)
+    if target_rows <= 0:
+        logger.warning(
+            "qwen2_5_omni.thinker2talker_full_payload: target_rows<=0 "
+            "(all_token_ids=%d, stop_drop=%d) for req=%s; nothing to ship.",
+            len(all_token_ids),
+            stop_emission_drop,
+            getattr(request, "request_id", "?"),
+        )
+        return None
+    if h.dim() >= 1 and h.shape[0] > target_rows:
+        logger.warning(
+            "qwen2_5_omni.thinker2talker_full_payload: excess hidden rows "
+            "(got %d, target %d, stop_drop %d) for req=%s; trimming",
+            int(h.shape[0]),
+            target_rows,
+            stop_emission_drop,
+            getattr(request, "request_id", None),
+        )
+        h = h[:target_rows]
+
+    prompt_len = len(prompt_token_ids)
+    if h.shape[0] < prompt_len:
+        # Under-captured prefill -- defensively skip rather than ship a
+        # truncated payload that would confuse the talker's prefill path.
+        logger.warning(
+            "qwen2_5_omni.thinker2talker_full_payload: hidden rows=%d < prompt_len=%d "
+            "for req=%s; under-captured prefill, skipping payload.",
+            int(h.shape[0]),
+            prompt_len,
+            getattr(request, "request_id", "?"),
+        )
+        return None
+
+    prefill_hidden = h[:prompt_len]
+    decode_hidden = h[prompt_len:]
+
+    payload: OmniPayload = to_dict(
+        OmniPayloadStruct(
+            hidden_states=HiddenStatesStruct(
+                output=decode_hidden,
+                output_shape=list(decode_hidden.shape),
+            ),
+            embed=EmbeddingsStruct(
+                prefill=prefill_hidden,
+                prefill_shape=list(prefill_hidden.shape),
+            ),
+            ids=IdsStruct(
+                prompt=list(prompt_token_ids),
+                output=list(output_token_ids),
+            ),
+        )
+    )
+    # Intentionally omit payload["meta"]: the thinker->talker transition
+    # carries no scheduler-relevant metadata (next_stage_prompt_len /
+    # left_context_size are not set on this edge).
+    return payload
diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
index fd7cfd2aa..aa98c2141 100644
--- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
+++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
@@ -35,6 +35,12 @@ logger = logging.getLogger(__name__)
 # Pooling output layer keys: "0" = word embedding, "24" = accept_hidden_layer
 _EMBED_LAYER_KEY = "0"
 _HIDDEN_LAYER_KEY = "24"
+# Per-model REPLACE-keys for the full-payload accumulator.  Keys in this
+# set use REPLACE semantics (subsequent emissions discard prior chunks)
+# instead of CONCAT.  qwen3-omni currently has none — model_outputs is
+# not emitted by the thinker/talker forward.
+_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset()
+
 _QWEN3_CODEC_CODEBOOK_SIZE = 2048
 _QWEN3_CODEC_PAD_TOKEN_ID = 4196
 _QWEN3_CODEC_BOS_TOKEN_ID = 4197
@@ -119,19 +125,6 @@ def _is_valid_qwen3_codec_token_id(token_id: Any) -> bool:
     return 0 <= token_id < _QWEN3_CODEC_CODEBOOK_SIZE
 
 
-def should_accumulate_qwen3_omni_full_payload_output(
-    model_config: Any,
-    custom_process_func: Any,
-) -> bool:
-    """Return whether Qwen3-Omni should accumulate full-payload outputs."""
-    return (
-        custom_process_func is not None
-        and not getattr(model_config, "async_chunk", False)
-        and getattr(model_config, "model_arch", None) == "Qwen3OmniMoeForConditionalGeneration"
-        and getattr(model_config, "model_stage", None) in {"thinker", "talker"}
-    )
-
-
 def _extract_qwen3_full_payload_codec_rows(
     code_predictor_codes: torch.Tensor,
     output_token_ids: list[int],
@@ -268,6 +261,89 @@ def _resolve_tts_token_embedding(
 # =========================
 
 
+def _construct_thinker2talker_streaming_input_async_chunk(
+    is_finished: bool,
+    request,
+    thinker_emb,
+    thinker_hid,
+    transfer_manager,
+) -> OmniPayloadStruct | None:
+    """Build Thinker -> Talker payloads for realtime streaming input chunks.
+
+    A resumable realtime request reuses the same logical request id across
+    audio segments. The first streaming prefill chunk is cached and returns ``None`` so the
+    connector does not emit an incomplete downstream chunk. The following
+    decode chunk flushes that cached prefill together with the current Thinker
+    output, keeping Talker ids and tensor rows aligned.
+    """
+    request_id = request.external_req_id
+    output_token_ids = request.output_token_ids
+    # Convert ConstantList to regular list for OmniSerializer serialization
+    output_token_ids = _ensure_list(output_token_ids)
+    speaker = extract_speaker_from_request(request)
+    language = extract_language_from_request(request)
+    finished = torch.tensor(is_finished, dtype=torch.bool)
+    emb_cpu = thinker_emb.detach().cpu()
+    hid_cpu = thinker_hid.detach().cpu()
+
+    if output_token_ids:
+        if thinker_emb.shape[0] > 1:
+            # if thinker_emb.shape[0] > 1, new streaming input segment is added
+            # and will transfer prefill embeddings and hidden states to talker.
+            new_prompt_len = thinker_emb.shape[0]
+            payload = OmniPayloadStruct(
+                meta=MetaStruct(finished=finished),
+                embed=EmbeddingsStruct(prefill=emb_cpu),
+                hidden_states=HiddenStatesStruct(output=hid_cpu),
+                ids=IdsStruct(
+                    all=_ensure_list(request.all_token_ids[-new_prompt_len - 1 :]),
+                    prompt=_ensure_list(request.prompt_token_ids[-new_prompt_len:]),
+                ),
+                speaker=speaker,
+                language=language,
+            )
+            transfer_manager._pending_streaming_prefills[request_id] = to_dict(payload)
+            return None
+        else:
+            save_payload = transfer_manager._pending_streaming_prefills.pop(request_id, None)
+            if save_payload is not None:
+                saved_prefill = save_payload.get("embed", {}).get("prefill")
+                saved_output = save_payload.get("hidden_states", {}).get("output")
+                if isinstance(saved_prefill, torch.Tensor) and isinstance(saved_output, torch.Tensor):
+                    return OmniPayloadStruct(
+                        meta=MetaStruct(finished=finished),
+                        embed=EmbeddingsStruct(prefill=torch.cat((saved_prefill, emb_cpu), dim=0)),
+                        hidden_states=HiddenStatesStruct(output=torch.cat((saved_output, hid_cpu), dim=0)),
+                        ids=IdsStruct(
+                            all=save_payload.get("ids", {}).get("all"),
+                            prompt=save_payload.get("ids", {}).get("prompt"),
+                        ),
+                        speaker=speaker,
+                        language=language,
+                    )
+            return OmniPayloadStruct(
+                meta=MetaStruct(
+                    finished=finished,
+                ),
+                embed=EmbeddingsStruct(decode=emb_cpu),
+                hidden_states=HiddenStatesStruct(output=hid_cpu),
+                ids=IdsStruct(output=output_token_ids),
+                speaker=speaker,
+                language=language,
+            )
+    else:
+        if not is_finished:
+            # do not send async chunk mode placeholder token or embedding/hidden of the stop token
+            return None
+        return OmniPayloadStruct(
+            meta=MetaStruct(finished=finished),
+            embed=EmbeddingsStruct(decode=emb_cpu),
+            hidden_states=HiddenStatesStruct(output=hid_cpu),
+            speaker=speaker,
+            language=language,
+        )
+
+
 @dataclass
 class _Thinker2TalkerStreamingState:
     last_prompt_len: int = 0
@@ -302,16 +378,16 @@ def _get_streaming_talker_tokens(
     streaming_context: Any | None = None,
     *,
     clear_state: bool = False,
-) -> tuple[list[int], list[int], list[int], list[int]]:
-    """Return streaming token slices and merged token views for thinker->talker.
-       e.g. For the second streaming input request:
-       merged_sequences: [input_prompt 1, output_tokens 1[:-1], input_prompt 2, output_tokens 2]
-      thinker_input_ids: [input_prompt 1, output_tokens 1[:-1], input_prompt 2]
+) -> tuple[list[int], list[int]]:
+    """Return prompt/output token deltas for the current streaming segment.
+
+    In non-async-chunk streaming, Thinker's prompt may already include the
+    next input segment. Remove that new prompt tail before building the Talker
+    delta for the previous segment.
+
     Returns:
         inc_prompt: prompt token delta for this segment.
         inc_output: output token delta for this segment.
-        merged_sequences: full thinker_sequences to send downstream.
-        thinker_input_ids: full thinker_input_ids paired with merged_sequences.
     """
     state = _get_qwen3_streaming_state(request_id, streaming_context).thinker2talker
     if new_prompt_len_snapshot:
@@ -321,15 +397,6 @@ def _get_streaming_talker_tokens(
 
     inc_prompt = prompt_token_ids[state.last_prompt_len :]
     inc_output = output_token_ids[state.last_output_len :]
-    delta_sequences = inc_prompt + inc_output
-    cached_sequences = state.merged_sequences
-
-    merged_sequences = cached_sequences + delta_sequences
-    thinker_input_ids = cached_sequences + inc_prompt
-
-    # Persist history for next segment. Drop the latest sampled token to keep
-    # thinker_input_ids / thinker_sequences alignment with next-step append.
-    cached_sequences.extend(delta_sequences[:-1])
 
     state.last_prompt_len = cur_prompt_len
     state.last_output_len = cur_output_len
@@ -339,7 +406,7 @@ def _get_streaming_talker_tokens(
         state.last_output_len = 0
         state.merged_sequences.clear()
 
-    return inc_prompt, inc_output, merged_sequences, thinker_input_ids
+    return inc_prompt, inc_output
 
 
 def _get_streaming_codec_delta_len(
@@ -436,6 +503,10 @@ def thinker2talker_async_chunk(
                 transfer_manager.request_payload[request_id] = to_dict(payload)
                 return None
     else:
+        if request.resumable:
+            return _construct_thinker2talker_streaming_input_async_chunk(
+                is_finished, request, thinker_emb, thinker_hid, transfer_manager
+            )
         if thinker_emb.shape[0] > 1:
             logger.warning(
                 "Unexpected multiple embeddings in thinker2talker_async_chunk for chunk_id %d: "
@@ -462,7 +533,13 @@ def thinker2talker_full_payload(
     request: OmniEngineCoreRequest,
 ) -> dict[str, Any] | None:
     """Pack complete thinker output for the non-async connector path."""
+    rid = getattr(request, "request_id", None)
     if not isinstance(pooling_output, dict):
+        logger.warning(
+            "thinker2talker_full_payload: pooling_output not a dict (type=%s) for req=%s; consumer wait gate may hang.",
+            type(pooling_output).__name__,
+            rid,
+        )
         return None
 
     layers = {
@@ -475,11 +552,13 @@ def thinker2talker_full_payload(
         hidden = pooling_output.get("hidden")
         thinker_emb = hidden if isinstance(hidden, torch.Tensor) else None
     if thinker_emb is None or thinker_hid is None:
-        logger.debug(
-            "thinker2talker_full_payload: missing thinker tensors for req=%s (embed=%s hidden=%s)",
-            getattr(request, "request_id", None),
+        logger.warning(
+            "thinker2talker_full_payload: missing thinker tensors for req=%s "
+            "(embed=%s hidden=%s keys=%s); consumer wait gate may hang.",
+            rid,
             thinker_emb is not None,
             thinker_hid is not None,
+            list(pooling_output.keys()),
         )
         return None
 
@@ -493,7 +572,7 @@ def thinker2talker_full_payload(
     # The accumulator captures one hidden-state row per executed thinker
     # forward (prefill + every decode step including the one that emitted
     # the stop_token), so for a finished request thinker_emb has exactly one
-    # row more than the rows the talker should consume.  async_chunk's
+    # row more than the rows the talker should consume. async_chunk's
     # chunk-0 path naturally captures only the prefill / non-stop portion,
     # which is why the [async_chunk] parametrization passes while [default]
     # over-generates one codec frame on short outputs (e.g.
@@ -566,7 +645,7 @@ def thinker2talker(
         output_ids = _ensure_list(output.cumulative_token_ids)
         is_streaming_session = bool(getattr(streaming_context, "enabled", False))
         if is_streaming_session:
-            prompt_token_ids, output_ids, thinker_sequences, thinker_input_ids = _get_streaming_talker_tokens(
+            prompt_token_ids, output_ids = _get_streaming_talker_tokens(
                 req_id,
                 prompt_token_ids,
                 output_ids,
@@ -574,9 +653,8 @@ def thinker2talker(
                 streaming_context,
                 clear_state=bool(getattr(thinker_output, "finished", False)),
             )
-        else:
-            thinker_sequences = prompt_token_ids + output_ids
-            thinker_input_ids = prompt_token_ids
+        thinker_sequences = prompt_token_ids + output_ids
+        thinker_input_ids = prompt_token_ids
         new_seq_length = len(prompt_token_ids + output_ids) - 1
         thinker_mm_raw = getattr(output, "multimodal_output", None)
         if not isinstance(thinker_mm_raw, dict):
@@ -682,7 +760,7 @@ def thinker2talker_token_only(
         output_ids = _ensure_list(output.cumulative_token_ids)
         is_streaming_session = bool(getattr(streaming_context, "enabled", False))
         if is_streaming_session:
-            prompt_token_ids, output_ids, thinker_sequences, thinker_input_ids = _get_streaming_talker_tokens(
+            prompt_token_ids, output_ids = _get_streaming_talker_tokens(
                 req_id,
                 prompt_token_ids,
                 output_ids,
@@ -690,9 +768,8 @@ def thinker2talker_token_only(
                 streaming_context,
                 clear_state=bool(getattr(thinker_output, "finished", False)),
             )
-        else:
-            thinker_sequences = prompt_token_ids + output_ids
-            thinker_input_ids = prompt_token_ids
+        thinker_sequences = prompt_token_ids + output_ids
+        thinker_input_ids = prompt_token_ids
         info_for_len = {"ids": {"all": thinker_sequences, "prompt": thinker_input_ids}}
         prompt_len = _compute_talker_prompt_ids_length(info_for_len, device="cpu")
 
@@ -758,6 +835,15 @@ def talker2code2wav_async_chunk(
     if isinstance(code_predictor_codes, torch.Tensor):
         if not code_predictor_codes.any():
             return None
+        sampling_params = getattr(request, "sampling_params", None)
+        stop_token_ids = set(getattr(sampling_params, "stop_token_ids", None) or [])
+        stop_token_id = getattr(sampling_params, "stop_token_id", None)
+        if stop_token_id is not None:
+            stop_token_ids.add(stop_token_id)
+        first_codebook = int(code_predictor_codes[0, 0].item())
+        if first_codebook in stop_token_ids:
+            logger.debug("skip stop-token codec frame: first_codebook=%s", first_codebook)
+            return None
     else:
         code_tensor = torch.tensor(code_predictor_codes, dtype=torch.long)
         if not code_tensor.any():
@@ -797,7 +883,14 @@ def talker2code2wav_full_payload(
     request: OmniEngineCoreRequest,
 ) -> dict[str, Any] | None:
     """Pack complete talker codec output for the non-async connector path."""
+    rid = getattr(request, "request_id", None)
     if not isinstance(pooling_output, dict):
+        logger.warning(
+            "talker2code2wav_full_payload: pooling_output not a dict "
+            "(type=%s) for req=%s; consumer wait gate may hang.",
+            type(pooling_output).__name__,
+            rid,
+        )
         return None
     code_predictor_codes = pooling_output.get("codes.audio")
     if code_predictor_codes is None:
@@ -805,10 +898,19 @@ def talker2code2wav_full_payload(
         if isinstance(codes, dict):
             code_predictor_codes = codes.get("audio")
     if code_predictor_codes is None:
+        logger.warning(
+            "talker2code2wav_full_payload: missing codes.audio (keys=%s) for req=%s; consumer wait gate may hang.",
+            list(pooling_output.keys()),
+            rid,
+        )
         return None
     if not isinstance(code_predictor_codes, torch.Tensor):
         code_predictor_codes = torch.as_tensor(code_predictor_codes)
     if code_predictor_codes.numel() == 0:
+        logger.warning(
+            "talker2code2wav_full_payload: empty codes.audio for req=%s; consumer wait gate may hang.",
+            rid,
+        )
         return None
 
     output_token_ids = _ensure_list(getattr(request, "output_token_ids", []) or [])
@@ -818,6 +920,16 @@ def talker2code2wav_full_payload(
         list(output_token_ids),
     )
     if code_predictor_codes.numel() == 0:
+        logger.warning(
+            "talker2code2wav_full_payload: no valid codec rows after filtering "
+            "(raw_shape=%s output_ids_len=%d aligned_rows=%s valid_rows=%s) for req=%s; "
+            "consumer wait gate may hang.",
+            raw_shape,
+            len(output_token_ids),
+            codec_stats["aligned_rows"],
+            codec_stats["valid_rows"],
+            rid,
+        )
         return None
 
     codec_codes = code_predictor_codes.transpose(0, 1).cpu().reshape(-1).tolist()
diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
index faa7e4cc4..1ffbdb931 100644
--- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
+++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
@@ -53,10 +53,13 @@ def talker2code2wav(
         # audio_codes may still contain zero-padded / invalid rows, so trim only
         # after filtering valid frames instead of trying to align EOS indices.
         seq_len = max(len(token_ids) - 1, 0)
-        # Filter invalid frames: zero-padded (EOS) and frames containing
-        # out-of-range values (e.g. stop_token_id=2150 exceeds codebook_size=2048).
+        # Filter invalid frames: zero-padded (EOS), out-of-range values (e.g.
+        # stop_token_id=2150 exceeds codebook_size=2048), and negative
+        # sentinels (e.g. -1 padding).
         _CODEBOOK_SIZE = 2048
-        valid_mask = audio_codes.any(dim=1) & (audio_codes.max(dim=1).values < _CODEBOOK_SIZE)
+        valid_mask = (
+            (audio_codes >= 0).all(dim=1) & audio_codes.any(dim=1) & (audio_codes.max(dim=1).values < _CODEBOOK_SIZE)
+        )
         audio_codes = audio_codes[valid_mask]
         if seq_len > 0 and audio_codes.ndim == 2 and int(audio_codes.shape[0]) > seq_len:
             audio_codes = audio_codes[-seq_len:]
@@ -280,3 +283,220 @@ def talker2code2wav_async_chunk(
         speaker=extract_speaker_from_request(request),
         language=extract_language_from_request(request),
     )
+
+
+# ============================================================================
+# Worker-connector data plane (non-async-chunk path).
+# AR runner's `flatten_payload` converts the model emit
+# `multimodal_outputs={"codes": {"audio": ..., "ref": ...},
+# "meta": {"ref_code_len": ..., "codec_streaming": ...}}` to flat dotted
+# keys (`codes.audio`, `codes.ref`, `meta.ref_code_len`,
+# `meta.codec_streaming`) before the full-payload accumulator runs.
+# - codes.audio is 2-D so default CONCAT across steps builds the full sequence.
+# - codes.ref is a list (not Tensor with dim>=2) so accumulator LATEST-wins
+#   keeps the prefill-emitted ref tensor across decode steps (which don't emit
+#   ref again).
+# - meta.ref_code_len is 1-D so LATEST-wins; consumer reads [-1].
+# ============================================================================
+
+# Per-model REPLACE-keys for the full-payload accumulator.  qwen3_tts's
+# producer side emits codec frames that should CONCAT (codes.audio) plus
+# scalars/lists that are correctly handled by default LATEST-wins, so this
+# stays empty.
+_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset()
+
+_CODEBOOK_SIZE = 2048
+_NUM_QUANTIZERS_DEFAULT = 16
+
+
+def _filter_audio_codes_qwen3_tts(audio_codes: torch.Tensor) -> torch.Tensor:
+    """Filter zero-padded, out-of-range, and negative-padded codec frames.
+
+    Mirrors the orchestrator-path body in `talker2code2wav` above.
+    """
+    if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
+        return audio_codes
+    if audio_codes.ndim != 2:
+        return audio_codes
+    valid_mask = (
+        (audio_codes >= 0).all(dim=1) & audio_codes.any(dim=1) & (audio_codes.max(dim=1).values < _CODEBOOK_SIZE)
+    )
+    return audio_codes[valid_mask]
+
+
+def _coerce_ref_code_len(raw) -> int:
+    """Coerce mm["meta"]["ref_code_len"] / pooling_output["meta.ref_code_len"]
+    raw value (Tensor | int | None) into a non-negative int.  Mirrors the
+    extraction inlined in the legacy ``talker2code2wav`` path; clamps any
+    negative input to 0 since downstream code treats this as a non-negative
+    frame count."""
+    if isinstance(raw, torch.Tensor):
+        value = int(raw.reshape(-1)[-1].item()) if raw.numel() > 0 else 0
+    elif raw is None:
+        value = 0
+    else:
+        value = int(raw)
+    return max(value, 0)
+
+
+def _normalize_ref_code(ref_code, num_quantizers: int, ref_code_len: int):
+    """Coerce ref_code into a [ref_len, Q] tensor or None.  Mirrors orchestrator path."""
+    if isinstance(ref_code, list):
+        ref_code = ref_code[0] if ref_code else None
+    if not isinstance(ref_code, torch.Tensor) or ref_code.numel() == 0:
+        return None, 0
+    ref_code = ref_code.to(torch.long).cpu().contiguous()
+    if ref_code.ndim == 1:
+        if ref_code.numel() % num_quantizers != 0:
+            return None, 0
+        ref_code = ref_code.reshape(-1, num_quantizers)
+    elif ref_code.ndim != 2:
+        return None, 0
+    if ref_code_len > 0 and int(ref_code.shape[0]) > ref_code_len:
+        ref_code = ref_code[:ref_code_len]
+    return ref_code, int(ref_code.shape[0])
+
+
+def talker2code2wav_token_only(
+    source_outputs: list,
+    prompt=None,
+    _requires_multimodal_data: bool = False,
+) -> list:
+    """Sync-side placeholder for the non-async-chunk Stage-1 (code2wav) input.
+
+    Sized to the expected codec token count (codebook-major flat:
+    Q * (ref_frames + audio_frames)).  Speaker / language metadata are
+    extracted from `prompt` and threaded via `additional_information`
+    (orchestrator-style; same as the legacy `talker2code2wav` builder).
+    Actual codec ids are delivered via the worker connector payload built
+    by `talker2code2wav_full_payload`.
+    """
+    from vllm_omni.inputs.data import OmniTokensPrompt
+
+    code2wav_inputs: list = []
+    for i, talker_output in enumerate(source_outputs):
+        if not talker_output.finished:
+            continue
+        output = talker_output.outputs[0]
+        mm = output.multimodal_output if hasattr(output, "multimodal_output") else None
+        mm = mm if isinstance(mm, dict) else {}
+        mm_codes = mm.get("codes", {}) if isinstance(mm, dict) else {}
+        token_ids = getattr(output, "cumulative_token_ids", []) or []
+        seq_len = max(len(token_ids) - 1, 0)
+
+        audio = mm_codes.get("audio") if isinstance(mm_codes, dict) else None
+        if isinstance(audio, torch.Tensor) and audio.numel() > 0:
+            audio = audio.to(torch.long)
+            audio = _filter_audio_codes_qwen3_tts(audio)
+            if seq_len > 0 and audio.ndim == 2 and int(audio.shape[0]) > seq_len:
+                audio = audio[-seq_len:]
+            num_audio_frames = int(audio.shape[0]) if audio.ndim == 2 else 0
+            num_quantizers = int(audio.shape[1]) if audio.ndim == 2 and audio.shape[1] > 0 else _NUM_QUANTIZERS_DEFAULT
+        else:
+            num_audio_frames = 0
+            num_quantizers = _NUM_QUANTIZERS_DEFAULT
+
+        ref_code_raw = mm_codes.get("ref") if isinstance(mm_codes, dict) else None
+        ref_code_len_raw = mm.get("meta", {}).get("ref_code_len") if isinstance(mm.get("meta"), dict) else None
+        ref_code_len = _coerce_ref_code_len(ref_code_len_raw)
+        _, ref_frames = _normalize_ref_code(ref_code_raw, num_quantizers, ref_code_len)
+
+        # Codebook-major flat: Q * (ref_frames + audio_frames)
+        prompt_len = num_quantizers * (ref_frames + num_audio_frames)
+
+        additional_info = to_dict(
+            OmniPayloadStruct(
+                meta=MetaStruct(left_context_size=ref_frames) if ref_frames > 0 else None,
+                speaker=extract_speaker_from_prompt(prompt, index=i),
+                language=extract_language_from_prompt(prompt, index=i),
+            )
+        )
+        code2wav_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[0] * prompt_len,
+                additional_information=additional_info if additional_info else None,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+            )
+        )
+    return code2wav_inputs
+
+
+def talker2code2wav_full_payload(
+    transfer_manager,
+    pooling_output,
+    request,
+):
+    """Producer-side payload builder.
+
+    Reads accumulated codec from `pooling_output["codes.audio"]` (CONCAT
+    across steps via flatten_payload), latest `pooling_output["codes.ref"]`
+    (prefill-emitted), and latest `pooling_output["meta.ref_code_len"]`.
+    Replicates the orchestrator-path body of `talker2code2wav` (filter,
+    crop to seq_len, prepend ref, codebook-major flatten).
+    """
+    del transfer_manager
+    rid = getattr(request, "request_id", "?")
+    if not isinstance(pooling_output, dict):
+        logger.warning(
+            "qwen3_tts.talker2code2wav_full_payload: pooling_output not a dict "
+            "(type=%s) for req=%s; consumer wait gate may hang.",
+            type(pooling_output).__name__,
+            rid,
+        )
+        return None
+
+    # codes.audio — try flat dotted first (flatten_payload), then nested fallback.
+    audio = pooling_output.get("codes.audio")
+    if audio is None:
+        codes_nested = pooling_output.get("codes")
+        if isinstance(codes_nested, dict):
+            audio = codes_nested.get("audio")
+    if not isinstance(audio, torch.Tensor) or audio.numel() == 0:
+        logger.warning(
+            "qwen3_tts.talker2code2wav_full_payload: missing/empty codes.audio "
+            "(keys=%s) for req=%s; consumer wait gate may hang.",
+            list(pooling_output.keys()),
+            rid,
+        )
+        return None
+    audio = audio.to(torch.long)
+    audio = _filter_audio_codes_qwen3_tts(audio)
+    if audio.numel() == 0:
+        logger.warning(
+            "qwen3_tts.talker2code2wav_full_payload: audio empty after codec "
+            "filter (negative/all-zero/out-of-range rows dropped) for req=%s.",
+            rid,
+        )
+        return None
+
+    output_token_ids = list(getattr(request, "output_token_ids", None) or [])
+    seq_len = max(len(output_token_ids) - 1, 0)
+    if seq_len > 0 and audio.ndim == 2 and int(audio.shape[0]) > seq_len:
+        audio = audio[-seq_len:]
+
+    num_quantizers = int(audio.shape[1]) if audio.ndim == 2 and audio.shape[1] > 0 else _NUM_QUANTIZERS_DEFAULT
+
+    # meta.ref_code_len — flat dotted then nested fallback.
+    ref_code_len_raw = pooling_output.get("meta.ref_code_len")
+    if ref_code_len_raw is None:
+        meta_nested = pooling_output.get("meta")
+        if isinstance(meta_nested, dict):
+            ref_code_len_raw = meta_nested.get("ref_code_len")
+    ref_code_len = _coerce_ref_code_len(ref_code_len_raw)
+
+    # codes.ref — flat dotted then nested fallback.
+    ref_code_raw = pooling_output.get("codes.ref")
+    if ref_code_raw is None:
+        codes_nested = pooling_output.get("codes")
+        if isinstance(codes_nested, dict):
+            ref_code_raw = codes_nested.get("ref")
+    ref_code, ref_frames = _normalize_ref_code(ref_code_raw, num_quantizers, ref_code_len)
+    if ref_code is not None:
+        audio = torch.cat([ref_code.to(audio.device), audio], dim=0)
+
+    codec_codes = audio.transpose(0, 1).to(device="cpu", dtype=torch.long).reshape(-1).contiguous()
+    return {
+        "codes": {"audio": codec_codes},
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+    }
diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py
index a3cc32c52..63002a98d 100644
--- a/vllm_omni/platforms/cuda/platform.py
+++ b/vllm_omni/platforms/cuda/platform.py
@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import importlib
+
 import torch
 import vllm.envs as envs
 from vllm.config import VllmConfig
@@ -58,14 +60,32 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
     ) -> str:
         from vllm_omni.diffusion.envs import PACKAGES_CHECKER
 
-        # Check compute capability for Flash Attention support
-        # Flash Attention requires compute capability >= 8.0 and < 10.0
+        # Check compute capability for Flash Attention support.
+        # FA requires sm_80+. Blackwell (sm_10x/sm_12x) only works with FA builds
+        # that include the Blackwell CUTE kernel — plain FA2 will crash there.
+        #
+        # Known Blackwell SKUs:
+        #   sm_100 = B200 / GB200 (datacenter)
+        #   sm_103 = B300 / GB300 (Blackwell Ultra)
+        #   sm_120 = RTX Pro 6000, RTX 50-series (consumer)
+        #   sm_121 = consumer Blackwell refresh
+        _known_blackwell_sms = {(10, 0), (10, 3), (12, 0), (12, 1)}
         compute_capability = cls.get_device_capability()
         compute_supported = False
+        is_blackwell = False
+        sm_str = ""
         if compute_capability is not None:
             major, minor = compute_capability
             capability = major * 10 + minor
-            compute_supported = 80 <= capability < 100
+            compute_supported = capability >= 80
+            sm_str = f"sm_{major}{minor}"
+            # Accept major in {10, 11, 12} to cover future Blackwell refreshes.
+            is_blackwell = major in (10, 11, 12)
+            if is_blackwell and (major, minor) not in _known_blackwell_sms:
+                logger.info(
+                    "Detected Blackwell-class GPU %s (untested variant); routing to CUDNN_ATTN with SDPA fallback.",
+                    sm_str,
+                )
 
         # Check if FA packages are available
         packages_info = PACKAGES_CHECKER.get_packages_info()
@@ -74,22 +94,86 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
         # Both compute capability and packages must be available for FA
         flash_attn_supported = compute_supported and packages_available
 
+        # cuDNN 9.5+ ships Blackwell FMHA kernels. If the runtime is older,
+        # the CUDNN_ATTN default would still work via internal fallback but
+        # without the tuned Blackwell path, so we skip routing there.
+        cudnn_version = torch.backends.cudnn.version() or 0
+        cudnn_blackwell_ready = cudnn_version >= 90500
+
+        # FlashInfer edges cuDNN by ~4% at the kernel level on sm_120 but
+        # regresses ~2x at e2e on HV-1.5 because its dense-prefill path can't
+        # take 2D attn_masks and the SDPA fallback dispatches to
+        # EFFICIENT_ATTENTION (~25 ms) instead of the cuDNN mask path (~11 ms).
+        # CUDNN_ATTN pins sdpa_kernel([CUDNN_ATTENTION]) directly so masked
+        # calls keep the cuDNN path. Blackwell default prefers CUDNN_ATTN;
+        # users can opt into FLASHINFER_ATTN explicitly for no-mask workloads.
+        flashinfer_available = False
+        try:
+            import flashinfer  # noqa: F401
+
+            flashinfer_available = True
+        except Exception as e:
+            # A partially installed / ABI-mismatched wheel can raise OSError or
+            # RuntimeError from extension loading, not just ImportError. This
+            # runs during default backend selection, so a probe failure must
+            # not abort startup — just treat FlashInfer as unavailable.
+            logger.debug("FlashInfer probe failed (%s); treating as unavailable", e)
+
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
             if backend_upper == "FLASH_ATTN" and not flash_attn_supported:
                 if not compute_supported:
                     logger.warning(
-                        "Flash Attention requires GPU with compute capability >= 8.0 "
-                        "and < 10.0. Falling back to TORCH_SDPA backend."
+                        "Flash Attention requires GPU with compute capability >= 8.0. "
+                        "Falling back to TORCH_SDPA backend."
                     )
                 elif not packages_available:
                     logger.warning("Flash Attention packages not available. Falling back to TORCH_SDPA backend.")
                 logger.debug("Defaulting to diffusion attention backend SDPA")
                 return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+            if backend_upper == "SAGE_ATTN_3":
+                sage_attn3_supported = compute_capability is not None and compute_capability.major >= 10
+                if not sage_attn3_supported:
+                    logger.warning(
+                        "SageAttention3 requires a Blackwell-class GPU with compute capability >= 10.0. "
+                        "Falling back to TORCH_SDPA backend."
+                    )
+                    return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+                try:
+                    importlib.import_module("sageattn3")
+                except ImportError:
+                    logger.warning(
+                        "SageAttention3 package not available. Install it from "
+                        "SageAttention/sageattention3_blackwell. Falling back to TORCH_SDPA backend."
+                    )
+                    return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
             backend = DiffusionAttentionBackendEnum[backend_upper]
             logger.debug("Using diffusion attention backend '%s'", backend_upper)
             return backend.get_path()
 
+        if is_blackwell and cudnn_blackwell_ready:
+            logger.info(
+                "Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell %s, cuDNN %d)",
+                sm_str,
+                cudnn_version,
+            )
+            return DiffusionAttentionBackendEnum.CUDNN_ATTN.get_path()
+
+        if is_blackwell and flashinfer_available:
+            logger.info(
+                "Defaulting to diffusion attention backend FLASHINFER_ATTN (Blackwell %s, cuDNN unavailable)",
+                sm_str,
+            )
+            return DiffusionAttentionBackendEnum.FLASHINFER_ATTN.get_path()
+
+        if is_blackwell and not cudnn_blackwell_ready:
+            logger.warning(
+                "Detected Blackwell %s but cuDNN %d < 9.5 — no tuned Blackwell FMHA. "
+                "Falling through to FLASH_ATTN / SDPA.",
+                sm_str,
+                cudnn_version,
+            )
+
         if flash_attn_supported:
             logger.debug("Defaulting to diffusion attention backend FLASH_ATTN")
             return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py
index f9ec84300..48c513da3 100644
--- a/vllm_omni/platforms/interface.py
+++ b/vllm_omni/platforms/interface.py
@@ -7,6 +7,8 @@ from typing import Any
 
 import torch
 import torch.nn as nn
+from vllm.config import CUDAGraphMode, VllmConfig
+from vllm.forward_context import BatchDescriptor
 from vllm.logger import init_logger
 from vllm.platforms import Platform
 
@@ -200,6 +202,40 @@ class OmniPlatform(Platform):
         """
         return "vllm_omni.profiler.omni_torch_profiler.OmniTorchProfilerWrapper"
 
+    @classmethod
+    def get_graph_wrapper_cls(cls) -> type:
+        """Return the platform's full-graph wrapper class.
+
+        Defaults to vLLM's CUDAGraphWrapper; NPU overrides with ACLGraphWrapper.
+        """
+        from vllm.compilation.cuda_graph import CUDAGraphWrapper
+
+        return CUDAGraphWrapper
+
+    @classmethod
+    def set_forward_context(
+        cls,
+        attn_metadata: Any,
+        vllm_config: VllmConfig,
+        *,
+        cudagraph_runtime_mode: CUDAGraphMode,
+        batch_descriptor: BatchDescriptor,
+    ):
+        """Platform-neutral wrapper around the device's set_forward_context.
+
+        Defaults to vLLM's ``set_forward_context``; NPU overrides to dispatch
+        to ``set_ascend_forward_context`` (renaming ``cudagraph_runtime_mode``
+        to ``aclgraph_runtime_mode``).
+        """
+        from vllm.forward_context import set_forward_context
+
+        return set_forward_context(
+            attn_metadata,
+            vllm_config,
+            cudagraph_runtime_mode=cudagraph_runtime_mode,
+            batch_descriptor=batch_descriptor,
+        )
+
 
 class UnspecifiedOmniPlatform(OmniPlatform):
     _omni_enum = OmniPlatformEnum.UNSPECIFIED
diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py
index e371c73a6..faf69c1b8 100644
--- a/vllm_omni/platforms/npu/platform.py
+++ b/vllm_omni/platforms/npu/platform.py
@@ -153,3 +153,27 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
     @classmethod
     def get_profiler_cls(cls) -> str:
         return "vllm_omni.platforms.npu.profiler.NPUTorchProfilerWrapper"
+
+    @classmethod
+    def get_graph_wrapper_cls(cls) -> type:
+        from vllm_ascend.compilation.acl_graph import ACLGraphWrapper
+
+        return ACLGraphWrapper
+
+    @classmethod
+    def set_forward_context(
+        cls,
+        attn_metadata,
+        vllm_config,
+        *,
+        cudagraph_runtime_mode,
+        batch_descriptor,
+    ):
+        from vllm_ascend.ascend_forward_context import set_ascend_forward_context
+
+        return set_ascend_forward_context(
+            attn_metadata,
+            vllm_config,
+            aclgraph_runtime_mode=cudagraph_runtime_mode,
+            batch_descriptor=batch_descriptor,
+        )
diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
index 82d628ae9..3a8556273 100644
--- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
+++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
@@ -155,6 +155,11 @@ class NPUARModelRunner(OmniNPUModelRunner):
         finally:
             set_cudagraph_capturing_enabled(False)
 
+    def _model_needs_full_prefix_hidden_states(self) -> bool:
+        """See gpu_ar_model_runner._model_needs_full_prefix_hidden_states."""
+        model = getattr(self, "model", None)
+        return bool(getattr(model, "requires_full_prefix_cached_hidden_states", True))
+
     def _maybe_update_prefix_cache(
         self,
         hidden_states: torch.Tensor,
@@ -169,8 +174,9 @@ class NPUARModelRunner(OmniNPUModelRunner):
                     type(multimodal_outputs),
                 )
 
+            hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None
             self.omni_prefix_cache.update_omni_tensor_prefix_cache(
-                hidden_states=hidden_states,
+                hidden_states=hs_for_cache,
                 multimodal_outputs=flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs,
                 num_tokens_unpadded=num_tokens_unpadded,
                 slot_mapping=self.input_batch.block_table[0].slot_mapping.cpu,
@@ -185,12 +191,13 @@ class NPUARModelRunner(OmniNPUModelRunner):
     ) -> tuple[dict[str, torch.Tensor] | None, dict | None]:
         combined_hidden_states, combined_multimodal_outputs = None, None
         if self.omni_prefix_cache is not None:
-            combined_hidden_states = self.omni_prefix_cache.get_merged_hidden_states(
-                query_start_loc=self.query_start_loc.cpu,
-                input_batch=self.input_batch,
-                hidden_states=hidden_states,
-                num_scheduled_tokens=num_scheduled_tokens,
-            )
+            if self._model_needs_full_prefix_hidden_states():
+                combined_hidden_states = self.omni_prefix_cache.get_merged_hidden_states(
+                    query_start_loc=self.query_start_loc.cpu,
+                    input_batch=self.input_batch,
+                    hidden_states=hidden_states,
+                    num_scheduled_tokens=num_scheduled_tokens,
+                )
             combined_multimodal_outputs = self.omni_prefix_cache.get_merged_multimodal_states(
                 query_start_loc=self.query_start_loc.cpu,
                 input_batch=self.input_batch,
@@ -730,6 +737,39 @@ class NPUARModelRunner(OmniNPUModelRunner):
 
         return None
 
+    def _sample(
+        self,
+        logits: torch.Tensor | None,
+        spec_decode_metadata: Any,
+    ):
+        sampling_metadata = self.input_batch.sampling_metadata
+        if spec_decode_metadata is None:
+            model_sample = getattr(self.model, "sample", None)
+            self.input_batch.update_async_output_token_ids()
+            if logits is not None and callable(model_sample) and getattr(self.model, "prefer_model_sampler", False):
+                # Apply logit bias (min_tokens, allowed_token_ids) before
+                # the custom model sampler — the standard GPU sampler does
+                # this internally, but prefer_model_sampler bypasses it.
+                if hasattr(self.sampler, "logit_bias_state"):
+                    self.sampler.logit_bias_state.apply_logit_bias(
+                        logits,
+                        self.input_batch.expanded_idx_mapping,
+                        self.input_batch.idx_mapping_np,
+                        self.input_batch.positions[self.input_batch.logits_indices],
+                    )
+                sampler_output = model_sample(
+                    logits,
+                    self._sampling_metadata_for_model_sampler(sampling_metadata),
+                )
+                if sampler_output is not None:
+                    return sampler_output
+            return self.sampler(
+                logits=logits,
+                sampling_metadata=sampling_metadata,
+            )
+
+        return super()._sample(logits, spec_decode_metadata)
+
     @torch.inference_mode()
     def sample_tokens(
         self, grammar_output: GrammarOutput | None
diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py
index c6963ac67..4137ffe92 100644
--- a/vllm_omni/platforms/npu/worker/npu_model_runner.py
+++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py
@@ -16,7 +16,6 @@ from vllm.v1.worker.gpu_model_runner import PerLayerAttnMetadata
 from vllm_ascend.ascend_forward_context import get_forward_context, set_ascend_forward_context
 from vllm_ascend.attention.attention_v1 import AscendAttentionState
 from vllm_ascend.attention.utils import using_paged_attention
-from vllm_ascend.compilation.acl_graph import ACLGraphWrapper
 from vllm_ascend.ops.rotary_embedding import update_cos_sin
 from vllm_ascend.utils import enable_sp, lmhead_tp_enable
 from vllm_ascend.worker.model_runner_v1 import SEQ_LEN_WITH_MAX_PA_WORKSPACE, NPUModelRunner
@@ -34,30 +33,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
         # in _pad_for_sequence_parallelism during execute_model.
         # This is a workaround for vllm-ascend not passing vllm_config to enable_sp().
         enable_sp(self.vllm_config)
-        # TODO move this model specific logic to a separate class
-        # TTS model IS the talker (no .talker sub-attr); use getattr to support both Omni and TTS.
-        self.has_talker_mtp = False
-        talker_mtp = getattr(self.model, "talker_mtp", None)
-        if talker_mtp is not None:
-            self.talker_mtp = talker_mtp  # type: ignore[assignment]
-            self.has_talker_mtp = True
-            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", False)
-            if cudagraph_mode.has_full_cudagraphs() and (has_separate_talker or talker_mtp_graph_safe):
-                self.talker_mtp = ACLGraphWrapper(talker_mtp, self.vllm_config, runtime_mode=CUDAGraphMode.FULL)
-            # TTS exposes mtp_hidden_size; Omni uses hf_text_config.hidden_size.
-            hidden_size = int(
-                getattr(self.model, "mtp_hidden_size", 0) or getattr(self.model_config.hf_text_config, "hidden_size")
-            )
-            max_batch_size = max(self.max_num_reqs, self.compilation_config.max_cudagraph_capture_size)
-            self.talker_mtp_input_ids = self._make_buffer(max_batch_size, dtype=torch.int32)
-            self.talker_mtp_inputs_embeds = self._make_buffer(
-                max_batch_size, hidden_size, dtype=self.dtype, numpy=False
-            )
-            self.last_talker_hidden = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
-            self.text_step = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
+        self._init_talker_mtp()
 
     @torch.inference_mode()
     def _dummy_run(
@@ -397,75 +373,3 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
             model_output = self._all_gather_hidden_states_and_aux(model_output)
 
         return model_output
-
-    def _talker_mtp_forward(self, decode_req_ids: list[str], inputs_embeds: torch.Tensor) -> None:
-        decode_batch_size = len(decode_req_ids)
-        if decode_batch_size == 0:
-            return
-        _cudagraph_mode, batch_desc, _, _, _ = self._determine_batch_execution_and_padding(
-            num_tokens=decode_batch_size,
-            num_reqs=decode_batch_size,
-            num_scheduled_tokens_np=np.ones(decode_batch_size, dtype=np.int32),
-            max_num_scheduled_tokens=1,
-            use_cascade_attn=False,
-        )
-        # Force eager for unwrapped code predictors (AR loops / multinomial).
-        if not isinstance(self.talker_mtp, ACLGraphWrapper):
-            _cudagraph_mode = CUDAGraphMode.NONE
-        num_tokens_padded = batch_desc.num_tokens
-        req_input_ids = self.talker_mtp_input_ids.gpu[:num_tokens_padded]
-        req_embeds = self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded]
-        last_talker_hidden = self.last_talker_hidden.gpu[:num_tokens_padded]
-        text_step = self.text_step.gpu[:num_tokens_padded]
-        subtalker_params = getattr(self.vllm_config.model_config, "subtalker_sampling_params", None)
-        if not isinstance(subtalker_params, dict):
-            subtalker_params = {}
-        # Extract seed from the first request's sampling params for Fast AR
-        # determinism.  NOTE: when batch_size > 1, all requests share the first
-        # request's seed.  Per-request Fast AR seeding requires row-by-row
-        # torch.multinomial calls and is left as a follow-up optimisation.
-        _seed = None
-        if decode_req_ids:
-            _first_sp = getattr(self.requests[decode_req_ids[0]], "sampling_params", None)
-            if _first_sp is not None and getattr(_first_sp, "seed", None) is not None:
-                _seed = _first_sp.seed
-            # Warn when batched requests have different seeds.
-            if len(decode_req_ids) > 1 and _seed is not None:
-                _other_seeds = {
-                    getattr(getattr(self.requests[rid], "sampling_params", None), "seed", None)
-                    for rid in decode_req_ids[1:]
-                }
-                if _other_seeds != {_seed}:
-                    logger.warning(
-                        "Fast AR seed: batch has mixed seeds; using first request's seed=%d for all %d requests.",
-                        _seed,
-                        len(decode_req_ids),
-                    )
-        talker_kwargs = {
-            "do_sample": subtalker_params.get("do_sample"),
-            "temperature": subtalker_params.get("temperature"),
-            "top_k": subtalker_params.get("top_k"),
-            "top_p": subtalker_params.get("top_p"),
-        }
-        if _seed is not None:
-            talker_kwargs["seed"] = _seed
-        with set_ascend_forward_context(
-            None, self.vllm_config, aclgraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc
-        ):
-            req_embeds, code_predictor_codes = self.talker_mtp(
-                req_input_ids,
-                req_embeds,
-                last_talker_hidden,
-                text_step,
-                **talker_kwargs,
-            )
-        # update the inputs_embeds and code_predictor_codes
-        out_key = getattr(self.model, "talker_mtp_output_key", ("codes", "audio"))
-        if not isinstance(out_key, tuple) or len(out_key) != 2:
-            raise TypeError(f"talker_mtp_output_key must be a 2-tuple, got {type(out_key).__name__}: {out_key!r}")
-        for idx, req_id in enumerate(decode_req_ids):
-            req_index = self.input_batch.req_ids.index(req_id)
-            start_offset = int(self.query_start_loc.cpu[req_index])
-            inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1]
-            update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}}
-            self._merge_additional_information_update(req_id, update_dict)
diff --git a/vllm_omni/platforms/rocm/patch/__init__.py b/vllm_omni/platforms/rocm/patch/__init__.py
new file mode 100644
index 000000000..ecf416ee3
--- /dev/null
+++ b/vllm_omni/platforms/rocm/patch/__init__.py
@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+
+
+def apply_patches():
+    """Apply all ROCm-specific patches."""
+    from vllm_omni.platforms.rocm.patch import worker  # noqa: F401
diff --git a/vllm_omni/platforms/rocm/patch/worker/__init__.py b/vllm_omni/platforms/rocm/patch/worker/__init__.py
new file mode 100644
index 000000000..824d33c5b
--- /dev/null
+++ b/vllm_omni/platforms/rocm/patch/worker/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from vllm_omni.platforms.rocm.patch.worker import patch_groupnorm  # noqa: F401
diff --git a/vllm_omni/platforms/rocm/patch/worker/patch_groupnorm.py b/vllm_omni/platforms/rocm/patch/worker/patch_groupnorm.py
new file mode 100644
index 000000000..7de866e73
--- /dev/null
+++ b/vllm_omni/platforms/rocm/patch/worker/patch_groupnorm.py
@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Patch ``initialize_model`` to replace VAE GroupNorm with AITER GroupNorm on ROCm."""
+
+import torch.nn as nn
+from vllm.logger import init_logger
+
+import vllm_omni.diffusion.registry as _registry_mod
+
+logger = init_logger(__name__)
+
+_original_initialize_model = _registry_mod.initialize_model
+
+
+def _replace_groupnorm_with_aiter(vae: nn.Module) -> bool:
+    from aiter.ops.groupnorm import GroupNorm as AiterGroupNorm
+
+    targets = [
+        (parent, name, child)
+        for parent in vae.modules()
+        for name, child in parent.named_children()
+        if isinstance(child, nn.GroupNorm) and child.affine
+    ]
+
+    for parent, name, child in targets:
+        new_group_norm = AiterGroupNorm(
+            num_groups=child.num_groups,
+            num_channels=child.num_channels,
+            eps=child.eps,
+            affine=True,
+            device=child.weight.device,
+            dtype=child.weight.dtype,
+        )
+        new_group_norm.weight = child.weight
+        new_group_norm.bias = child.bias
+        setattr(parent, name, new_group_norm)
+
+    return len(targets) > 0
+
+
+def _patched_initialize_model(od_config):
+    model = _original_initialize_model(od_config)
+
+    if hasattr(model, "vae"):
+        try:
+            from vllm._aiter_ops import is_aiter_found_and_supported
+
+            if is_aiter_found_and_supported() and _replace_groupnorm_with_aiter(model.vae):
+                logger.info("AITER GroupNorm is enabled for VAE.")
+        except Exception:
+            logger.warning("Failed to apply AITER GroupNorm to VAE.")
+
+    return model
+
+
+_registry_mod.initialize_model = _patched_initialize_model
diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py
index d7ad2b0f7..c2e3fae67 100644
--- a/vllm_omni/platforms/rocm/platform.py
+++ b/vllm_omni/platforms/rocm/platform.py
@@ -10,6 +10,7 @@ from vllm.platforms.rocm import RocmPlatform
 
 from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
 from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
+from vllm_omni.platforms.rocm.patch import apply_patches
 
 logger = init_logger(__name__)
 
@@ -51,6 +52,10 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform):
 
     _omni_enum = OmniPlatformEnum.ROCM
 
+    def __init__(self):
+        super().__init__()
+        apply_patches()
+
     @classmethod
     def get_omni_ar_worker_cls(cls) -> str:
         return "vllm_omni.worker.gpu_ar_worker.GPUARWorker"
diff --git a/vllm_omni/quantization/component_config.py b/vllm_omni/quantization/component_config.py
index f9286079b..25d1abf1e 100644
--- a/vllm_omni/quantization/component_config.py
+++ b/vllm_omni/quantization/component_config.py
@@ -23,10 +23,10 @@ if TYPE_CHECKING:
     )
 
 
-# Pre-quantized checkpoints (modelopt FP8/FP4/MXFP8) only quantize the
+# Pre-quantized checkpoints (modelopt FP8/FP4/MXFP8/mixed) only quantize the
 # Thinker LM.  Vision and audio encoder weights remain in BF16 with no
 # corresponding scale tensors in the checkpoint.
-PRE_QUANTIZED_METHODS: frozenset[str] = frozenset({"modelopt", "modelopt_fp4", "modelopt_mxfp8"})
+PRE_QUANTIZED_METHODS: frozenset[str] = frozenset({"modelopt", "modelopt_fp4", "modelopt_mxfp8", "modelopt_mixed"})
 
 
 def resolve_encoder_quant_config(
diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py
index 3766e4596..2a0326c2e 100644
--- a/vllm_omni/quantization/factory.py
+++ b/vllm_omni/quantization/factory.py
@@ -138,11 +138,19 @@ def _normalize_quant_method_alias(method: str | None) -> str | None:
 
 _MODEL_OPT_METHODS = {
     "modelopt",
+    "modelopt_fp4",
+    "modelopt_mixed",
 }
 _MODEL_OPT_FP8_ALGOS = {
     "FP8",
     "FP8_PER_CHANNEL_PER_TOKEN",
 }
+_MODEL_OPT_NVFP4_ALGOS = {
+    "NVFP4",
+}
+_MODEL_OPT_MIXED_ALGOS = {
+    "MIXED_PRECISION",
+}
 
 
 def _normalize_method_name(method: Any) -> str:
@@ -170,6 +178,10 @@ def _detect_modelopt_method(config: Mapping[str, Any]) -> str | None:
     if quant_algo:
         if quant_algo in _MODEL_OPT_FP8_ALGOS:
             return "modelopt"
+        if quant_algo in _MODEL_OPT_NVFP4_ALGOS:
+            return "modelopt_fp4"
+        if quant_algo in _MODEL_OPT_MIXED_ALGOS:
+            return "modelopt_mixed"
         return None
 
     if method is not None:
@@ -387,6 +399,12 @@ def resolve_quant_config_from_disk(
         )
         return build_quant_config(qc_method, **qc_kwargs)
 
+    # AutoRound MXFP8 checkpoints use data_type="mx_fp" instead of
+    # is_checkpoint_*_serialized; rebuild so the offline path is selected.
+    if qc_kwargs.get("data_type") == "mx_fp":
+        logger.info("config.json declares data_type='mx_fp'; rebuilding as offline AutoRound MXFP8.")
+        return build_quant_config(qc_method, **qc_kwargs)
+
     if (
         "ignored_layers" in qc_kwargs
         and hasattr(quant_config, "ignored_layers")
diff --git a/vllm_omni/quantization/inc_config.py b/vllm_omni/quantization/inc_config.py
index fb4851cca..d7dad3d71 100644
--- a/vllm_omni/quantization/inc_config.py
+++ b/vllm_omni/quantization/inc_config.py
@@ -7,8 +7,13 @@ from __future__ import annotations
 from os.path import commonprefix
 from typing import TYPE_CHECKING, Any
 
+import torch
+from torch.nn import Module
+from vllm.model_executor.layers.linear import LinearMethodBase
 from vllm.model_executor.layers.quantization.inc import INCConfig
 from vllm.model_executor.models.utils import WeightsMapper
+from vllm.model_executor.parameter import ModelWeightParameter
+from vllm.model_executor.utils import replace_parameter
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import (
@@ -45,12 +50,39 @@ def _map_with_stage_prefix(
 
 
 class OmniINCConfig(INCConfig):
-    """INCConfig extended with multi-stage prefix remapping."""
+    """INCConfig extended with multi-stage prefix remapping and MXFP8 support.
+
+    AutoRound MXFP8 checkpoints declare quant_method="auto-round" with data_type="mx_fp".
+    This config detects that case and dispatches to IncMxfp8OfflineLinearMethod instead
+    of the standard INT quantization path.
+
+    Architecture:
+      - AutoRound INT4/INT8 → standard INCConfig path (xpu_w4a16_quant_layer, etc.)
+      - AutoRound MXFP8 (data_type="mx_fp") → IncMxfp8OfflineLinearMethod
+      - Native MXFP8 (quant_method="mxfp8") → DiffusionMXFP8Config → NPUMxfp8LinearMethod (NPU only)
+    """
+
+    # Extend supported data types and formats to include MXFP8
+    SUPPORTED_DTYPES = {"int", "mx_fp"}
+    SUPPORTED_FORMATS = {"auto_round:auto_gptq", "auto_round:auto_awq", "auto_round:llm_compressor"}
 
     # ------------------------------------------------------------------
     # Core integration: called by vLLM's configure_quant_config()
     # ------------------------------------------------------------------
 
+    def get_quant_method(self, layer, prefix: str):
+        """Get quantization method, handling AutoRound MXFP8 as a special case."""
+        # Check if this is an AutoRound MXFP8 checkpoint (data_type="mx_fp")
+        if hasattr(self, "data_type") and self.data_type == "mx_fp" and self.weight_bits == 8:
+            from vllm.model_executor.layers.linear import LinearBase
+
+            if isinstance(layer, LinearBase):
+                return IncMxfp8OfflineLinearMethod()
+            return None
+
+        # Otherwise, use parent INCConfig logic
+        return super().get_quant_method(layer, prefix)
+
     def apply_vllm_mapper(self, hf_to_vllm_mapper: WeightsMapper) -> None:
         """Remap HF checkpoint names to vLLM runtime prefixes."""
         prefix_map = getattr(hf_to_vllm_mapper, "orig_to_new_prefix", None) or {}
@@ -138,3 +170,108 @@ class OmniINCConfig(INCConfig):
         if isinstance(quant_config, INCConfig):
             return cls.from_inc_config(quant_config)
         return quant_config
+
+
+# ---------------------------------------------------------------------------
+# IncMxfp8OfflineLinearMethod - used by AutoRound MXFP8 checkpoints
+# ---------------------------------------------------------------------------
+
+
+class IncMxfp8OfflineLinearMethod(LinearMethodBase):
+    """Offline MXFP8 linear method for AutoRound MXFP8 checkpoints.
+
+    For pre-quantized AutoRound MXFP8 checkpoints (data_type="mx_fp") where weights
+    are stored as BF16 (losslessly representing FP8 E4M3 values) with uint8 MX block scales.
+    Delegates to the platform-appropriate kernel (XPUMxFp8LinearKernel, etc.).
+
+    Used by OmniINCConfig when data_type="mx_fp".
+    """
+
+    def __init__(self) -> None:
+        from vllm.model_executor.kernels.linear import init_mxfp8_linear_kernel
+        from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+            MXFP8_BLOCK_SIZE,
+            MXFP8_SCALE_DTYPE,
+        )
+
+        self.kernel = init_mxfp8_linear_kernel()
+        self.block_size = MXFP8_BLOCK_SIZE
+        self.scale_dtype = MXFP8_SCALE_DTYPE
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: list[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ) -> None:
+        if input_size_per_partition % self.block_size != 0:
+            raise ValueError(
+                f"MXFP8 requires input_size_per_partition "
+                f"({input_size_per_partition}) to be divisible by "
+                f"{self.block_size}."
+            )
+
+        output_size_per_partition = sum(output_partition_sizes)
+        weight_loader = extra_weight_attrs.get("weight_loader")
+
+        layer.logical_widths = output_partition_sizes
+        layer.input_size_per_partition = input_size_per_partition
+        layer.output_size_per_partition = output_size_per_partition
+        layer.orig_dtype = params_dtype
+
+        # Allocate as params_dtype (BF16) to match checkpoint dtype.
+        # Will be cast to FP8 in process_weights_after_loading.
+        layer.register_parameter(
+            "weight",
+            ModelWeightParameter(
+                data=torch.empty(output_size_per_partition, input_size_per_partition, dtype=params_dtype),
+                input_dim=1,
+                output_dim=0,
+                weight_loader=weight_loader,
+            ),
+        )
+
+        num_groups = input_size_per_partition // self.block_size
+        layer.register_parameter(
+            "weight_scale",
+            ModelWeightParameter(
+                data=torch.empty(output_size_per_partition, num_groups, dtype=self.scale_dtype),
+                input_dim=1,
+                output_dim=0,
+                weight_loader=weight_loader,
+            ),
+        )
+
+    def process_weights_after_loading(self, layer: Module) -> None:
+        if getattr(layer, "_already_called_process_weights_after_loading", False):
+            return
+
+        # Cast BF16 weight to FP8.
+        # Checkpoint stores weights as BF16 (losslessly representing FP8 E4M3 values).
+        w = layer.weight.data
+        if w.dtype != torch.float8_e4m3fn:
+            # Cast to FP8 E4M3. AutoRound stores FP8 values as BF16, so this is lossless.
+            w = w.to(torch.float8_e4m3fn).contiguous()
+            replace_parameter(layer, "weight", w)
+
+        # Delegate layout transforms (transpose, scale reinterpret) to the platform kernel.
+        self.kernel.process_weights_after_loading(layer)
+        layer._already_called_process_weights_after_loading = True
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        ori_shape = x.shape
+        if x.dim() > 2:
+            x = x.reshape(-1, ori_shape[-1])
+        output = self.kernel.apply_weights(layer, x, bias)
+        if len(ori_shape) > 2:
+            output = output.reshape(*ori_shape[:-1], -1)
+        return output
diff --git a/vllm_omni/quantization/mxfp8_config.py b/vllm_omni/quantization/mxfp8_config.py
index 544c6a4ca..9d3216000 100644
--- a/vllm_omni/quantization/mxfp8_config.py
+++ b/vllm_omni/quantization/mxfp8_config.py
@@ -133,10 +133,17 @@ class DiffusionMXFP8Config(QuantizationConfig):
                 if self.is_checkpoint_mxfp8_serialized:
                     return NPUMxfp8LinearMethod(self)
                 return NPUMxfp8OnlineLinearMethod(self)
-            # Placeholder for future platforms: add `elif is_cuda(): ...` here.
+            if current_omni_platform.is_xpu():
+                if self.is_checkpoint_mxfp8_serialized:
+                    raise NotImplementedError(
+                        "Native MXFP8 offline mode is not supported on XPU. "
+                        "Use AutoRound MXFP8 checkpoints (quant_method='auto-round', data_type='mx_fp') instead, "
+                        "or use online MXFP8 mode without is_checkpoint_mxfp8_serialized."
+                    )
+                return VllmMxfp8OnlineLinearMethod()
             raise NotImplementedError(
                 "DiffusionMXFP8Config (W8A8 MXFP8) is currently only supported "
-                "on NPU (Ascend) platforms. CUDA support is not yet implemented."
+                "on NPU (Ascend) and XPU (Intel) platforms."
             )
         return None
 
@@ -450,3 +457,64 @@ class NPUMxfp8OnlineLinearMethod(_LazyWeightMixin, NPUMxfp8LinearMethod):
         replace_parameter(layer, "weight", weight_fp8)
         replace_parameter(layer, "weight_scale", weight_scale)
         layer._already_called_process_weights_after_loading = True
+
+
+# ---------------------------------------------------------------------------
+# vLLM kernel-based methods (XPU, CUDA, ROCm — any platform with Mxfp8LinearKernel)
+# ---------------------------------------------------------------------------
+# Note: VllmMxfp8OfflineLinearMethod removed - XPU only supports AutoRound MXFP8 offline
+#       (see IncMxfp8OfflineLinearMethod in inc_config.py)
+
+try:
+    from vllm.model_executor.layers.quantization.online.mxfp8 import (
+        Mxfp8OnlineLinearMethod as _VllmMxfp8OnlineBase,
+    )
+except ImportError as e:
+    raise ImportError(
+        "vLLM MXFP8 support is not available. "
+        "Online MXFP8 quantization requires vLLM with MXFP8 kernel support. "
+        "Please upgrade vLLM or use a compatible build."
+    ) from e
+
+
+class VllmMxfp8OnlineLinearMethod(_VllmMxfp8OnlineBase):
+    """Online MXFP8 linear method extending vLLM's implementation with 3D tensor support.
+
+    Loads BF16 weights, quantizes to MXFP8 at load time using vLLM's kernel.
+    Adds reshape logic to handle 3D tensors from diffusion models.
+
+    Inherits from vLLM's Mxfp8OnlineLinearMethod and only overrides __init__ and apply().
+    __init__ directly initializes the kernel without calling super().__init__() to avoid
+    vLLM config dependency in unit tests. All other methods (create_weights,
+    process_weights_after_loading) are inherited directly from vLLM.
+    """
+
+    def __init__(self) -> None:
+        """Initialize kernel directly without calling parent __init__ to avoid config dependency.
+
+        TODO: skipping super().__init__() couples this to vLLM internals. If the parent
+        adds required initialization, consider a version-aware assertion or calling
+        super().__init__() with a fallback config.
+        """
+        from vllm.model_executor.kernels.linear import init_mxfp8_linear_kernel
+
+        self.kernel = init_mxfp8_linear_kernel()
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Apply with 3D tensor support for diffusion models.
+
+        vLLM's kernel expects 2D tensors (M, K), but diffusion models pass
+        3D tensors (batch, seq, hidden). Flatten before kernel, reshape after.
+        """
+        ori_shape = x.shape
+        if x.dim() > 2:
+            x = x.reshape(-1, ori_shape[-1])
+        output = super().apply(layer, x, bias)
+        if len(ori_shape) > 2:
+            output = output.reshape(*ori_shape[:-1], -1)
+        return output
diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py
index 708d9c2b6..ba6d329ea 100644
--- a/vllm_omni/transformers_utils/configs/__init__.py
+++ b/vllm_omni/transformers_utils/configs/__init__.py
@@ -17,12 +17,14 @@ _CLASS_TO_MODULE: dict[str, str] = {
     "FishSpeechConfig": "vllm_omni.transformers_utils.configs.fish_speech",
     "FishSpeechSlowARConfig": "vllm_omni.transformers_utils.configs.fish_speech",
     "FishSpeechFastARConfig": "vllm_omni.transformers_utils.configs.fish_speech",
+    "GLMTTSConfig": "vllm_omni.transformers_utils.configs.glm_tts",
     "VoxCPM2Config": "vllm_omni.transformers_utils.configs.voxcpm2",
     "VoxtralTTSConfig": "vllm_omni.transformers_utils.configs.voxtral_tts",
     "CosyVoice3Config": "vllm_omni.transformers_utils.configs.cosyvoice3",
     "OmniVoiceConfig": "vllm_omni.transformers_utils.configs.omnivoice",
     "BailingMoeV2Config": "vllm_omni.transformers_utils.configs.ming_flash_omni",
     "BailingMM2Config": "vllm_omni.transformers_utils.configs.ming_flash_omni",
+    "MingImageGenConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni",
     "MingFlashOmniConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni",
     "Qwen3VLMoeVisionConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni",
     "WhisperEncoderConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni",
@@ -36,12 +38,14 @@ __all__ = [
     "FishSpeechConfig",
     "FishSpeechSlowARConfig",
     "FishSpeechFastARConfig",
+    "GLMTTSConfig",
     "VoxCPM2Config",
     "VoxtralTTSConfig",
     "CosyVoice3Config",
     "OmniVoiceConfig",
     "BailingMoeV2Config",
     "BailingMM2Config",
+    "MingImageGenConfig",
     "MingFlashOmniConfig",
     "Qwen3VLMoeVisionConfig",
     "WhisperEncoderConfig",
@@ -65,6 +69,7 @@ def __dir__():
 # run as soon as `vllm_omni.transformers_utils.configs` is imported.
 from vllm_omni.transformers_utils.configs import cosyvoice3 as _cosyvoice3  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech  # noqa: F401, E402
+from vllm_omni.transformers_utils.configs import glm_tts as _glm_tts  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import mammoth_moda2 as _mammoth_moda2  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import ming_flash_omni as _ming_flash_omni  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import omnivoice as _omnivoice  # noqa: F401, E402
diff --git a/vllm_omni/transformers_utils/configs/glm_tts.py b/vllm_omni/transformers_utils/configs/glm_tts.py
new file mode 100644
index 000000000..9a7f8dfbf
--- /dev/null
+++ b/vllm_omni/transformers_utils/configs/glm_tts.py
@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GLM-TTS config registration with transformers AutoConfig.
+
+Registers GLMTTSConfig (model_type="glm_tts") so that
+``AutoConfig.from_pretrained("path/to/glm-tts")`` returns the correct config class.
+
+Note: GLM-TTS uses a Llama backbone, but we register a custom config
+to handle the special token IDs and flow model parameters.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from transformers import AutoConfig, PretrainedConfig
+
+
+class GLMTTSConfig(PretrainedConfig):
+    """Llama-based AR model for text-to-speech token generation.
+
+    Special token IDs are loaded dynamically from the tokenizer at init time.
+    """
+
+    model_type: str = "glm_tts"
+
+    def __init__(
+        self,
+        vocab_size: int = 98304,
+        hidden_size: int = 2048,
+        intermediate_size: int = 6144,
+        num_hidden_layers: int = 28,
+        num_attention_heads: int = 16,
+        num_key_value_heads: int = 4,
+        hidden_act: str = "silu",
+        max_position_embeddings: int = 8192,
+        rms_norm_eps: float = 1e-5,
+        rope_theta: float = 10000.0,
+        audio_token_start: int = -1,
+        audio_token_end: int = -1,
+        boa_token_id: int = -1,
+        eoa_token_id: int = -1,
+        pad_token_id: int = -1,
+        bos_token_id: int = -1,
+        speech_token_vocab_size: int = 32768,
+        speech_token_dim: int = 512,
+        spk_embed_dim: int = 192,
+        mel_dim: int = 80,
+        input_frame_rate: float = 25.0,
+        mel_framerate: int = 50,
+        max_token_text_ratio: float = 20.0,
+        min_token_text_ratio: float = 2.0,
+        sample_method: str = "ras",
+        ras_top_p: float = 0.8,
+        ras_top_k: int = 25,
+        ras_win_size: int = 10,
+        ras_tau_r: float = 0.1,
+        **kwargs: Any,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_theta = rope_theta
+        self.audio_token_start = audio_token_start
+        self.audio_token_end = audio_token_end
+        self.boa_token_id = boa_token_id
+        self.eoa_token_id = eoa_token_id
+        self.speech_token_vocab_size = speech_token_vocab_size
+        self.speech_token_dim = speech_token_dim
+        self.spk_embed_dim = spk_embed_dim
+        self.mel_dim = mel_dim
+        self.input_frame_rate = input_frame_rate
+        self.mel_framerate = mel_framerate
+        self.max_token_text_ratio = max_token_text_ratio
+        self.min_token_text_ratio = min_token_text_ratio
+        self.sample_method = sample_method
+        self.ras_top_p = ras_top_p
+        self.ras_top_k = ras_top_k
+        self.ras_win_size = ras_win_size
+        self.ras_tau_r = ras_tau_r
+        super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, **kwargs)
+
+
+AutoConfig.register("glm_tts", GLMTTSConfig)
+
+__all__ = [
+    "GLMTTSConfig",
+]
diff --git a/vllm_omni/transformers_utils/configs/ming_flash_omni.py b/vllm_omni/transformers_utils/configs/ming_flash_omni.py
index 408b20868..376d902e3 100644
--- a/vllm_omni/transformers_utils/configs/ming_flash_omni.py
+++ b/vllm_omni/transformers_utils/configs/ming_flash_omni.py
@@ -260,6 +260,61 @@ class BailingMM2Config(PretrainedConfig):
         return self.llm_config
 
 
+class MingImageGenConfig(PretrainedConfig):
+    """Configuration for Ming-flash-omni-2.0 image generation stage.
+
+    Mirrors the layout of the HF checkpoint at
+    https://huggingface.co/inclusionAI/Ming-flash-omni-2.0 where image-gen
+    components live in sibling subfolders (``connector/``, ``transformer/``,
+    ``vae/``, ``scheduler/``, ``mlp/``).
+    """
+
+    model_type = "ming_flash_omni_imagegen"
+
+    def __init__(
+        self,
+        connector_subfolder: str = "connector",
+        transformer_subfolder: str = "transformer",
+        vae_subfolder: str = "vae",
+        scheduler_subfolder: str = "scheduler",
+        mlp_subfolder: str = "mlp",
+        diffusion_c_input_dim: int = 2560,
+        thinker_hidden_size: int = 4096,
+        img_gen_scales: list[int] | None = None,
+        text_encoder_norm: bool = True,
+        # Default sampling parameters — match Ming's defaults in
+        # ``modeling_bailingmm2.py::generate`` (``image_gen_steps=30``,
+        # ``image_gen_cfg=2.0``).
+        num_inference_steps: int = 30,
+        guidance_scale: float = 2.0,
+        default_height: int = 1024,
+        default_width: int = 1024,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.connector_subfolder = connector_subfolder
+        self.transformer_subfolder = transformer_subfolder
+        self.vae_subfolder = vae_subfolder
+        self.scheduler_subfolder = scheduler_subfolder
+        self.mlp_subfolder = mlp_subfolder
+        self.diffusion_c_input_dim = diffusion_c_input_dim
+        self.thinker_hidden_size = thinker_hidden_size
+        self.img_gen_scales = img_gen_scales if img_gen_scales is not None else [16]
+        self.text_encoder_norm = text_encoder_norm
+        self.num_inference_steps = num_inference_steps
+        self.guidance_scale = guidance_scale
+        self.default_height = default_height
+        self.default_width = default_width
+
+    @property
+    def num_query_tokens(self) -> int:
+        """Total number of learnable query tokens appended to the thinker input.
+
+        For ``img_gen_scales=[16]`` this yields 256 tokens (a single 16x16 grid).
+        """
+        return sum(s * s for s in self.img_gen_scales)
+
+
 class MingFlashOmniTalkerConfig(PretrainedConfig):
     """Configuration class for Ming-flash-omni-2.0 talker (TTS) stage.
 
@@ -309,13 +364,14 @@ class MingFlashOmniConfig(PretrainedConfig):
     is_composition = True
     sub_configs: ClassVar = {
         "thinker_config": BailingMM2Config,
+        "image_gen_config": MingImageGenConfig,
         "talker_config": MingFlashOmniTalkerConfig,
     }
 
     def __init__(
         self,
         thinker_config: BailingMM2Config | dict[str, Any] | None = None,
-        image_gen_config: dict[str, Any] | None = None,
+        image_gen_config: MingImageGenConfig | dict[str, Any] | None = None,
         talker_config: MingFlashOmniTalkerConfig | dict[str, Any] | None = None,
         **kwargs,
     ):
@@ -326,8 +382,13 @@ class MingFlashOmniConfig(PretrainedConfig):
         else:
             self.thinker_config = thinker_config or BailingMM2Config()
 
-        # Image generation config (for future implementation)
-        self.image_gen_config = image_gen_config
+        # Image generation config
+        if isinstance(image_gen_config, dict):
+            self.image_gen_config = MingImageGenConfig(**image_gen_config)
+        elif image_gen_config is None:
+            self.image_gen_config = MingImageGenConfig()
+        else:
+            self.image_gen_config = image_gen_config
 
         # Talker config
         if isinstance(talker_config, dict):
@@ -342,6 +403,7 @@ class MingFlashOmniConfig(PretrainedConfig):
 # Register model_type -> config class for AutoConfig
 AutoConfig.register(BailingMoeV2Config.model_type, BailingMoeV2Config)
 AutoConfig.register(BailingMM2Config.model_type, BailingMM2Config)
+AutoConfig.register(MingImageGenConfig.model_type, MingImageGenConfig)
 AutoConfig.register(MingFlashOmniTalkerConfig.model_type, MingFlashOmniTalkerConfig)
 AutoConfig.register(MingFlashOmniConfig.model_type, MingFlashOmniConfig)
 
diff --git a/vllm_omni/transformers_utils/configs/omnivoice.py b/vllm_omni/transformers_utils/configs/omnivoice.py
index d1bd3527f..1bd7207b6 100644
--- a/vllm_omni/transformers_utils/configs/omnivoice.py
+++ b/vllm_omni/transformers_utils/configs/omnivoice.py
@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """OmniVoice configuration for vLLM-Omni two-stage pipeline."""
 
+import os
+
 from transformers import AutoConfig
 from transformers.configuration_utils import PretrainedConfig
 
@@ -80,6 +82,14 @@ class OmniVoiceConfig(PretrainedConfig):
 
         # Serving
         self.speculative_config = None
+        # Precedence: config.json > OMNIVOICE_CUDA_GRAPH env var > default (enabled).
+        # Read once at init; changing OMNIVOICE_CUDA_GRAPH at runtime has no effect.
+        self.enable_cuda_graph = getattr(self, "enable_cuda_graph", os.environ.get("OMNIVOICE_CUDA_GRAPH", "1") != "0")
+        self.cuda_graph_capture_sizes = getattr(
+            self,
+            "cuda_graph_capture_sizes",
+            [128, 192, 256, 320, 384, 448, 512, 640, 768, 1024],
+        )
 
 
 AutoConfig.register("omnivoice", OmniVoiceConfig)
diff --git a/vllm_omni/transformers_utils/parsers/voxtral_tts.py b/vllm_omni/transformers_utils/parsers/voxtral_tts.py
index d4669258a..52241fd74 100644
--- a/vllm_omni/transformers_utils/parsers/voxtral_tts.py
+++ b/vllm_omni/transformers_utils/parsers/voxtral_tts.py
@@ -17,6 +17,7 @@ logger = init_logger(__name__)
 
 _VOXTRAL_TTS_ARCHS = frozenset({"VoxtralTTSForConditionalGeneration"})
 _VOXTRAL_TTS_MODEL_TYPE = "voxtral_tts"
+_VOXTRAL_TTS_DTYPE = "bfloat16"
 
 
 def _is_voxtral_tts_params(config_dict: dict) -> bool:
@@ -55,6 +56,10 @@ def _parse_voxtral_tts(config_dict: dict) -> tuple[dict, PretrainedConfig]:
         _remap_mistral_quantization_args,
     )
 
+    if config_dict.get("dtype") is None:
+        config_dict["dtype"] = _VOXTRAL_TTS_DTYPE
+    dtype = config_dict["dtype"]
+
     audio_config: dict[str, Any] = {}
     if (config_dict.get("multimodal") or {}).get("audio_model_args"):
         audio_config = _remap_voxtral_tts_audio_args(config_dict)
@@ -69,6 +74,7 @@ def _parse_voxtral_tts(config_dict: dict) -> tuple[dict, PretrainedConfig]:
         text_config=PretrainedConfig.from_dict(text_config),
         audio_config=audio_config,
         architectures=config_dict.get("architectures", ["VoxtralTTSForConditionalGeneration"]),
+        dtype=dtype,
     )
     return config_dict, config
 
diff --git a/vllm_omni/utils/mm_outputs.py b/vllm_omni/utils/mm_outputs.py
index cf2837a7c..43e8fa110 100644
--- a/vllm_omni/utils/mm_outputs.py
+++ b/vllm_omni/utils/mm_outputs.py
@@ -68,12 +68,13 @@ def to_payload_element(
             if we need to avoid splitting nonempty lists prior to calling
             postprocess, which is the case for prefix cache.
         seq_len: Optional sequence length (i.e., dim 0 of hidden states).
-            This should be set to None in the prefix caching case, because
-            the condition that would be executed here is the same as the
-            criteria for being added to the multimodal outputs cache.
+            When set, a tensor whose first dimension equals seq_len is
+            sliced per request. The prefix cache passthrough also passes
+            the total scheduled token count here so 1D (seq_len,) metadata
+            that is intentionally not cached is still split per request.
     """
-    # Prefix cache won't hit this case because this is the condition
-    # for being a mm_cache_key in the multimodal outputs tensor.
+    # Cached per-token tensors are merged elsewhere; here a first dim
+    # equal to seq_len means a per-request slice is required.
     if seq_len is not None and isinstance(element, torch.Tensor) and element.shape[0] == seq_len:
         return element[start:end].contiguous()
     # Every other case is shared between prefix cache (passthrough data)
diff --git a/vllm_omni/worker/base.py b/vllm_omni/worker/base.py
index 67b6f77b3..dd431d16b 100644
--- a/vllm_omni/worker/base.py
+++ b/vllm_omni/worker/base.py
@@ -232,10 +232,10 @@ class OmniGPUWorkerBase(GPUWorker):
                 if hasattr(self.model_runner, "graph_runners"):
                     self.model_runner.graph_runners.clear()
                     logger.info(f"[LLM Worker {self.rank}] LLM CUDA Graphs cleared.")
-            mem_before = current_omni_platform.get_current_memory_usage(self.device)
+            mem_before = current_omni_platform.get_free_memory(self.device)
             self.sleep(level=task.level)
-            mem_after = current_omni_platform.get_current_memory_usage(self.device)
-            rank_freed = max(0, mem_before - mem_after)
+            mem_after = current_omni_platform.get_free_memory(self.device)
+            rank_freed = max(0, mem_after - mem_before)
             if torch.distributed.is_initialized():
                 t_freed = torch.tensor([float(rank_freed)], device=self.device)
                 torch.distributed.all_reduce(t_freed)
@@ -246,6 +246,11 @@ class OmniGPUWorkerBase(GPUWorker):
             if self.rank != 0:
                 return None
             current_stage_id = getattr(self.vllm_config.model_config, "stage_id", 0)
+            try:
+                total_mem = current_omni_platform.get_device_total_memory()
+            except (NotImplementedError, AttributeError):
+                total_mem = torch.cuda.get_device_properties(self.device).total_memory
+            residual_gib = (total_mem - mem_after) / 1024**3
             ack = OmniACK(
                 task_id=task.task_id,
                 status="SUCCESS",
@@ -255,7 +260,7 @@ class OmniGPUWorkerBase(GPUWorker):
                 metadata={
                     "source": "omni_platform_audit",
                     "total_freed_gib": f"{total_freed / 1024**3:.2f}",
-                    "rank_residual_gib": f"{mem_after / 1024**3:.2f}",
+                    "rank_residual_gib": f"{residual_gib:.2f}",
                 },
             )
             if hasattr(self, "result_mq") and self.result_mq:
@@ -281,7 +286,12 @@ class OmniGPUWorkerBase(GPUWorker):
                 torch.distributed.barrier()
             gc.collect()
             current_omni_platform.synchronize()
-            usage_now = current_omni_platform.get_current_memory_usage(self.device)
+            free_now = current_omni_platform.get_free_memory(self.device)
+            try:
+                total_mem = current_omni_platform.get_device_total_memory()
+            except (NotImplementedError, AttributeError):
+                total_mem = torch.cuda.get_device_properties(self.device).total_memory
+            current_used_gib = (total_mem - free_now) / 1024**3
             if self.rank != 0:
                 return None
             current_stage_id = getattr(self.vllm_config.model_config, "stage_id", 0)
@@ -290,7 +300,7 @@ class OmniGPUWorkerBase(GPUWorker):
                 status="SUCCESS",
                 stage_id=current_stage_id,
                 rank=self.rank,
-                metadata={"state": "WARM", "current_vram_gib": f"{usage_now / 1024**3:.2f}"},
+                metadata={"state": "WARM", "current_vram_gib": f"{current_used_gib:.2f}"},
             )
             if hasattr(self, "result_mq") and self.result_mq:
                 self.result_mq.put(ack)
diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py
index c4ba29a35..78e2acdb5 100644
--- a/vllm_omni/worker/gpu_ar_model_runner.py
+++ b/vllm_omni/worker/gpu_ar_model_runner.py
@@ -84,11 +84,23 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False)
         # Initialize KV cache manager (preserve vllm_config fallback behavior)
         self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config)
-        # Only Qwen3-Omni currently consumes the connector-based full-payload
-        # handoff added in this PR. Other model architectures (e.g. Bagel
-        # diffusion) retain their pre-existing runner behavior so this PR
-        # does not perturb them.
-        if getattr(self.model_config, "model_arch", None) == "Qwen3OmniMoeForConditionalGeneration":
+        # Worker-connector init is gated by a per-`model_arch` allowlist
+        # (covers both producer-side and consumer-side runners for the
+        # arches below).  Consumer-wait stages must be registered
+        # separately as `(model_arch, model_stage)` tuples in
+        # `omni_scheduling_coordinator._FULL_PAYLOAD_INPUT_STAGES`;
+        # forgetting that produces a Stage-1 hang on the consumer.
+        _OMNI_CONNECTOR_INIT_ARCHS = {
+            "Qwen3OmniMoeForConditionalGeneration",
+            "Qwen2_5OmniForConditionalGeneration",
+            "CovoAudioForConditionalGeneration",
+            "MiMoAudioModel",
+            "Qwen3TTSTalkerForConditionalGeneration",
+            "Qwen3TTSCode2Wav",
+            "CosyVoice3Model",
+            "DyninOmniForConditionalGeneration",
+        }
+        if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS:
             self.init_omni_connectors(
                 vllm_config=self.vllm_config,
                 model_config=self.model_config,
@@ -148,6 +160,10 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             num_to_replace = min(num_sampled_ids, num_placeholders)
             req_history[first_placeholder : first_placeholder + num_to_replace] = new_ids[:num_to_replace]
 
+        for index, req_history in enumerate(output_token_ids):
+            if -1 in req_history:
+                output_token_ids[index] = req_history[: req_history.index(-1)]
+
         return output_token_ids
 
     def _sampling_metadata_for_model_sampler(self, sampling_metadata):
@@ -194,7 +210,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         return result
 
     def _capture_talker_mtp_graphs(self) -> None:
-        from vllm_omni.worker.gpu_model_runner import CUDAGraphWrapper
+        from vllm.compilation.cuda_graph import CUDAGraphWrapper
 
         if not self.has_talker_mtp or not isinstance(self.talker_mtp, CUDAGraphWrapper):
             return
@@ -250,6 +266,24 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         finally:
             set_cudagraph_capturing_enabled(False)
 
+    def _model_needs_full_prefix_hidden_states(self) -> bool:
+        """Opt-out hook for models whose postprocess only consumes the tail.
+
+        When False, we skip both the per-step GPU->CPU hidden-state write into
+        OmniTensorPrefixCache and the merged-tensor reconstruction on hits;
+        postprocess receives the normal scheduled-token slice instead. Models
+        that need the full cached_prefix + new_tail span (default) are not
+        affected.
+        """
+        model = getattr(self, "model", None)
+        return bool(getattr(model, "requires_full_prefix_cached_hidden_states", True))
+
+    def _deferred_prefix_cache_mm_keys(self) -> set[str]:
+        """Model-declared multimodal keys whose prefix-cache writes are deferred."""
+        model = getattr(self, "model", None)
+        keys = getattr(model, "deferred_prefix_cache_mm_keys", ())
+        return set(keys or ())
+
     def _maybe_update_prefix_cache(
         self,
         hidden_states: torch.Tensor,
@@ -272,12 +306,19 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                     type(multimodal_outputs),
                 )
 
+            hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None
+            # FIX: The .cpu attribute of slot_mapping is stale (not updated by the Triton
+            # _compute_slot_mapping_kernel which only writes to .gpu). We must use .gpu and
+            # sync back to CPU to get the correctly computed slot mapping.
+            slot_mapping_gpu = self.input_batch.block_table[0].slot_mapping.gpu
+            slot_mapping_cpu = slot_mapping_gpu[:num_tokens_padded].cpu()
             self.omni_prefix_cache.update_omni_tensor_prefix_cache(
-                hidden_states=hidden_states,
+                hidden_states=hs_for_cache,
                 multimodal_outputs=flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs,
                 num_tokens_unpadded=num_tokens_unpadded,
-                slot_mapping=self.input_batch.block_table[0].slot_mapping.cpu,
+                slot_mapping=slot_mapping_cpu,
                 num_tokens_padded=num_tokens_padded,
+                skip_mm_cache_keys=self._deferred_prefix_cache_mm_keys(),
             )
 
     def _maybe_get_combined_prefix_cache_tensors(
@@ -293,12 +334,18 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         # the prefix cached hidden states and multimodal states.
         combined_hidden_states, combined_multimodal_outputs = None, None
         if self.omni_prefix_cache is not None:
-            combined_hidden_states = self.omni_prefix_cache.get_merged_hidden_states(
-                query_start_loc=self.query_start_loc.cpu,
-                input_batch=self.input_batch,
-                hidden_states=hidden_states,
-                num_scheduled_tokens=num_scheduled_tokens,
-            )
+            if (
+                not self._model_needs_full_prefix_hidden_states()
+                and not self.omni_prefix_cache.has_prefix_cached_new_req_ids()
+            ):
+                return None, None
+            if self._model_needs_full_prefix_hidden_states():
+                combined_hidden_states = self.omni_prefix_cache.get_merged_hidden_states(
+                    query_start_loc=self.query_start_loc.cpu,
+                    input_batch=self.input_batch,
+                    hidden_states=hidden_states,
+                    num_scheduled_tokens=num_scheduled_tokens,
+                )
             combined_multimodal_outputs = self.omni_prefix_cache.get_merged_multimodal_states(
                 query_start_loc=self.query_start_loc.cpu,
                 input_batch=self.input_batch,
@@ -360,6 +407,12 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 if flush_ids:
                     self.flush_full_payload_outputs(flush_ids)
 
+        if self.omni_prefix_cache is not None and scheduler_output.finished_req_ids:
+            self.omni_prefix_cache.commit_deferred_mm_outputs(
+                set(scheduler_output.finished_req_ids),
+                self.input_batch,
+            )
+
         if self.routed_experts_initialized:
             capturer = get_global_experts_capturer()
             if capturer is not None and hasattr(capturer, "finalize_pending_copy"):
@@ -524,6 +577,14 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 cascade_attn_prefix_lens=cascade_attn_prefix_lens,
                 slot_mappings=slot_mappings_by_group,
             )
+            self._maybe_attach_attention_metadata_extensions(
+                attn_metadata=attn_metadata,
+                num_reqs=num_reqs,
+                num_reqs_padded=num_reqs_padded,
+                max_query_len=max_num_scheduled_tokens,
+                pad_attn=pad_attn,
+                num_scheduled_tokens_np=num_scheduled_tokens_np,
+            )
 
             (
                 input_ids,
@@ -709,6 +770,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         sampling_metadata = self.input_batch.sampling_metadata
         if spec_decode_metadata is None:
             model_sample = getattr(self.model, "sample", None)
+            self.input_batch.update_async_output_token_ids()
             if logits is not None and callable(model_sample) and getattr(self.model, "prefer_model_sampler", False):
                 # Apply logit bias (min_tokens, allowed_token_ids) before
                 # the custom model sampler — the standard GPU sampler does
@@ -726,7 +788,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 )
                 if sampler_output is not None:
                     return sampler_output
-            self.input_batch.update_async_output_token_ids()
             return self.sampler(
                 logits=logits,
                 sampling_metadata=sampling_metadata,
@@ -931,9 +992,22 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 dtype=np.int32,
             )
         query_start_loc_cpu = self.query_start_loc.cpu
+        if self.omni_prefix_cache is not None:
+            deferred_mm_cache_keys = self._deferred_prefix_cache_mm_keys()
+            if deferred_mm_cache_keys:
+                self.omni_prefix_cache.stage_deferred_mm_outputs(
+                    query_start_loc=query_start_loc_cpu,
+                    input_batch=self.input_batch,
+                    multimodal_outputs=(
+                        flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs
+                    ),
+                    num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
+                    deferred_mm_cache_keys=deferred_mm_cache_keys,
+                )
 
         pooler_output: list[dict[str, object]] | None = None
         if needs_pooler_payload:
+            mm_cpu = None
             # Prior to applying the post-processing func, extract
             # the prefix cached hidden states and multimodal states.
             if self.omni_prefix_cache is not None:
@@ -946,8 +1020,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                     scheduler_output.num_scheduled_tokens,
                 )
             # Otherwise we don't have the mm CPU data yet, so we still need to build it
-            if self.omni_prefix_cache is None:
-                mm_cpu = build_mm_cpu(flatten_payload(multimodal_outputs))
+            if self.omni_prefix_cache is None or combined_multimodal_outputs is None:
+                mm_cpu = build_mm_cpu(flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs)
 
             self._process_additional_information_updates(
                 hidden_states,
@@ -968,7 +1042,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                     req_hidden_states_cpu[rid] = hidden_states[start:end].detach().to("cpu").contiguous()
 
             pooler_output = []
-            for rid in req_ids_output_copy:
+            for out_idx, rid in enumerate(req_ids_output_copy):
                 if rid not in downstream_req_id_set:
                     pooler_output.append({})
                     continue
diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py
index 5a5175dc0..9f1060ed1 100644
--- a/vllm_omni/worker/gpu_generation_model_runner.py
+++ b/vllm_omni/worker/gpu_generation_model_runner.py
@@ -59,10 +59,18 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        # Scope full-payload connector init to Qwen3-Omni: other generation
-        # models (e.g. Bagel DiT) retain their pre-existing runner setup
-        # so this refactor does not perturb them.
-        if getattr(self.model_config, "model_arch", None) == "Qwen3OmniMoeForConditionalGeneration":
+        # Mirrors the init allowlist in gpu_ar_model_runner.py.
+        _OMNI_CONNECTOR_INIT_ARCHS = {
+            "Qwen3OmniMoeForConditionalGeneration",
+            "Qwen2_5OmniForConditionalGeneration",
+            "CovoAudioForConditionalGeneration",
+            "MiMoAudioModel",
+            "Qwen3TTSTalkerForConditionalGeneration",
+            "Qwen3TTSCode2Wav",
+            "CosyVoice3Model",
+            "DyninOmniForConditionalGeneration",
+        }
+        if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS:
             self.init_omni_connectors(
                 vllm_config=self.vllm_config,
                 model_config=self.model_config,
@@ -278,6 +286,14 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
                 cascade_attn_prefix_lens=cascade_attn_prefix_lens,
                 slot_mappings=slot_mappings_by_group,
             )
+            self._maybe_attach_attention_metadata_extensions(
+                attn_metadata=attn_metadata,
+                num_reqs=num_reqs,
+                num_reqs_padded=num_reqs_padded,
+                max_query_len=max_num_scheduled_tokens,
+                pad_attn=pad_attn,
+                num_scheduled_tokens_np=num_scheduled_tokens_np,
+            )
 
             (
                 input_ids,
@@ -291,7 +307,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
                 num_tokens_padded,
                 intermediate_tensors,
             )
-
             # [Omni] Pass token counts per request for code2wav output slicing
             model_kwargs["seq_token_counts"] = tokens
 
@@ -730,6 +745,14 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
                     slot_mappings=slot_mappings_by_group,
                     use_spec_decode=self.speculative_config is not None,
                 )
+                self._maybe_attach_attention_metadata_extensions(
+                    attn_metadata=attn_metadata,
+                    num_reqs=num_reqs_padded,
+                    num_reqs_padded=num_reqs_padded,
+                    max_query_len=max_query_len,
+                    pad_attn=True,
+                    for_cudagraph_capture=is_graph_capturing,
+                )
 
         with self.maybe_dummy_run_with_lora(
             self.lora_config,
diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py
index eae7c0a11..0663673d2 100644
--- a/vllm_omni/worker/gpu_model_runner.py
+++ b/vllm_omni/worker/gpu_model_runner.py
@@ -1,10 +1,10 @@
 import contextlib
 from collections.abc import Callable
+from dataclasses import replace
 from typing import TYPE_CHECKING, Any, cast
 
 import numpy as np
 import torch
-from vllm.compilation.cuda_graph import CUDAGraphWrapper
 from vllm.config import CUDAGraphMode
 from vllm.distributed.parallel_state import get_pp_group
 from vllm.forward_context import set_forward_context
@@ -32,6 +32,7 @@ from vllm_omni.core.prefix_cache import OmniTensorPrefixCache
 from vllm_omni.engine.serialization import deserialize_additional_information
 from vllm_omni.model_executor.layers.rotary_embedding.mrope import OmniMRotaryEmbedding as MRotaryEmbedding
 from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.platforms import current_omni_platform
 
 if TYPE_CHECKING:
     from vllm.v1.core.sched.output import SchedulerOutput
@@ -95,31 +96,121 @@ class OmniGPUModelRunner(GPUModelRunner):
     @instrument(span_name="Loading (GPU)")
     def load_model(self, *args, **kwargs) -> None:
         super().load_model(*args, **kwargs)
+        self._init_talker_mtp()
+        self._prewarm_attention_capture_workspaces()
 
+    def _init_talker_mtp(self) -> None:
         # TODO move this model specific logic to a separate class
         # TTS model IS the talker (no .talker sub-attr); use getattr to support both Omni and TTS.
         self.has_talker_mtp = False
         talker_mtp = getattr(self.model, "talker_mtp", None)
-        if talker_mtp is not None:
-            self.talker_mtp = talker_mtp  # type: ignore[assignment]
-            self.has_talker_mtp = True
-            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", False)
-            if cudagraph_mode.has_full_cudagraphs() and (has_separate_talker or talker_mtp_graph_safe):
-                self.talker_mtp = CUDAGraphWrapper(talker_mtp, self.vllm_config, runtime_mode=CUDAGraphMode.FULL)
-            # TTS exposes mtp_hidden_size; Omni uses hf_text_config.hidden_size.
-            hidden_size = int(
-                getattr(self.model, "mtp_hidden_size", 0) or getattr(self.model_config.hf_text_config, "hidden_size")
-            )
-            max_batch_size = max(self.max_num_reqs, self.compilation_config.max_cudagraph_capture_size)
-            self.talker_mtp_input_ids = self._make_buffer(max_batch_size, dtype=torch.int32)
-            self.talker_mtp_inputs_embeds = self._make_buffer(
-                max_batch_size, hidden_size, dtype=self.dtype, numpy=False
-            )
-            self.last_talker_hidden = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
-            self.text_step = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
+        if talker_mtp is None:
+            return
+        self.talker_mtp = talker_mtp  # type: ignore[assignment]
+        self.has_talker_mtp = True
+        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", False)
+        if cudagraph_mode.has_full_cudagraphs() and (has_separate_talker or 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.
+        hidden_size = int(
+            getattr(self.model, "mtp_hidden_size", 0) or getattr(self.model_config.hf_text_config, "hidden_size")
+        )
+        max_batch_size = max(self.max_num_reqs, self.compilation_config.max_cudagraph_capture_size)
+        self.talker_mtp_input_ids = self._make_buffer(max_batch_size, dtype=torch.int32)
+        self.talker_mtp_inputs_embeds = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
+        self.last_talker_hidden = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
+        self.text_step = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
+
+    def _prewarm_attention_capture_workspaces(self) -> None:
+        capture_sizes = getattr(self.compilation_config, "cudagraph_capture_sizes", None)
+        if not capture_sizes:
+            return
+        from vllm_omni.attention.fish_kvcache_backend import prewarm_fish_kvcache_attn_capture_workspaces
+
+        prewarm_fish_kvcache_attn_capture_workspaces(
+            model_config=self.model_config,
+            device=self.device,
+            dtype=self.dtype,
+            capture_sizes=capture_sizes,
+        )
+
+    def _maybe_attach_attention_metadata_extensions(
+        self,
+        *,
+        attn_metadata: Any,
+        num_reqs: int,
+        num_reqs_padded: int,
+        max_query_len: int,
+        pad_attn: bool,
+        for_cudagraph_capture: bool = False,
+        num_scheduled_tokens_np: np.ndarray | None = None,
+    ) -> None:
+        from vllm_omni.attention.fish_kvcache_backend import maybe_attach_fish_kvcache_seq_lens_upper_bound
+
+        maybe_attach_fish_kvcache_seq_lens_upper_bound(
+            model_config=self.model_config,
+            attn_metadata=attn_metadata,
+            input_batch=self.input_batch,
+            optimistic_seq_lens_cpu=self.optimistic_seq_lens_cpu,
+            num_reqs=num_reqs,
+            num_reqs_padded=num_reqs_padded,
+            max_query_len=max_query_len,
+            pad_attn=pad_attn,
+            for_cudagraph_capture=for_cudagraph_capture,
+            num_scheduled_tokens_np=num_scheduled_tokens_np,
+        )
+
+    def _build_model_sampler_output_token_ids(self) -> list[list[int]]:
+        """Build decoded-token history for ``prefer_model_sampler`` models.
+
+        vLLM only populates ``sampling_metadata.output_token_ids`` when penalties
+        or logits processors require it. Models that opt into a custom sampler
+        (e.g. CosyVoice3 RAS sampler, HunyuanImage3 stage-transition sampler)
+        also depend on this history, so we reconstruct it directly from the
+        input batch. Shared by GPU and NPU AR runners.
+        """
+        req_output_token_ids = getattr(self.input_batch, "req_output_token_ids", [])
+        req_ids = list(getattr(self.input_batch, "req_ids", []))
+        output_token_ids = [list(req_output_token_ids[idx] or []) for idx in range(len(req_ids))]
+
+        sampled_token_ids_cpu = getattr(self.input_batch, "sampled_token_ids_cpu", None)
+        async_copy_ready_event = getattr(self.input_batch, "async_copy_ready_event", None)
+        prev_req_id_to_index = getattr(self.input_batch, "prev_req_id_to_index", None)
+        if sampled_token_ids_cpu is None or not output_token_ids or prev_req_id_to_index is None:
+            return output_token_ids
+
+        sampled_token_ids: list[list[int]] | None = None
+        for index, req_id in enumerate(req_ids):
+            prev_index = prev_req_id_to_index.get(req_id)
+            if prev_index is None:
+                continue
+            req_history = output_token_ids[index]
+            if not req_history or req_history[-1] != -1:
+                continue
+            if sampled_token_ids is None:
+                assert async_copy_ready_event is not None
+                async_copy_ready_event.synchronize()
+                sampled_token_ids = sampled_token_ids_cpu.tolist()
+            new_ids = list(sampled_token_ids[prev_index])
+            if not new_ids:
+                continue
+            num_sampled_ids = len(new_ids) if new_ids[-1] != -1 else new_ids.index(-1)
+            first_placeholder = req_history.index(-1)
+            num_placeholders = len(req_history) - first_placeholder
+            num_to_replace = min(num_sampled_ids, num_placeholders)
+            req_history[first_placeholder : first_placeholder + num_to_replace] = new_ids[:num_to_replace]
+
+        return output_token_ids
+
+    def _sampling_metadata_for_model_sampler(self, sampling_metadata):
+        output_token_ids = self._build_model_sampler_output_token_ids()
+        if output_token_ids == sampling_metadata.output_token_ids:
+            return sampling_metadata
+        return replace(sampling_metadata, output_token_ids=output_token_ids)
 
     def _init_mrope_positions(self, req_state: CachedRequestState):
         """Initialize M-RoPE positions for multimodal inputs.
@@ -272,17 +363,22 @@ class OmniGPUModelRunner(GPUModelRunner):
 
         # Remove finished requests from the cached states.
         # cleanup_finished_request lives on OmniConnectorModelRunnerMixin and
-        # is only safe to call once init_omni_connectors() has populated the
-        # mixin state. Archs that inherit the method via MRO without running
-        # that init must be skipped, so probe a mixin-owned attribute as the
-        # "state initialized" gate.
+        # is only safe to call once init_omni_connectors() has finished
+        # populating mixin state (it sets ``_omni_connector_initialized = True``
+        # at the very end).  Archs that inherit the method via MRO without
+        # running that init must be skipped, so gate on the explicit flag
+        # rather than probing private attribute names.
         cleanup_finished_request = (
-            getattr(self, "cleanup_finished_request", None) if hasattr(self, "_request_ids_mapping") else None
+            getattr(self, "cleanup_finished_request", None)
+            if getattr(self, "_omni_connector_initialized", False)
+            else None
         )
         for req_id in scheduler_output.finished_req_ids:
             self.requests.pop(req_id, None)
             self.model_intermediate_buffer.pop(req_id, None)
             self.num_prompt_logprobs.pop(req_id, None)
+            if self.omni_prefix_cache is not None:
+                self.omni_prefix_cache.discard_deferred_mm_outputs(req_id)
             if hasattr(self, "_downstream_payload_cache"):
                 self._downstream_payload_cache.pop(req_id, None)
             if hasattr(self, "_talker_mtp_generators"):
@@ -341,7 +437,7 @@ class OmniGPUModelRunner(GPUModelRunner):
         for new_req_data in scheduler_output.scheduled_new_reqs:
             req_id = new_req_data.req_id
             if req_id in self.requests:
-                self._update_streaming_input_additional_info(new_req_data, req_id)
+                self._update_streaming_input_additional_info(req_id)
                 req_state = self._update_streaming_request(req_id, new_req_data)
                 reqs_to_add.append(req_state)
                 continue
@@ -847,6 +943,14 @@ class OmniGPUModelRunner(GPUModelRunner):
                     slot_mappings=slot_mappings_by_group,
                     use_spec_decode=self.speculative_config is not None,
                 )
+                self._maybe_attach_attention_metadata_extensions(
+                    attn_metadata=attn_metadata,
+                    num_reqs=num_reqs_padded,
+                    num_reqs_padded=num_reqs_padded,
+                    max_query_len=max_query_len,
+                    pad_attn=True,
+                    for_cudagraph_capture=is_graph_capturing,
+                )
 
         with self.maybe_dummy_run_with_lora(
             self.lora_config,
@@ -1543,9 +1647,10 @@ class OmniGPUModelRunner(GPUModelRunner):
             use_cascade_attn=False,
         )
         # Force eager for unwrapped code predictors (AR loops / multinomial).
-        # When talker_mtp is not a CUDAGraphWrapper, it manages its own CUDA
-        # graphs internally (code_predictor has its own bucket sizes).
-        if not isinstance(self.talker_mtp, CUDAGraphWrapper):
+        # When talker_mtp is not wrapped by the platform's full-graph wrapper,
+        # it manages its own device graphs internally (code_predictor has its
+        # own bucket sizes).
+        if not isinstance(self.talker_mtp, current_omni_platform.get_graph_wrapper_cls()):
             _cudagraph_mode = CUDAGraphMode.NONE
             num_tokens_padded = decode_batch_size
         else:
@@ -1609,7 +1714,7 @@ class OmniGPUModelRunner(GPUModelRunner):
         }
         if generator is not None:
             talker_kwargs["generator"] = generator
-        with set_forward_context(
+        with current_omni_platform.set_forward_context(
             None, self.vllm_config, cudagraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc
         ):
             req_embeds, code_predictor_codes = self.talker_mtp(
@@ -1663,7 +1768,12 @@ class OmniGPUModelRunner(GPUModelRunner):
             if key in gpu_keys:
                 dest[key] = value.detach().clone()
             else:
-                dest[key] = value.detach().to("cpu").contiguous()
+                t = value.detach()
+                if t.is_cuda:
+                    dest[key] = t.to("cpu").contiguous()
+                else:
+                    # If the tensor is already on the CPU, there is no need to unload it to the CPU.
+                    dest[key] = t.contiguous()
         elif isinstance(value, list):
             dest[key] = [
                 (item.detach().to("cpu").contiguous() if isinstance(item, torch.Tensor) else item) for item in value
@@ -1696,34 +1806,12 @@ class OmniGPUModelRunner(GPUModelRunner):
         logger.warning_once("_merge_additional_information_update is deprecated, use _update_intermediate_buffer")
         return self._update_intermediate_buffer(req_id, upd)
 
-    def _update_streaming_input_additional_info(self, new_req_data, req_id):
-        # For streaming input prefill case only. Update buffer from last segment input
+    def _update_streaming_input_additional_info(self, req_id):
+        # For streaming input prefill case only. Set num processed tokens = 0 for new segment input
         cached_additional_info = self.model_intermediate_buffer.get(req_id, {})
         if cached_additional_info:
-            payload_info = getattr(new_req_data, "additional_information", None)
-            inc_info = deserialize_additional_information(payload_info)
-            if isinstance(inc_info, dict) and inc_info:
-                accumulated_keys: set[tuple[str, str]] = set()
-                if hasattr(self, "model") and hasattr(self.model, "streaming_accumulated_keys"):
-                    accumulated_keys = self.model.streaming_accumulated_keys
-                merged_info = dict(cached_additional_info)
-                for key, value in inc_info.items():
-                    if isinstance(value, dict):
-                        existing_sub = merged_info.get(key)
-                        merged_sub = dict(existing_sub) if isinstance(existing_sub, dict) else {}
-                        for sk, sv in value.items():
-                            if (key, sk) in accumulated_keys and isinstance(sv, torch.Tensor):
-                                inc_tensor = sv.detach().to("cpu").contiguous()
-                                old_tensor = merged_sub.get(sk)
-                                if old_tensor is None:
-                                    merged_sub[sk] = inc_tensor
-                                else:
-                                    merged_sub[sk] = torch.cat((old_tensor, inc_tensor), dim=0)
-                            else:
-                                merged_sub[sk] = sv
-                        merged_info[key] = merged_sub
-                    else:
-                        merged_info[key] = value
-                merged_info.setdefault("meta", {})["num_processed_tokens"] = 0
-                self.model_intermediate_buffer[req_id] = merged_info
-                setattr(self.requests[req_id], "additional_information_cpu", merged_info)
+            merged_info = dict(cached_additional_info)
+            merged_info.setdefault("meta", {})["num_processed_tokens"] = 0
+            merged_info.setdefault("meta", {})["resumable"] = True
+            self.model_intermediate_buffer[req_id] = merged_info
+            setattr(self.requests[req_id], "additional_information_cpu", merged_info)
diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py
index eb5b47b53..6d991d0f2 100644
--- a/vllm_omni/worker/omni_connector_model_runner_mixin.py
+++ b/vllm_omni/worker/omni_connector_model_runner_mixin.py
@@ -46,6 +46,27 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+def should_accumulate_full_payload_output(model_config, custom_process_func) -> bool:
+    """Producer-side structural gate.
+
+    Fires iff the stage explicitly declares a downstream full-payload
+    producer hook via ``custom_process_next_stage_input_func``.  Consumer
+    stages may have ``custom_process_input_func`` values that can be
+    mechanically derived to ``*_full_payload`` helper names in the same
+    module; those are intentionally not enough to make the stage a producer.
+    """
+    if custom_process_func is None:
+        return False
+    if getattr(model_config, "async_chunk", False):
+        return False
+    if getattr(model_config, "final_output", False):
+        return False
+    next_stage_func = getattr(model_config, "custom_process_next_stage_input_func", None)
+    if not isinstance(next_stage_func, str) or not next_stage_func:
+        return False
+    return getattr(model_config, "model_stage", None) is not None
+
+
 class OmniConnectorModelRunnerMixin:
     """Unified data-plane communication mixin for Model Runners.
 
@@ -189,6 +210,13 @@ class OmniConnectorModelRunnerMixin:
             )
             self._save_thread.start()
 
+        # Explicit "fully initialised" marker so other parts of the runner
+        # (e.g. _update_states cleanup) can branch on a stable contract
+        # instead of probing for private mixin attribute names.  Must be set
+        # only after every field above has been bound, so a partially
+        # constructed mixin is never observable as initialised.
+        self._omni_connector_initialized = True
+
     def shutdown_omni_connectors(self) -> None:
         """Stop background threads and release connector resources."""
         self._stop_event.set()
@@ -217,6 +245,30 @@ class OmniConnectorModelRunnerMixin:
         saves is added to ``_deferred_send_cleanup`` so the bg save's
         decrement path drains it without leaving orphans.
         """
+        # Force-flush any pending full-payload accumulator entry before
+        # cleanup proceeds.  Without this, finished requests with no
+        # downstream consumer (e.g. text-only on multi-modal arch) leave
+        # the entry orphaned in _pending_full_payload_send across requests,
+        # which empirically destabilises subsequent thinker forwards by
+        # making prefix-cache reuse observe stale accumulator state.  The
+        # flush is idempotent when the entry has already been flushed by the
+        # scheduler-driven path, but this cleanup path runs for every request,
+        # so skip it entirely when the request never accumulated a payload.
+        if req_id in self._pending_full_payload_send:
+            try:
+                self.flush_full_payload_outputs({req_id})
+            except Exception:
+                # Cleanup must still proceed regardless of flush errors here --
+                # we already gated on ``_omni_connector_initialized`` upstream,
+                # so any exception here reflects a real connector-side issue
+                # (shared memory corruption, background thread crash) worth
+                # surfacing rather than silently swallowing.
+                logger.warning(
+                    "flush_full_payload_outputs(%s) raised during cleanup; continuing tear-down.",
+                    req_id,
+                    exc_info=True,
+                )
+
         ext_id = self._request_ids_mapping.pop(req_id, None)
         keys_to_clean: list[str] = [req_id]
         if ext_id is not None and ext_id != req_id:
@@ -684,6 +736,27 @@ class OmniConnectorModelRunnerMixin:
         _custom_process_func, both of which are set at init time. Avoid
         the per-step dynamic import inside the model decode loop.
         """
+        if getattr(self, "_omni_connector", None) is None:
+            # No connector at all: send_full_payload_outputs would no-op.
+            # Skip the per-step accumulator+build that would otherwise be
+            # silently discarded.  Defends against a terminal stage whose
+            # custom_process_input_func has a *_full_payload derivative in
+            # the same module (e.g. dynin stage 2 token2image_to_token2audio
+            # in pipelines that don't configure any connector at all).
+            #
+            # Known limitation: a *terminal-consumer* stage that has a
+            # connector configured for receiving upstream input is NOT
+            # caught here -- ``_omni_connector`` is non-None for it, and
+            # ``_load_custom_func`` may still resolve a ``*_full_payload``
+            # derivative from this stage's ``custom_process_input_func``.
+            # In that case the accumulator builds payloads that
+            # ``send_full_payload_outputs`` later drops via its own
+            # connector-side checks (wasted CPU, not a functional bug).
+            # A topology-aware gate (explicit producer field or pipeline
+            # is_terminal info) would close the gap; that change is out
+            # of scope for this PR.
+            self._should_accumulate_full_payload_output_cached = False
+            return False
         cached = getattr(self, "_should_accumulate_full_payload_output_cached", None)
         if cached is not None:
             return cached
@@ -691,19 +764,12 @@ class OmniConnectorModelRunnerMixin:
         if model_config is None:
             self._should_accumulate_full_payload_output_cached = False
             return False
-        if getattr(model_config, "model_arch", None) == "Qwen3OmniMoeForConditionalGeneration":
-            from vllm_omni.model_executor.stage_input_processors.qwen3_omni import (
-                should_accumulate_qwen3_omni_full_payload_output,
-            )
-
-            result = should_accumulate_qwen3_omni_full_payload_output(
-                model_config,
-                getattr(self, "_custom_process_func", None),
-            )
-            self._should_accumulate_full_payload_output_cached = result
-            return result
-        self._should_accumulate_full_payload_output_cached = False
-        return False
+        result = should_accumulate_full_payload_output(
+            model_config,
+            getattr(self, "_custom_process_func", None),
+        )
+        self._should_accumulate_full_payload_output_cached = result
+        return result
 
     @staticmethod
     def _new_full_payload_accumulator(output: dict[str, Any]):
@@ -729,6 +795,58 @@ class OmniConnectorModelRunnerMixin:
                 output[k] = tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=0)
         return output, request
 
+    def _resolve_full_payload_replace_keys(self) -> frozenset:
+        """Per-model REPLACE-key set for the full-payload accumulator.
+
+        Looked up from the stage-input-processor module that ships the model's sync builder
+        (`model_config.custom_process_input_func.__module__`).  The module
+        declares ``_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str]``; if absent,
+        returns the empty set.
+
+        Cached per instance.  Keys in this set use REPLACE semantics in the
+        accumulator (subsequent emissions discard prior chunks) instead of
+        the default CONCAT semantics.  Use for tensors that carry the full
+        result so far rather than per-step deltas (e.g. ``model_outputs``).
+        """
+        cached = getattr(self, "_full_payload_replace_keys_cached", None)
+        if cached is not None:
+            return cached
+        proc = getattr(self, "_custom_process_func", None)
+        if proc is None:
+            self._full_payload_replace_keys_cached = frozenset()
+            return self._full_payload_replace_keys_cached
+        module_name = getattr(proc, "__module__", None)
+        if module_name is None:
+            self._full_payload_replace_keys_cached = frozenset()
+            return self._full_payload_replace_keys_cached
+        try:
+            import sys as _sys
+
+            mod = _sys.modules.get(module_name) or importlib.import_module(module_name)
+            keys = getattr(mod, "_FULL_PAYLOAD_REPLACE_KEYS", frozenset())
+        except ImportError:
+            logger.debug(
+                "Could not import stage input processor module %s while resolving "
+                "_FULL_PAYLOAD_REPLACE_KEYS; using CONCAT semantics for all keys.",
+                module_name,
+                exc_info=True,
+            )
+            keys = frozenset()
+        if not isinstance(keys, (frozenset, set)):
+            logger.debug(
+                "Ignoring non-set _FULL_PAYLOAD_REPLACE_KEYS from %s: %s",
+                module_name,
+                type(keys).__name__,
+            )
+            keys = frozenset()
+        self._full_payload_replace_keys_cached = frozenset(keys)
+        logger.debug(
+            "Resolved _FULL_PAYLOAD_REPLACE_KEYS for %s: %s",
+            module_name,
+            sorted(self._full_payload_replace_keys_cached),
+        )
+        return self._full_payload_replace_keys_cached
+
     def accumulate_full_payload_output(
         self,
         req_id: str,
@@ -751,6 +869,7 @@ class OmniConnectorModelRunnerMixin:
         The data is actually sent when ``flush_full_payload_outputs`` is called
         with the finished request IDs from the next scheduler cycle.
         """
+        replace_keys = self._resolve_full_payload_replace_keys()
         existing = self._pending_full_payload_send.get(req_id)
 
         if existing is None:
@@ -766,6 +885,19 @@ class OmniConnectorModelRunnerMixin:
         for k, v in pooler_output.items():
             if v is None:
                 continue
+            if k in replace_keys:
+                # Explicit REPLACE semantics: the new value supersedes any
+                # prior chunks (e.g. `model_outputs` carries the full result
+                # so far, not an appendable per-step delta).
+                latest.pop(k, None)
+                if isinstance(v, torch.Tensor) and v.dim() >= 2:
+                    chunks[k] = [v]
+                    rows[k] = int(v.shape[0])
+                else:
+                    chunks.pop(k, None)
+                    rows.pop(k, None)
+                    latest[k] = v
+                continue
             if isinstance(v, torch.Tensor) and v.dim() >= 2:
                 if k in chunks and chunks[k] and v.shape[1:] == chunks[k][0].shape[1:]:
                     chunks[k].append(v)
@@ -783,6 +915,10 @@ class OmniConnectorModelRunnerMixin:
 
     def flush_full_payload_outputs(self, finished_req_ids: set[str]) -> None:
         """Send accumulated full_payload outputs for requests that just finished."""
+        pending_req_ids = set(self._pending_full_payload_send.keys())
+        if not (finished_req_ids & pending_req_ids):
+            return
+
         logger.info(
             "[Stage-%s] flush_full_payload_outputs: finished_req_ids=%s, pending=%s",
             self._stage_id,
@@ -927,11 +1063,11 @@ class OmniConnectorModelRunnerMixin:
         if self._stage_id == 0:
             return
         request_id = request.request_id
-        self._request_ids_mapping[request_id] = getattr(
-            request,
-            "external_req_id",
-            request_id,
-        )
+        # Explicit external_req_id=None must fall back to request_id;
+        # otherwise recv keys become `None_<stage>_<chunk>` and collide
+        # across requests.
+        ext = getattr(request, "external_req_id", None)
+        self._request_ids_mapping[request_id] = ext if ext is not None else request_id
         with self._lock:
             if request_id in self._stage_recv_req_ids:
                 return
@@ -1109,9 +1245,7 @@ class OmniConnectorModelRunnerMixin:
         if self._kv_transfer_manager is None:
             return False
 
-        request_id = getattr(req, "request_id", None) or (
-            req.request_ids[0] if hasattr(req, "request_ids") and req.request_ids else None
-        )
+        request_id = getattr(req, "request_id", None)
         if not request_id:
             logger.warning("Request has no ID, cannot receive KV cache")
             return False
@@ -2115,7 +2249,10 @@ class OmniConnectorModelRunnerMixin:
         if mapped is not None:
             return mapped
         if request is not None:
-            return getattr(request, "external_req_id", fallback_req_id)
+            # external_req_id may be explicitly None; fall back.
+            ext = getattr(request, "external_req_id", None)
+            if ext is not None:
+                return ext
         return fallback_req_id
 
     def _resolve_next_stage_id(self, model_config: Any) -> int: