[Fix] Put VoxCPM2 hf_overrides under engine_extras

Signed-off-by: Sy03 <1370724210@qq.com>
2026-07-04 16:58:48 +08:00
208 changed files with 2296 additions and 12386 deletions
@@ -108,10 +108,7 @@ steps:
    depends_on: upload-merge-pipeline
    steps:
      - label: "Entrypoint Test with H100"
-        # The suite relaunches an OmniServer per test class (~30 model loads,
-        # incl. ten ~27.5 GiB Bagel loads); 60 min passed only ~1 in 3 runs
-        # even on main, and vLLM 0.25 weight loading needs ~10 min more.
-        timeout_in_minutes: 90
+        timeout_in_minutes: 60
        commands:
          - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
        agents:
@@ -149,9 +146,7 @@ steps:
                      type: DirectoryOrCreate

      - label: "Entrypoint Test with L4"
-        # 30 min left no headroom for a cold image pull (~11 min for new base
-        # layers) plus per-class engine relaunches; main passed at 26/30.
-        timeout_in_minutes: 45
+        timeout_in_minutes: 30
        commands:
          - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4" --run-level "advanced_model"
        agents:
@@ -632,10 +627,11 @@ steps:
      - label: "Diffusion · FLUX.2 Klein Test"
        source_file_dependencies:
          - tests/e2e/offline_inference/test_flux2_klein.py
+          - tests/e2e/online_serving/test_flux2_klein.py
          - vllm_omni/diffusion/models/flux2_klein/
        timeout_in_minutes: 30
        commands:
-          - pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
+          - pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py tests/e2e/online_serving/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
        agents:
          queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
        plugins:
@@ -598,38 +598,6 @@ steps:
                      path: /mnt/hf-cache
                      type: DirectoryOrCreate

-  - group: ":card_index_dividers: Tiny Model Tests [Multi-GPU]"
-    key: nightly-tiny-model-test-group
-    depends_on: upload-nightly-pipeline
-    if: >-
-      build.env("NIGHTLY") == "1" ||
-      build.pull_request.labels includes "nightly-test"
-    steps:
-      - label: ":full_moon: Tiny Model Tests · Diffusion (Multi-GPU accelerations)"
-        timeout_in_minutes: 30
-        # NOTE: Single-GPU (core_model) tiny model tests run on the ready label,
-        # so nightly only adds the multi-GPU parallelism configurations.
-        #
-        # This is a bit confusing, but `full_model` is used to mark the multigpu tests, and
-        # --run-level `core_model` is used to indicate that we should use tiny model weights
-        # for the CI, since these tests should pass irrespective of whether or not we use
-        # the full weights or not.
-        commands:
-          - timeout 25m pytest -sv tests/model_tests/diffusion -m 'full_model and cuda' --run-level "core_model"
-        agents:
-          queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-        plugins:
-          - docker#v5.2.0:
-              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-              always-pull: true
-              propagate-environment: true
-              shm-size: "8gb"
-              environment:
-                - "HF_HOME=/fsx/hf_cache"
-                - "HF_TOKEN"
-              volumes:
-                - "/fsx/hf_cache:/fsx/hf_cache"
-
  # 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
@@ -1319,44 +1287,6 @@ steps:
                      path: /mnt/hf-cache
                      type: DirectoryOrCreate

-      - label: ":full_moon: Diffusion X2V · Doc Test"
-        timeout_in_minutes: 60
-        commands:
-          - pytest -s -v tests/examples/offline_inference/test_image_to_video.py -m "full_model and example and H100" --run-level "full_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: ":full_moon: Diffusion X2V · Other Function Test"
        timeout_in_minutes: 90
        commands:
@@ -157,27 +157,6 @@ steps:
                - "/fsx/hf_cache:/fsx/hf_cache"


-  - group: ":card_index_dividers: Tiny Model Tests [Single GPU]"
-    depends_on: upload-ready-pipeline
-    steps:
-      - label: "Tiny Model Tests · Diffusion (base)"
-        commands:
-          - timeout 15m pytest -sv tests/model_tests/diffusion -m 'core_model and cuda' --run-level "core_model"
-        agents:
-          # Ready label only runs the cases that don't need parallelism for tiny
-          # model tests, which is usually just a smoke test for TTI and I2I
-          queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
-        plugins:
-          - docker#v5.2.0:
-              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-              always-pull: true
-              propagate-environment: true
-              environment:
-                - "HF_HOME=/fsx/hf_cache"
-                - "HF_TOKEN"
-              volumes:
-                - "/fsx/hf_cache:/fsx/hf_cache"
-
  - label: "Engine&Model Executor Test"
    depends_on: upload-ready-pipeline
    commands:
@@ -208,7 +208,6 @@ cache/
 .cache/
 diffusion_cache/
 kv_cache/
-!vllm_omni/experimental/ar_diffusion/kv_cache/

 # Logs
 logs/
@@ -275,8 +274,3 @@ tests/.ci_generated/

 # Local Lance precision-matching dump scripts (not part of the repo)
 lance_compare/
-
-# FP8 quantization test sample images: tests/diffusion/quantization/
-# test_quantization_fp8.py saves these into the repo root for manual
-# inspection; they are byproducts and must never be committed.
-test_*_fp8.png
@@ -14,71 +14,8 @@ RUN apt-get update && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

-
-# Install vLLM (CUDA 13.0 build) from precompiled wheel at the selected commit.
-# Must use direct URL because the wheel has a PEP 440 local version identifier
-# (e.g. +gaeee7ef93.cu130) which pip/uv refuse to install from a PEP 503 package index.
-# The cu130 wheel index lives at /<commit>/cu130/vllm/ but the actual wheel files
-# are stored at the top level /<commit>/<wheel>.
-ENV VLLM_PRECOMPILED_WHEEL_COMMIT=dd10e03f95f94edbea1975c67ace3a35ec9a8a40
-# Keep the base image's CUDA-12 toolkit libs on the loader path. NOTE: this does
-# NOT resolve libnvJitLink.so.13 (that CUDA-13 lib is not in /usr/local/cuda here) —
-# the .so.13 JIT linker is handled separately via ldconfig after the cu130 wheels
-# are installed (see the nvidia-cuda-nvjitlink-cu13 step below).
-ENV LD_LIBRARY_PATH=/usr/local/cuda/targets/x86_64-linux/lib:${LD_LIBRARY_PATH}
-RUN VLLM_WHEEL_URL=$(python3 -c "import urllib.request,re; \
-    from urllib.parse import urljoin, quote; \
-    index_url='https://wheels.vllm.ai/${VLLM_PRECOMPILED_WHEEL_COMMIT}/cu130/vllm/'; \
-    html=urllib.request.urlopen(index_url).read().decode(); \
-    m=re.search(r'href=\"([^\"]*x86_64\.whl)\"',html); \
-    assert m, f'No x86_64 vLLM wheel found at {index_url}'; \
-    print(quote(urljoin(index_url,m.group(1)), safe=':/+%'))") && \
-    echo "Installing vLLM from: ${VLLM_WHEEL_URL}" && \
-    uv pip install --system --force-reinstall "${VLLM_WHEEL_URL}"
-
-
-
 RUN uv pip install --system ".[dev]"

-RUN uv pip install --system --upgrade \
-        "flashinfer-python==0.6.13" \
-        "flashinfer-cubin==0.6.13"
-
-RUN uv pip install --system --upgrade \
-    "flashinfer-jit-cache==0.6.13" \
-    --index-url https://flashinfer.ai/whl/cu130
-
-# Force-reinstall torch AFTER flashinfer, because flashinfer --upgrade pulls in
-# a newer torch (2.12.x) whose libtorch_cuda.so lacks
-# at::cuda::getCurrentCUDABlasHandle, causing "undefined symbol" at import time.
-RUN uv pip install --system --force-reinstall \
-    "torch==2.11.0" "torchaudio==2.11.0" "torchvision==0.26.0" \
-    --torch-backend cu130
-
-# CUDA 13 JIT linker: libnvJitLink.so.13 is provided by the nvidia-nvjitlink
-# package that the cu130 torch/vllm wheels above pull into site-packages (the
-# v0.24.0 base image's /usr/local/cuda toolkit is CUDA 12, so the LD_LIBRARY_PATH
-# above never resolves the .so.13). Its lib dir is not on the dynamic-loader path,
-# so bitsandbytes / torch.compile worker subprocesses fail to dlopen it. Locate the
-# already-installed lib and register its dir with ldconfig (system-wide, survives
-# subprocess env resets). Only fall back to the pip wheel (correct name:
-# nvidia-nvjitlink-cu13, NOT nvidia-cuda-nvjitlink-cu13, which does not exist) if
-# the torch deps did not bring it in.
-RUN NVJITLINK_LIB="$(find / -name 'libnvJitLink.so.13*' 2>/dev/null | head -1)" && \
-    if [ -z "$NVJITLINK_LIB" ]; then \
-        uv pip install --system "nvidia-nvjitlink-cu13" && \
-        NVJITLINK_LIB="$(find / -name 'libnvJitLink.so.13*' 2>/dev/null | head -1)"; \
-    fi && \
-    test -n "$NVJITLINK_LIB" && \
-    dirname "$NVJITLINK_LIB" > /etc/ld.so.conf.d/nvjitlink-cu13.conf && \
-    ldconfig && \
-    echo "Registered ${NVJITLINK_LIB} with ldconfig"
-
-# Re-pin numpy LAST. numba requires numpy <= 2.4, but the flashinfer and torch cu130
-# reinstalls above resolve numpy back up to 2.5, which breaks `import numba` at test
-# time ("Numba needs NumPy 2.4 or less"). This must be the final dependency op so
-# nothing can clobber it again.
-RUN uv pip install --system "numpy==2.2.6"

 RUN ln -sf /usr/bin/python3 /usr/bin/python

@@ -1,239 +0,0 @@
-# Composable parallel strategies
-
-In vLLM-Omni, the `composable_parallel` system layers a declarative *parallel strategy* on top of the existing stage configs. A strategy YAML maps each pipeline stage (by its `model_stage` name) to a stack of axis declarations — one per parallelism dimension — that the engine translates into per-stage parallel sizing (`tensor_parallel_size`, `data_parallel_size`, `pipeline_parallel_size`, expert-parallel toggle, replica count) and, when relevant, an orchestrator-wide load-balancer policy.
-
-Strategies are intentionally narrow: each one only writes the axes it declares, leaving every other engine arg untouched. This makes a strategy file a small, reusable overlay you can swap independently of the deploy YAML (see `configuration/stage_configs.md`).
-
-!!! note
-    Composable parallel is **opt-in**. Passing `--strategy-config <path/to/strategy.yaml>` overlays the derived sizing onto the registry-merged stages *before* any CLI overrides are applied. When `--strategy-config` is omitted, stages remain exactly as their deploy YAML produces them.
-
-## Top-level schema reference
-
-A strategy file has a single top-level mapping, `strategies`, keyed by the `model_stage` name of each stage you want to overlay. Each value is a list of axis declarations (see [StrategySpec fields](#strategyspec-fields) below).
-
-| Field | Type | Required | Default | Description |
-|-------|------|----------|---------|-------------|
-| `strategies` | `dict[str, list]` | required | — | Maps `model_stage` name (e.g. `thinker`, `talker`, `code2wav`, `dit`) to a list of axis declarations. The key MUST be the stage name, not the integer `stage_id`. See `vllm_omni/config/composable_parallel/strategy_loader.py`. |
-
-### StrategySpec fields
-
-Each entry in a stage's list becomes one `StrategySpec` (see `vllm_omni/config/composable_parallel/spec.py`). The loader fills in `routing` and `aggregation` from per-kind defaults, so the file stays compact.
-
-| Field | Type | Required | Default | Description |
-|-------|------|----------|---------|-------------|
-| `axis` | str | required | — | Mesh-axis kind. One of `tp`, `dp`, `pp`, `ep`, `stage_replica` (wired today) or `sp_ulysses`, `sp_ring`, `cfg`, `vae_pp`, `hsdp`, `stage_pp`, `cp` (reserved; declaring one raises `AxisTranslationError` at translate time). See `MeshAxisKind` in `vllm_omni/config/composable_parallel/spec.py`. |
-| `size` | int | required | — | Axis degree. Must be `> 0`. |
-| `routing` | str | optional | per-kind default | Stateless routing policy. Accepted only for `dp` / `stage_replica`; other kinds reject it. Values: `random`, `round_robin`, `least_queue`. (`hash` is recognised but currently raises `NotImplementedError`.) See `RouteByStage` in `vllm_omni/config/composable_parallel/routing.py`. |
-| `l1_owner` | str | optional | per-kind default | Overrides which layer owns request routing for this axis. `engine` for `tp`/`dp`/`pp`/`ep`, `delegated` for `stage_replica`. Recorded into `shard_extension`. |
-| `name` | str | optional | axis kind | Human-readable label for this spec, used in error messages. |
-
-Two additional `StrategySpec` slots — `layer_hook_specs` and `kernel_specs` — let advanced authors bind layer-walker hook IDs (`LayerHookSpec`) or kernel IDs (`KernelSpec`) onto an axis. These are not exposed by the YAML loader today and are intended for programmatic construction. See `vllm_omni/config/composable_parallel/spec.py` for the dataclasses.
-
-### CLI flags
-
-| Flag | Description |
-|------|-------------|
-| `--strategy-config PATH` | Path to a `strategy.yaml`. Only takes effect on registry-based deploy paths (`--deploy-config` or the bundled default deploy YAML). Its derived sizing is overlaid onto the merged stages *before* CLI overrides. |
-| `--stage-configs-path PATH` | **Legacy.** Loads the old `stage_args` YAML schema for un-migrated models. Passing `--strategy-config` together with this path is silently inapplicable — vLLM-Omni emits a warning and the strategy is ignored (see `vllm_omni/entrypoints/utils.py`, lines 372–380). |
-| `--omni-lb-policy POLICY` | Orchestrator-wide load-balancer policy. Strategy-derived from any `stage_replica` axis; an explicit CLI value wins and must match the derived one or `AsyncOmniEngine` raises `Conflicting load-balancer policy` at construction. |
-
-!!! warning
-    `--strategy-config` cannot be used with `--stage-configs-path` (the legacy deploy path). If your model has not been migrated to the registry-based pipeline yet, the strategy is silently dropped with a warning. Migrate to a registry-based model or use `--deploy-config` to apply a strategy.
-
-## Sub-schemas
-
-### `axis`
-
-The set of recognised axis kinds is the `MeshAxisKind` enum in `vllm_omni/config/composable_parallel/spec.py`:
-
-```yaml
-# Translatable today (see `_SUPPORTED_KINDS` in translator.py).
-axis: tp              # tensor parallel — sets tensor_parallel_size
-axis: dp              # engine data parallel — sets data_parallel_size (vLLM's intra-engine DP LB)
-axis: pp              # pipeline parallel — sets pipeline_parallel_size
-axis: ep              # dense expert parallel — sets enable_expert_parallel=True; size must equal tp*dp
-axis: stage_replica   # omni-coordinator-level replica fan-out — sets num_replicas + omni_lb_policy
-
-# Reserved (raise AxisTranslationError today; declared in the type system for forward compat).
-axis: sp_ulysses
-axis: sp_ring
-axis: cfg
-axis: vae_pp
-axis: hsdp
-axis: stage_pp
-axis: cp
-```
-
-### `routing`
-
-The full closed hierarchy lives in `vllm_omni/config/composable_parallel/routing.py` (`Broadcast`, `PartitionByHash`, `ShardSequence`, `ShardSpatial`, `RouteByExpert`, `PipelineMicrobatch`, `DuplicateWithCondTag`, `RouteByStage`). For YAML authors only `RouteByStage` is exposed as a `routing:` string, and only for `dp` / `stage_replica` axes (other axes have fixed routing — `tp`/`ep` broadcast, `pp` microbatch). Recognised policies:
-
-```yaml
-routing: random         # uniform-random across replicas/DP ranks
-routing: round_robin    # cycle through replicas/DP ranks (the default for dp/stage_replica)
-routing: least_queue    # send to the replica with the smallest backlog
-# routing: hash         # recognised but raises NotImplementedError; no key-stable balancer yet
-```
-
-### `aggregation`
-
-Aggregation is filled in automatically by the loader (`_default_aggregation` in `vllm_omni/config/composable_parallel/strategy_loader.py`) from the axis kind, so users do not author it:
-
-| Axis kind | Aggregation pattern |
-|-----------|---------------------|
-| `tp`, `ep` | `TakeRank` (all ranks agree; pick one) |
-| `dp`, ep (in some flows) | `Union` (disjoint replica results) |
-| `pp` | `StitchPipeline` |
-| `stage_replica` | `FanInByStage` |
-
-Full pattern hierarchy: `TakeRank`, `Union`, `GatherDim`, `AllGather`, `StitchSpatial`, `StitchPipeline`, `Combine`, `FanInByStage` (see `vllm_omni/config/composable_parallel/aggregation.py`).
-
-## Precedence
-
-From highest to lowest, when the same axis is touched by multiple layers:
-
-1. **CLI overrides** (`--stage_N_tensor_parallel_size`, `--stage-overrides`, `--omni-lb-policy`, etc.). A CLI value always wins; a warning is emitted whenever it overrides a strategy-declared axis (see `_reconcile_strategy_with_cli` in `vllm_omni/config/config_factory.py`).
-2. **Strategy YAML** (`--strategy-config`). Derived sizing is written onto the merged stages with *conflict-on-explicit* semantics — if the deploy YAML already set a value to something different, application raises `StrategyApplyError`.
-3. **Deploy YAML** (`--deploy-config` or the bundled `vllm_omni/deploy/<model_type>.yaml`).
-4. **Parser defaults.**
-
-After CLI overrides land, the device-layout guard re-runs (`check_device_layout`, called from `_reconcile_strategy_with_cli`) against the *effective* `tp * dp * pp * num_replicas` and the resolved `devices` string, so a `--devices` or `--tensor-parallel-size` that conflicts with the strategy's world size cannot slip past the pre-spawn check.
-
-## Worked example
-
-### Example 1 — `strategy_tp2.yaml`
-
-Shard the `thinker` stage with TP=2 on top of the bundled Qwen2.5-Omni deploy YAML. Routing is fixed (broadcast) for TP, so no `routing:` field is needed.
-
-```yaml
-# examples/offline_inference/qwen2_5_omni/strategy_tp2.yaml
-strategies:
-  thinker:
-    - axis: tp
-      size: 2
-```
-
-Launch:
-
-```bash
-python3 end2end.py --query-type text --num-prompts 6 \
-    --strategy-config strategy_tp2.yaml \
-    --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}'
-```
-
-Effective `OmniParallelConfig` per stage after the overlay (see `OmniParallelConfig` in `vllm_omni/config/composable_parallel/translator.py`):
-
-| Stage (`model_stage`) | `tensor_parallel_size` | `data_parallel_size` | `pipeline_parallel_size` | `stage_replica_size` | `omni_lb_policy` |
-|-----------------------|------------------------|----------------------|--------------------------|----------------------|------------------|
-| `thinker`             | **2** (from strategy)  | 1                    | 1                        | 1                    | —                |
-| `talker`              | 1                      | 1                    | 1                        | 1                    | —                |
-| `code2wav`            | 1                      | 1                    | 1                        | 1                    | —                |
-
-### Example 2 — `strategy_stage_replica.yaml`
-
-Replicate the `thinker` stage across two independent engines and let omni's StagePool round-robin requests across them.
-
-```yaml
-# examples/offline_inference/qwen2_5_omni/strategy_stage_replica.yaml
-strategies:
-  thinker:
-    - axis: stage_replica
-      size: 2
-      routing: round_robin
-```
-
-Launch:
-
-```bash
-python3 end2end.py --query-type text --num-prompts 6 \
-    --strategy-config strategy_stage_replica.yaml \
-    --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}'
-```
-
-Effective per-stage config and the derived orchestrator knob:
-
-| Stage (`model_stage`) | `tensor_parallel_size` | `stage_replica_size` | `omni_lb_policy` (derived) |
-|-----------------------|------------------------|----------------------|----------------------------|
-| `thinker`             | 1                      | **2**                | **`round-robin`**          |
-| `talker`              | 1                      | 1                    | —                          |
-| `code2wav`            | 1                      | 1                    | —                          |
-
-The derived `omni_lb_policy = round-robin` is **not** written into any stage's engine args. Instead it is surfaced on `StrategyApplyResult.omni_lb_policy` and applied once at orchestrator construction by `AsyncOmniEngine._apply_strategy_lb_policy` (see `vllm_omni/engine/async_omni_engine.py`, lines 1107–1134). If you also pass `--omni-lb-policy round-robin` the values must match; passing a different value raises `Conflicting load-balancer policy`.
-
-## Per-field deep-dive
-
-### `name`
-
-Default: the axis kind (e.g. `"tp"`).
-
-Human-readable label that appears in `AxisTranslationError` and `StrategyApplyError` messages. Set it to disambiguate two stages declaring the same axis kind when scanning logs.
-
-### `mesh_axis` (`axis` + `size`)
-
-Default: no default — both fields are required.
-
-`axis` chooses the parallelism dimension; `size` is its degree. Five kinds are wired end-to-end today (`tp`, `dp`, `pp`, `ep`, `stage_replica`); the remaining kinds in `MeshAxisKind` are reserved and raise `AxisTranslationError` at translate time so a typo cannot silently become a no-op (see `_SUPPORTED_KINDS` in `vllm_omni/config/composable_parallel/translator.py`). For `ep`, the size MUST equal `tensor_parallel_size * data_parallel_size`.
-
-### `routing`
-
-Default: per-kind — `RouteByStage(round_robin)` for `dp` and `stage_replica`; `PipelineMicrobatch()` for `pp`; `Broadcast()` for `tp`/`ep` and any other kind.
-
-User-facing string only on `dp` and `stage_replica` axes. The four recognised values are `random`, `round_robin`, `least_queue`, `hash`; `hash` is rejected today with `NotImplementedError`. See `RouteByStage` in `vllm_omni/config/composable_parallel/routing.py`.
-
-### `l1_owner`
-
-Default: per-kind — `engine` for `tp`/`dp`/`pp`/`ep`, `delegated` for `stage_replica`.
-
-Overrides which layer owns request routing along this axis. `engine` means vLLM realizes routing intra-engine; `delegated` means the omni coordinator owns routing across independent replicas. Set this only when overriding the per-kind default in `_DEFAULT_L1_OWNER` (`vllm_omni/config/composable_parallel/translator.py`); invalid values fail fast in `_resolve_l1_owner`.
-
-### `stage_replica` and `omni_lb_policy`
-
-`stage_replica` is *not* a vLLM world dimension; it spins up `size` independent engine replicas of one stage, coordinated by the omni `StagePool` load balancer. The axis writes both:
-
- The stage's `num_replicas` (a per-stage deploy knob).
- The pipeline-wide `omni_lb_policy` string (an orchestrator-level knob).
-
-Because `omni_lb_policy` is read once at `AsyncOmniEngine` construction — not from per-stage engine args — only one value is allowed across the whole pipeline. If two stages declare conflicting policies, application raises (see `apply_strategy_specs` in `vllm_omni/config/composable_parallel/apply.py`, lines 253–260). To set this knob manually instead, use `--omni-lb-policy` on the CLI; see [the model stage docs](stage_configs.md) for how stage replicas plug into the orchestrator at runtime.
-
-### `model_stage` keys
-
-Default: no default — required for every entry under `strategies:`.
-
-Strategy keys MUST be `model_stage` *names* (strings), the same labels stage configs use for `engine_args.model_stage` (see [`model_stage` in `stage_configs.md`](stage_configs.md#engine_argsmodel_stage)). Integer `stage_id` values are not accepted; the resolver matches on `model_stage` exactly (`_resolve_stage` in `vllm_omni/config/composable_parallel/apply.py`).
-
-## Common pitfalls
-
-!!! warning "Using `stage_id` integers as keys"
-    Strategy keys are `model_stage` names (str), not integer stage IDs. A file that uses an integer index will raise at apply time:
-
-    ```text
-    StrategyApplyError: strategy role model_stage=0 did not match any stage;
-    available stages: 'thinker'(id=0), 'talker'(id=1), 'code2wav'(id=2)
-    ```
-
-    Use the `model_stage` name (e.g. `thinker`) instead.
-
-!!! warning "Mixing CLI overrides with a strategy on the same axis"
-    A strategy is meant to be the single writer for the axes it declares. If a CLI override (`--stage_0_tensor_parallel_size`, `--stage-overrides '{"0": {...}}'`) sets the same axis, the **CLI value wins** and a warning is emitted from `_reconcile_strategy_with_cli`. The device-layout guard then re-runs against the effective layout, so a CLI override that breaks `tp * dp * pp * num_replicas` will still fail fast at config time.
-
-!!! warning "Passing `--strategy-config` against the legacy `--stage-configs-path`"
-    Composable parallel only applies on the registry-based deploy path. If a model still resolves through the legacy `stage_args` YAML (loaded via `--stage-configs-path`), `--strategy-config` is silently inapplicable and emits:
-
-    ```text
-    --strategy-config (<path>) was provided but model '<model>' resolves via the
-    legacy stage_configs YAML path, which does not support composable-parallel
-    strategies; the strategy is ignored. Use a registry-based model to apply it.
-    ```
-
-    Migrate the model to the new deploy schema (see `configuration/stage_configs.md`) before applying a strategy.
-
-!!! warning "Conflicting `omni_lb_policy` across stages"
-    `omni_lb_policy` is a single pipeline-wide knob. If two `stage_replica` axes in different stages derive different policies, `apply_strategy_specs` raises:
-
-    ```text
-    StrategyApplyError: role 'talker' derives omni_lb_policy='least-queue-length'
-    but role 'thinker' already derived 'round-robin'. The load-balancer policy is
-    pipeline-wide; only one value is allowed.
-    ```
-
-    Make the policies match, drop one, or supply the value explicitly via `--omni-lb-policy`.
@@ -123,7 +123,7 @@ After inserting the extra thinker stage, renumber the remaining stages:
    runtime:
      devices: "2"
    engine_input_source: [1]
-    sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker

  - stage_id: 3
    runtime:
@@ -131,7 +131,7 @@ After inserting the extra thinker stage, renumber the remaining stages:
    engine_args:
      max_num_seqs: 1
    engine_input_source: [2]
-    sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_token_only
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
 ```

 Compared with the default Qwen3-Omni config:
@@ -475,64 +475,92 @@ thinker_hidden_states = output.multimodal_output["24"]

 ### Implementation Example

-Create stage transition processors in `vllm_omni/model_executor/stage_input_processors/your_model_name.py`. Each inter-stage edge should provide a **coherent processor set** rather than a single monolithic function:
-
-| Suffix | Role | Runs when |
-|--------|------|-----------|
-| `*_full_payload` | Worker-side payload producer | `async_chunk=false`; accumulates tensors and ships via connector |
-| `*_async_chunk` | Scheduler-side streaming producer | `async_chunk=true`; emits per-chunk payloads |
-| `*_token_only` | Orchestrator placeholder builder | `async_chunk=false`; allocates downstream prompt slots only |
+Create stage transition processors in `vllm_omni/model_executor/stage_input_processors/your_model_name.py`:

 ```python
-# qwen3_omni.py (Thinker → Talker, non-async path)
+# qwen3_omni.py

-def thinker2talker_token_only(
-    source_outputs: list[Any],
+def thinker2talker(
+    stage_list: list[Any],
+    engine_input_source: list[int],
    prompt: OmniTokensPrompt | TextPrompt | None = None,
    requires_multimodal_data: bool = False,
-    streaming_context: Any | None = None,
 ) -> list[OmniTokensPrompt]:
-    """Allocate talker prefill slots; bulk tensors arrive via the connector."""
-    ...
+    """
+    Process thinker outputs to create talker inputs.
+
+    Args:
+        stage_list: List of stage objects
+        engine_input_source: Source stage IDs (typically [0] for thinker)
+        prompt: Original prompt data
+
+    Returns:
+        List of OmniTokensPrompt for talker stage
+    """
+    source_stage_id = engine_input_source[0]
+    thinker_outputs = stage_list[source_stage_id].engine_outputs
+    talker_inputs = []
+
+    for thinker_output in thinker_outputs:
+        output = thinker_output.outputs[0]
+        # Extract thinker embeddings and hidden states
+        thinker_prefill_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
+        thinker_hidden_states = output.multimodal_output["24"].float().clone().detach().cuda()
+
+        info = {
+            "thinker_prefill_embeddings": thinker_prefill_embeddings,
+            "thinker_hidden_states": thinker_hidden_states,
+            "thinker_sequences": thinker_output.prompt_token_ids + output.token_ids,
+            "thinker_input_ids": thinker_output.prompt_token_ids,
+        }
+
+        talker_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[0] * computed_length,
+                additional_information=info,
+                multi_modal_data=None,
+            )
+        )
+
+    return talker_inputs


-def thinker2talker_full_payload(
-    transfer_manager: Any,
-    pooling_output: dict[str, Any],
-    request: OmniEngineCoreRequest,
-) -> dict[str, Any] | None:
-    """Pack accumulated thinker hidden states into OmniPayload for stage-1."""
-    ...
+def talker2code2wav(
+    stage_list: list[Any],
+    engine_input_source: list[int],
+    prompt: OmniTokensPrompt | TextPrompt | None = None,
+    requires_multimodal_data: bool = False,
+) -> list[OmniTokensPrompt]:
+    """
+    Process talker outputs to create code2wav inputs.
+    """
+    source_stage_id = engine_input_source[0]
+    talker_outputs = stage_list[source_stage_id].engine_outputs
+    code2wav_inputs = []

+    for talker_output in talker_outputs:
+        output = talker_output.outputs[0]
+        # Extract codec codes
+        codec_codes = (
+            output.multimodal_output["code_predictor_codes"]
+            .to(torch.long)
+            .transpose(0, 1)
+            .cpu()
+            .to(torch.long)
+            .reshape(-1)
+            .tolist()
+        )

-def thinker2talker_async_chunk(
-    transfer_manager: Any,
-    multimodal_output: OmniPayload | dict[str, Any],
-    request: OmniEngineCoreRequest,
-    is_finished: bool = False,
-) -> OmniPayloadStruct | None:
-    """Stream thinker rows to talker while async_chunk is enabled."""
-    ...
+        code2wav_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=codec_codes,
+                multi_modal_data=None,
+            )
+        )
+
+    return code2wav_inputs
 ```

-Wire these in `pipeline.py`:
-
-```python
-StagePipelineConfig(
-    stage_id=0,
-    custom_process_next_stage_input_func=f"{_PROC}.thinker2talker_full_payload",
-    async_chunk_process_next_stage_input_func=f"{_PROC}.thinker2talker_async_chunk",
-    ...
-),
-StagePipelineConfig(
-    stage_id=1,
-    sync_process_input_func=f"{_PROC}.thinker2talker_token_only",
-    ...
-),
-```
-
-Do **not** add a no-suffix `thinker2talker` when a `sync_process_input_func` is already declared — `_select_processor_funcs()` always prefers the `*_token_only` hook in non-async mode, so the bare function would never run. See `docs/design/rfc_stage_input_processors_refactor.md` for the full contract.
-
 ## Testing

 For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
@@ -20,8 +20,6 @@ The full set of backends and their platform defaults is in the **Backend Options
 | `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. |
-| `FLASH_ATTN_HUB` | FlashAttention 2 from HuggingFace `kernels` library. Useful for train/rollout alignment. |
-| `FLASH_ATTN_3_HUB` | FlashAttention 3 from HuggingFace `kernels` library. CUDA Hopper (sm_90+) only; falls back to `FLASH_ATTN_HUB` on older GPUs. |


 ## Configuration
@@ -212,36 +210,6 @@ 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.

-## HuggingFace Kernels Hub Backends
-
-To achieve perfect numerical consistency between **training** (typically using HuggingFace Diffusers with Hub kernels) and **serving/rollout** (using vLLM-Omni), you can use the Hub-based attention backends. This eliminates numerical drift / sampling divergence caused by executing different local kernel versions during rollout.
-
-The following backend options are supported:
- `FLASH_ATTN_HUB` (HuggingFace `kernels-community/flash-attn2`)
- `FLASH_ATTN_3_HUB` (HuggingFace `kernels-community/flash-attn3`, Hopper sm_90+ only)
-
-### Installation
-
-To use these backends, you must install the `kernels` library:
-
-```bash
-pip install kernels==0.14.1
-```
-
-If the `kernels` library is not available in the environment, vLLM-Omni will log a warning and fall back gracefully to the corresponding local backend implementations (`FLASH_ATTN`). On CUDA GPUs below Hopper (compute capability < 9.0), `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`.
-
-### Usage
-
-Select a Hub backend using the global CLI flag or environment variables:
-
-```bash
-# Environment variable
-export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN_HUB
-
-# CLI flag
-vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN_HUB
-```
-
 ## Usage Examples

 ### Default (auto-route)
@@ -126,7 +126,7 @@ The following tables show which models support each feature:
 | **MagiHuman**            |     ❌     |     ❌      |           ❌           |       ❓        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **MammothModa2(T2I)**    |     ❌     |     ❌      |           ❌           |       ❌        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Nextstep_1(T2I)**      |     ❓     |     ❓      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **OmniGen2**             |     ❌     |     ✅      |           ✅           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **OmniGen2**             |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Ovis-Image**           |     ❌     |     ✅      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Qwen-Image**           |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ✅         |
 | **Qwen-Image-2512**      |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ✅         |
@@ -22,7 +22,7 @@ Text-to-image now runs through the shared offline image example
 ```bash
 python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
  --model ./MammothModa2-Preview \
-  --deploy-config ./vllm_omni/deploy/mammoth_moda2_ar.yaml \
+  --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
  --question "Summarize this image." \
  --image ./image.png
 ```
@@ -138,15 +138,6 @@ def parse_args() -> argparse.Namespace:
        required=True,
        help="Path(s) to input image file(s) (PNG, JPG, etc.). Can specify multiple images.",
    )
-    parser.add_argument(
-        "--deploy-config",
-        type=str,
-        default=None,
-        help=(
-            "Path to a deploy YAML. Required for multi-stage image-edit pipelines "
-            "whose deploy config is not auto-loaded."
-        ),
-    )
    parser.add_argument(
        "--prompt",
        type=str,
@@ -358,12 +349,7 @@ def parse_args() -> argparse.Namespace:
    parser.add_argument(
        "--enforce-eager",
        action="store_true",
-        default=None,
-        help=(
-            "Disable torch.compile and force eager execution. Left unset (None) "
-            "so it is only forwarded when explicitly given; "
-            "otherwise the per-stage deploy YAML value wins."
-        ),
+        help="Disable torch.compile and force eager execution.",
    )
    parser.add_argument(
        "--vae-use-slicing",
@@ -482,7 +468,7 @@ def main():
        }

    # Initialize Omni with appropriate pipeline
-    omni_kwargs: dict[str, Any] = dict(
+    omni = Omni(
        model=args.model,
        enable_layerwise_offload=args.enable_layerwise_offload,
        vae_use_slicing=args.vae_use_slicing,
@@ -490,15 +476,11 @@ def main():
        cache_backend=args.cache_backend,
        cache_config=cache_config,
        parallel_config=parallel_config,
+        enforce_eager=args.enforce_eager,
        enable_cpu_offload=args.enable_cpu_offload,
        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
        profiler_config=args.profiler_config,
    )
-    if args.enforce_eager is not None:
-        omni_kwargs["enforce_eager"] = args.enforce_eager
-    if args.deploy_config:
-        omni_kwargs["deploy_config"] = args.deploy_config
-    omni = Omni(**omni_kwargs)
    model_class_name = get_model_class_name(omni)
    declared_extra_body_params = get_extra_body_params(model_class_name)
    print("Pipeline loaded")
@@ -3,140 +3,14 @@
 This shared example generates videos from images with VACE, Wan2.2,
 LTX-2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, and other compatible pipelines.

- `image_to_video.py`: command-line script for single video generation with advanced options.
+## Local CLI Usage

-## Table of Contents
-
- [Overview](#overview)
- [Prerequisites](#prerequisites)
- [Quick Start](#quick-start)
- [Key Arguments](#key-arguments)
- [More CLI Examples](#more-cli-examples)
- [Advanced Features](#advanced-features)
- [FAQ](#faq)
-
-## Overview
-
-This folder provides a unified CLI script for image-to-video generation using vLLM-Omni diffusion/video pipelines. The script selects practical defaults for supported model families while still exposing common sampling, memory, and parallelism options.
-
-### Supported Models
-
-| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM Notes |
-| ----- | ------------------ | -------------- | ------------- | -------- | ---------- |
-| `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480 x 832 | 81 | 50 | 5.0 | Around 60 GiB BF16 for basic single-card usage |
-| `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 480 x 832 | 81 | 50 | 4.0 | Around 20–25 GiB BF16, smallest I2V model |
-| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v` | 480 x 832 | 121 | 50 | 6.0 | Around 100 GiB at default settings; the example enables `--enable-cpu-offload` + VAE tiling/slicing to fit an 80 GiB card |
-| LTX2 (local path + `--model-class-name LTX2ImageToVideoPipeline`) | 512 x 768 | 121 | 40 | 4.0 | Memory use depends on frame count and tensor parallelism |
-
-!!! info
-    Peak VRAM: based on basic single-card usage, batch size = 1, without any acceleration/optimization features. Some model weights cannot fit into one card with 80 GiB VRAM, which may need to use CPU offloading.
-
-Default model: `Wan-AI/Wan2.2-I2V-A14B-Diffusers`.
-
-## Prerequisites
-
-Download the example image used in the snippets below:
+Download the example image:

 ```bash
-curl -L -o cherry_blossom.jpg https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
+wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
 ```

-## Quick Start
-
-### Python API
-
-Single-prompt generation using TI2V-5B (lightweight):
-
-```python
-import PIL.Image
-import torch
-
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-if __name__ == "__main__":
-    image = PIL.Image.open("cherry_blossom.jpg").convert("RGB")
-    image = image.resize((576, 320))
-
-    omni = Omni(
-        model="Wan-AI/Wan2.2-TI2V-5B-Diffusers",
-        flow_shift=12.0,
-    )
-    outputs = omni.generate(
-        {
-            "prompt": "Cherry blossoms swaying gently in the breeze, petals falling",
-            "multi_modal_data": {"image": image},
-        },
-        OmniDiffusionSamplingParams(
-            height=320,
-            width=576,
-            num_frames=17,
-            num_inference_steps=20,
-            guidance_scale=4.0,
-            generator=torch.Generator(device="cuda").manual_seed(42),
-        ),
-    )
-    from diffusers.utils import export_to_video
-    frames = outputs[0].request_output.images
-    export_to_video(frames, "quick_test_i2v.mp4", fps=16)
-```
-
-### Local CLI Usage
-
-Quick test using TI2V-5B with a small resolution and few frames:
-
-```bash
-python image_to_video.py \
-  --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
-  --image cherry_blossom.jpg \
-  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
-  --height 320 \
-  --width 576 \
-  --num-frames 17 \
-  --guidance-scale 4.0 \
-  --num-inference-steps 20 \
-  --flow-shift 12.0 \
-  --fps 16 \
-  --output quick_test_i2v.mp4
-```
-
-## Key Arguments
-
-| Argument | Type | Default | Description |
-| -------- | ---- | ------- | ----------- |
-| `--model` | str | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | Diffusers I2V model ID or local path |
-| `--model-class-name` | str | `None` | Override model class name (e.g., `LTX2ImageToVideoPipeline`) |
-| `--image` | str | (required) | Path to input image |
-| `--prompt` | str | `""` | Text description of desired motion/animation |
-| `--negative-prompt` | str | `""` | Optional list of artifacts to suppress |
-| `--seed` | int | `42` | Random seed for deterministic sampling |
-| `--guidance-scale` | float | `5.0` | CFG scale |
-| `--guidance-scale-high` | float | `None` | Separate CFG for high-noise stage (MoE only) |
-| `--height` | int | auto | Video height (auto-calculated from image if not set). Multiples of 16 |
-| `--width` | int | auto | Video width (auto-calculated from image if not set). Multiples of 16 |
-| `--num-frames` | int | `81` | Number of frames |
-| `--num-inference-steps` | int | `50` | Number of denoising steps |
-| `--boundary-ratio` | float | `0.875` | Boundary split ratio for two-stage MoE models |
-| `--flow-shift` | float | `5.0` | Scheduler flow shift (5.0 for 720p, 12.0 for 480p) |
-| `--sample-solver` | str | `unipc` | Wan2.2 sampling solver (`unipc` or `euler` for Lightning/Distill) |
-| `--fps` | int | `None` | Frames per second for the saved MP4 |
-| `--frame-rate` | float | `None` | Generation frame rate for pipelines that require it (e.g., LTX2) |
-| `--output` | str | `i2v_output.mp4` | Path to save the generated video |
-| `--vae-use-slicing` | flag | off | Enable VAE slicing for memory optimization |
-| `--vae-use-tiling` | flag | off | Enable VAE tiling for memory optimization |
-| `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
-| `--enable-layerwise-offload` | flag | off | Enable layerwise offloading on DiT modules |
-| `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
-| `--tensor-parallel-size` | int | `1` | Tensor parallel size (effective for models that support TP, e.g. LTX2) |
-| `--ulysses-degree` | int | `1` | Ulysses sequence parallel degree |
-| `--ring-degree` | int | `1` | Ring sequence parallel degree |
-| `--cache-backend` | str | `None` | Cache backend: `cache_dit` or `tea_cache` |
-| `--use-hsdp` | flag | off | Enable Hybrid Sharded Data Parallel |
-| `--hsdp-shard-size` | int | `-1` | GPUs per shard group (-1 auto-calculates) |
-| `--hsdp-replicate-size` | int | `1` | Number of replica groups for HSDP |
-
-## More CLI Examples
-
 ### Wan2.2-I2V-A14B-Diffusers (MoE)

 ```bash
@@ -144,7 +18,7 @@ python image_to_video.py \
  --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
  --image cherry_blossom.jpg \
  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
-  --negative-prompt "low quality, blurry" \
+  --negative-prompt "<optional quality filter>" \
  --height 480 \
  --width 832 \
  --num-frames 48 \
@@ -154,7 +28,7 @@ python image_to_video.py \
  --boundary-ratio 0.875 \
  --flow-shift 12.0 \
  --fps 16 \
-  --output i2v_wan_moe.mp4
+  --output i2v_output.mp4
 ```

 ### Wan2.2-TI2V-5B-Diffusers (Unified)
@@ -164,7 +38,7 @@ python image_to_video.py \
  --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
  --image cherry_blossom.jpg \
  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
-  --negative-prompt "low quality, blurry" \
+  --negative-prompt "<optional quality filter>" \
  --height 480 \
  --width 832 \
  --num-frames 48 \
@@ -172,65 +46,7 @@ python image_to_video.py \
  --num-inference-steps 40 \
  --flow-shift 12.0 \
  --fps 16 \
-  --output i2v_wan_ti2v.mp4
-```
-
-### HunyuanVideo-1.5 I2V (480p)
-
-```bash
-python image_to_video.py \
-  --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v \
-  --image cherry_blossom.jpg \
-  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
-  --height 480 \
-  --width 832 \
-  --num-frames 121 \
-  --guidance-scale 6.0 \
-  --flow-shift 5.0 \
-  --num-inference-steps 50 \
-  --fps 24 \
-  --enable-cpu-offload \
-  --vae-use-tiling \
-  --vae-use-slicing \
-  --output hunyuan_i2v.mp4
-```
-
-### LTX2 Image-to-Video
-
-```bash
-python image_to_video.py \
-  --model /path/to/LTX-2 \
-  --model-class-name LTX2ImageToVideoPipeline \
-  --image cherry_blossom.jpg \
-  --prompt "A cinematic dolly shot of cherry blossoms" \
-  --height 512 \
-  --width 768 \
-  --num-frames 121 \
-  --num-inference-steps 40 \
-  --guidance-scale 4.0 \
-  --frame-rate 24 \
-  --fps 24 \
-  --output ltx2_i2v.mp4
-```
-
-## Advanced Features
-
-### CFG Parallel
-
-Set `--cfg-parallel-size 2` to enable CFG Parallel for faster inference on multi-GPU setups.
-See more examples in the [cfg_parallel user guide](../../../docs/user_guide/diffusion/parallelism/cfg_parallel.md).
-
-### Cache Acceleration
-
-Use `--cache-backend cache_dit` for Cache-DiT acceleration or `--cache-backend tea_cache` for Timestep Embedding Aware Cache:
-
-```bash
-python image_to_video.py \
-  --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
-  --image cherry_blossom.jpg \
-  --prompt "Cherry blossoms swaying gently in the breeze" \
-  --cache-backend cache_dit \
-  --output i2v_cached.mp4
+  --output i2v_output.mp4
 ```

 ### LTX-2.3
@@ -404,11 +220,3 @@ The VACE T2V command is documented in the shared

 For Wan2.2 LightX2V-converted local Diffusers directories and related LoRA
 assets, see the [LoRA guide](../../../docs/user_guide/diffusion/lora.md#wan22-lightx2v-offline-assembly).
-
-## FAQ
-
-**OOM errors**: Try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage. For very large models, add `--enable-cpu-offload` or `--enable-layerwise-offload`.
-
-**Auto-calculated resolution**: If `--height` and `--width` are not provided, the script calculates output dimensions from the input image while maintaining aspect ratio and targeting 480 x 832 area (or 512 x 768 for LTX2).
-
-**Wan2.2 MoE vs unified**: `I2V-A14B` is a larger Mixture-of-Experts model with separate low/high-noise DiT stages (use `--guidance-scale-high` and `--boundary-ratio`). `TI2V-5B` is a smaller unified T2V+I2V model that does not need these extra arguments.
@@ -19,7 +19,7 @@ hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-P
 ```bash
 python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
  --model ./MammothModa2-Preview \
-  --deploy-config ./vllm_omni/deploy/mammoth_moda2_ar.yaml \
+  --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
  --question "Summarize this image." \
  --image ./image.png
 ```
@@ -4,7 +4,7 @@ Offline inference example: MammothModa2 image summarization (single AR stage).
 Example:
  uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
    --model path/to/MammothModa2-Preview \
-    --deploy-config vllm_omni/deploy/mammoth_moda2_ar.yaml \
+    --stage-config vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
    --image /path/to/input.jpg \
    --question "Please summarize the content of this image."
 """
@@ -27,7 +27,9 @@ DEFAULT_QUESTION = "Please summarize the content of this image."
 def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).")
    parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.")
-    parser.add_argument("--deploy-config", type=str, required=True, help="Path to deploy config yaml.")
+    parser.add_argument(
+        "--stage-config", type=str, required=True, help="Path to stage config yaml (single-stage AR->text)."
+    )
    parser.add_argument("--image", type=str, required=True, help="Path to input image.")
    parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.")
    parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.")
@@ -71,7 +73,7 @@ def main() -> None:

    omni = Omni(
        model=args.model,
-        deploy_config=args.deploy_config,
+        stage_configs_path=args.stage_config,
        trust_remote_code=args.trust_remote_code,
        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
    )
@@ -90,40 +90,14 @@ Pass `--deploy-config /path/to/your_deploy.yaml` to any of the commands above to

 ### Image generation (text-to-image / img2img)

-Image generation is served through the standard task examples.
-The diffusion-stage knobs are declared centrally in `vllm_omni/model_extras/ming_flash_omni.py` and routed via `--extra-body`. With deploy yaml assigned properly, the model can run through the shared example scripts like any other diffusion model:
-
-Text-to-image (offline):
-
-```bash
-python examples/offline_inference/text_to_image/text_to_image.py \
-    --model Jonathan1909/Ming-flash-omni-2.0 \
-    --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml \
-    --prompt "Please draw a cute cat." \
-    --height 1024 \
-    --width 1024 \
-    --extra-body '{"steps": 30, "cfg": 2.0, "seed": 42}' \
-    --output ming_flash_omni_t2i.png
-```
-
-Image-to-image (offline):
-```bash
-# Reference image: figures/cases/person_gen_05.png from the upstream Ming repo
-# https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/examples/vllm_demo.py
-wget https://raw.githubusercontent.com/inclusionAI/Ming/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/figures/cases/person_gen_05.png
-
-python examples/offline_inference/image_to_image/image_edit.py \
-    --model Jonathan1909/Ming-flash-omni-2.0 \
-    --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml \
-    --image person_gen_05.png \
-    --prompt "Put a pair of sunglasses on the person." \
-    --extra-args '{"steps": 30, "cfg": 2.0, "seed": 42}' \
-    --output ming_flash_omni_i2i.png
-```
-
-For the online path and the full knob list
-(`steps`/`cfg`/`seed`/`byte5_text`/`negative_prompt`/`height`/`width`), see the
-[image-generation section in the recipe](../../../recipes/inclusionAI/Ming-flash-omni-2.0.md#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

@@ -60,24 +60,3 @@ If media file paths are not provided, the script will use default assets. Suppor
 - `mixed_modalities`: Audio + image + video
 - `use_audio_in_video`: Extract audio from video
 - `text`: Text-only query
-
-### Composable parallelism (strategy configs)
-
-You can shard or replicate a stage with a small composable-parallel
-`strategy.yaml` overlaid onto the **bundled default deploy config** -- no bespoke
-deploy YAML required. Supply the strategy with `--strategy-config` and any matching
-device layout via `--stage-overrides` (explicit kwargs forwarded to `Omni`).
-
-Both examples below require >= 3 GPUs (the thinker uses 2, talker + code2wav share a third):
-
-```bash
-# Tensor-parallel the thinker (TP=2)
-python end2end.py --query-type text --num-prompts 6 \
-    --strategy-config strategy_tp2.yaml \
-    --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}'
-
-# Replicate the thinker across 2 engines (round-robin load balancing)
-python end2end.py --query-type text --num-prompts 6 \
-    --strategy-config strategy_stage_replica.yaml \
-    --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}'
-```
@@ -553,39 +553,6 @@ def parse_args():
        default=False,
        help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
    )
-    parser.add_argument(
-        "--deploy-config",
-        type=str,
-        default=None,
-        help="Optional explicit deploy YAML (otherwise the bundled default is used).",
-    )
-    parser.add_argument(
-        "--strategy-config",
-        type=str,
-        default=None,
-        help=(
-            "Optional composable-parallel strategy.yaml mapping role -> parallel axes "
-            "(e.g. tp / stage_replica). Overlaid onto the merged stage configs."
-        ),
-    )
-    parser.add_argument(
-        "--stage-overrides",
-        type=str,
-        default=None,
-        help=(
-            "Optional JSON of per-stage overrides applied on top of the default "
-            'deploy config, e.g. \'{"0": {"devices": "0,1"}}\' to give the '
-            "thinker a 2-GPU pool. Lets you run a strategy on the bundled default "
-            "deploy config without writing a bespoke deploy YAML."
-        ),
-    )
-    parser.add_argument(
-        "--omni-lb-policy",
-        type=str,
-        default=None,
-        choices=["random", "round-robin", "least-queue-length"],
-        help="StagePool load-balancer policy for replicated stages (orchestrator-level knob).",
-    )
    return parser.parse_args()


@@ -1,22 +0,0 @@
-# Composable-parallel strategy for Qwen2.5-Omni.
-#
-# Replicate the thinker stage across 2 independent engines and let the omni
-# StagePool balance requests round-robin across them. The thinker's `devices`
-# entry acts as a per-replica template, so omni expands it to one GPU per replica.
-# The `routing: round_robin` below auto-derives the StagePool load-balancer policy,
-# so --omni-lb-policy is not needed.
-#
-# Rather than hand-write a bespoke deploy YAML, run on the bundled default deploy
-# config and pass the matching 2-GPU device template via --stage-overrides (explicit
-# kwargs): the thinker gets a two-GPU pool while talker + code2wav share a third
-# card. Requires >= 3 GPUs.
-#
-# Usage:
-#   python3 end2end.py --query-type text --num-prompts 6 \
-#       --strategy-config strategy_stage_replica.yaml \
-#       --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}'
-strategies:
-  thinker:
-    - axis: stage_replica
-      size: 2
-      routing: round_robin
@@ -1,19 +0,0 @@
-# Composable-parallel strategy for Qwen2.5-Omni: tensor-parallel the thinker.
-#
-# Shard the thinker stage across 2 GPUs with tensor parallelism (TP=2). TP routing
-# is fixed (broadcast), so no `routing` / load-balancer policy is involved -- this
-# is the simplest axis and a good starting point for the mental model.
-#
-# The strategy sets tensor_parallel_size=2 on the thinker. Rather than hand-write a
-# bespoke deploy YAML, run on the bundled default deploy config and pass a matching
-# 2-GPU device layout via --stage-overrides (explicit kwargs), so the thinker gets a
-# two-GPU pool while talker + code2wav share a third card. Requires >= 3 GPUs.
-#
-# Usage:
-#   python3 end2end.py --query-type text --num-prompts 6 \
-#       --strategy-config strategy_tp2.yaml \
-#       --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}'
-strategies:
-  thinker:
-    - axis: tp
-      size: 2
@@ -69,16 +69,7 @@ def parse_args() -> argparse.Namespace:
        "--stage-configs-path",
        type=str,
        default=None,
-        help="[Deprecated] Path to a legacy stage_args-format YAML. Prefer --deploy-config.",
-    )
-    parser.add_argument(
-        "--deploy-config",
-        type=str,
-        default=None,
-        help=(
-            "Path to a deploy YAML (new pipeline/stages format). Required for multi-stage "
-            "text-to-image pipelines whose deploy config is not auto-loaded."
-        ),
+        help="Path to a YAML file containing stage configurations for Omni.",
    )
    parser.add_argument("--prompt", default="a cup of coffee on the table", help="Text prompt for image generation.")
    parser.add_argument(
@@ -164,12 +155,7 @@ def parse_args() -> argparse.Namespace:
    parser.add_argument(
        "--enforce-eager",
        action="store_true",
-        default=None,
-        help=(
-            "Disable torch.compile and force eager execution. Left unset (None) "
-            "so it is only forwarded when explicitly given; "
-            "otherwise the per-stage deploy YAML value wins."
-        ),
+        help="Disable torch.compile and force eager execution.",
    )
    parser.add_argument(
        "--enable-cpu-offload",
@@ -235,13 +221,8 @@ def parse_args() -> argparse.Namespace:
    parser.add_argument(
        "--tensor-parallel-size",
        type=int,
-        default=None,
-        help=(
-            "Number of GPUs used for tensor parallelism (TP) inside the DiT. "
-            "Left unset so it is only forwarded when explicitly given; "
-            "otherwise the per-stage deploy YAML (or engine default of 1) wins. "
-            "Passing it always overrides the deploy YAML."
-        ),
+        default=1,
+        help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
    )
    parser.add_argument(
        "--enable-expert-parallel",
@@ -429,8 +410,10 @@ def main():
        "ring_degree": args.ring_degree,
        "ulysses_mode": args.ulysses_mode,
        "cfg_parallel_size": args.cfg_parallel_size,
+        "tensor_parallel_size": args.tensor_parallel_size,
        "vae_patch_parallel_size": args.vae_patch_parallel_size,
        "enable_expert_parallel": args.enable_expert_parallel,
+        "enforce_eager": args.enforce_eager,
        "enable_cpu_offload": args.enable_cpu_offload,
        "mode": "text-to-image",
        "log_stats": args.log_stats,
@@ -442,14 +425,8 @@ def main():
        **lora_args,
        **quant_kwargs,
    }
-    if args.tensor_parallel_size is not None:
-        omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size
-    if args.enforce_eager is not None:
-        omni_kwargs["enforce_eager"] = args.enforce_eager
    if args.stage_configs_path:
        omni_kwargs["stage_configs_path"] = args.stage_configs_path
-    if args.deploy_config:
-        omni_kwargs["deploy_config"] = args.deploy_config
    if use_nextstep:
        # NextStep-1.1 requires explicit pipeline class
        omni_kwargs["model_class_name"] = "NextStep11Pipeline"
@@ -474,9 +451,8 @@ def main():
    print(f"  Quantization: {args.quantization if args.quantization else 'None (BF16)'}")
    if ignored_layers:
        print(f"  Ignored layers: {ignored_layers}")
-    tp_display = args.tensor_parallel_size if args.tensor_parallel_size is not None else "deploy/default"
    print(
-        f"  Parallel configuration: tensor_parallel_size={tp_display}, "
+        f"  Parallel configuration: tensor_parallel_size={args.tensor_parallel_size}, "
        f"ulysses_degree={args.ulysses_degree}, ulysses_mode={args.ulysses_mode}, "
        f"ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, "
        f"vae_patch_parallel_size={args.vae_patch_parallel_size}, "
@@ -546,8 +522,8 @@ def main():
        diffusion_params.extra_args.update({k: v for k, v in user_extra.items() if v is not None})

    if lora_request:
-        diffusion_params.lora_request = lora_request
-        diffusion_params.lora_scale = args.lora_scale
+        diffusion_params.extra_args["lora_request"] = lora_request
+        diffusion_params.extra_args["lora_scale"] = args.lora_scale

    # Build per-stage sampling params for multi-stage models (e.g. BAGEL),
    # or wrap single diffusion params for single-stage models.
@@ -34,7 +34,7 @@ def run_e2e():
        help="Model name or path (HuggingFace or local)",
    )
    parser.add_argument(
-        "--deploy-config",
+        "--stage-config",
        type=str,
        default="vllm_omni/deploy/omnivoice.yaml",
    )
@@ -87,14 +87,14 @@ def run_e2e():
    )
    args = parser.parse_args()

-    if not os.path.exists(args.deploy_config):
-        raise FileNotFoundError(f"Deploy config not found: {args.deploy_config}")
+    if not os.path.exists(args.stage_config):
+        raise FileNotFoundError(f"Stage config not found: {args.stage_config}")

    print(f"Initializing OmniVoice with model={args.model}")

    omni = Omni(
        model=args.model,
-        deploy_config=args.deploy_config,
+        stage_configs_path=args.stage_config,
        log_stats=True,
    )

@@ -49,14 +49,22 @@ python examples/online_serving/openai_chat_completion_client_for_multimodal_gene
    --modalities text
 ```

+Parameterized curl wrapper in this directory:

-## Image generation (text-to-image / img2img)
+```bash
+bash run_curl_multimodal_generation.sh text
+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
+```

-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.
+## Image generation (text-to-image)

-The image-generation stage is a standard vLLM-Omni diffusion pipeline (`MingImagePipeline`); its request knobs are declared in `vllm_omni/model_extras/ming_flash_omni.py` and routed through `extra_body`, so they no longer need a bespoke `sampling_params_list` recipe (that form is still available for per-stage thinker sampling — see below).
-
-### Launch
+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 \
@@ -66,58 +74,32 @@ vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \
    --port 8091
 ```

-
-### Text-to-image
-
-Request image output with `"modalities": ["image"]`:
+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
+    -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
 ```

-Pass generation knobs under a literal `extra_body` object (the OpenAI Python client's `extra_body=` kwarg produces the same request). Keys are filtered against the declared set and routed into every stage's `extra_args`:
+### Optional 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"],
-    "messages": [
-      {
-        "role": "user",
-        "content": "Draw a poster."
-      }
-    ],
-    "extra_body": {
-      "steps": 6,
-      "cfg": 1.5,
-      "height": 512,
-      "width": 512,
-      "seed": 123,
-      "byte5_text": ["理解与生成统一"],
-      "negative_prompt": "ugly, blurry, distorted"
-    }
-  }' \
-  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
-  | base64 -d > ming_imagegen_extra_body.png
-```
+Pass image-gen overrides as flat keys on the diffusion-stage `sampling_params_list[1].extra_args`:

-NOTE: `extra_body` does **not** set the thinker's own sampling params. To tune the thinker (stage-0) sampling (`temperature` / `top_p` / `top_k` / `max_tokens`) or place knobs explicitly per stage, use `sampling_params_list` (`[thinker, imagegen]`).
-`negative_prompt` must sit on the **stage-0** entry to trigger the real-CFG companion; the imagegen knobs go on **stage-1**:
+| 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 (set on **stage-0 thinker** `extra_args` so `expand_cfg_prompts` spawns the companion). |
+
+Example with all the knobs:

 ```bash
 curl http://127.0.0.1:8091/v1/chat/completions \
@@ -161,48 +143,27 @@ curl http://127.0.0.1:8091/v1/chat/completions \

 ### img2img (reference image)

-Add an `image_url` content part (a base64 data URL) to the user message; it is routed into the DiT stage as `extra[reference_image]`. base64-encode a local file and stream it through `jq` via stdin - piping `base64 -> jq -> curl` so that avoids the shell `ARG_MAX` limit that inlining a large base64 string in `-d '…'` hits.
+Add an `image_url` content part to the user message; the parser routes it
+into the diffusion stage as `extra[reference_image]`:

-```bash
-# Reference image: figures/cases/person_gen_05.png from the upstream Ming repo
-# Check https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/examples/vllm_demo.py
-wget https://raw.githubusercontent.com/inclusionAI/Ming/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/figures/cases/person_gen_05.png
-
-base64 -w0 person_gen_05.png \
-| jq -Rs --arg prompt "Put a pair of sunglasses on the person." '{
-    model: "Jonathan1909/Ming-flash-omni-2.0",
-    modalities: ["image"],
-    messages: [
-      {
-        role: "user",
-        content: [
-          { type: "text", text: $prompt },
-          { type: "image_url", image_url: { url: ("data:image/png;base64," + (. | rtrimstr("\n"))) } }
-        ]
-      }
-    ]
-  }' \
-| curl http://127.0.0.1:8091/v1/chat/completions \
-    -H "Content-Type: application/json" \
-    -d @- \
-| jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
-| base64 -d > ming_img2img.png
+```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>"}}
+  ]
+}]
 ```

-The reference image can also be a public URL (`"url": "https://…/photo.jpg"`) or the simplified `{"image": "<base64>"}` content-part form — see the [image-to-image request formats](../image_to_image/README.md#request-format).
+### GPU layout

-### Knobs (declared `extra_body` params)
-
-| Key | Default | Description |
-| --- | --- | --- |
-| `height` / `width` | 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. |
-| `byte5_text` | (auto) | Glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "…". ` format. Auto-extracted from quoted spans in the prompt when omitted. |
-| `negative_prompt` | (empty) | Real CFG negative conditioning (text-to-image only). |
-
-For the offline `text_to_image.py` / `image_edit.py` scripts and the full knob reference, see the [image-generation section in the recipe](../../../recipes/inclusionAI/Ming-flash-omni-2.0.md#image-generation-text-to-image--img2img).
+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

@@ -0,0 +1,169 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Server port
+PORT="${PORT:-8091}"
+# Default query type
+QUERY_TYPE="${1:-text}"
+
+# Validate query type
+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|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
+
+# Define URLs for assets
+MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg"
+CHERRY_BLOSSOM_IMAGE_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
+SAMPLE_VIDEO_URL="https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
+
+# Build user content based on query type
+case "$QUERY_TYPE" in
+  text)
+    user_content='[
+      {
+        "type": "text",
+        "text": "请详细介绍鹦鹉的生活习性。"
+      }
+    ]'
+    ;;
+  use_image)
+    user_content='[
+        {
+          "type": "image_url",
+          "image_url": {
+            "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'"
+          }
+        },
+        {
+          "type": "text",
+          "text": "Describe this image in detail."
+        }
+      ]'
+    ;;
+  use_audio)
+    user_content='[
+        {
+          "type": "audio_url",
+          "audio_url": {
+            "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'"
+          }
+        },
+        {
+          "type": "text",
+          "text": "Please recognize the language of this speech and transcribe it. Format: oral."
+        }
+      ]'
+    ;;
+  use_video)
+    user_content='[
+        {
+          "type": "video_url",
+          "video_url": {
+            "url": "'"$SAMPLE_VIDEO_URL"'"
+          }
+        },
+        {
+          "type": "text",
+          "text": "Describe what is happening in this video."
+        }
+      ]'
+    ;;
+  use_mixed_modalities)
+    user_content='[
+        {
+          "type": "image_url",
+          "image_url": {
+            "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'"
+          }
+        },
+        {
+          "type": "audio_url",
+          "audio_url": {
+            "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'"
+          }
+        },
+        {
+          "type": "text",
+          "text": "Describe the image, and recognize the language of this speech and transcribe it. Format: oral"
+        }
+      ]'
+    ;;
+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",
+  "modalities": ["text"],
+  "messages": [
+    {
+      "role": "system",
+      "content": [
+        {
+          "type": "text",
+          "text": "你是一个友好的AI助手。\n\ndetailed thinking off"
+        }
+      ]
+    },
+    {
+      "role": "user",
+      "content": $user_content
+    }
+  ]
+}
+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")
+
+echo "Output of request: $(echo "$output" | jq '.choices[0].message.content')"
@@ -41,7 +41,6 @@ dev = [
    "pytest-cov>=4.0.0",
    "pytest-mock>=3.10.0",
    "websockets>=15.0.1",
-    "pytest-subtests>=0.11.0",
    "datasets>=2.14.0",
    "mypy==1.11.1",
    "pre-commit==4.0.1",
@@ -32,12 +32,12 @@ flags. Image size uses the standard `--height` / `--width` flags.
  [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py)
 - Declared parameters:
  [`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)
- Deploy config:
-  [`vllm_omni/deploy/mammoth_moda2.yaml`](../../vllm_omni/deploy/mammoth_moda2.yaml)
+- Stage config:
+  [`vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml`](../../vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml)

 ## Hardware Support

-The default deploy config runs both the AR and DiT stages on a single GPU
+The default stage config runs both the AR and DiT stages on a single GPU
 (`devices: "0"`). The committed `gpu_memory_utilization` split (stage-0 AR `0.5`,
 stage-1 DiT `0.3`) is sized for an ~80 GB GPU. The model also fits on a 48 GB GPU
 after rebalancing the split so the AR weights (~23 GB) leave room for the KV
@@ -48,7 +48,7 @@ cache — see the note under *1x L40S 48GB*.
 ### 1x L40S 48GB

 > **48 GB config adjustment:** the committed
-> `vllm_omni/deploy/mammoth_moda2.yaml` uses
+> `vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml` uses
 > `gpu_memory_utilization` 0.5 / 0.3 (sized for ~80 GB). To fit on a 48 GB L40S,
 > set the stage-0 (AR) value to `0.8` and the stage-1 (DiT) value to `0.16`
 > before running. (On an ~80 GB GPU, leave the defaults unchanged.)
@@ -71,13 +71,13 @@ hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-P
 ```

 Run text-to-image with the shared offline example from the repository root. The
-deploy config sets `trust_remote_code`, so no extra flag is needed. Forward the
+stage config sets `trust_remote_code`, so no extra flag is needed. Forward the
 MammothModa2 generation parameters as a JSON object through `--extra-body`:

 ```bash
 python examples/offline_inference/text_to_image/text_to_image.py \
  --model ./MammothModa2-Preview \
-  --stage-configs-path vllm_omni/deploy/mammoth_moda2.yaml \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \
  --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
  --height 1024 \
  --width 1024 \
@@ -102,9 +102,8 @@ curl http://localhost:8099/v1/chat/completions \
 ```

 **Text + speech in one response** (the headline 4.5 feature). The TTS
-path is gated by a Jinja flag on the chat template. Pass
-`use_tts_template=true` via the **top-level** `chat_template_kwargs`
-field (curl does not flatten nested `extra_body`):
+path is gated by a Jinja flag on the chat template, passed through
+`extra_body.chat_template_kwargs.use_tts_template=true`:

 ```bash
 curl http://localhost:8099/v1/chat/completions \
@@ -113,18 +112,12 @@ curl http://localhost:8099/v1/chat/completions \
        "model": "openbmb/MiniCPM-o-4_5",
        "messages": [{"role": "user", "content": "Say hello, then introduce vLLM in one sentence."}],
        "modalities": ["text", "audio"],
-        "chat_template_kwargs": {"use_tts_template": true}
+        "extra_body": {"chat_template_kwargs": {"use_tts_template": true}}
    }'
 ```

-When using the OpenAI Python SDK, the same flag can also be sent as
-`extra_body={"chat_template_kwargs": {"use_tts_template": True}}`
-because the client merges `extra_body` into the request root.
-
-Response carries text in one choice's `message.content` and base64 WAV
-in another choice's `message.audio.data` (24 kHz mono, see Notes). With
-`modalities: ["text", "audio"]` you typically get two `choices` entries
-(one text, one audio).
+Response carries text in `choices[0].message.content` and base64 WAV
+in `choices[0].message.audio.data` (24 kHz mono, see Notes).

 **Gradio demo (text + image + audio + video UI)**:

@@ -216,13 +209,9 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \
  missing dep raises `ImportError` at first request with the same
  install hint instead of silently emitting empty audio.

- **TTS trigger**: speech output requires
-  `chat_template_kwargs.use_tts_template=true` so the chat template
-  appends `<|tts_bos|>` before generation. Without it, Stage-1 talker
-  receives no TTS token span and returns silent audio (not text-only).
-  For **curl**, put `chat_template_kwargs` at the request root; nested
-  `extra_body.chat_template_kwargs` is ignored. The OpenAI Python SDK
-  may use `extra_body` because it flattens those fields into the root.
+- **TTS trigger**: speech output is only emitted when the client passes
+  `extra_body.chat_template_kwargs.use_tts_template=true`. Without it,
+  the response is text-only (which is also faster).

 - **Output audio**: 24 kHz mono WAV inside the OpenAI-style
  `message.audio.data` (base64). The Gradio demo's WAV player decodes
@@ -18,7 +18,6 @@ Use this recipe when you want a known-good starting point for serving

 - **Thinker only** — multimodal understanding with text output.
 - **Thinker + Talker (omni-speech)** — multimodal understanding with text and spoken output.
- **Thinker + Imagegen (DiT)** — text-to-image / img2img
 - **Talker only (TTS)** — standalone text-to-speech via the OpenAI `/v1/audio/speech` endpoint.

 ## References
@@ -161,129 +160,6 @@ deltas appear at `choices[0].delta.content`.
 - Memory usage: size depends on output modalities and multimodal input; leave
  headroom for video frames and audio caches.

-## Image generation (text-to-image / img2img)
-
-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.
-
-The image-generation stage is a standard vLLM-Omni diffusion pipeline (`MingImagePipeline`); its request knobs are declared in `vllm_omni/model_extras/ming_flash_omni.py` and routed through `extra_body`, so they no longer need a bespoke `sampling_params_list` recipe (that form is still available for per-stage thinker sampling — see below).
-
-```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
-```
-
-With fewer GPUs, copy the YAML and drop the thinker to TP=2 with the DiT on a
-free card.
-
-### Online (text-to-image)
-
-Request image output with `"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
-```
-
-**Quick form — `extra_body`** (keys are filtered against the declared set and routed into every stage's `extra_args`). Convenient, but it does **not** let you set the thinker (stage-0) sampling params. Wrap the knobs in a literal `"extra_body"` object for raw curl; the OpenAI Python client's `extra_body=` kwarg produces the same request:
-
-```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"],
-    "messages": [
-      {
-        "role": "user",
-        "content": "Draw a poster."
-      }
-    ],
-    "extra_body": {
-      "steps": 6,
-      "cfg": 1.5,
-      "height": 512,
-      "width": 512,
-      "seed": 123,
-      "byte5_text": ["理解与生成统一"],
-      "negative_prompt": "ugly, blurry, distorted"
-    }
-  }' \
-  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
-  | base64 -d > ming_imagegen_extra_body.png
-```
-
-**Full control — `sampling_params_list`** (one entry per stage: `[thinker, imagegen]`).
-Use this when you need to tune the thinker's own sampling (`temperature` / `top_p` / `top_k` / `max_tokens`), or to place knobs explicitly per stage.
-Note `negative_prompt` must sit on the **stage-0 thinker** `extra_args`; the imagegen-stage knobs (`steps` / `cfg` / `height` / `width` / `seed` / `byte5_text`) go on the **stage-1** entry:
-
-```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": {
-          "negative_prompt": "ugly, blurry, distorted"
-        }
-      },
-      {
-        "seed": 42,
-        "extra_args": {
-          "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
-```
-
-### Knobs (declared `extra_body` params)
-
-| Key | Default | Description |
-| --- | --- | --- |
-| `height` / `width` | 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. |
-| `byte5_text` | (auto) | Glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "…". ` format. Auto-extracted from quoted spans in the prompt when omitted. |
-| `negative_prompt` | (empty) | Real CFG negative conditioning. Spawns a CFG-text companion via `expand_cfg_prompts`; **online / text-to-image only** (offline uses Ming's default zero-negative). |
-
-For img2img, add an `image_url` content part to the user message (online) or pass `--image` (offline); the reference image is routed into the DiT stage as `extra[reference_image]`.
-
 ### 1x H100 80GB — standalone TTS (talker only)

 The bundled `ming_flash_omni_tts.yaml` runs the talker on a single GPU and exposes the OpenAI `/v1/audio/speech` endpoint.
@@ -13,7 +13,7 @@ transformers >= 5.5.3
 av>=14.0.0
 omegaconf>=2.3.0
 diffusers==0.38.0
-safetensors>=0.8.0
+safetensors>=0.8.0rc0
 accelerate==1.12.0
 soundfile>=0.13.1
 cache-dit==1.3.0
@@ -1 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
@@ -1,136 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Tests for the generic engine-backend dispatcher + AR-Diffusion runner wiring.
-
-``DiffusionEngine.resolve_engine_class`` is a generic dispatcher (``"default"`` / a
-``DiffusionEngine`` subclass / an import-path string). DreamZero selects the
-AR-Diffusion engine via its deploy config's ``engine_backend`` qualname — no
-DreamZero/ar_diffusion knowledge lives in the public base, so the routing check
-here is simply that the ``ARDiffusionEngine`` qualname resolves correctly.
-"""
-
-from dataclasses import fields
-from types import SimpleNamespace
-
-import pytest
-from vllm.utils.import_utils import resolve_obj_by_qualname
-
-from vllm_omni.diffusion.data import OmniDiffusionConfig
-from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
-from vllm_omni.experimental.ar_diffusion.engine import (
-    AR_DIFFUSION_MODEL_RUNNER_CLS,
-    ARDiffusionEngine,
-)
-
-
-def _cfg(backend):
-    return SimpleNamespace(engine_backend=backend)
-
-
-def test_default_resolves_to_diffusion_engine():
-    assert DiffusionEngine.resolve_engine_class(_cfg("default")) is DiffusionEngine
-
-
-def test_subclass_type_is_returned():
-    assert DiffusionEngine.resolve_engine_class(_cfg(ARDiffusionEngine)) is ARDiffusionEngine
-
-
-def test_ar_diffusion_qualname_resolves():
-    # How DreamZero's deploy config selects the engine: a full import-path string.
-    cls = DiffusionEngine.resolve_engine_class(_cfg("vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"))
-    assert cls is ARDiffusionEngine
-
-
-def test_non_engine_type_raises():
-    with pytest.raises(TypeError):
-        DiffusionEngine.resolve_engine_class(_cfg(dict))
-
-
-def test_qualname_not_an_engine_raises():
-    with pytest.raises(TypeError):
-        DiffusionEngine.resolve_engine_class(_cfg("builtins.dict"))
-
-
-def test_bad_qualname_raises():
-    with pytest.raises(ValueError):
-        DiffusionEngine.resolve_engine_class(_cfg("not.a.real.module.NoSuchEngine"))
-
-
-def test_ar_diffusion_engine_is_diffusion_engine_subclass():
-    assert issubclass(ARDiffusionEngine, DiffusionEngine)
-
-
-def test_config_engine_backend_field_defaults_to_default():
-    field = {f.name: f for f in fields(OmniDiffusionConfig)}["engine_backend"]
-    assert field.default == "default"
-
-
-# --- AR-Diffusion worker/runner wiring -----------------------------------------------
-
-
-def test_ar_diffusion_model_runner_cls_resolves():
-    cls = resolve_obj_by_qualname(AR_DIFFUSION_MODEL_RUNNER_CLS)
-    assert cls.__name__ == "ARDiffusionModelRunner"
-
-
-def test_engine_class_declares_its_runner():
-    """Routing lives on the engine class (review: hsliuustc0106) — engines
-    never mutate od_config; the worker resolves the runner off the class."""
-    assert ARDiffusionEngine.default_diffusion_model_runner_cls == AR_DIFFUSION_MODEL_RUNNER_CLS
-    assert DiffusionEngine.default_diffusion_model_runner_cls is None
-
-
-def _select_runner(od_config, platform_default="PLATFORM"):
-    """Mirror of the worker's runner-selection chain (override > engine > platform)."""
-    runner_override = getattr(od_config, "diffusion_model_runner_cls", None)
-    if isinstance(runner_override, str) and runner_override:
-        return runner_override
-    engine_cls = DiffusionEngine.resolve_engine_class(od_config)
-    engine_runner = getattr(engine_cls, "default_diffusion_model_runner_cls", None)
-    if isinstance(engine_runner, str) and engine_runner:
-        return engine_runner
-    return platform_default
-
-
-def test_runner_selection_chain():
-    ar = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
-    # engine-declared runner wins over platform default
-    assert (
-        _select_runner(SimpleNamespace(diffusion_model_runner_cls=None, engine_backend=ar))
-        == AR_DIFFUSION_MODEL_RUNNER_CLS
-    )
-    # explicit override wins over the engine's declaration
-    assert (
-        _select_runner(SimpleNamespace(diffusion_model_runner_cls="my.custom.Runner", engine_backend=ar))
-        == "my.custom.Runner"
-    )
-    # default engine -> platform default
-    assert _select_runner(SimpleNamespace(diffusion_model_runner_cls=None, engine_backend="default")) == "PLATFORM"
-
-
-def test_dreamzero_pipeline_rejects_non_ar_diffusion_engine():
-    """Review (hsliuustc0106): a stale config with engine_backend='default'
-    must fail at init, not mid-forward on the first KV access."""
-    from types import SimpleNamespace
-
-    import pytest
-
-    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
-
-    with pytest.raises(ValueError, match="requires the AR-Diffusion engine"):
-        DreamZeroPipeline(od_config=SimpleNamespace(engine_backend="default"))
-    with pytest.raises(ValueError, match="requires the AR-Diffusion engine"):
-        DreamZeroPipeline(od_config=SimpleNamespace())
-
-
-def test_runner_rejects_multi_sequence_config():
-    """Review (hsliuustc0106): the paged path writes batch index 0 only, so
-    max_num_seqs > 1 must be rejected at init."""
-    from types import SimpleNamespace
-
-    import pytest
-
-    from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
-
-    fake = SimpleNamespace(od_config=SimpleNamespace(max_num_seqs=4))
-    with pytest.raises(ValueError, match="max_num_seqs=1"):
-        ARDiffusionModelRunner._preallocate_kv_cache(fake)
@@ -1,254 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Unit tests for the AR-Diffusion KV cache helpers (Phase 1, PR-2).
-
-Covers the request adapter, the chunk-window spec/manager (registration + the
-eviction policy), and the pool builder — exercised against the installed vLLM
-KV stack on CPU (block bookkeeping only, no GPU tensors).
-"""
-
-import pytest
-import torch
-from vllm.v1.kv_cache_interface import KVCacheSpecKind, get_kv_cache_spec_kind
-from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
-from vllm.v1.request import RequestStatus
-
-from vllm_omni.experimental.ar_diffusion.kv_cache import (
-    ARDiffusionKVCache,
-    ARDiffusionKVConfig,
-    ARDiffusionRequestAdapter,
-    ChunkWindowManager,
-    ChunkWindowSpec,
-    allocate_kv_pool_with_views,
-    build_kv_manager,
-    compute_num_blocks,
-)
-from vllm_omni.experimental.ar_diffusion.kv_cache.paged import chunk_window_skipped_tokens
-
-BLOCK = 16
-
-
-def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False):
-    return ChunkWindowSpec(
-        block_size=BLOCK,
-        num_kv_heads=4,
-        head_size=64,
-        dtype=torch.float16,
-        sliding_window=window_chunks * chunk_size,
-        chunk_size=chunk_size,
-        window_chunks=window_chunks,
-        sink_chunks=sink_chunks,
-        reset_at_boundary=reset_at_boundary,
-    )
-
-
-# --- ChunkWindowSpec registration -------------------------------------------
-
-
-def test_spec_registration_resolves_to_chunk_window_manager():
-    # Without explicit registration the MRO walk would fall back to the parent
-    # SlidingWindowManager; assert the subclass manager wins.
-    spec = make_spec()
-    assert KVCacheSpecRegistry.get_manager_class(spec) is ChunkWindowManager
-
-
-def test_spec_kind_is_sliding_window():
-    assert get_kv_cache_spec_kind(make_spec()) == KVCacheSpecKind.SLIDING_WINDOW
-
-
-def test_spec_rejects_inconsistent_window():
-    with pytest.raises(ValueError):
-        ChunkWindowSpec(
-            block_size=BLOCK,
-            num_kv_heads=4,
-            head_size=64,
-            dtype=torch.float16,
-            sliding_window=99,  # != window_chunks * chunk_size
-            chunk_size=BLOCK,
-            window_chunks=2,
-        )
-
-
-# --- eviction policy (pure) -------------------------------------------------
-
-
-def test_sliding_replace_keeps_window():
-    # window = 2 chunks * 16 = 32. Base sliding formula keeps `window` tokens;
-    # the snap is to chunk boundaries.
-    def skip(n):
-        return chunk_window_skipped_tokens(n, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False)
-
-    assert skip(32) == 0  # nothing past the window yet
-    assert skip(48) == 16  # one chunk fell out of the window
-    assert skip(64) == 32
-
-
-def test_sliding_replace_snaps_to_chunk_boundary():
-    # A non-chunk-aligned overflow must snap down so a chunk is never half-evicted.
-    skip = chunk_window_skipped_tokens(50, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False)
-    assert skip % 16 == 0 and skip == 16
-
-
-def test_sink_chunks_protected():
-    # sink = 1 chunk (16 tokens) is never skipped.
-    skip = chunk_window_skipped_tokens(80, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=False)
-    # raw overflow = 80 - 32 - 16 = 32 -> snapped 32
-    assert skip == 32
-
-
-def test_reset_at_boundary_drops_completed_past_sink():
-    skip = chunk_window_skipped_tokens(48, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=True)
-    # completed = 48; sink = 16 -> drop 32
-    assert skip == 32
-
-
-# --- ARDiffusionKVConfig ------------------------------------------------------------
-
-
-def test_kv_config_sliding_window_property():
-    assert ARDiffusionKVConfig(chunk_size=16, window_chunks=3).sliding_window == 48
-    assert ARDiffusionKVConfig(chunk_size=16, window_chunks=None).sliding_window is None
-
-
-# --- ARDiffusionRequestAdapter ------------------------------------------------------
-
-
-def test_adapter_advances_per_chunk_not_per_step():
-    a = ARDiffusionRequestAdapter("r0", chunk_size=16)
-    assert a.num_computed_tokens == 0
-    assert a.num_tokens == 16  # in-flight chunk
-    a.on_chunk_committed()
-    assert a.num_computed_tokens == 16
-    assert a.num_tokens == 32
-
-
-def test_adapter_accounts_for_prefill_prefix():
-    a = ARDiffusionRequestAdapter("r0", chunk_size=16, prefill_prefix_tokens=4)
-    assert a.num_computed_tokens == 4
-    assert a.num_prompt_tokens == 4
-    assert a.num_tokens == 20
-
-
-def test_adapter_status_is_vllm_enum():
-    assert isinstance(ARDiffusionRequestAdapter("r0", chunk_size=16).status, RequestStatus)
-
-
-# --- pool / manager ---------------------------------------------------------
-
-
-def test_compute_num_blocks():
-    # 1 MiB budget at fraction 0.5 with 16 KiB pages -> 32 blocks.
-    assert compute_num_blocks(1 << 20, 0.5, 16 << 10) == 32
-
-
-def test_paged_pool_layout_exposes_flat_slot_views():
-    kv_pools, k_pools, v_pools = allocate_kv_pool_with_views(
-        num_blocks=4,
-        block_size=BLOCK,
-        num_layers=1,
-        num_kv_heads=4,
-        head_dim=64,
-        dtype=torch.float32,
-        device=torch.device("cpu"),
-    )
-    assert kv_pools[0].shape == (2, 4, BLOCK, 4, 64)
-    assert k_pools[0].shape == (4 * BLOCK, 4, 64)
-    assert v_pools[0].shape == (4 * BLOCK, 4, 64)
-
-    k_pools[0][BLOCK + 3].fill_(7)
-    v_pools[0][2 * BLOCK + 5].fill_(11)
-    assert torch.equal(kv_pools[0][0, 1, 3], k_pools[0][BLOCK + 3])
-    assert torch.equal(kv_pools[0][1, 2, 5], v_pools[0][2 * BLOCK + 5])
-
-
-def test_build_manager_allocate_free_roundtrip():
-    """End-to-end: a ARDiffusionRequestAdapter drives a real KVCacheManager.
-
-    This is the adapter conformance check — if the adapter were missing an
-    attribute the manager reads, allocate_slots/free would raise here.
-    """
-    spec = make_spec()
-    mgr = build_kv_manager(spec, ["layer0"], num_blocks=16, max_model_len=1024)
-    free_before = mgr.block_pool.get_num_free_blocks()
-
-    adapter = ARDiffusionRequestAdapter("req-0", chunk_size=BLOCK)
-    blocks = mgr.allocate_slots(adapter, num_new_tokens=BLOCK, full_sequence_must_fit=True)
-    assert blocks is not None
-    assert mgr.block_pool.get_num_free_blocks() < free_before
-
-    mgr.free(adapter)
-    assert mgr.block_pool.get_num_free_blocks() == free_before
-
-
-def test_cross_attn_pool_deducted_from_self_attn_budget():
-    """The cross-attn pool is allocated directly; its bytes are subtracted from
-    the self-attn paged-pool budget so the two together stay within the free
-    memory budget (review: zwhzzz0821)."""
-    L = 512
-    avail = 1 << 30  # 1 GiB
-    kv = ARDiffusionKVCache(
-        ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2, gpu_memory_fraction=0.5),
-        num_layers=2,
-        num_kv_heads=4,
-        head_size=64,
-        dtype=torch.float16,
-        block_size=BLOCK,
-        max_model_len=4096,
-        available_bytes=avail,
-        cross_attn_length=L,
-        device=torch.device("cpu"),
-    )
-    cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2  # K+V, pos+neg, layers
-    expected = compute_num_blocks(avail - cross_bytes, 0.5, kv.spec.page_size_bytes * 2)
-    assert kv.num_blocks == expected
-    assert expected > 2 * (2 + 1) + 2  # above the min-blocks floor, so the deduction is what's tested
-
-
-def _make_kv(*, local_branches, num_frame_per_block=2, window_chunks=9):
-    return ARDiffusionKVCache(
-        ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks),
-        num_layers=1,
-        num_kv_heads=4,
-        head_size=64,
-        dtype=torch.float32,
-        block_size=BLOCK,
-        max_model_len=4096,
-        available_bytes=1 << 16,  # tiny -> the floor binds
-        device=torch.device("cpu"),
-        local_branches=local_branches,
-        num_frame_per_block=num_frame_per_block,
-    )
-
-
-def test_pool_floor_is_branch_aware():
-    """CFG-parallel rank (one local branch) sizes for one window + in-flight chunk;
-    a single-process run (both branches) sizes for two. Scratch scales the same way."""
-    one = _make_kv(local_branches=1)
-    two = _make_kv(local_branches=2)
-
-    assert one.managed_num_blocks == 1 * (9 + 2) + 2  # 13
-    assert two.managed_num_blocks == 2 * (9 + 2) + 2  # 24
-    assert one.scratch_num_blocks == one.scratch_blocks_per_branch
-    assert two.scratch_num_blocks == 2 * two.scratch_blocks_per_branch
-    assert one.num_blocks_total == 13 + one.scratch_blocks_per_branch
-
-
-def test_scratch_maps_to_slot_zero_with_one_local_branch():
-    """A CFG-parallel rank runs exactly one branch: whichever CFG side it is,
-    its scratch lands in the rank's single slot (no dead second slot)."""
-    one = _make_kv(local_branches=1)
-    assert one.scratch_block_ids(True, 0, 2) == one.scratch_block_ids(False, 0, 2)
-
-    two = _make_kv(local_branches=2)
-    assert two.scratch_block_ids(True, 0, 2) != two.scratch_block_ids(False, 0, 2)
-
-
-def test_scratch_exhaustion_still_raises():
-    one = _make_kv(local_branches=1)
-    cap = one.scratch_blocks_per_branch
-    with pytest.raises(RuntimeError, match="scratch blocks exhausted"):
-        one.scratch_block_ids(False, 0, cap + 1)
-
-
-def test_invalid_local_branches_rejected():
-    with pytest.raises(ValueError, match="local_branches"):
-        _make_kv(local_branches=3)
@@ -1,93 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Tests for ARDiffusionKVCache — the engine-level KV orchestration body (Phase 1)."""
-
-import pytest
-import torch
-
-from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVCache, ARDiffusionKVConfig
-
-DIMS = dict(num_layers=2, num_kv_heads=4, head_size=64, dtype=torch.float16, block_size=16)
-
-
-def make_cache(*, chunk_size=16, window_chunks=2, available_bytes=1 << 24):
-    cfg = ARDiffusionKVConfig(enable=True, chunk_size=chunk_size, window_chunks=window_chunks)
-    return ARDiffusionKVCache(cfg, max_model_len=4096, available_bytes=available_bytes, **DIMS)
-
-
-def test_requires_enabled_config():
-    with pytest.raises(ValueError):
-        ARDiffusionKVCache(ARDiffusionKVConfig(enable=False), max_model_len=256, available_bytes=1 << 20, **DIMS)
-
-
-def test_requires_bounded_window():
-    cfg = ARDiffusionKVConfig(enable=True, chunk_size=16, window_chunks=None)
-    with pytest.raises(ValueError):
-        ARDiffusionKVCache(cfg, max_model_len=256, available_bytes=1 << 20, **DIMS)
-
-
-def test_full_request_lifecycle_and_eviction():
-    """begin -> per-chunk allocate/slots/commit over a long rollout -> free.
-
-    Exercises the orchestrator end-to-end and asserts the chunk window bounds
-    memory (pool plateaus) and frees cleanly.
-    """
-    kv = make_cache(chunk_size=16, window_chunks=2)
-    free_total = kv.manager.block_pool.get_num_free_blocks()
-
-    adapter = kv.begin_request("req-0")
-    free_after = []
-    for k in range(10):
-        block_table = kv.allocate_chunk(adapter)
-        slots = kv.chunk_write_slots(adapter)
-        # Slots target real blocks in this chunk's table; length == chunk_size.
-        assert len(slots) == kv.spec.chunk_size
-        used = {int(s) // kv.block_size for s in slots}
-        assert kv.null_block_id not in used
-        assert used <= set(block_table)
-        # Resident window stays bounded by window + the in-flight chunk.
-        assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks + 1
-        free_after.append(kv.manager.block_pool.get_num_free_blocks())
-        kv.commit_chunk(adapter)
-
-    # Pool memory plateaus once the window is full (eviction recycles blocks).
-    assert free_after[-1] == free_after[kv.spec.window_chunks]
-
-    kv.end_request(adapter)
-    assert kv.manager.block_pool.get_num_free_blocks() == free_total
-
-
-def test_num_computed_advances_per_chunk():
-    kv = make_cache(chunk_size=16, window_chunks=2)
-    a = kv.begin_request("r")
-    assert a.num_computed_tokens == 0
-    kv.allocate_chunk(a)
-    kv.commit_chunk(a)
-    assert a.num_computed_tokens == 16
-    kv.end_request(a)
-
-
-def test_state_close_frees_both_branch_blocks():
-    """ARDiffusionKVState.close() returns both CFG branches' pool blocks to the pool.
-
-    This is the primitive the runner's LRU eviction relies on: when a session is
-    evicted, close() must free the blocks both adapters hold, or session churn
-    leaks pool ownership (review P1).
-    """
-    from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
-
-    kv = make_cache(chunk_size=16, window_chunks=2)
-    free_total = kv.manager.block_pool.get_num_free_blocks()
-
-    pos = kv.begin_request("bde__s")
-    neg = kv.begin_request("bde__s__neg")
-    state = ARDiffusionKVState(kv, pos, neg, num_layers=kv.num_layers)
-    for _ in range(3):
-        for adapter in (pos, neg):
-            kv.allocate_chunk(adapter)
-            kv.commit_chunk(adapter)
-    # Both branches hold resident blocks now.
-    assert kv.manager.block_pool.get_num_free_blocks() < free_total
-
-    state.close()
-    # All blocks returned — no leak across an evicted session.
-    assert kv.manager.block_pool.get_num_free_blocks() == free_total
@@ -1,394 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Tests for AR-Diffusion paged self-attention contexts."""
-
-from __future__ import annotations
-
-import subprocess
-from importlib.util import find_spec
-from pathlib import Path
-from unittest.mock import MagicMock
-
-import pytest
-import torch
-
-from vllm_omni.experimental.ar_diffusion.kv_cache import (
-    ARDiffusionKVCache,
-    ARDiffusionKVConfig,
-    ARDiffusionPagedLayerContext,
-    ARDiffusionPagedLayerInputs,
-    ar_diffusion_paged_attention,
-    paged_write_attn,
-)
-from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
-
-BLOCK = 16
-N_HEADS = 4
-HEAD_DIM = 64
-
-
-def make_state(*, num_layers=1, window_chunks=2, dtype=torch.float32, device=torch.device("cpu")):
-    cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
-    kv = ARDiffusionKVCache(
-        cfg,
-        num_layers=num_layers,
-        num_kv_heads=N_HEADS,
-        head_size=HEAD_DIM,
-        dtype=dtype,
-        block_size=BLOCK,
-        max_model_len=4096,
-        available_bytes=1 << 26,
-        device=device,
-    )
-    pos = kv.begin_request("r-pos")
-    neg = kv.begin_request("r-neg")
-    return kv, ARDiffusionKVState(kv, pos, neg, num_layers=num_layers)
-
-
-def _dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
-    scores = torch.einsum("bqhd,bkhd->bhqk", query.float(), key.float()) * (HEAD_DIM**-0.5)
-    probs = torch.softmax(scores, dim=-1).to(value.dtype)
-    return torch.einsum("bhqk,bkhd->bqhd", probs, value)
-
-
-def _cuda_flash_attn_usable() -> bool:
-    if not torch.cuda.is_available():
-        return False
-    try:
-        spec = find_spec("vllm.vllm_flash_attn")
-        if spec is None or spec.origin is None:
-            return True
-        fa2_so = Path(spec.origin).parent / "_vllm_fa2_C.abi3.so"
-        linked = subprocess.check_output(["ldd", str(fa2_so)], text=True, timeout=5)
-    except Exception:
-        return True
-    if "libcudart.so.13" not in linked:
-        return True
-    try:
-        out = subprocess.check_output(
-            ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"],
-            text=True,
-            timeout=5,
-        )
-        driver_major = int(out.splitlines()[0].split(".")[0])
-    except Exception:
-        return True
-    return driver_major >= 580
-
-
-def _commit_video_span(
-    kv: ARDiffusionKVCache,
-    st: ARDiffusionKVState,
-    *,
-    is_negative: bool,
-    n_chunks: int,
-    dtype: torch.dtype,
-    device: torch.device,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    ctx = st.get_kv_caches(is_negative, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
-    ctx.ensure_video_slots(device)
-    k = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-    v = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-    kv._k_pools[0][ctx.current_video_slot_mapping] = k[0]
-    kv._v_pools[0][ctx.current_video_slot_mapping] = v[0]
-    st.commit_paged_context(is_negative)
-    return k, v
-
-
-def test_paged_context_allocates_lazily_and_commits_after_forward():
-    _, st = make_state()
-
-    contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)
-    ctx = contexts[0].forward_ctx
-    assert isinstance(contexts[0], ARDiffusionPagedLayerContext)
-    assert st.pos.completed_chunks == 0
-    assert ctx.current_video_slot_mapping is None
-
-    ctx.ensure_video_slots(torch.device("cpu"))
-    assert st.pos.completed_chunks == 0
-    assert len(ctx.current_video_block_ids) == 1
-
-    st.commit_paged_context(False)
-    assert st.pos.completed_chunks == 1
-    assert st._committed[False] == BLOCK
-
-
-def test_scratch_video_and_action_blocks_do_not_commit():
-    kv, st = make_state()
-
-    ctx = st.get_kv_caches(False, seq_len=2 * BLOCK, commit_current=False)[0].forward_ctx
-    ctx.ensure_video_slots(torch.device("cpu"))
-    ctx.ensure_action_slots(3, torch.device("cpu"))
-
-    assert ctx.current_video_block_ids == kv.scratch_block_ids(False, 0, 2)
-    assert ctx.action_scratch_block_ids == kv.scratch_block_ids(False, 2, 1)
-    st.commit_paged_context(False)
-    assert st.pos.completed_chunks == 0
-    assert st._committed[False] == 0
-
-
-def test_pipeline_kv_get_paged_path_has_no_gather_backend():
-    kv, st = make_state()
-    assert not hasattr(kv, "gather_window_all_layers")
-
-    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
-
-    pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline)
-    pipeline._ar_diffusion_kv_state = st
-    contexts = pipeline._kv_get(MagicMock(), False, seq_len=BLOCK, update_kv_cache=False)
-
-    assert len(contexts) == 1
-    assert isinstance(contexts[0], ARDiffusionPagedLayerContext)
-
-
-@pytest.mark.parametrize("history_chunks", [0, 1, 3])
-@pytest.mark.parametrize("action_len", [0, 3])
-@pytest.mark.parametrize("commit_current", [False, True])
-def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len, commit_current):
-    torch.manual_seed(0)
-    device = torch.device("cpu")
-    dtype = torch.float32
-    kv, st = make_state(dtype=dtype, device=device, window_chunks=2)
-
-    history_k_parts: list[torch.Tensor] = []
-    history_v_parts: list[torch.Tensor] = []
-    if history_chunks:
-        k, v = _commit_video_span(
-            kv,
-            st,
-            is_negative=False,
-            n_chunks=history_chunks,
-            dtype=dtype,
-            device=device,
-        )
-        history_k_parts.append(k)
-        history_v_parts.append(v)
-
-    ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=commit_current)[0].forward_ctx
-    ctx.ensure_video_slots(device)
-    current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-    current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-    kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0]
-    kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0]
-
-    action_k = action_v = None
-    if action_len:
-        ctx.ensure_action_slots(action_len, device)
-        action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-        action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-        kv._k_pools[0][ctx.action_slot_mapping] = action_k[0]
-        kv._v_pools[0][ctx.action_slot_mapping] = action_v[0]
-
-    query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-    block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
-        action_len=action_len,
-        query_len=query.shape[1],
-        device=device,
-    )
-    paged = ar_diffusion_paged_attention(
-        query,
-        kv.key_cache(0),
-        kv.value_cache(0),
-        block_table=block_table,
-        query_start_loc=query_start_loc,
-        seq_lens=seq_lens,
-        max_query_len=max_query_len,
-        max_seq_len=max_seq_len,
-        softmax_scale=HEAD_DIM**-0.5,
-        causal=False,
-    )
-
-    if history_k_parts:
-        history_k = torch.cat(history_k_parts, dim=1)
-        history_v = torch.cat(history_v_parts, dim=1)
-    else:
-        history_k = torch.empty(1, 0, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-        history_v = torch.empty(1, 0, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-    new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
-    new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
-    if action_len:
-        new_k = torch.cat([new_k, action_k], dim=1)
-        new_v = torch.cat([new_v, action_v], dim=1)
-    ref = _dense_attention(query, new_k, new_v)
-
-    torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)
-
-    before = st.pos.completed_chunks
-    st.commit_paged_context(False)
-    assert st.pos.completed_chunks == before + (1 if commit_current else 0)
-
-
-@pytest.mark.skipif(not _cuda_flash_attn_usable(), reason="usable CUDA FlashAttention is required")
-@pytest.mark.parametrize("history_chunks", [1, 3])
-@pytest.mark.parametrize("action_len", [0, 3])
-@pytest.mark.parametrize("commit_current", [False, True])
-def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len, commit_current):
-    pytest.importorskip("vllm.vllm_flash_attn")
-    torch.manual_seed(0)
-    device = torch.device("cuda")
-    dtype = torch.float16
-    kv, st = make_state(dtype=dtype, device=device, window_chunks=2)
-
-    history_k, history_v = _commit_video_span(
-        kv,
-        st,
-        is_negative=False,
-        n_chunks=history_chunks,
-        dtype=dtype,
-        device=device,
-    )
-
-    layer_ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=commit_current)[0]
-    ctx = layer_ctx.forward_ctx
-    current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-    current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-    action_k = action_v = None
-    if action_len:
-        action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-        action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-
-    query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-
-    # The production path: once-per-forward host prep, then the fused
-    # write+attend custom op consuming the NamedTuple payload.
-    ctx.prepare(device=device, action_len=action_len, query_len=query.shape[1])
-    inputs = layer_ctx.to_layer_inputs()
-    assert isinstance(inputs, ARDiffusionPagedLayerInputs)
-    paged = paged_write_attn(
-        inputs,
-        query[0],
-        current_k[0],
-        current_v[0],
-        action_k[0] if action_len else None,
-        action_v[0] if action_len else None,
-        HEAD_DIM**-0.5,
-    ).unsqueeze(0)
-
-    # Direct python-fn call on the same (already written) pools must be
-    # bit-exact: identical kernel, identical inputs.
-    direct = ar_diffusion_paged_attention(
-        query,
-        kv.key_cache(0),
-        kv.value_cache(0),
-        block_table=ctx.block_table,
-        query_start_loc=ctx.query_start_loc,
-        seq_lens=ctx.seq_lens,
-        max_query_len=ctx.max_query_len,
-        max_seq_len=ctx.max_seq_len,
-        softmax_scale=HEAD_DIM**-0.5,
-        causal=False,
-    )
-    assert torch.equal(paged, direct)
-
-    new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
-    new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
-    if action_len:
-        new_k = torch.cat([new_k, action_k], dim=1)
-        new_v = torch.cat([new_v, action_v], dim=1)
-    ref = _dense_attention(query, new_k, new_v)
-
-    torch.testing.assert_close(paged, ref, rtol=2e-2, atol=2e-2)
-
-
-def test_block_table_padded_to_fixed_width():
-    """Shapes must be constant across window growth: only values change."""
-    device = torch.device("cpu")
-    kv, st = make_state(window_chunks=2)
-
-    ctx1 = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx
-    ctx1.prepare(device=device, action_len=0, query_len=BLOCK)
-    st.commit_paged_context(False)
-
-    ctx2 = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx
-    ctx2.prepare(device=device, action_len=0, query_len=BLOCK)
-    st.commit_paged_context(False)
-
-    # 1-block vs 2-block visible history: same table width, same max_seq_len.
-    assert ctx1.block_table.shape == ctx2.block_table.shape
-    assert ctx1.max_seq_len == ctx2.max_seq_len
-    expected_width = kv.spec.sliding_window // kv.block_size + 1
-    assert ctx1.block_table.shape == (1, expected_width)
-    # Real lengths live in seq_lens, not the padded table.
-    assert int(ctx1.seq_lens[0]) == BLOCK
-    assert int(ctx2.seq_lens[0]) == 2 * BLOCK
-
-
-def test_prepare_is_idempotent_and_layers_share_metadata():
-    device = torch.device("cpu")
-    _, st = make_state(num_layers=2)
-    contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)
-    fctx = contexts[0].forward_ctx
-    fctx.prepare(device=device, action_len=0, query_len=BLOCK)
-    table = fctx.block_table
-    fctx.prepare(device=device, action_len=0, query_len=BLOCK)
-    assert fctx.block_table is table  # memoized, not rebuilt
-
-    i0, i1 = contexts[0].to_layer_inputs(), contexts[1].to_layer_inputs()
-    # 0-dim tensors (NOT python ints) so dynamo doesn't install per-layer
-    # value guards on the shared block code object.
-    assert isinstance(i0.layer_idx, torch.Tensor) and int(i0.layer_idx) == 0
-    assert isinstance(i1.layer_idx, torch.Tensor) and int(i1.layer_idx) == 1
-    # All layers share the same metadata tensor objects.
-    assert i0.block_table is i1.block_table
-    assert i0.seq_lens is i1.seq_lens
-    assert i0.video_slots is i1.video_slots
-
-
-def test_layer_inputs_before_prepare_raises():
-    _, st = make_state()
-    layer_ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)[0]
-    with pytest.raises(RuntimeError, match="before prepare"):
-        layer_ctx.to_layer_inputs()
-
-
-def test_custom_op_registration_idempotent():
-    import importlib
-    import sys
-
-    assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn")
-    mod = "vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention"
-    saved = sys.modules.pop(mod)
-    try:
-        importlib.import_module(mod)  # re-registration must not raise
-    finally:
-        sys.modules[mod] = saved
-    assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn")
-
-
-def test_custom_op_compiles_fullgraph_without_recompile_on_value_change():
-    """The op must trace as one opaque node: fullgraph OK, and changed tensor
-    VALUES (new slots / block ids) must not trigger recompilation."""
-    import torch._dynamo
-
-    device = torch.device("cpu")
-    kv, st = make_state(num_layers=2, window_chunks=2)
-
-    def run_one_forward(commit):
-        contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=commit)
-        fctx = contexts[0].forward_ctx
-        fctx.prepare(device=device, action_len=0, query_len=BLOCK)
-        q = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
-        k = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
-        v = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
-        # Both layers through ONE compiled fn: layer_idx is a tensor, so a
-        # different layer must NOT recompile (all 40 DiT blocks share the
-        # block-forward code object in production).
-        for layer_ctx in contexts:
-            out = compiled(layer_ctx.to_layer_inputs(), q, k, v)
-        st.commit_paged_context(False)
-        return out
-
-    torch._dynamo.reset()
-    from torch._dynamo.testing import CompileCounter
-
-    counter = CompileCounter()
-
-    def fn(inputs, q, k, v):
-        return paged_write_attn(inputs, q, k, v, None, None, HEAD_DIM**-0.5) * 1.0
-
-    compiled = torch.compile(fn, backend=counter, fullgraph=True)
-
-    run_one_forward(commit=True)  # history grows between calls ->
-    run_one_forward(commit=True)  # block-table VALUES change, shapes don't
-    run_one_forward(commit=False)
-
-    assert counter.frame_count == 1, f"recompiled: frame_count={counter.frame_count}"
@@ -1,222 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Tests for the ARDiffusionKVState paged-attention pipeline bridge."""
-
-import pytest
-import torch
-
-from vllm_omni.experimental.ar_diffusion.kv_cache import (
-    ARDiffusionKVCache,
-    ARDiffusionKVConfig,
-    ARDiffusionPagedLayerContext,
-)
-from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
-
-BLOCK = 16
-N_HEADS = 4
-HEAD_DIM = 64
-
-
-def make_state(num_layers=1, window_chunks=4, cross_attn_length=0):
-    cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
-    kv = ARDiffusionKVCache(
-        cfg,
-        num_layers=num_layers,
-        num_kv_heads=N_HEADS,
-        head_size=HEAD_DIM,
-        dtype=torch.float32,
-        block_size=BLOCK,
-        max_model_len=4096,
-        available_bytes=1 << 24,
-        cross_attn_length=cross_attn_length,
-        device=torch.device("cpu"),
-    )
-    pos = kv.begin_request("r-pos")
-    neg = kv.begin_request("r-neg")
-    return kv, ARDiffusionKVState(kv, pos, neg, num_layers=num_layers)
-
-
-def _prepare_and_commit(st: ARDiffusionKVState, is_negative: bool, n_chunks: int) -> None:
-    ctx = st.get_kv_caches(is_negative, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
-    ctx.ensure_video_slots(torch.device("cpu"))
-    st.commit_paged_context(is_negative)
-
-
-def test_get_returns_paged_layer_contexts_when_nothing_committed():
-    _, st = make_state(num_layers=2)
-    for neg in (False, True):
-        contexts = st.get_kv_caches(neg, seq_len=BLOCK, commit_current=False)
-        assert len(contexts) == 2
-        assert all(isinstance(ctx, ARDiffusionPagedLayerContext) for ctx in contexts)
-        assert contexts[0].history_block_ids == []
-
-
-def test_managed_paged_context_commits_only_after_forward():
-    kv, st = make_state(num_layers=1)
-    contexts = st.get_kv_caches(False, seq_len=2 * BLOCK, commit_current=True)
-    ctx = contexts[0].forward_ctx
-
-    assert st.pos.completed_chunks == 0
-    assert kv.window_block_ids(st.pos) == []
-    ctx.ensure_video_slots(torch.device("cpu"))
-    assert st.pos.completed_chunks == 0
-    assert len(ctx.current_video_block_ids) == 2
-
-    st.commit_paged_context(False)
-
-    assert st.pos.completed_chunks == 2
-    assert st._committed[False] == 2 * BLOCK
-    assert len(kv.window_block_ids(st.pos)) == 2
-
-
-def test_scratch_paged_context_does_not_commit():
-    kv, st = make_state(num_layers=1)
-    contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)
-    ctx = contexts[0].forward_ctx
-    ctx.ensure_video_slots(torch.device("cpu"))
-
-    assert ctx.current_video_block_ids == kv.scratch_block_ids(False, 0, 1)
-    st.commit_paged_context(False)
-
-    assert st.pos.completed_chunks == 0
-    assert st._committed[False] == 0
-    assert kv.window_block_ids(st.pos) == []
-
-
-def test_eviction_bounds_resident_window():
-    kv, st = make_state(num_layers=1, window_chunks=3)
-    _prepare_and_commit(st, False, 3)
-    assert len(kv.window_block_ids(st.pos)) == 3
-
-    _prepare_and_commit(st, False, 2)
-    ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)[0].forward_ctx
-    visible_blocks, video_len = ctx.video_block_table(torch.device("cpu"))
-    assert len(visible_blocks) == 3
-    assert video_len == 3 * BLOCK
-
-
-def test_branches_are_independent():
-    kv, st = make_state()
-    _prepare_and_commit(st, False, 1)
-    assert len(kv.window_block_ids(st.pos)) == 1
-    assert kv.window_block_ids(st.neg) == []
-
-
-def test_reset_clears_session_window():
-    kv, st = make_state(num_layers=1, window_chunks=4)
-    _prepare_and_commit(st, False, 2)
-    _prepare_and_commit(st, True, 1)
-    assert st._committed[False] == 2 * BLOCK and st._committed[True] == BLOCK
-    free_before = kv.manager.block_pool.get_num_free_blocks()
-
-    st.reset()
-
-    assert st._committed == {False: 0, True: 0}
-    assert st._paged_pending == {False: None, True: None}
-    assert kv.window_block_ids(st.pos) == []
-    assert kv.window_block_ids(st.neg) == []
-    assert kv.manager.block_pool.get_num_free_blocks() > free_before
-    _prepare_and_commit(st, False, 1)
-    assert len(kv.window_block_ids(st.pos)) == 1
-
-
-def test_get_cross_kv_caches_raises_before_populate():
-    """Engine ownership guard: a cross-attn read before _kv_populate_cross must fail loud."""
-    _, st = make_state()  # cross_attn_length == 0, nothing populated
-    for neg in (False, True):
-        with pytest.raises(RuntimeError, match="cross-attn read before"):
-            st.get_cross_kv_caches(neg)
-    # Guards the full AND: the populated flag alone is not enough without a pool.
-    st._cross_text_populated[False] = True
-    with pytest.raises(RuntimeError, match="cross-attn read before"):
-        st.get_cross_kv_caches(False)
-
-
-def test_get_cross_kv_caches_returns_pool_dicts_when_populated():
-    L = 8
-    kv, st = make_state(num_layers=2, cross_attn_length=L)
-    written = []
-    for i in range(2):
-        k = torch.randn(1, L, N_HEADS, HEAD_DIM)
-        v = torch.randn(1, L, N_HEADS, HEAD_DIM)
-        kv.write_cross_kv(i, False, k, v)
-        written.append((k, v))
-    st._cross_text_populated[False] = True
-
-    out = st.get_cross_kv_caches(False)
-    assert len(out) == 2
-    for i, (k, v) in enumerate(written):
-        assert out[i]["is_init"] is True
-        assert out[i]["k"].shape == (1, L, N_HEADS, HEAD_DIM)
-        assert torch.equal(out[i]["k"], k)
-        assert torch.equal(out[i]["v"], v)
-
-
-def test_kv_create_is_noop_engine_owns_allocation():
-    from unittest.mock import MagicMock
-
-    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
-
-    p = DreamZeroPipeline.__new__(DreamZeroPipeline)
-    p._ar_diffusion_kv_state = object()
-    state = MagicMock()
-    p._kv_create(state, 1, "float32", "cpu", 24, 4, 64)
-    state.create_kv_caches.assert_not_called()
-
-
-def test_get_kv_cache_requires_frame_aligned_seqlen():
-    _, st = make_state()
-    with pytest.raises(AssertionError, match="frame-aligned"):
-        st.get_kv_caches(False, seq_len=BLOCK + 1, commit_current=True)
-
-
-def test_prepare_one_branch_allocates_nothing_for_the_other():
-    """CFG-parallel laziness: a rank prepares only the branch it runs."""
-    kv, st = make_state(num_layers=1, window_chunks=4)
-    pos_ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx
-    neg_ctx = st.get_kv_caches(True, seq_len=BLOCK, commit_current=True)[0].forward_ctx
-
-    pos_ctx.prepare(device=torch.device("cpu"), action_len=0, query_len=BLOCK)
-
-    assert pos_ctx._allocated_video
-    assert not neg_ctx._allocated_video
-    assert kv.window_block_ids(st.neg) == []
-
-
-def test_failed_forward_tears_down_session(monkeypatch):
-    """Review follow-up (hsliuustc0106): a forward that dies partway must not
-    leave allocated-but-uncommitted KV behind — the session is torn down (pool
-    blocks freed, model-local state dropped) and the error propagates."""
-    from types import SimpleNamespace
-
-    from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
-    from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
-
-    kv, st = make_state(num_layers=1, window_chunks=4)
-    _prepare_and_commit(st, False, 2)  # session owns pool blocks
-    free_before_failure = kv.manager.block_pool.get_num_free_blocks()
-
-    runner = object.__new__(ARDiffusionModelRunner)
-    runner.kv_cache = kv
-    runner._ar_diffusion_states = __import__("collections").OrderedDict({"s1": st})
-    runner._max_ar_diffusion_states = 4
-    runner.pipeline = SimpleNamespace(_states={"s1": object()}, _ar_diffusion_kv_state=None)
-    runner.device = None
-    runner._perf_e2e_times = []
-
-    def boom(self, req):
-        raise RuntimeError("layer 17 exploded")
-
-    monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
-    req = SimpleNamespace(
-        request_id="r0",
-        sampling_params=SimpleNamespace(extra_args={"session_id": "s1"}),
-    )
-
-    with pytest.raises(RuntimeError, match="layer 17 exploded"):
-        ARDiffusionModelRunner.execute_model(runner, req)
-
-    assert "s1" not in runner._ar_diffusion_states
-    assert "s1" not in runner.pipeline._states
-    assert runner.pipeline._ar_diffusion_kv_state is None
-    # close() freed the session's resident blocks.
-    assert kv.manager.block_pool.get_num_free_blocks() > free_before_failure
@@ -1,76 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Slot-mapping / block-table tests for the AR-Diffusion engine (Phase 1, Step 3)."""
-
-import pytest
-import torch
-
-from vllm_omni.experimental.ar_diffusion.kv_cache import (
-    ARDiffusionRequestAdapter,
-    ChunkWindowSpec,
-    build_kv_manager,
-    chunk_slot_mapping,
-    compute_slot_mapping,
-    resident_block_ids,
-)
-
-BLOCK = 16
-
-
-def make_spec(*, window_chunks=2):
-    return ChunkWindowSpec(
-        block_size=BLOCK,
-        num_kv_heads=4,
-        head_size=64,
-        dtype=torch.float16,
-        sliding_window=window_chunks * BLOCK,
-        chunk_size=BLOCK,
-        window_chunks=window_chunks,
-    )
-
-
-def test_slot_mapping_matches_block_offsets():
-    # block table [5, 2, 9], block_size 4. slot = block*4 + (pos % 4).
-    slots = compute_slot_mapping([5, 2, 9], [0, 1, 4, 5, 8, 9], block_size=4)
-    assert slots.tolist() == [20, 21, 8, 9, 36, 37]
-
-
-def test_slot_mapping_rejects_bad_block_size():
-    with pytest.raises(ValueError):
-        compute_slot_mapping([1], [0], block_size=0)
-
-
-def test_chunk_slot_mapping_targets_current_chunk():
-    # After 2 committed chunks (num_computed=32), the in-flight chunk maps to the
-    # block covering positions 32..47 — i.e. block table index 2.
-    block_ids = [7, 4, 3]  # block index 2 -> physical block 3
-    slots = chunk_slot_mapping(block_ids, num_computed_tokens=32, chunk_size=4, block_size=16)
-    # positions 32..35 -> block_index 2 -> physical 3 -> slots 48..51
-    assert slots.tolist() == [48, 49, 50, 51]
-
-
-def test_resident_block_ids_excludes_null():
-    assert resident_block_ids([0, 1, 0, 2, 3], null_block_id=0) == [1, 2, 3]
-
-
-def test_blocktable_build_single_request():
-    """Integration: a real KVCacheManager block table -> a valid slot mapping.
-
-    Every slot for the in-flight chunk must land inside a real (non-null) block
-    that the manager actually allocated for the request.
-    """
-    spec = make_spec()
-    mgr = build_kv_manager(spec, ["l0"], num_blocks=16, max_model_len=1024)
-    null_id = mgr.block_pool.null_block.block_id
-    adapter = ARDiffusionRequestAdapter("req", chunk_size=BLOCK)
-
-    mgr.allocate_slots(adapter, num_new_tokens=BLOCK)
-    block_ids = mgr.get_block_ids(adapter.request_id)[0]
-    real = set(resident_block_ids(block_ids, null_id))
-    assert real, "expected at least one real block after allocation"
-
-    slots = chunk_slot_mapping(block_ids, adapter.num_computed_tokens, spec.chunk_size, spec.block_size)
-    # Each slot resolves to a real allocated block, none to null_block.
-    blocks_used = {int(s) // spec.block_size for s in slots}
-    assert blocks_used <= real
-    assert null_id not in blocks_used
-    assert len(slots) == spec.chunk_size
@@ -1,154 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for applying strategy specs onto merged stage configs."""
-
-from __future__ import annotations
-
-import pytest
-
-from vllm_omni.config.composable_parallel import (
-    Broadcast,
-    FanInByStage,
-    MeshAxisSpec,
-    RouteByStage,
-    StrategyApplyError,
-    StrategySpec,
-    TakeRank,
-    apply_strategy_specs,
-)
-from vllm_omni.config.stage_config import StageConfig
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def _tp(size: int) -> StrategySpec:
-    return StrategySpec("tp", MeshAxisSpec("tp", size), Broadcast(), TakeRank())
-
-
-def _stage_replica(size: int, policy: str = "round_robin") -> StrategySpec:
-    return StrategySpec("stage_replica", MeshAxisSpec("stage_replica", size), RouteByStage(policy), FanInByStage())
-
-
-def _stage(stage_id: int, model_stage: str, engine_args=None, runtime=None) -> StageConfig:
-    return StageConfig(
-        stage_id=stage_id,
-        model_stage=model_stage,
-        yaml_engine_args=dict(engine_args or {}),
-        yaml_runtime=dict(runtime or {"num_replicas": 1}),
-    )
-
-
-def _qwen_stages() -> list[StageConfig]:
-    return [
-        _stage(0, "thinker"),
-        _stage(1, "talker"),
-        _stage(2, "code2wav"),
-    ]
-
-
-def test_apply_tp_by_role():
-    stages = _qwen_stages()
-    apply_strategy_specs(stages, {"thinker": [_tp(2)]})
-    assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
-    # untouched roles keep their config
-    assert "tensor_parallel_size" not in stages[1].yaml_engine_args
-
-
-def test_apply_by_model_stage():
-    stages = _qwen_stages()
-    apply_strategy_specs(stages, {"thinker": [_tp(2)]})
-    assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
-
-
-def test_apply_stage_replica_sets_num_replicas_and_surfaces_lb():
-    stages = _qwen_stages()
-    result = apply_strategy_specs(stages, {"talker": [_stage_replica(2, "round_robin")]})
-    assert stages[1].yaml_runtime["num_replicas"] == 2
-    assert result.omni_lb_policy == "round-robin"
-
-
-def test_only_declared_axes_are_written():
-    stages = _qwen_stages()
-    # strategy declares only stage_replica -> tp must not be forced.
-    apply_strategy_specs(stages, {"talker": [_stage_replica(2)]})
-    assert "tensor_parallel_size" not in stages[1].yaml_engine_args
-
-
-def test_conflict_on_explicit_tp():
-    stages = _qwen_stages()
-    stages[0].yaml_engine_args["tensor_parallel_size"] = 4
-    with pytest.raises(StrategyApplyError):
-        apply_strategy_specs(stages, {"thinker": [_tp(2)]})
-
-
-def test_equal_explicit_value_is_noop():
-    stages = _qwen_stages()
-    stages[0].yaml_engine_args["tensor_parallel_size"] = 2
-    apply_strategy_specs(stages, {"thinker": [_tp(2)]})
-    assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
-
-
-def test_explicit_none_conflicts_with_derived_value():
-    # An explicit YAML ``null`` (``tensor_parallel_size: null``) is a *present*
-    # value, not a missing key, so a strategy deriving a non-None size must raise
-    # rather than silently clobber the explicit None.
-    stages = _qwen_stages()
-    stages[0].yaml_engine_args["tensor_parallel_size"] = None
-    with pytest.raises(StrategyApplyError):
-        apply_strategy_specs(stages, {"thinker": [_tp(2)]})
-
-
-def test_missing_key_is_filled_without_conflict():
-    # A genuinely absent key (never set in the YAML) is filled by the strategy
-    # and must NOT raise — the contrast case to an explicit None.
-    stages = _qwen_stages()
-    assert "tensor_parallel_size" not in stages[0].yaml_engine_args
-    apply_strategy_specs(stages, {"thinker": [_tp(2)]})
-    assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
-
-
-def test_num_replicas_conflict():
-    stages = _qwen_stages()
-    stages[1].yaml_runtime["num_replicas"] = 3
-    with pytest.raises(StrategyApplyError):
-        apply_strategy_specs(stages, {"talker": [_stage_replica(2)]})
-
-
-def test_device_count_ok_template():
-    stages = _qwen_stages()
-    stages[0].yaml_runtime["devices"] = "0,1"
-    apply_strategy_specs(stages, {"thinker": [_tp(2)]})
-    assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
-
-
-def test_device_count_ok_pool():
-    # tp=2 -> world=2; 2 replicas -> pool of 4 device ids is valid.
-    stages = _qwen_stages()
-    stages[1].yaml_runtime["devices"] = "0,1,2,3"
-    apply_strategy_specs(stages, {"talker": [_tp(2), _stage_replica(2)]})
-    assert stages[1].yaml_runtime["num_replicas"] == 2
-
-
-def test_device_count_mismatch():
-    stages = _qwen_stages()
-    stages[0].yaml_runtime["devices"] = "0,1,2"
-    with pytest.raises(StrategyApplyError):
-        apply_strategy_specs(stages, {"thinker": [_tp(2)]})
-
-
-def test_unknown_role_raises():
-    stages = _qwen_stages()
-    with pytest.raises(StrategyApplyError):
-        apply_strategy_specs(stages, {"nonexistent": [_tp(2)]})
-
-
-def test_conflicting_lb_policy_across_roles():
-    stages = _qwen_stages()
-    with pytest.raises(StrategyApplyError):
-        apply_strategy_specs(
-            stages,
-            {
-                "talker": [_stage_replica(2, "round_robin")],
-                "code2wav": [_stage_replica(2, "least_queue")],
-            },
-        )
@@ -1,141 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Integration tests: strategy overlay on the real merge_pipeline_deploy seam."""
-
-from __future__ import annotations
-
-import logging
-from pathlib import Path
-
-import pytest
-
-from vllm_omni.config.composable_parallel import (
-    Broadcast,
-    FanInByStage,
-    MeshAxisSpec,
-    RouteByStage,
-    StrategyApplyError,
-    StrategySpec,
-    TakeRank,
-    apply_strategy_specs,
-)
-from vllm_omni.config.config_factory import StageConfigFactory
-from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
-from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-_DEPLOY = Path(__file__).parents[3] / "vllm_omni" / "deploy" / "qwen2_5_omni.yaml"
-
-
-def _tp(size: int) -> StrategySpec:
-    return StrategySpec("tp", MeshAxisSpec("tp", size), Broadcast(), TakeRank())
-
-
-def _stage_replica(size: int, policy: str = "round_robin") -> StrategySpec:
-    return StrategySpec("stage_replica", MeshAxisSpec("stage_replica", size), RouteByStage(policy), FanInByStage())
-
-
-def _qwen_stages():
-    pipeline = OMNI_PIPELINES["qwen2_5_omni"]
-    deploy = load_deploy_config(_DEPLOY)
-    return merge_pipeline_deploy(pipeline, deploy)
-
-
-def _stage(stages, role):
-    return next(s for s in stages if s.model_stage == role)
-
-
-def test_overlay_tp_on_thinker():
-    stages = _qwen_stages()
-    # The bundled deploy pins the thinker to one GPU; a TP=2 strategy needs a
-    # matching 2-GPU layout (mirrors what a TP2 deploy would declare).
-    _stage(stages, "thinker").yaml_runtime["devices"] = "0,1"
-    result = apply_strategy_specs(stages, {"thinker": [_tp(2)]})
-    assert _stage(result.stages, "thinker").yaml_engine_args["tensor_parallel_size"] == 2
-
-
-def test_device_guard_rejects_tp_on_single_gpu_deploy():
-    # The bundled deploy pins the thinker to a single GPU, so the pre-spawn
-    # device check must refuse a TP=2 strategy on it.
-    stages = _qwen_stages()
-    with pytest.raises(StrategyApplyError):
-        apply_strategy_specs(stages, {"thinker": [_tp(2)]})
-
-
-def test_overlay_stage_replica_on_talker():
-    stages = _qwen_stages()
-    result = apply_strategy_specs(stages, {"talker": [_stage_replica(2, "round_robin")]})
-    assert _stage(result.stages, "talker").yaml_runtime["num_replicas"] == 2
-    assert result.omni_lb_policy == "round-robin"
-
-
-def test_overlay_mixed_roles():
-    stages = _qwen_stages()
-    _stage(stages, "thinker").yaml_runtime["devices"] = "0,1"
-    result = apply_strategy_specs(
-        stages,
-        {
-            "thinker": [_tp(2)],
-            "talker": [_stage_replica(2, "round_robin")],
-            "code2wav": [_stage_replica(2, "round_robin")],
-        },
-    )
-    by_role = {s.model_stage: s for s in result.stages}
-    assert by_role["thinker"].yaml_engine_args["tensor_parallel_size"] == 2
-    assert by_role["talker"].yaml_runtime["num_replicas"] == 2
-    assert by_role["code2wav"].yaml_runtime["num_replicas"] == 2
-    assert result.omni_lb_policy == "round-robin"
-
-
-def _resolved(stages, role):
-    """Return a role's fully-resolved (post-CLI) OmegaConf stage config."""
-    stage = next(s for s in stages if s.model_stage == role)
-    return stage.to_omegaconf()
-
-
-def test_device_check_survives_cli_override():
-    # Strategy replicates the talker (1-GPU template -> valid at apply time),
-    # but a CLI --stage_1_devices with 3 ids must NOT slip past the device
-    # guard: effective world=1, replicas=2 admits only 1 or 2 device ids.
-    pipeline_cfg = OMNI_PIPELINES["qwen2_5_omni"]
-    with pytest.raises(StrategyApplyError):
-        StageConfigFactory._create_from_registry(
-            "qwen2_5_omni",
-            pipeline_cfg,
-            cli_overrides={"stage_1_devices": "0,1,2"},
-            strategy_specs={"talker": [_stage_replica(2, "round_robin")]},
-        )
-
-
-def test_cli_overrides_strategy_with_warning():
-    # Strategy derives num_replicas=2 for the talker; a CLI override to 3 wins
-    # (it is the most explicit user action) but must be surfaced loudly rather
-    # than silently. The resulting layout (1 template device) stays valid.
-    #
-    # vLLM's logger sets propagate=False, so attach a handler directly to it
-    # rather than relying on pytest's caplog (which listens on the root logger).
-    messages: list[str] = []
-
-    class _Capture(logging.Handler):
-        def emit(self, record: logging.LogRecord) -> None:
-            messages.append(record.getMessage())
-
-    log = logging.getLogger("vllm_omni.config.config_factory")
-    handler = _Capture(level=logging.WARNING)
-    log.addHandler(handler)
-    try:
-        pipeline_cfg = OMNI_PIPELINES["qwen2_5_omni"]
-        stages, _ = StageConfigFactory._create_from_registry(
-            "qwen2_5_omni",
-            pipeline_cfg,
-            cli_overrides={"stage_1_num_replicas": 3},
-            strategy_specs={"talker": [_stage_replica(2, "round_robin")]},
-        )
-    finally:
-        log.removeHandler(handler)
-
-    # CLI value wins in the resolved config.
-    assert _resolved(stages, "talker").runtime.num_replicas == 3
-    # ...and the override was warned about, naming the conflicting field.
-    assert any("num_replicas" in m and "overrides the strategy-derived" in m for m in messages)
@@ -1,42 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for the vendored spec/pattern data types."""
-
-from __future__ import annotations
-
-import pytest
-
-from vllm_omni.config.composable_parallel.aggregation import AggregationPattern, TakeRank
-from vllm_omni.config.composable_parallel.routing import RoutingPattern
-from vllm_omni.config.composable_parallel.spec import MeshAxisSpec
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_mesh_axis_rejects_bad_kind():
-    with pytest.raises(ValueError):
-        MeshAxisSpec(kind="not_a_kind", size=2)  # type: ignore[arg-type]
-
-
-def test_mesh_axis_rejects_nonpositive_size():
-    with pytest.raises(ValueError):
-        MeshAxisSpec(kind="tp", size=0)
-
-
-def test_routing_pattern_is_closed():
-    with pytest.raises(TypeError):
-
-        class RogueRouting(RoutingPattern):  # noqa: F811
-            pass
-
-
-def test_aggregation_pattern_is_closed():
-    with pytest.raises(TypeError):
-
-        class RogueAggregation(AggregationPattern):  # noqa: F811
-            pass
-
-
-def test_take_rank_type_guard():
-    with pytest.raises(TypeError):
-        TakeRank(rank=True)  # bool is not a valid rank
@@ -1,97 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for parsing/loading strategy files."""
-
-from __future__ import annotations
-
-import pytest
-
-from vllm_omni.config.composable_parallel.routing import Broadcast, RouteByStage
-from vllm_omni.config.composable_parallel.strategy_loader import (
-    StrategyLoadError,
-    load_strategy_specs,
-    parse_strategy_specs,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_parse_basic():
-    data = {
-        "strategies": {
-            "thinker": [{"axis": "tp", "size": 2}],
-            "talker": [{"axis": "stage_replica", "size": 2, "routing": "round_robin"}],
-        }
-    }
-    specs = parse_strategy_specs(data)
-    assert set(specs) == {"thinker", "talker"}
-
-    (tp_spec,) = specs["thinker"]
-    assert tp_spec.mesh_axis.kind == "tp"
-    assert tp_spec.mesh_axis.size == 2
-    assert isinstance(tp_spec.routing, Broadcast)
-
-    (sr_spec,) = specs["talker"]
-    assert sr_spec.mesh_axis.kind == "stage_replica"
-    assert isinstance(sr_spec.routing, RouteByStage)
-    assert sr_spec.routing.routing_policy == "round_robin"
-
-
-def test_parse_without_strategies_key():
-    # A bare role mapping (no top-level "strategies") is also accepted.
-    specs = parse_strategy_specs({"thinker": [{"axis": "tp", "size": 4}]})
-    assert specs["thinker"][0].mesh_axis.size == 4
-
-
-def test_l1_owner_goes_to_shard_extension():
-    specs = parse_strategy_specs({"talker": [{"axis": "stage_replica", "size": 2, "l1_owner": "delegated"}]})
-    assert specs["talker"][0].shard_extension["l1_owner"] == "delegated"
-
-
-def test_missing_axis_raises():
-    with pytest.raises(StrategyLoadError):
-        parse_strategy_specs({"thinker": [{"size": 2}]})
-
-
-def test_missing_size_raises():
-    with pytest.raises(StrategyLoadError):
-        parse_strategy_specs({"thinker": [{"axis": "tp"}]})
-
-
-def test_routing_on_non_policy_axis_raises():
-    with pytest.raises(StrategyLoadError):
-        parse_strategy_specs({"thinker": [{"axis": "tp", "size": 2, "routing": "random"}]})
-
-
-def test_bad_size_raises():
-    with pytest.raises(StrategyLoadError):
-        parse_strategy_specs({"thinker": [{"axis": "tp", "size": "two"}]})
-
-
-def test_entries_must_be_list():
-    with pytest.raises(StrategyLoadError):
-        parse_strategy_specs({"thinker": {"axis": "tp", "size": 2}})
-
-
-def test_non_mapping_entry_raises():
-    # A list whose elements are not mappings (e.g. a bare string) must raise a
-    # StrategyLoadError, not an opaque TypeError from dict(entry).
-    with pytest.raises(StrategyLoadError):
-        parse_strategy_specs({"thinker": ["tp"]})
-
-
-def test_load_from_file(tmp_path):
-    path = tmp_path / "strategy.yaml"
-    path.write_text(
-        "strategies:\n"
-        "  thinker:\n"
-        "    - axis: tp\n"
-        "      size: 2\n"
-        "  talker:\n"
-        "    - axis: stage_replica\n"
-        "      size: 2\n"
-        "      routing: least_queue\n"
-    )
-    specs = load_strategy_specs(str(path))
-    assert specs["thinker"][0].mesh_axis.size == 2
-    assert specs["talker"][0].routing.routing_policy == "least_queue"
@@ -1,195 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for StrategySpec -> OmniParallelConfig translation."""
-
-from __future__ import annotations
-
-import dataclasses
-
-import pytest
-
-from vllm_omni.config.composable_parallel import (
-    AxisTranslationError,
-    Broadcast,
-    FanInByStage,
-    MeshAxisSpec,
-    PartitionByHash,
-    PipelineMicrobatch,
-    RouteByStage,
-    StrategySpec,
-    TakeRank,
-    Union,
-    translate_strategy_stack,
-)
-from vllm_omni.config.composable_parallel.aggregation import StitchPipeline
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def _tp(size: int) -> StrategySpec:
-    return StrategySpec("tp", MeshAxisSpec("tp", size), Broadcast(), TakeRank())
-
-
-def _dp(size: int, policy: str = "round_robin") -> StrategySpec:
-    return StrategySpec("dp", MeshAxisSpec("dp", size), RouteByStage(policy), Union())
-
-
-def _pp(size: int) -> StrategySpec:
-    return StrategySpec("pp", MeshAxisSpec("pp", size), PipelineMicrobatch(), StitchPipeline())
-
-
-def _ep(size: int) -> StrategySpec:
-    return StrategySpec("ep", MeshAxisSpec("ep", size), Broadcast(), Union())
-
-
-def _stage_replica(size: int, policy: str = "round_robin") -> StrategySpec:
-    return StrategySpec("stage_replica", MeshAxisSpec("stage_replica", size), RouteByStage(policy), FanInByStage())
-
-
-def test_tp_only():
-    cfg = translate_strategy_stack([_tp(2)])
-    assert cfg.tensor_parallel_size == 2
-    assert cfg.data_parallel_size == 1
-    assert cfg.world_size == 2
-    assert cfg.stage_replica_size == 1
-    assert cfg.l1_owners == {"tp": "engine"}
-
-
-def test_dp_is_engine_owned():
-    cfg = translate_strategy_stack([_dp(2)])
-    assert cfg.data_parallel_size == 2
-    assert cfg.world_size == 2
-    # engine DP does not emit an omni LB policy (vLLM's internal DP LB is used).
-    assert cfg.omni_lb_policy is None
-    assert cfg.as_engine_kwargs().get("omni_lb_policy") is None
-    assert cfg.l1_owners["dp"] == "engine"
-
-
-def test_tp_times_dp_world_size():
-    cfg = translate_strategy_stack([_tp(2), _dp(2)])
-    assert cfg.world_size == 4
-    assert cfg.tensor_parallel_size == 2
-    assert cfg.data_parallel_size == 2
-
-
-def test_stage_replica_is_delegated():
-    cfg = translate_strategy_stack([_stage_replica(3, "least_queue")])
-    # stage_replica is NOT a world dimension.
-    assert cfg.world_size == 1
-    assert cfg.stage_replica_size == 3
-    assert cfg.omni_lb_policy == "least-queue-length"
-    assert cfg.l1_owners["stage_replica"] == "delegated"
-    # stage_replica_size is not an engine kwarg; it's a per-stage deploy knob.
-    assert "stage_replica_size" not in cfg.as_engine_kwargs()
-
-
-def test_stage_replica_lb_policy_only_when_replicated():
-    cfg = translate_strategy_stack([_stage_replica(1)])
-    assert cfg.stage_replica_size == 1
-    assert cfg.as_engine_kwargs().get("omni_lb_policy") is None
-
-
-def test_tp_with_stage_replica():
-    cfg = translate_strategy_stack([_tp(2), _stage_replica(2)])
-    assert cfg.world_size == 2  # tp only
-    assert cfg.tensor_parallel_size == 2
-    assert cfg.stage_replica_size == 2
-    assert cfg.omni_lb_policy == "round-robin"
-
-
-def test_ep_must_match_tp_dp_product():
-    cfg = translate_strategy_stack([_tp(2), _ep(2)])
-    assert cfg.enable_expert_parallel is True
-    assert cfg.as_engine_kwargs()["enable_expert_parallel"] is True
-
-
-def test_ep_mismatch_raises():
-    with pytest.raises(Exception) as exc:
-        translate_strategy_stack([_tp(2), _ep(4)])
-    assert "must equal" in str(exc.value)
-
-
-def test_duplicate_kind_rejected():
-    with pytest.raises(AxisTranslationError):
-        translate_strategy_stack([_tp(2), _tp(2)])
-
-
-def test_unsupported_kind_rejected():
-    spec = StrategySpec("sp", MeshAxisSpec("sp_ulysses", 2), Broadcast(), TakeRank())
-    with pytest.raises(AxisTranslationError):
-        translate_strategy_stack([spec])
-
-
-def test_dp_hash_routing_rejected():
-    spec = StrategySpec("dp", MeshAxisSpec("dp", 2), PartitionByHash(), Union())
-    with pytest.raises(NotImplementedError):
-        translate_strategy_stack([spec])
-
-
-def test_dp_invalid_routing_policy_rejected():
-    spec = StrategySpec("dp", MeshAxisSpec("dp", 2), RouteByStage("bogus"), Union())
-    with pytest.raises(AxisTranslationError):
-        translate_strategy_stack([spec])
-
-
-def test_lb_policy_never_in_engine_kwargs_even_when_replicated():
-    # omni_lb_policy is pipeline-wide (applied at engine construction), not a
-    # per-stage engine arg, so it must never leak into as_engine_kwargs().
-    cfg = translate_strategy_stack([_tp(2), _stage_replica(2, "least_queue")])
-    assert cfg.omni_lb_policy == "least-queue-length"
-    assert "omni_lb_policy" not in cfg.as_engine_kwargs()
-
-
-def test_stage_replica_hash_routing_rejected():
-    spec = StrategySpec("sr", MeshAxisSpec("stage_replica", 2), RouteByStage("hash"), FanInByStage())
-    with pytest.raises(NotImplementedError):
-        translate_strategy_stack([spec])
-
-
-def test_dp_wrong_owner_rejected():
-    spec = StrategySpec(
-        "dp",
-        MeshAxisSpec("dp", 2),
-        RouteByStage("round_robin"),
-        Union(),
-        shard_extension={"l1_owner": "delegated"},
-    )
-    with pytest.raises(AxisTranslationError):
-        translate_strategy_stack([spec])
-
-
-def test_stage_replica_wrong_owner_rejected():
-    spec = StrategySpec(
-        "sr",
-        MeshAxisSpec("stage_replica", 2),
-        RouteByStage("round_robin"),
-        FanInByStage(),
-        shard_extension={"l1_owner": "engine"},
-    )
-    with pytest.raises(AxisTranslationError):
-        translate_strategy_stack([spec])
-
-
-@pytest.mark.parametrize(
-    "specs",
-    [
-        [_tp(2)],
-        [_dp(2)],
-        [_tp(2), _dp(2)],
-        [_tp(2), _ep(2)],
-        [_tp(2), _stage_replica(2)],
-        [_pp(2)],
-    ],
-)
-def test_as_engine_kwargs_are_valid_omni_engine_args(specs):
-    # Guard against desync: every key emitted by as_engine_kwargs() must be a
-    # real OmniEngineArgs field, otherwise constructing the per-stage engine args
-    # would explode at runtime. Keep this CPU-only by checking field names rather
-    # than instantiating the engine.
-    from vllm_omni.engine.arg_utils import OmniEngineArgs
-
-    valid_fields = {f.name for f in dataclasses.fields(OmniEngineArgs)}
-    cfg = translate_strategy_stack(specs)
-    kwargs = cfg.as_engine_kwargs()
-    unknown = set(kwargs) - valid_fields
-    assert not unknown, f"as_engine_kwargs emitted non-OmniEngineArgs fields: {unknown}"
@@ -1,87 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for endpoint restrictions logic."""
-
-import pytest
-from fastapi import FastAPI
-from fastapi.testclient import TestClient
-
-from vllm_omni.config.endpoint_policy import (
-    EndpointRestriction,
-    OmniServingCapability,
-    shutdown_unsupported_routes,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-SERVER_ROUTE = "/v1/completions"
-REJECTION_REASON = "This route is banned!"
-
-
-def _make_app_with_server_route():
-    """Create a FastAPI app with a route to reject."""
-    app = FastAPI()
-
-    @app.post(SERVER_ROUTE)
-    async def existing_handler():
-        """Returns a response; we want to disable this."""
-        return {"ok": True}
-
-    return app
-
-
-def test_restricted_completions_returns_400():
-    """Ensure that simple route rejection works."""
-    app = _make_app_with_server_route()
-    restrictions = (EndpointRestriction(OmniServingCapability.COMPLETIONS, REJECTION_REASON),)
-    shutdown_unsupported_routes(app, restrictions)
-
-    client = TestClient(app)
-    resp = client.post(SERVER_ROUTE, json={"prompt": "hello", "model": "m"})
-    assert resp.status_code == 400
-    body = resp.json()
-    assert body["error"]["message"] == REJECTION_REASON
-    assert body["error"]["type"] == "BadRequestError"
-
-
-def test_unrestricted_completions_not_blocked():
-    """Ensure that if we don't shutdown any routes, the route stays."""
-    app = _make_app_with_server_route()
-    shutdown_unsupported_routes(app, ())
-
-    client = TestClient(app)
-    resp = client.post(SERVER_ROUTE)
-    assert resp.status_code == 200
-    assert resp.json() == {"ok": True}
-
-
-def test_route_rejection_is_idempotent():
-    """Ensure that calling rejection twice doesn't break things."""
-    app = _make_app_with_server_route()
-    restrictions = (EndpointRestriction(OmniServingCapability.COMPLETIONS, REJECTION_REASON),)
-    shutdown_unsupported_routes(app, restrictions)
-    shutdown_unsupported_routes(app, restrictions)
-
-    client = TestClient(app)
-    resp = client.post(SERVER_ROUTE, json={"prompt": "hello", "model": "m"})
-    assert resp.status_code == 400
-    body = resp.json()
-    assert body["error"]["message"] == REJECTION_REASON
-    assert body["error"]["type"] == "BadRequestError"
-
-
-def test_invalid_route_rejection():
-    """Ensure that a restriction won't crash if a route doesn't exist."""
-    # NOTE: This test is for safety since the route organization is still messy.
-    # Since the restriction needs a valid OmniServingCapability, its implied the route
-    # is a valid vLLM route, but was already removed from the app by earlier processing.
-    app = FastAPI()
-    restrictions = (EndpointRestriction(OmniServingCapability.COMPLETIONS, REJECTION_REASON),)
-    shutdown_unsupported_routes(app, restrictions)
-
-    client = TestClient(app)
-    resp = client.post(SERVER_ROUTE, json={"prompt": "hello", "model": "m"})
-    assert resp.status_code == 400
-    body = resp.json()
-    assert body["error"]["message"] == REJECTION_REASON
-    assert body["error"]["type"] == "BadRequestError"
@@ -6,7 +6,6 @@ import pytest
 from vllm.v1.request import RequestStatus

 import vllm_omni.core.sched.omni_ar_scheduler as scheduler_mod
-import vllm_omni.model_executor.models.voxcpm2.scheduler as voxcpm2_scheduler_mod
 from vllm_omni.model_executor.models.voxcpm2.scheduler import VoxCPM2OmniARAsyncScheduler

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -52,30 +51,14 @@ class _MockRequest:
        return RequestStatus.is_finished(self.status)


-@pytest.fixture(autouse=True)
-def _mock_cuda_graph_platform(monkeypatch) -> None:
-    monkeypatch.setattr(voxcpm2_scheduler_mod.current_omni_platform, "is_cuda", lambda: True)
-
-
-def _make_scheduler(
-    *,
-    enable_unified_decode_graph: bool | None = True,
-    deterministic_cfm_noise: bool = False,
-    runtime_config_on_model_config: bool = False,
-) -> VoxCPM2OmniARAsyncScheduler:
+def _make_scheduler(*, enable_unified_decode_graph: bool = True) -> VoxCPM2OmniARAsyncScheduler:
    sched = VoxCPM2OmniARAsyncScheduler.__new__(VoxCPM2OmniARAsyncScheduler)
-    hf_config = SimpleNamespace()
-    model_config = SimpleNamespace(hf_config=hf_config)
-    if enable_unified_decode_graph is not None or deterministic_cfm_noise:
-        runtime_config = SimpleNamespace(
-            enable_unified_decode_graph=True if enable_unified_decode_graph is None else enable_unified_decode_graph,
-            deterministic_cfm_noise=deterministic_cfm_noise,
+    runtime_config = SimpleNamespace(enable_unified_decode_graph=enable_unified_decode_graph)
+    sched.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_config=SimpleNamespace(voxcpm2_runtime_config=runtime_config),
        )
-        if runtime_config_on_model_config:
-            model_config.voxcpm2_runtime_config = runtime_config
-        else:
-            hf_config.voxcpm2_runtime_config = runtime_config
-    sched.vllm_config = SimpleNamespace(model_config=model_config)
+    )
    return sched


@@ -103,39 +86,6 @@ def test_voxcpm2_unified_decode_graph_does_not_defer_when_disabled() -> None:
    assert not scheduler._should_defer_waiting_for_unified_decode_graph()


-def test_voxcpm2_unified_decode_graph_uses_model_runtime_defaults() -> None:
-    scheduler = _make_scheduler(enable_unified_decode_graph=None)
-    scheduler.running = [_MockRequest("decode")]
-    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
-
-    assert scheduler._should_defer_waiting_for_unified_decode_graph()
-
-
-def test_voxcpm2_unified_decode_graph_reads_model_config_runtime_config() -> None:
-    scheduler = _make_scheduler(enable_unified_decode_graph=True, runtime_config_on_model_config=True)
-    scheduler.running = [_MockRequest("decode")]
-    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
-
-    assert scheduler._should_defer_waiting_for_unified_decode_graph()
-
-
-def test_voxcpm2_unified_decode_graph_does_not_defer_with_deterministic_noise() -> None:
-    scheduler = _make_scheduler(deterministic_cfm_noise=True)
-    scheduler.running = [_MockRequest("decode")]
-    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
-
-    assert not scheduler._should_defer_waiting_for_unified_decode_graph()
-
-
-def test_voxcpm2_unified_decode_graph_does_not_defer_without_cuda_graph(monkeypatch) -> None:
-    monkeypatch.setattr(voxcpm2_scheduler_mod.current_omni_platform, "is_cuda", lambda: False)
-    scheduler = _make_scheduler()
-    scheduler.running = [_MockRequest("decode")]
-    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
-
-    assert not scheduler._should_defer_waiting_for_unified_decode_graph()
-
-
 def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> None:
    scheduler = _make_scheduler()
    scheduler.running = [_MockRequest("decode")]
@@ -3,10 +3,7 @@
    "test_name": "test_voxcpm2",
    "server_params": {
      "model": "openbmb/VoxCPM2",
-      "trust_remote_code": true,
-      "extra_cli_args": [
-        "--trust-remote-code"
-      ]
+      "trust_remote_code": true
    },
    "benchmark_params": [
      {
@@ -1,104 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import sys
-import types
-
-import pytest
-import torch
-
-from tests.helpers.mark import hardware_test
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
-from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
-from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
-from vllm_omni.platforms import current_omni_platform
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_kernels_hub_platform_fallback(monkeypatch: pytest.MonkeyPatch):
-    """Test that when kernels package is not available, the platform fallback logic
-
-    routes to local FLASH_ATTN.
-    """
-    from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
-
-    # Temporarily hide/remove kernels module if present
-    monkeypatch.setitem(sys.modules, "kernels", None)
-
-    # Use monkeypatch to mock CudaOmniPlatform capability and package checker to allow FLASH_ATTN
-    from vllm.platforms.interface import DeviceCapability
-
-    from vllm_omni.diffusion.envs import PACKAGES_CHECKER
-
-    monkeypatch.setattr(
-        CudaOmniPlatform,
-        "get_device_capability",
-        classmethod(lambda cls, device_id=0: DeviceCapability(8, 0)),
-    )
-    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True})
-
-    # Test FLASH_ATTN_HUB falls back to FLASH_ATTN when kernels is unavailable
-    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64)
-    assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
-
-    # Test FLASH_ATTN_3_HUB falls back to FLASH_ATTN when kernels is unavailable
-    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
-    assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
-
-    # Test FLASH_ATTN_3_HUB falls back to FLASH_ATTN_HUB on pre-Hopper GPUs
-    kernels_module = types.ModuleType("kernels")
-    monkeypatch.setitem(sys.modules, "kernels", kernels_module)
-    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
-    assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN_HUB.get_path()
-
-
-@pytest.mark.core_model
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_kernels_hub_execution():
-    """Verify basic forward of flash_attn_hub and flash_attn_3_hub, comparing with SDPA reference."""
-    device = torch.device(current_omni_platform.device_type)
-    dtype = torch.bfloat16
-
-    num_heads = 8
-    head_dim = 64
-    seq_len = 32
-    batch_size = 1
-
-    torch.manual_seed(42)
-    q = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
-    k = q.clone()
-    v = q.clone()
-
-    # 1. Test PyTorch SDPA reference
-    sdpa_impl = SDPAImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False)
-    attn_metadata_sdpa = AttentionMetadata(attn_mask=None)
-    output_ref = sdpa_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa)
-
-    # 2. Test FlashAttentionHubBackend (FlashAttention 2)
-    from vllm_omni.diffusion.attention.backends.flash_attn_hub import FlashAttentionHubImpl
-
-    fa_hub_impl = FlashAttentionHubImpl(
-        num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False
-    )
-    output_fa_hub = fa_hub_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa)
-    assert output_fa_hub.shape == q.shape
-    assert not torch.isnan(output_fa_hub).any()
-    max_diff = torch.max(torch.abs(output_ref - output_fa_hub)).item()
-    assert max_diff < 1e-2, f"FlashAttentionHub output differs too much from SDPA reference: {max_diff}"
-
-    # 3. Test FlashAttention3HubBackend (FlashAttention 3, Hopper+ only)
-    major, _minor = torch.cuda.get_device_capability()
-    if major < 9:
-        pytest.skip("FLASH_ATTN_3_HUB execution requires Hopper-class GPU (compute capability >= 9.0)")
-
-    from vllm_omni.diffusion.attention.backends.flash_attn_hub import FlashAttention3HubImpl
-
-    fa3_hub_impl = FlashAttention3HubImpl(
-        num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False
-    )
-    output_fa3_hub = fa3_hub_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa)
-    assert output_fa3_hub.shape == q.shape
-    assert not torch.isnan(output_fa3_hub).any()
-    max_diff = torch.max(torch.abs(output_ref - output_fa3_hub)).item()
-    assert max_diff < 1e-2, f"FlashAttention3Hub output differs too much from SDPA reference: {max_diff}"
@@ -19,10 +19,6 @@ class _FakeGroup:
        self.group_ranks = group_ranks
        self.parallel_mode = parallel_mode
        self.device_group = object()
-        self.device_communicator = kwargs.get("device_communicator")
-        reduce_scatter = kwargs.get("reduce_scatter")
-        if reduce_scatter is not None:
-            self.reduce_scatter = reduce_scatter
        self.ulysses_group = kwargs.get("ulysses_group")
        self.ring_group = kwargs.get("ring_group")
        self.local_group = next(group for group in group_ranks if local_rank in group)
@@ -59,23 +55,6 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
        created_groups.append(group)
        return group

-    def fake_init_vllm_model_parallel_group(
-        group_ranks,
-        local_rank,
-        backend,
-        group_name,
-    ):
-        del backend
-        group = _FakeGroup(
-            [list(ranks) for ranks in group_ranks],
-            local_rank,
-            f"vllm_{group_name}",
-            device_communicator=object(),
-            reduce_scatter=lambda tensor, **kwargs: tensor,
-        )
-        created_groups.append(group)
-        return group
-
    fake_world_group = SimpleNamespace(
        rank_in_group=local_rank,
        local_rank=local_rank,
@@ -89,7 +68,6 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
    monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
    monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
    monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
-    monkeypatch.setattr(parallel_state, "init_vllm_model_parallel_group", fake_init_vllm_model_parallel_group)
    monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)

    for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
@@ -109,7 +87,7 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
        backend="gloo",
    )

-    assert parallel_state.vllm_parallel_state._PCP is not parallel_state._SP
+    assert parallel_state.vllm_parallel_state._PCP is parallel_state._SP
    assert parallel_state.vllm_parallel_state._PCP.world_size == 2
    assert parallel_state._DP.world_size == 2
    assert parallel_state.vllm_parallel_state._DP is not parallel_state._DP
@@ -118,13 +96,6 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
    assert parallel_state.vllm_parallel_state._TP.world_size == 2
    assert parallel_state._PP.world_size == 2

-    assert parallel_state.vllm_parallel_state._PCP.device_communicator is not None
-    assert parallel_state.vllm_parallel_state._DP.device_communicator is not None
-    assert parallel_state.vllm_parallel_state._EP.device_communicator is not None
-    assert hasattr(parallel_state.vllm_parallel_state._PCP, "reduce_scatter")
-    assert hasattr(parallel_state.vllm_parallel_state._DP, "reduce_scatter")
-    assert hasattr(parallel_state.vllm_parallel_state._EP, "reduce_scatter")
-
    assert parallel_state.vllm_parallel_state._PCP.local_group == [0, 2]
    assert parallel_state.vllm_parallel_state._DP.local_group == [0, 8, 16, 24]
    assert parallel_state.vllm_parallel_state._EP.local_group == [
@@ -184,9 +155,7 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
        ],
    ]

-    vllm_group_names = [group.parallel_mode for group in created_groups if group.parallel_mode.startswith("vllm_")]
-    assert vllm_group_names == ["vllm_pcp", "vllm_tp", "vllm_dp", "vllm_ep"]
-    ep_groups = [group.local_group for group in created_groups if group.parallel_mode == "vllm_ep"]
+    ep_groups = [group.local_group for group in created_groups if group.parallel_mode == "expert"]
    assert ep_groups == [parallel_state.vllm_parallel_state._EP.local_group]


@@ -1,236 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Shape broadcasting regression tests for RotaryEmbedding.forward_native.
-
-Covers the batch-dimension shape mismatch fix: forward_native now strips the
-batch dim from 3D cos/sin before applying RoPE, consistent with forward_cuda,
-forward_hip, and apply_rotary_emb_mindiesd paths.
-"""
-
-from __future__ import annotations
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.layers.rope import RotaryEmbedding
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def _make_inputs(
-    batch: int,
-    seq_len: int,
-    num_heads: int,
-    head_dim: int,
-    cos_batch: int | None = None,
-    dtype: torch.dtype = torch.float32,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Generate x, cos, sin tensors for RoPE testing.
-
-    Args:
-        batch: Batch dimension of x.
-        seq_len: Sequence length.
-        num_heads: Number of attention heads.
-        head_dim: Head dimension (must be even).
-        cos_batch: Batch dimension of cos/sin. Defaults to batch if None.
-        dtype: Data type for the tensors.
-    """
-    if cos_batch is None:
-        cos_batch = batch
-
-    torch.manual_seed(42)
-    x = torch.randn(batch, seq_len, num_heads, head_dim, dtype=dtype)
-    cos = torch.randn(cos_batch, seq_len, head_dim // 2, dtype=dtype)
-    sin = torch.randn(cos_batch, seq_len, head_dim // 2, dtype=dtype)
-    return x, cos, sin
-
-
-class TestRotaryEmbeddingNativeShapeRegression:
-    """Regression tests for forward_native shape handling (#8e297d81)."""
-
-    @pytest.fixture
-    def rope_neox(self) -> RotaryEmbedding:
-        """NeoX-style (non-interleaved) RoPE."""
-        return RotaryEmbedding(is_neox_style=True)
-
-    @pytest.fixture
-    def rope_gptj(self) -> RotaryEmbedding:
-        """GPT-J style (interleaved) RoPE."""
-        return RotaryEmbedding(is_neox_style=False)
-
-    # ── 2D cos/sin (no batch dim) ──────────────────────────────────────
-
-    def test_2d_cos_sin_unchanged(self, rope_neox: RotaryEmbedding) -> None:
-        """2D cos/sin [S, D/2] pass through unchanged (no-op path)."""
-        x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, cos_batch=2)
-        # Manually drop batch dim to create 2D cos/sin
-        cos_2d = cos[0]
-        sin_2d = sin[0]
-        assert cos_2d.dim() == 2
-        output = rope_neox.forward_native(x, cos_2d, sin_2d)
-        assert output.shape == x.shape
-        assert not torch.isnan(output).any()
-        assert not torch.isinf(output).any()
-
-    # ── 3D cos/sin — batch-dim stripping ──────────────────────────────
-
-    @pytest.mark.parametrize("cos_batch", [1, 2, 4])
-    def test_3d_cos_sin_no_error(self, rope_neox: RotaryEmbedding, cos_batch: int) -> None:
-        """3D cos/sin with batch dim should not raise errors."""
-        x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, cos_batch=cos_batch)
-        output = rope_neox.forward_native(x, cos, sin)
-        assert output.shape == x.shape
-
-    @pytest.mark.parametrize("interleaved", [False, True])
-    def test_3d_cos_identical_to_2d(self, interleaved: bool) -> None:
-        """3D cos/sin with identical batch slices == 2D cos/sin result."""
-        rope = RotaryEmbedding(is_neox_style=not interleaved)
-        B, S, H, D = 2, 16, 8, 64
-        x = torch.randn(B, S, H, D)
-        cos_2d = torch.randn(S, D // 2)
-        sin_2d = torch.randn(S, D // 2)
-
-        # 3D variant: expand the same 2D data across batch
-        cos_3d = cos_2d.unsqueeze(0).expand(B, -1, -1)
-        sin_3d = sin_2d.unsqueeze(0).expand(B, -1, -1)
-
-        out_2d = rope.forward_native(x, cos_2d, sin_2d)
-        out_3d = rope.forward_native(x, cos_3d, sin_3d)
-
-        torch.testing.assert_close(out_3d, out_2d, atol=0, rtol=0)
-
-    # ── Batch size mismatch — the regression scenario ─────────────────
-
-    @pytest.mark.parametrize(
-        "x_batch,cos_batch",
-        [
-            (2, 1),  # cos has fewer batch elements
-            (1, 2),  # cos has more batch elements (takes [0])
-            (4, 1),  # large x batch, single cos
-        ],
-    )
-    def test_batch_size_mismatch_regression(
-        self,
-        rope_neox: RotaryEmbedding,
-        x_batch: int,
-        cos_batch: int,
-    ) -> None:
-        """forward_native works when cos batch dim differs from x batch dim.
-
-        This is the exact regression scenario — before the fix, a 3D cos/sin
-        with a different batch size than x would cause a shape mismatch error
-        inside apply_rotary_emb_torch during broadcast.
-        """
-        x, cos, sin = _make_inputs(batch=x_batch, seq_len=16, num_heads=8, head_dim=64, cos_batch=cos_batch)
-        output = rope_neox.forward_native(x, cos, sin)
-        assert output.shape == x.shape
-        assert output.dtype == x.dtype
-
-    # ── Batch consistency — all elements get same rotation ────────────
-
-    def test_2d_cos_broadcasts_identically_across_batch(self, rope_neox: RotaryEmbedding) -> None:
-        """With 2D cos, identical x slices → identical output slices."""
-        B, S, H, D = 4, 16, 8, 64
-        x_template = torch.randn(1, S, H, D)
-        x = x_template.expand(B, -1, -1, -1)  # all batch elements identical
-        cos = torch.randn(S, D // 2)
-        sin = torch.randn(S, D // 2)
-
-        output = rope_neox.forward_native(x, cos, sin)
-
-        # All output batch slices must be identical
-        for i in range(1, B):
-            torch.testing.assert_close(output[i], output[0], atol=0, rtol=0)
-
-    def test_3d_identical_cos_broadcasts_identically_across_batch(self, rope_neox: RotaryEmbedding) -> None:
-        """With 3D cos (identical slices), identical x → identical output."""
-        B, S, H, D = 4, 16, 8, 64
-        x_template = torch.randn(1, S, H, D)
-        x = x_template.expand(B, -1, -1, -1)
-        cos = torch.randn(1, S, D // 2).expand(B, -1, -1)
-        sin = torch.randn(1, S, D // 2).expand(B, -1, -1)
-
-        output = rope_neox.forward_native(x, cos, sin)
-
-        for i in range(1, B):
-            torch.testing.assert_close(output[i], output[0], atol=0, rtol=0)
-
-    # ── Output shape / dtype invariants ───────────────────────────────
-
-    def test_output_shape_preserved(self, rope_neox: RotaryEmbedding) -> None:
-        """Output shape must match input shape for all cos/sin dimensionalities."""
-        B, S, H, D = 2, 32, 4, 128
-        x = torch.randn(B, S, H, D)
-
-        # 2D cos/sin
-        cos_2d = torch.randn(S, D // 2)
-        sin_2d = torch.randn(S, D // 2)
-        assert rope_neox.forward_native(x, cos_2d, sin_2d).shape == (B, S, H, D)
-
-        # 3D cos/sin, matching batch
-        cos_3d = torch.randn(B, S, D // 2)
-        sin_3d = torch.randn(B, S, D // 2)
-        assert rope_neox.forward_native(x, cos_3d, sin_3d).shape == (B, S, H, D)
-
-        # 3D cos/sin, different batch
-        cos_3d_1 = torch.randn(1, S, D // 2)
-        sin_3d_1 = torch.randn(1, S, D // 2)
-        assert rope_neox.forward_native(x, cos_3d_1, sin_3d_1).shape == (B, S, H, D)
-
-    def test_output_dtype_preserved(self, rope_neox: RotaryEmbedding) -> None:
-        """Output dtype must match input dtype."""
-        for dtype in [torch.float32, torch.float16, torch.bfloat16]:
-            x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, dtype=dtype)
-            output = rope_neox.forward_native(x, cos, sin)
-            assert output.dtype == dtype, f"dtype mismatch: {output.dtype} != {dtype}"
-
-    def test_output_no_nan_inf(
-        self,
-        rope_neox: RotaryEmbedding,
-        rope_gptj: RotaryEmbedding,
-    ) -> None:
-        """Output must be finite for both RoPE styles."""
-        for rope in (rope_neox, rope_gptj):
-            x, cos, sin = _make_inputs(batch=2, seq_len=32, num_heads=4, head_dim=64)
-            output = rope.forward_native(x, cos, sin)
-            assert not torch.isnan(output).any(), f"NaN in {rope}"
-            assert not torch.isinf(output).any(), f"Inf in {rope}"
-
-    def test_output_contiguous(self, rope_neox: RotaryEmbedding) -> None:
-        """Output should be contiguous."""
-        x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64)
-        output = rope_neox.forward_native(x, cos, sin)
-        assert output.is_contiguous()
-
-    # ── CFG-style doubled batch ───────────────────────────────────────
-
-    def test_cfg_style_doubled_batch(self, rope_neox: RotaryEmbedding) -> None:
-        """CFG (classifier-free guidance) doubles batch — cos may have
-        original batch size while x has 2× batch."""
-        B, S, H, D = 2, 16, 8, 64
-        # Simulate CFG: x has 2× batch (cond + uncond concatenated)
-        x = torch.randn(B * 2, S, H, D)
-        # cos/sin only have the original batch
-        cos = torch.randn(B, S, D // 2)
-        sin = torch.randn(B, S, D // 2)
-
-        output = rope_neox.forward_native(x, cos, sin)
-        assert output.shape == (B * 2, S, H, D)
-
-    # ── RotaryEmbeddingWan forward_native is unchanged ────────────────
-
-    def test_wan_forward_native_unchanged(self) -> None:
-        """RotaryEmbeddingWan.forward_native still uses its own implementation
-        (unflatten + stack + flatten) and should not crash with 3D cos/sin."""
-        from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan
-
-        rope = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True)
-        B, S, H, D = 2, 16, 8, 64
-        x = torch.randn(B, S, H, D)
-        cos = torch.randn(B, S, 1, D // 2)
-        sin = torch.randn(B, S, 1, D // 2)
-
-        output = rope.forward_native(x, cos, sin)
-        assert output.shape == x.shape
-        assert not torch.isnan(output).any()
@@ -50,48 +50,6 @@ class TestSetForwardContextNumTokens:
        mock_get.assert_not_called()


-class TestSetForwardContextDPMetadata:
-    """Test diffusion MoE DP metadata setup for vLLM 0.24 all2all dispatch."""
-
-    def test_sets_dp_metadata_from_vllm_dp_group(self, mocker):
-        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
-
-        mock_ctx = mocker.MagicMock()
-        mock_ctx.dp_metadata = None
-        fake_group = mocker.MagicMock()
-        fake_group.world_size = 2
-        fake_group.cpu_group = object()
-
-        def _all_gather_object(output, obj, group):
-            assert group is fake_group.cpu_group
-            output[:] = [obj, obj + 1]
-
-        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
-        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
-        mocker.patch.object(hunyuan_moe._vllm_ps, "_DP", fake_group, create=True)
-        mocker.patch.object(hunyuan_moe.torch.distributed, "all_gather_object", side_effect=_all_gather_object)
-
-        hunyuan_moe._set_forward_context_dp_metadata(1024)
-
-        assert mock_ctx.dp_metadata is not None
-        assert mock_ctx.dp_metadata.num_tokens_across_dp_cpu.tolist() == [1024, 1025]
-
-    def test_keeps_existing_dp_metadata(self, mocker):
-        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
-
-        existing_metadata = object()
-        mock_ctx = mocker.MagicMock()
-        mock_ctx.dp_metadata = existing_metadata
-        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
-        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
-        mock_all_gather = mocker.patch.object(hunyuan_moe.torch.distributed, "all_gather_object")
-
-        hunyuan_moe._set_forward_context_dp_metadata(1024)
-
-        assert mock_ctx.dp_metadata is existing_metadata
-        mock_all_gather.assert_not_called()
-
-
 class TestHunyuanFusedMoEPlatformDispatch:
    """Test platform dispatch via platform qualname hooks."""

@@ -4,6 +4,7 @@
 from collections import OrderedDict

 import pytest
+import torch

 from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
 from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState
@@ -45,11 +46,14 @@ def test_dreamzero_pipeline_state_lru_caps_retained_sessions() -> None:
    assert "session-b" not in pipeline._states


-def test_dreamzero_state_owns_no_kv_caches() -> None:
-    """KV (self- and cross-attention) is engine-owned since the AR-Diffusion
-    paged backend: the model-local cache accessors must be gone so nothing can
-    silently bypass the engine pool."""
+def test_dreamzero_state_cache_access_requires_initialization() -> None:
    state = DreamZeroState()

-    for removed in ("get_kv_caches", "create_kv_caches", "update_kv_cache", "get_crossattn_caches"):
-        assert not hasattr(state, removed)
+    with pytest.raises(RuntimeError, match="KV caches not initialized"):
+        state.get_kv_caches()
+
+    with pytest.raises(RuntimeError, match="Cross-attn caches not initialized"):
+        state.get_crossattn_caches()
+
+    with pytest.raises(RuntimeError, match="create_kv_caches first"):
+        state.update_kv_cache(0, torch.empty(0))
@@ -27,15 +27,8 @@ WIDTH = 1024
 PROMPT = "A brown and white dog is running on the grass."
 MEAN_THRESHOLD = 3e-2
 P99_THRESHOLD = 3e-1
-# NOTE: re-baselined for the vLLM rebase (dev/vllm-align). HunyuanImage-3 is an
-# expert-parallel MoE model; the newer upstream vLLM (65a7b4628409, 585 commits
-# past the v0.24.0 tag that main tracks) shifted MoE kernel numerics, so the
-# structural-similarity metrics against the committed baseline drift to a
-# deterministic SSIM 0.9490 / PSNR 26.21 dB (the mean/p99 pixel-diff checks above
-# still pass — the image renders correctly). Thresholds relaxed to accept this
-# drift as the new bar. On main (older vLLM) the stricter 0.97 / 30.0 hold.
-SSIM_THRESHOLD = 0.94
-PSNR_THRESHOLD = 25.0
+SSIM_THRESHOLD = 0.97
+PSNR_THRESHOLD = 30.0

 _REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent
 BASELINE_PATH = _REPO_ROOT / "tests" / "assets" / "hunyuan" / "hunyuan_baseline.png"
@@ -18,8 +18,6 @@ from tests.helpers.runtime import DiffusionResponse, OmniRunnerHandler
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform

-# NOTE: Most of the flux2klein tests have migrated to the common diffusion tests.
-# This model will also use the tiny model builder when run at the core model level.
 MODEL = "black-forest-labs/FLUX.2-klein-4B"

 _OMNI_RUNNER_PARAM = (MODEL, None)
@@ -72,6 +70,41 @@ def _send_inpaint_with_generator(
    )


+# Regression test for https://github.com/vllm-project/vllm-omni/issues/3097
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+def test_flux2_klein_can_accept_text_inputs(omni_runner_handler: OmniRunnerHandler):
+    omni_runner_handler.send_diffusion_request(
+        {
+            "model": MODEL,
+            "prompt": "a cup of coffee on the table",
+            "sampling_params": OmniDiffusionSamplingParams(num_inference_steps=2, seed=42),
+        }
+    )
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_flux2_klein_inpaint_basic(omni_runner_handler: OmniRunnerHandler):
+    input_image, mask_image = _create_test_inputs()
+    omni_runner_handler.send_diffusion_request(
+        {
+            "model": MODEL,
+            "prompt": "Fill in the masked area with a beautiful garden",
+            "multi_modal_data": {"image": input_image, "mask_image": mask_image},
+            "sampling_params": OmniDiffusionSamplingParams(
+                height=_HEIGHT,
+                width=_WIDTH,
+                num_inference_steps=_NUM_INFERENCE_STEPS,
+                guidance_scale=0.0,
+                generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
+                num_outputs_per_prompt=1,
+            ),
+        }
+    )
+
+
@pytest.mark.advanced_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@@ -22,7 +22,6 @@ _OMNI_RUNNER_PARAM = (
    {
        "init_timeout": 1200,
        "tensor_parallel_size": 2,
-        "vae_use_tiling": True,
    },
 )

@@ -33,7 +32,6 @@ _OMNI_RUNNER_PARAM_OFFLOAD = (
        "init_timeout": 1200,
        "tensor_parallel_size": 2,
        "enable_cpu_offload": True,
-        "vae_use_tiling": True,
    },
 )

@@ -47,6 +45,10 @@ pytestmark = [
    ),
 ]

+_SKIP_ISSUE_3236 = pytest.mark.skip(
+    reason="https://github.com/vllm-project/vllm-omni/issues/3236",
+)
+

 def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
    """Validate that the MP4 contains meaningful video and audio tracks."""
@@ -72,6 +74,7 @@ def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
    container.close()


+@_SKIP_ISSUE_3236
@hardware_test(res={"cuda": "H100"}, num_cards=2)
 def test_magi_human_e2e(omni_runner: OmniRunner):
    """End-to-end test for MagiHuman generating video and audio."""
@@ -5,7 +5,7 @@ Verifies that the AR->DiT pipeline produces an image tensor whose pixel values
 match a golden reference.

 Model Hub repo id: ``bytedance-research/MammothModa2-Preview``.
-Deploy config: ``get_deploy_config_path("mammoth_moda2.yaml")`` -> ``vllm_omni/deploy/mammoth_moda2.yaml``
+Stage config: ``get_deploy_config_path("mammoth_moda2.yaml")`` → ``vllm_omni/deploy/mammoth_moda2.yaml``

 Golden pixel file: ``tests/e2e/offline_inference/fixtures/mammoth_moda2_t2i_golden.json``
  Regenerate with: ``UPDATE_GOLDEN=1 pytest tests/e2e/offline_inference/test_mammoth_moda2_expansion.py``
@@ -36,9 +36,9 @@ _VISION_END_TOKEN_ID = 151653  # "<|vision_end|>"
 _AR_PATCH_SIZE = 16

 MODEL_PATH = "bytedance-research/MammothModa2-Preview"
-T2I_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")
+T2I_STAGE_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")

-_OMNI_RUNNER_PARAM = (MODEL_PATH, None, {"deploy_config": T2I_DEPLOY_CONFIG})
+_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_STAGE_CONFIG)

 # Golden pixel reference file.  Set UPDATE_GOLDEN=1 to regenerate.
 _GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json"
@@ -18,15 +18,13 @@ from tests.helpers.runtime import OmniRunner
 from tests.helpers.stage_config import get_deploy_config_path

 MODEL = "k2-fsa/OmniVoice"
-DEPLOY_CONFIG = get_deploy_config_path("omnivoice.yaml")
+STAGE_CONFIG = get_deploy_config_path("omnivoice.yaml")

-# OmniRunner tuple: model, legacy stage config path, extra Omni kwargs.
-# The migrated test passes deploy_config through extra Omni kwargs.
+# (model, stage_config_path, extra_omni_kwargs) — see ``omni_runner`` in tests.helpers.fixtures.runtime
 _OMNI_RUNNER_PARAM = (
    MODEL,
-    None,
+    STAGE_CONFIG,
    {
-        "deploy_config": DEPLOY_CONFIG,
        "trust_remote_code": True,
        "log_stats": True,
    },
@@ -34,10 +34,12 @@ def get_cuda_graph_config():
            "stages": {
                0: {
                    "max_num_seqs": 1,
+                    "gpu_memory_utilization": 0.2,
                    "enforce_eager": True,
                    "async_scheduling": False,
                },
                1: {
+                    "gpu_memory_utilization": 0.2,
                    "enforce_eager": True,
                    "async_scheduling": False,
                },
@@ -0,0 +1,62 @@
+"""
+Tests for Flux2 Klein; currently Dev is implemented separately,
+but ideally these models will fold together in the future.
+"""
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import (
+    OmniServer,
+    OmniServerParams,
+    OpenAIClientHandler,
+    dummy_messages_from_mix_data,
+)
+
+FOUR_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=4)
+POSITIVE_PROMPT = "A cat sitting on a windowsill"
+NEGATIVE_PROMPT = "blurry, low quality"
+
+
+# For now, we only test a good configuration for this model to keep the CI
+# light, but in the near future may explore using random models to check
+# cross-feature compatibility more generally.
+def _get_diffusion_feature_cases(model: str):
+    return [
+        # FP8 / Hybrid sequence parallelism
+        pytest.param(
+            OmniServerParams(
+                model=model,
+            ),
+            id="default",
+            marks=FOUR_CARD_FEATURE_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_diffusion_feature_cases(
+        model="black-forest-labs/FLUX.2-klein-4B",
+    ),
+    indirect=True,
+)
+def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+    messages = dummy_messages_from_mix_data(content_text=POSITIVE_PROMPT)
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "height": 512,
+            "width": 512,
+            "num_inference_steps": 2,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "true_cfg_scale": 4.0,
+            "seed": 42,
+        },
+    }
+
+    openai_client.send_diffusion_request(request_config)
@@ -0,0 +1,163 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+End-to-end tests for Flux2 Klein inpainting in online serving mode.
+
+Uses /v1/images/edits endpoint which is the correct API for image inpainting.
+"""
+
+import base64
+from io import BytesIO
+
+import httpx
+import pytest
+from PIL import Image, ImageDraw
+
+from tests.helpers.runtime import OmniServer, OmniServerParams
+
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
+
+MODEL = "black-forest-labs/FLUX.2-klein-4B"
+
+_HEIGHT = 512
+_WIDTH = 512
+_NUM_INFERENCE_STEPS = 4
+
+
+def _get_diffusion_feature_cases(model: str):
+    return [
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=["--tensor-parallel-size", "2", "--enable-cpu-offload"],
+            ),
+            id="tp2_cpu_offload",
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=["--tensor-parallel-size", "2"],
+            ),
+            id="tp2_basic",
+        ),
+    ]
+
+
+def _image_to_base64_jpeg(image: Image.Image) -> str:
+    buffer = BytesIO()
+    image.save(buffer, format="JPEG")
+    buffer.seek(0)
+    return base64.b64encode(buffer.read()).decode("utf-8")
+
+
+def _create_test_mask_base64(width: int = _WIDTH, height: int = _HEIGHT) -> str:
+    mask = Image.new("L", (width, height), 0)
+    draw = ImageDraw.Draw(mask)
+    draw.rectangle([width // 4, height // 4, width * 3 // 4, height * 3 // 4], fill=255)
+    return _image_to_base64_jpeg(mask)
+
+
+def _compare_images(img1: Image.Image, img2: Image.Image) -> bool:
+    return list(img1.getdata()) == list(img2.getdata())
+
+
+def _send_edit_request(host: str, port: int, model: str, image_b64: str, mask_b64: str, prompt: str, **kwargs):
+    url = f"http://{host}:{port}/v1/images/edits"
+    files = {
+        "image": ("image.jpg", base64.b64decode(image_b64), "image/jpeg"),
+        "mask_image": ("mask.jpg", base64.b64decode(mask_b64), "image/jpeg"),
+    }
+    data = {"prompt": prompt, "model": model, **kwargs}
+    with httpx.Client(timeout=60.0) as client:
+        response = client.post(url, files=files, data=data)
+        response.raise_for_status()
+        return response.json()
+
+
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_flux2_klein_inpaint_basic(omni_server: OmniServer):
+    input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
+    mask_b64 = _create_test_mask_base64()
+
+    result = _send_edit_request(
+        host=omni_server.host,
+        port=omni_server.port,
+        model=MODEL,
+        image_b64=input_image_b64,
+        mask_b64=mask_b64,
+        prompt="Fill in the masked area with a beautiful garden",
+        guidance_scale=1.0,
+        num_inference_steps=_NUM_INFERENCE_STEPS,
+        n=1,
+        seed=42,
+    )
+
+    assert "data" in result and len(result["data"]) == 1
+    img_data = result["data"][0].get("b64_json") or result["data"][0].get("url", "").split(",")[-1]
+    img = Image.open(BytesIO(base64.b64decode(img_data)))
+    assert img.size == (_WIDTH, _HEIGHT)
+
+
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_flux2_klein_inpaint_deterministic(omni_server: OmniServer):
+    input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
+    mask_b64 = _create_test_mask_base64()
+    prompt = "A red flower in a field"
+
+    result1 = _send_edit_request(
+        host=omni_server.host,
+        port=omni_server.port,
+        model=MODEL,
+        image_b64=input_image_b64,
+        mask_b64=mask_b64,
+        prompt=prompt,
+        guidance_scale=1.0,
+        num_inference_steps=_NUM_INFERENCE_STEPS,
+        n=1,
+        seed=12345,
+    )
+
+    result2 = _send_edit_request(
+        host=omni_server.host,
+        port=omni_server.port,
+        model=MODEL,
+        image_b64=input_image_b64,
+        mask_b64=mask_b64,
+        prompt=prompt,
+        guidance_scale=1.0,
+        num_inference_steps=_NUM_INFERENCE_STEPS,
+        n=1,
+        seed=12345,
+    )
+
+    img1_data = result1["data"][0].get("b64_json") or result1["data"][0].get("url", "").split(",")[-1]
+    img2_data = result2["data"][0].get("b64_json") or result2["data"][0].get("url", "").split(",")[-1]
+
+    img1 = Image.open(BytesIO(base64.b64decode(img1_data)))
+    img2 = Image.open(BytesIO(base64.b64decode(img2_data)))
+
+    assert _compare_images(img1, img2), (
+        "Same input with same seed should produce identical output. This is critical for offline/online consistency."
+    )
+
+
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_flux2_klein_inpaint_multiple_outputs(omni_server: OmniServer):
+    input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
+    mask_b64 = _create_test_mask_base64()
+
+    result = _send_edit_request(
+        host=omni_server.host,
+        port=omni_server.port,
+        model=MODEL,
+        image_b64=input_image_b64,
+        mask_b64=mask_b64,
+        prompt="A beautiful landscape",
+        guidance_scale=1.0,
+        num_inference_steps=_NUM_INFERENCE_STEPS,
+        n=2,
+        seed=42,
+    )
+
+    assert "data" in result and len(result["data"]) == 2
@@ -127,7 +127,7 @@ class TestHiggsAudioV3OnlineHappyPath:
        windows stitch into a coherent PCM stream. The byte-count gate is the same as
        the sync paths; per-chunk audio content is verified offline against Whisper.

-        NOTE: ``min_hnr_db=0.0`` sit below the typical speech-noise floor so the
+        NOTE: ``min_hnr_db=0.0`` sits below the typical speech-noise floor so the
        check still catches catastrophic codec failure (silence, white noise, sample
        scramble all give HNR << 0) while allowing for the sliding-window codec's
        slightly-noisier-than-sync output. The default 1.0 dB threshold has only
@@ -226,27 +226,3 @@ def test_thinker_prefix_caching_audio_output(omni_server, openai_client) -> None
    }

    _run_prefix_cache_check(openai_client, request_config)
-
-
-@pytest.mark.advanced_model
-@pytest.mark.core_model
-@pytest.mark.omni
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_completions_rejected_for_thinker_talker(omni_server, openai_client) -> None:
-    """Ensure Thinker-talker models reject /v1/completions; we do this because the
-    thinker-talker handoff implementations currently use ChatML <|im_start|> and
-    <|im_end|> markers to segment the input sequence; when we don't have them,
-    the talker does not get any embeddings, which breaks the server.
-    """
-    responses = openai_client.send_completions_http_request(
-        {
-            "json": {
-                "model": omni_server.model,
-                "prompt": "Hello, how are you?",
-                "max_tokens": 10,
-            },
-        },
-        err_code=400,
-    )
-    assert not responses[0].success
@@ -410,44 +410,6 @@ def test_audio_in_video_002(omni_server, openai_client) -> None:
    openai_client.send_omni_request(request_config, request_num=get_max_batch_size())


-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_audio_in_video_default_loader_sampling_regression(omni_server, openai_client) -> None:
-    """
-    Regression: ``use_audio_in_video`` with default video loader sampling (no
-    ``media_io_kwargs.video.fps``) after a prior request that pinned fps.
-
-    Before the fix the second request could partial-hit mm-cache (audio cached,
-    video re-processed without audio) and fail with ``StopIteration`` in
-    ``replace_multimodal_special_tokens``.
-    """
-    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 128, embed_audio=True)['base64']}"
-    messages = dummy_messages_from_mix_data(
-        system_prompt=get_system_prompt(),
-        video_data_url=video_data_url,
-        content_text="Reply with one word: OK.",
-    )
-
-    base_config = {
-        "model": omni_server.model,
-        "messages": messages,
-        "stream": False,
-        "use_audio_in_video": True,
-        "modalities": ["text"],
-    }
-
-    # Prime mm-cache with explicit loader fps (historical failure mode).
-    openai_client.send_omni_request(
-        {
-            **base_config,
-            "extra_body": {"media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}}},
-        }
-    )
-
-    # Default loader sampling: must succeed without explicit media_io fps.
-    openai_client.send_omni_request(dict(base_config))
-
-
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_one_word_prompt_001(omni_server, openai_client) -> None:
@@ -739,7 +739,7 @@ def test_resolve_stage_configs_injects_global_diffusion_attention_when_missing(m
    monkeypatch.setattr(
        engine_mod,
        "load_and_resolve_stage_configs",
-        lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
+        lambda *args, **kwargs: ("dummy-config", [stage_cfg]),
    )

    _config_path, stage_configs = engine._resolve_stage_configs(
@@ -772,7 +772,7 @@ def test_resolve_stage_configs_preserves_stage_diffusion_attention(monkeypatch):
    monkeypatch.setattr(
        engine_mod,
        "load_and_resolve_stage_configs",
-        lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
+        lambda *args, **kwargs: ("dummy-config", [stage_cfg]),
    )

    _config_path, stage_configs = engine._resolve_stage_configs(
@@ -798,7 +798,7 @@ def test_resolve_stage_configs_does_not_inject_diffusion_attention_into_llm_stag
    monkeypatch.setattr(
        engine_mod,
        "load_and_resolve_stage_configs",
-        lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
+        lambda *args, **kwargs: ("dummy-config", [stage_cfg]),
    )

    _config_path, stage_configs = engine._resolve_stage_configs(
@@ -874,54 +874,3 @@ def test_build_engine_args_dict_uses_diffusion_attention_config_key():

    assert "attention_config" not in engine_args_dict
    assert engine_args_dict["diffusion_attention_config"].default.backend == "FLASH_ATTN"
-
-
-def test_omni_master_server_allocates_globally_unique_route_ports(monkeypatch):
-    """Regression: two stages must never draw the same ZMQ port.
-
-    ``get_open_ports_list`` only dedups within a single call, so per-route
-    allocation used to let a later stage reuse an earlier stage's port. The
-    second engine to ``bind()`` then died with ``zmq.error.ZMQError: Address
-    already in use`` (flaky multi-stage startup, e.g. Qwen3-Omni thinker/talker/
-    code2wav). ``OmniMasterServer`` now dedups every port it hands out.
-    """
-    import itertools
-
-    from vllm_omni.engine import stage_engine_startup as ses
-
-    # A colliding prefix (repeats + the master port 9000) followed by an endless
-    # fresh stream, so the only way to succeed is to redraw the collisions.
-    supply = itertools.chain([9000, 9000, 9001, 9001, 9000], itertools.count(9002))
-
-    def fake_get_open_ports_list(count):
-        return [next(supply) for _ in range(count)]
-
-    monkeypatch.setattr(ses, "get_open_ports_list", fake_get_open_ports_list)
-
-    server = ses.OmniMasterServer(
-        master_address="127.0.0.1",
-        master_port=9000,  # seed: a route must not reuse the registration port
-        stage_ids=[0, 1, 2],
-    )
-
-    ports = []
-    for sid in (0, 1, 2):
-        alloc = server.get_allocation(sid)
-        for addr in (
-            alloc.handshake_bind_address,
-            alloc.input_bind_address,
-            alloc.output_bind_address,
-        ):
-            ports.append(ses._port_from_zmq_address(addr))
-
-    assert len(ports) == len(set(ports)), f"duplicate route ports allocated: {ports}"
-    assert 9000 not in ports, "route reused the master registration port"
-
-
-def test_port_from_zmq_address_parsing():
-    from vllm_omni.engine.stage_engine_startup import _port_from_zmq_address
-
-    assert _port_from_zmq_address("tcp://127.0.0.1:34277") == 34277
-    assert _port_from_zmq_address(None) is None
-    assert _port_from_zmq_address("ipc:///tmp/sock") is None
-    assert _port_from_zmq_address("tcp://host:not-a-port") is None
@@ -1,63 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for strategy-derived omni_lb_policy threading into AsyncOmniEngine.
-
-``AsyncOmniEngine._apply_strategy_lb_policy`` only touches ``self._omni_lb_policy``
-and module-level logging, so we exercise it with a lightweight stub ``self`` to
-avoid standing up a full engine (which needs GPUs/model weights).
-"""
-
-from __future__ import annotations
-
-import pytest
-
-from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-class _Stub:
-    def __init__(self, current: str = "random") -> None:
-        self._omni_lb_policy = current
-
-
-_apply = AsyncOmniEngine._apply_strategy_lb_policy
-
-
-def test_no_derived_policy_is_noop():
-    stub = _Stub("random")
-    _apply(stub, None, {})
-    assert stub._omni_lb_policy == "random"
-
-
-def test_derived_overrides_default_random():
-    # "random" is the engine default == "unset", so the strategy value wins.
-    stub = _Stub("random")
-    _apply(stub, "round-robin", {})
-    assert stub._omni_lb_policy == "round-robin"
-
-
-def test_derived_overrides_when_flag_absent():
-    stub = _Stub("random")
-    _apply(stub, "least-queue-length", {"omni_lb_policy": None})
-    assert stub._omni_lb_policy == "least-queue-length"
-
-
-def test_explicit_matching_flag_is_fine():
-    stub = _Stub("round-robin")
-    _apply(stub, "round-robin", {"omni_lb_policy": "round-robin"})
-    assert stub._omni_lb_policy == "round-robin"
-
-
-def test_explicit_conflicting_flag_raises():
-    stub = _Stub("round-robin")
-    with pytest.raises(ValueError, match="Conflicting load-balancer policy"):
-        _apply(stub, "least-queue-length", {"omni_lb_policy": "round-robin"})
-
-
-def test_explicit_random_is_treated_as_unset():
-    # An explicit --omni-lb-policy=random is indistinguishable from the default,
-    # so a strategy value still overrides it (documented Phase-1 semantics).
-    stub = _Stub("random")
-    _apply(stub, "round-robin", {"omni_lb_policy": "random"})
-    assert stub._omni_lb_policy == "round-robin"
@@ -198,37 +198,16 @@ def _assert_realtime_smoke(result: dict) -> None:
    assert result["output_sample_rate"] > 0


-def _assert_realtime_accuracy(
-    result: dict,
-    threshold: float = 0.8,
-    whisper_model_size: str = "large-v3",
-) -> None:
-    """Assert that whisper transcription of audio output matches model text.
-
-    Args:
-        result: Roundtrip result dict from ``_run_realtime_audio_roundtrip``.
-        threshold: Minimum cosine similarity between whisper text and model text.
-                   Default 0.8 is appropriate for sync mode (exact match expected).
-                   Use 0.65 for async_chunk mode where chunked upload with pacing
-                   delays introduces inherent variability in audio codec quality.
-        whisper_model_size: Whisper model used to transcribe the generated audio
-                   for the accuracy check. Defaults to ``large-v3``: the default
-                   ``small`` model mishears short Chinese TTS clips (observed:
-                   北京→韦京 and a dropped leading sentence, sim=0.443), which
-                   caused spurious sim<threshold failures under async_chunk codec
-                   variability even though audio generation was correct. large-v3
-                   transcribes these clips reliably, so a failure here now points
-                   at the model, not the ASR grader.
-    """
+def _assert_realtime_accuracy(result: dict) -> None:
    final_text = (result["transcription_text"] or "").strip()
    assert final_text, "Expected non-empty transcription (model text stream)"

    wav_out = _wav_bytes_from_pcm16(result["output_pcm"], result["output_sample_rate"])
-    whisper_text = convert_audio_bytes_to_text(wav_out, model_size=whisper_model_size).strip()
+    whisper_text = convert_audio_bytes_to_text(wav_out).strip()
    assert whisper_text, "Whisper returned empty string for synthesized output audio"

    sim = cosine_similarity_text(whisper_text.lower(), final_text.lower())
-    assert sim > threshold, (
+    assert sim > 0.8, (
        f"Output audio transcript should match model text (sim={sim:.3f}): "
        f"whisper={whisper_text!r}, model_text={final_text!r}"
    )
@@ -255,12 +234,7 @@ class TestQwen3OmniRealtimeWebSocket:
        )

        _assert_realtime_smoke(result)
-        # Use threshold=0.65 for async_chunk: chunked upload with send_delay_ms
-        # introduces inherent variability in audio codec quality. Whisper
-        # transcription is nondeterministic and may miss 1-2 words in async
-        # mode even when model text generation is correct.
-        # Sync mode uses the default 0.8 (exact match expected).
-        _assert_realtime_accuracy(result, threshold=0.65)
+        _assert_realtime_accuracy(result)

    @pytest.mark.advanced_model
    @pytest.mark.omni
@@ -87,8 +87,8 @@ def test_create_chat_completion_converts_value_error_to_error_response(mocker: M
    serving_chat.chat_template = None
    serving_chat.chat_template_content_format = "string"
    serving_chat.default_chat_template_kwargs = {}
-    serving_chat.online_renderer = mocker.MagicMock()
-    serving_chat.online_renderer.validate_chat_template.return_value = None
+    serving_chat.openai_serving_render = mocker.MagicMock()
+    serving_chat.openai_serving_render.validate_chat_template.return_value = None
    serving_chat._effective_chat_template_kwargs = mocker.MagicMock(return_value={})
    serving_chat._preprocess_chat = mocker.AsyncMock(
        side_effect=ValueError("Invalid speaker 'uncle_fu'. Supported: ethan, vivian")
@@ -39,8 +39,7 @@ from vllm_omni.entrypoints.openai.serving_speech import (
    OmniOpenAIServingSpeech,
    _create_wav_header,
 )
-from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest, SpeechServingContext
-from vllm_omni.entrypoints.openai.tts_adapters.ming_tts import MingTTSAdapter
+from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest
 from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
    FISH_TEXT_ONLY_SYSTEM_PROMPT,
    build_fish_voice_clone_prompt_ids,
@@ -974,29 +973,6 @@ class TestTTSMethods:
                )
            )

-    def test_upload_ming_voice_embedding_wrong_dims_rejected_without_replacing_existing(self, speech_server):
-        """Ming embedding uploads must be 192-dim before replacing stored voices."""
-        speech_server._tts_model_type = "ming_tts"
-        existing = {
-            "name": "bad_emb_voice",
-            "file_path": "/tmp/voice_samples/bad_emb_voice.safetensors",
-            "mime_type": "application/x-safetensors",
-            "embedding_source": "direct",
-            "embedding_dim": 192,
-        }
-        speech_server.uploaded_speakers = {"bad_emb_voice": existing.copy()}
-
-        with pytest.raises(ValueError, match="Ming speaker embedding must have 192 dims, got 191"):
-            asyncio.run(
-                speech_server.upload_voice_embedding(
-                    embedding_json=json.dumps([0.0] * 191),
-                    consent="consent",
-                    name="bad_emb_voice",
-                )
-            )
-
-        assert speech_server.uploaded_speakers["bad_emb_voice"] == existing
-
    def test_base_task_requires_ref_audio_or_speaker_embedding(self, speech_server):
        """Base task without ref_audio or speaker_embedding is rejected."""
        req = OpenAICreateSpeechRequest(input="Hello", task_type="Base")
@@ -1581,86 +1557,6 @@ class TestTTSMethods:
        assert params["x_vector_only_mode"] == [True]
        assert "ref_audio" not in params

-    def test_ming_adapter_uploaded_direct_embedding_uses_embedding_not_audio(
-        self, speech_server, mocker: MockerFixture
-    ):
-        """Ming uploaded direct embeddings are loaded into prompt construction."""
-        speech_server.uploaded_speakers = {
-            "emb_voice": {
-                "name": "emb_voice",
-                "file_path": "/tmp/voice_samples/emb_voice.safetensors",
-                "mime_type": "application/x-safetensors",
-                "embedding_source": "direct",
-                "embedding_dim": 192,
-            }
-        }
-        fake_embedding = [0.1] * 192
-        mock_get_emb = mocker.patch.object(
-            speech_server, "_get_uploaded_speaker_embedding", return_value=fake_embedding
-        )
-        mock_get_audio = mocker.patch.object(speech_server, "_get_uploaded_audio_data")
-        mock_prompt = mocker.patch.object(
-            speech_server,
-            "_build_ming_dense_prompt",
-            return_value={"additional_information": {"speaker_count": 1}},
-        )
-
-        adapter = MingTTSAdapter(SpeechServingContext(server=speech_server))
-        req = OpenAICreateSpeechRequest(input="Hello", voice="emb_voice")
-        prepared = asyncio.run(adapter.build(req, [], False))
-
-        mock_get_emb.assert_called_once_with("emb_voice")
-        mock_get_audio.assert_not_called()
-        mock_prompt.assert_called_once_with(req, ref_audio_data=None)
-        prompt_request = mock_prompt.call_args.args[0]
-        assert prompt_request.speaker_embedding == fake_embedding
-        assert len(prompt_request.speaker_embedding) == 192
-        assert prepared.prompt["additional_information"]["speaker_count"] == 1
-        assert prepared.model_type == "ming_tts"
-
-    def test_ming_adapter_uploaded_audio_path_preserves_ref_text(self, speech_server, mocker: MockerFixture):
-        """Ming uploaded audio voices still resolve audio and stored ref_text."""
-        speech_server.uploaded_speakers = {
-            "audio_voice": {
-                "name": "audio_voice",
-                "file_path": "/tmp/voice_samples/audio_voice.wav",
-                "mime_type": "audio/wav",
-                "ref_text": "Reference transcript.",
-            }
-        }
-        ref_audio_source = "data:audio/wav;base64,ZmFrZWF1ZGlv"
-        ref_audio_data = ([0.0, 0.1], 16000)
-        fake_embedding = [0.2] * 192
-        mock_get_audio = mocker.patch.object(speech_server, "_get_uploaded_audio_data", return_value=ref_audio_source)
-        mock_get_emb = mocker.patch.object(speech_server, "_get_uploaded_speaker_embedding")
-        mocker.patch.object(
-            speech_server,
-            "_resolve_ref_audio",
-            new=mocker.AsyncMock(return_value=ref_audio_data),
-        )
-        mock_extract = mocker.patch.object(
-            speech_server,
-            "_extract_ming_speaker_embeddings_from_ref_audio",
-            return_value=[fake_embedding],
-        )
-        mock_prompt = mocker.patch.object(
-            speech_server,
-            "_build_ming_dense_prompt",
-            return_value={"additional_information": {"speaker_count": 1}},
-        )
-
-        adapter = MingTTSAdapter(SpeechServingContext(server=speech_server))
-        req = OpenAICreateSpeechRequest(input="Hello", voice="audio_voice")
-        prepared = asyncio.run(adapter.build(req, [], False))
-
-        mock_get_audio.assert_called_once_with("audio_voice")
-        mock_get_emb.assert_not_called()
-        mock_extract.assert_called_once_with([ref_audio_data])
-        mock_prompt.assert_called_once_with(req, ref_audio_data=ref_audio_data)
-        assert req.ref_text == "Reference transcript."
-        assert req.speaker_embedding == fake_embedding
-        assert prepared.model_type == "ming_tts"
-
    # ── regression: full flow from issue #1603 ──

    def test_regression_1603_speaker_key_with_uploaded_audio_voice(self, speech_server, mocker: MockerFixture):
@@ -3448,27 +3344,6 @@ class TestTTSAsyncOffloading:
        qwen3_tts_server._build_tts_params.assert_called_once()
        qwen3_tts_server._estimate_prompt_len_async.assert_awaited_once()

-    def test_prepare_speech_generation_qwen3_default_seed_sets_tts_local_seed(
-        self, qwen3_tts_server, mocker: MockerFixture
-    ):
-        """Deploy default seed should seed Qwen3 TTS residual MTP sampling."""
-        qwen3_tts_server.engine_client.default_sampling_params_list = [
-            SimpleNamespace(max_tokens=2048, seed=42, extra_args=None)
-        ]
-        qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None)
-        qwen3_tts_server._build_tts_params = mocker.MagicMock(
-            return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]}
-        )
-        qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512)
-        request = OpenAICreateSpeechRequest(input="hello")
-
-        asyncio.run(qwen3_tts_server._prepare_speech_generation(request))
-
-        stage0_params = qwen3_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0]
-        assert stage0_params.seed == 42
-        assert stage0_params.extra_args["tts_local_seed"] == 42
-        assert qwen3_tts_server.engine_client.default_sampling_params_list[0].extra_args is None
-
    def test_prepare_speech_generation_uses_adapter_model_type_label(
        self,
        voxtral_server,
@@ -3554,48 +3429,6 @@ class TestTTSAsyncOffloading:

        mock_coerce.assert_called_once_with(qwen3_tts_server.engine_client.default_sampling_params_list, True)

-    def test_prepare_speech_generation_no_async_chunk_stream_uses_final_only(
-        self, qwen3_tts_server, mocker: MockerFixture
-    ):
-        """Full-payload TTS streaming should not request delta multimodal outputs."""
-        qwen3_tts_server.engine_client.model_config.async_chunk = False
-        qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None)
-        qwen3_tts_server._build_tts_params = mocker.MagicMock(
-            return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]}
-        )
-        qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512)
-        mock_coerce = mocker.patch(
-            "vllm_omni.entrypoints.openai.serving_speech.coerce_param_message_types",
-            return_value=qwen3_tts_server.engine_client.default_sampling_params_list,
-        )
-        request = OpenAICreateSpeechRequest(input="hello", stream_format="audio", response_format="pcm")
-
-        asyncio.run(qwen3_tts_server._prepare_speech_generation(request))
-
-        mock_coerce.assert_called_once_with(qwen3_tts_server.engine_client.default_sampling_params_list, False)
-
-    def test_prepare_speech_generation_no_async_chunk_stream_keeps_delta_for_non_qwen3(
-        self, voxtral_server, mocker: MockerFixture
-    ):
-        """FINAL_ONLY streaming for async_chunk=False is scoped to qwen3_tts only."""
-        voxtral_server.engine_client.model_config.async_chunk = False
-        mocker.patch.object(voxtral_server._get_tts_adapter(), "validate", return_value=None)
-        voxtral_server._build_voxtral_prompt_async = mocker.AsyncMock(
-            return_value={
-                "prompt_token_ids": [1, 2, 3],
-                "additional_information": {"voice": ["test"]},
-            }
-        )
-        mock_coerce = mocker.patch(
-            "vllm_omni.entrypoints.openai.serving_speech.coerce_param_message_types",
-            return_value=voxtral_server.engine_client.default_sampling_params_list,
-        )
-        request = OpenAICreateSpeechRequest(input="hello", voice="test", stream_format="audio", response_format="pcm")
-
-        asyncio.run(voxtral_server._prepare_speech_generation(request))
-
-        mock_coerce.assert_called_once_with(voxtral_server.engine_client.default_sampling_params_list, True)
-
    def test_prepare_speech_generation_qwen3_voicedesign_non_streaming_mode_false(
        self, qwen3_tts_server, mocker: MockerFixture
    ):
@@ -343,7 +343,7 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo
    )
    mocker.patch(
        "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
-        return_value=("dummy.yaml", [fake_llm_stage, fake_diffusion_stage], None),
+        return_value=("dummy.yaml", [fake_llm_stage, fake_diffusion_stage]),
    )

    engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
@@ -107,7 +107,6 @@ class FakeAsyncOmniEngine:
        self.stage_vllm_configs = [None for _ in range(self.num_stages)]
        self.output_processors = [SimpleNamespace(tokenizer=None) for _ in range(self.num_stages)]
        self.input_processor = None
-        self.endpoint_restrictions = ()

        self.output_q: queue.Queue[Any] = queue.Queue()
        self.submitted: list[dict[str, Any]] = []
@@ -179,7 +179,7 @@ class TestOmniSleepMode:
        try:
            device_id = 0
            used_before = get_device_global_memory_used_gib(device_id)
-            acks = await llm_engine.sleep(stage_ids=[0], level=1)
+            acks = await llm_engine.sleep(stage_ids=[0], level=2)
            await asyncio.sleep(1.5)
            used_after = get_device_global_memory_used_gib(device_id)
            drop_gib = used_before - used_after
@@ -209,7 +209,7 @@ class TestOmniSleepMode:
        """Diffusion Worker stage signal loop"""
        try:
            logger.info("Starting Diffusion Worker Handshake Test")
-            acks = await diffusion_engine.sleep(stage_ids=[0], level=1)
+            acks = await diffusion_engine.sleep(stage_ids=[0], level=2)

            def _get_status(ack):
                return ack.status if hasattr(ack, "status") else ack.get("status")
@@ -264,9 +264,9 @@ class TestOmniSleepMode:
                base_output = output
            assert base_output is not None and len(base_output.images) > 0
            logger.info("Baseline Generation successful.")
-            # Sleep Level 1
+            # Sleep Level 2
            logger.info("Entering Deep Sleep (VRAM Scavenging)...")
-            await diffusion_engine.sleep(stage_ids=[0], level=1)
+            await diffusion_engine.sleep(stage_ids=[0], level=2)
            # Wake-up
            logger.info("Waking up (Reloading Weights)...")
            await diffusion_engine.wake_up(stage_ids=[0])
@@ -353,8 +353,8 @@ class TestOmniSleepMode:
            logger.info(f"Diffusion Initial VRAM: {vram_initial:.2f} GiB")

            # Sleep
-            logger.info("Triggering Level 1 Deep Sleep (Partial Weight Offloading)...")
-            acks = await diffusion_engine.sleep(stage_ids=[0], level=1)
+            logger.info("Triggering Level 2 Deep Sleep (Full Weight Offloading)...")
+            acks = await diffusion_engine.sleep(stage_ids=[0], level=2)

            reported_freed_bytes = sum(getattr(ack, "freed_bytes", 0) for ack in acks)
            reported_freed_gib = reported_freed_bytes / 1024**3
@@ -407,47 +407,6 @@ class TestOmniSleepMode:
            diffusion_engine.shutdown()
            await asyncio.sleep(1)

-    @pytest.mark.asyncio
-    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-    async def test_level2_sleep_wake_raises(self, llm_engine: AsyncOmni):
-        """Regression for #4473 Repro A: wake_up() after sleep(level=2) must raise
-        NotImplementedError instead of silently producing corrupted output."""
-        try:
-            await llm_engine.sleep(stage_ids=[0], level=2)
-            with pytest.raises(NotImplementedError, match="sleep\\(level=2\\)"):
-                await llm_engine.wake_up(stage_ids=[0])
-        finally:
-            llm_engine.shutdown()
-
-    @pytest.mark.asyncio
-    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-    async def test_partial_wake_blocks_generate(self, llm_engine: AsyncOmni):
-        """Regression for #4473 Repro B: generate() must be rejected if kv_cache
-        is still asleep after wake_up(tags=["weights"]), instead of crashing with
-        CUDA illegal memory access."""
-        try:
-            await llm_engine.sleep(stage_ids=[0], level=1)
-            await llm_engine.wake_up(stage_ids=[0], tags=["weights"])
-            with pytest.raises(RuntimeError, match="partially or fully asleep"):
-                async for _ in llm_engine.generate("test", sampling_params=SamplingParams(max_tokens=4)):
-                    pass
-        finally:
-            llm_engine.shutdown()
-
-    @pytest.mark.asyncio
-    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-    async def test_duplicate_wake_is_idempotent(self, llm_engine: AsyncOmni):
-        """Regression for #4473 Repro C: duplicate wake_up(tags=None) must be a
-        safe no-op instead of raising a cumem CUDA invalid argument error."""
-        try:
-            await llm_engine.sleep(stage_ids=[0], level=1)
-            first_acks = await llm_engine.wake_up(stage_ids=[0])
-            assert len(first_acks) > 0, "First wake_up() should return ACKs"
-            second_acks = await llm_engine.wake_up(stage_ids=[0])
-            assert second_acks == [], f"Duplicate wake_up() should return [] but got {second_acks}"
-        finally:
-            llm_engine.shutdown()
-

 # ---------------------------------------------------------------------------
 # H100 (or MI325) e2e: BAGEL / pure diffusion, multi-TP
@@ -485,7 +444,7 @@ async def test_diffusion_model_sleep_tp(tp_size: int):
        async for _ in engine.generate("test", sampling_params_list=[llm_sp, diff_sp]):
            pass

-        acks = await engine.sleep(level=1)
+        acks = await engine.sleep(level=2)
        statuses = [get_ack_info(ack, "status") for ack in acks]
        assert all(s == "SUCCESS" for s in statuses), f"Sleep failed. Statuses: {statuses}"

@@ -541,7 +500,7 @@ async def test_multistage_sleep_h100(tp_size: int):
        async for _ in engine.generate("warmup", sampling_params_list=[SamplingParams(), sp]):
            pass

-        acks = await engine.sleep(stage_ids=[0, 1], level=1)
+        acks = await engine.sleep(stage_ids=[0, 1], level=2)
        assert len(acks) == 2

        await engine.wake_up(stage_ids=[0, 1])
@@ -24,15 +24,15 @@ def test_dreamzero_vla_resolves_to_dreamzero_config(monkeypatch):

 def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch):
    monkeypatch.setattr(
-        "vllm_omni.config.config_factory.StageConfigFactory._try_infer_model_type",
-        classmethod(lambda _cls, model, trust_remote_code=True: "vla"),
+        "vllm_omni.config.config_factory.StageConfigFactory._auto_detect_model_type",
+        classmethod(lambda _cls, _model, trust_remote_code=True: ("vla", None)),
    )
    monkeypatch.setattr(
        "vllm_omni.diffusion.utils.hf_utils._looks_like_dreamzero",
        lambda _model: True,
    )

-    stage_configs, _ = load_stage_configs_from_model("GEAR-Dreams/DreamZero-DROID")
+    stage_configs = load_stage_configs_from_model("GEAR-Dreams/DreamZero-DROID")
    engine_args = stage_configs[0].engine_args

    assert engine_args.model_class_name == "DreamZeroPipeline"
@@ -9,7 +9,6 @@ import pytest
 from pytest_mock import MockerFixture

 from vllm_omni.entrypoints.cli.serve import OmniServeCommand, run_headless
-from vllm_omni.entrypoints.utils import parse_stage_overrides
 from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -32,20 +31,6 @@ def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None:
    assert not explicit["async_chunk"]


-def test_serve_parser_accepts_strategy_config() -> None:
-    """``--strategy-config`` must parse onto the ``strategy_config`` dest and be
-    forwarded as an explicit kwarg so the engine can overlay the strategy."""
-    parser = TrackingArgumentParser()
-    subparsers = parser.add_subparsers(dest="subcommand")
-    cmd = OmniServeCommand()
-    cmd.subparser_init(subparsers)
-
-    argv = ["serve", "fake-model", "--omni", "--strategy-config", "/tmp/strategy.yaml"]
-    args = parser.parse_args(argv)
-    assert args.strategy_config == "/tmp/strategy.yaml"
-    assert args.get_explicit_kwargs_dict()["strategy_config"] == "/tmp/strategy.yaml"
-
-
 def _make_headless_args(**kwargs) -> TrackingNamespace:
    defaults = {
        "model": "fake-model",
@@ -95,97 +80,13 @@ def test_run_headless_rejects_non_multiprocess_worker_backend() -> None:
        run_headless(args)


-# ---------------------------------------------------------------------------
-# --stage-overrides parsing parity (headless vs standard path)
-# ---------------------------------------------------------------------------
-
-
-def test_parse_stage_overrides_valid_json() -> None:
-    """A valid JSON string is parsed into the nested per-stage dict."""
-    parsed = parse_stage_overrides('{"0": {"devices": "0,1"}, "1": {"devices": "2"}}')
-    assert parsed == {"0": {"devices": "0,1"}, "1": {"devices": "2"}}
-
-
-def test_parse_stage_overrides_none_and_empty_return_none() -> None:
-    """No overrides (None / empty string) resolve to ``None``."""
-    assert parse_stage_overrides(None) is None
-    assert parse_stage_overrides("") is None
-
-
-def test_parse_stage_overrides_empty_dict_returns_none() -> None:
-    """An empty dict is falsy and must resolve to ``None``, locking in parity
-    with the original standard-path ``if stage_overrides_json:`` falsy check."""
-    assert parse_stage_overrides({}) is None
-
-
-def test_parse_stage_overrides_passes_through_non_str() -> None:
-    """An already-parsed mapping is returned unchanged (identity)."""
-    overrides = {"0": {"devices": "0"}}
-    assert parse_stage_overrides(overrides) is overrides
-
-
-def test_parse_stage_overrides_invalid_json_raises() -> None:
-    """Invalid JSON raises ValueError whose message matches the standard path
-    verbatim: the ``--stage-overrides is not valid JSON:`` prefix AND the
-    ``Got: <repr>`` suffix echoing the raw input."""
-    bad = "{not valid json}"
-    with pytest.raises(ValueError) as excinfo:
-        parse_stage_overrides(bad)
-    message = str(excinfo.value)
-    assert message.startswith("--stage-overrides is not valid JSON:")
-    assert f"Got: {bad!r}" in message
-
-
-def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture) -> None:
-    """Regression: the headless path must parse ``--stage-overrides`` (a JSON
-    string) and forward the parsed dict to ``load_and_resolve_stage_configs``,
-    mirroring the standard engine path. Previously it was dropped entirely,
-    silently producing a different per-stage device layout."""
-    captured: dict = {}
-
-    def _fake_resolve(*args, **kwargs):
-        captured.update(kwargs)
-        # Return a stage that does NOT match stage_id=0 so run_headless stops
-        # right after the resolver call (we only care about how it was called).
-        return ("/fake/stages.yaml", [SimpleNamespace(stage_id=99)], None)
-
-    mocker.patch(
-        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        side_effect=_fake_resolve,
-    )
-
-    args = _make_headless_args(
-        stage_id=0,
-        strategy_config="/tmp/strategy.yaml",
-        stage_overrides='{"0": {"devices": "0,1"}, "1": {"devices": "2"}}',
-    )
-    with pytest.raises(ValueError, match="No stage config found for stage_id=0"):
-        run_headless(args)
-
-    assert captured["stage_overrides"] == {"0": {"devices": "0,1"}, "1": {"devices": "2"}}
-    assert captured["strategy_config_path"] == "/tmp/strategy.yaml"
-
-
-def test_run_headless_invalid_stage_overrides_raises(mocker: MockerFixture) -> None:
-    """Invalid ``--stage-overrides`` JSON in headless mode fails fast with the
-    shared ValueError instead of being silently ignored."""
-    mocker.patch(
-        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [SimpleNamespace(stage_id=0)], None),
-    )
-
-    args = _make_headless_args(stage_id=0, stage_overrides="{not valid json}")
-    with pytest.raises(ValueError, match="--stage-overrides is not valid JSON"):
-        run_headless(args)
-
-
 def test_run_headless_raises_when_stage_id_not_in_configs(mocker: MockerFixture) -> None:
    """Headless looks up its assigned stage_id in the loaded deploy YAML and
    fails fast when the launcher's --stage-id doesn't match any entry."""
    other_stage = SimpleNamespace(stage_id=99)
    mocker.patch(
        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [other_stage], None),
+        return_value=("/fake/stages.yaml", [other_stage]),
    )

    args = _make_headless_args(stage_id=0)
@@ -233,7 +134,7 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker

    mocker.patch(
        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [stage_cfg], None),
+        return_value=("/fake/stages.yaml", [stage_cfg]),
    )
    mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
@@ -309,7 +210,7 @@ def test_run_headless_llm_launches_one_manager_per_omni_dp_size_local(mocker: Mo

    mocker.patch(
        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [stage_cfg], None),
+        return_value=("/fake/stages.yaml", [stage_cfg]),
    )
    mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
@@ -368,7 +269,7 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture)

    mocker.patch(
        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [stage_cfg], None),
+        return_value=("/fake/stages.yaml", [stage_cfg]),
    )
    mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
@@ -448,7 +349,7 @@ def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixtur

    mocker.patch(
        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [stage_cfg], None),
+        return_value=("/fake/stages.yaml", [stage_cfg]),
    )
    mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
@@ -144,7 +144,7 @@ def _build_serving_chat():
        engine_client=mock_engine,
        models=models,
        response_role="assistant",
-        online_renderer=mock_render,
+        openai_serving_render=mock_render,
        request_logger=None,
        chat_template=None,
        chat_template_content_format="auto",
@@ -318,7 +318,7 @@ class TestLoadAndResolveStageConfigs:
    def test_load_and_resolve_with_kwargs(self):
        """Ensure that dtype survives default stage creation."""
        kwargs = {"dtype": torch.float32}
-        config_path, stage_configs, _ = load_and_resolve_stage_configs(
+        config_path, stage_configs = load_and_resolve_stage_configs(
            model="black-forest-labs/FLUX.2-klein-4B",
            stage_configs_path=None,
            kwargs=kwargs,
@@ -349,10 +349,10 @@ class TestLoadAndResolveStageConfigs:
        ]
        load_stage_configs = mocker.patch(
            "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
-            return_value=(returned_stage_configs, None),
+            return_value=returned_stage_configs,
        )

-        config_path, stage_configs, _ = load_and_resolve_stage_configs(
+        config_path, stage_configs = load_and_resolve_stage_configs(
            model="dummy-model",
            stage_configs_path=str(deploy_path),
            kwargs={},
@@ -363,7 +363,6 @@ class TestLoadAndResolveStageConfigs:
            base_engine_args={},
            deploy_config_path=str(deploy_path),
            stage_overrides=None,
-            strategy_config_path=None,
        )
        assert config_path == str(deploy_path)
        assert len(stage_configs) == 2
@@ -1,81 +0,0 @@
-"""
-Offline inference tests: image-to-video.
-See examples/offline_inference/image_to_video/README.md
-"""
-
-import shutil
-from pathlib import Path
-
-import pytest
-import requests
-
-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"}),
-]
-
-I2V_SCRIPT = EXAMPLES / "offline_inference" / "image_to_video" / "image_to_video.py"
-README_PATH = I2V_SCRIPT.with_name("README.md")
-EXAMPLE_OUTPUT_SUBFOLDER = "example_offline_i2v"
-
-_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
-_IMAGE_NAME = "cherry_blossom.jpg"
-
-_SKIP_SECTIONS = {
-    "Prerequisites",
-    "Advanced Features",
-    "FAQ",
-    # VACE conditional tasks are covered by test_vace_video_generation.py, which
-    # runs the same README snippets with synthetic assets and smoke settings. The
-    # prep block here downloads inputs via wget (not portable, not collected), so
-    # exercising this section from the shared runner would only duplicate and fail.
-    "Wan2.1 VACE Conditional Tasks",
-}
-
-
-def _ensure_test_image(run_dir: Path) -> None:
-    """Download or copy the example image into run_dir so CLI snippets can find it."""
-    dest = run_dir / _IMAGE_NAME
-    if dest.exists():
-        return
-    src = I2V_SCRIPT.parent / _IMAGE_NAME
-    if src.exists():
-        shutil.copy2(src, dest)
-        return
-    response = requests.get(_IMAGE_URL, timeout=60)
-    response.raise_for_status()
-    dest.write_bytes(response.content)
-
-
-def _skip_readme_snippet(language: str, code: str, h2_title: str) -> tuple[bool, str]:
-    if h2_title in _SKIP_SECTIONS:
-        return True, f"README section '{h2_title}' is intentionally excluded for examples tests"
-    if language == "python":
-        return True, "Python API snippets produce video files that ExampleRunner does not auto-collect"
-    if "/path/to/" in code:
-        return True, "Snippet references a placeholder local model path"
-    return False, ""
-
-
-README_SNIPPETS = ReadmeSnippet.extract_readme_snippets(README_PATH, skipif=_skip_readme_snippet)
-
-
-@pytest.mark.parametrize("snippet", README_SNIPPETS, ids=lambda snippet: snippet.test_id)
-def test_image_to_video(snippet: ReadmeSnippet, example_runner: ExampleRunner):
-    should_skip, reason = snippet.skip
-    if should_skip:
-        pytest.skip(reason)
-
-    run_dir = example_runner.output_root / EXAMPLE_OUTPUT_SUBFOLDER / snippet.test_id
-    run_dir.mkdir(parents=True, exist_ok=True)
-    _ensure_test_image(run_dir)
-
-    result = example_runner.run(snippet, output_subfolder=Path(EXAMPLE_OUTPUT_SUBFOLDER))
-    for asset in result.assets:
-        assert_video_valid(asset)
@@ -19,20 +19,19 @@ from PIL import Image
 from tests.helpers.media import (
    convert_audio_bytes_to_text,
    cosine_similarity_text,
-    preprocess_text,
 )

 _GENDER_PIPELINE = None
 _GENDER_PIPELINE_LOCK = threading.Lock()
 # Transcript gates default to whisper ``small`` for speed. ``small`` mishears a
 # short TTS clip ~0.5% of the time (e.g. "Hello"->"fellow", or hallucinating a
-# leading SFX token), which flakes the deterministic similarity gate. Short
-# clips first get a conservative containment fallback for minor ASR repeats/noise.
-# A test can also opt in to ASR escalation by setting
-# ``transcript_escalation_model`` to a whisper model name (e.g. ``"large-v3"``)
-# in its request_config: on a failed fast pass the clip is re-transcribed with
-# that stronger ASR before the test fails, so a weak-ASR mishear is rescued while
-# a genuine model artifact still fails (the strong ASR mismatches too).
+# leading SFX token), which flakes the deterministic similarity gate. A test can
+# opt in to ASR escalation by setting ``transcript_escalation_model`` to a
+# whisper model name (e.g. ``"large-v3"``) in its request_config: on a failed
+# fast pass the clip is re-transcribed with that stronger ASR before the test
+# fails, so a weak-ASR mishear is rescued while a genuine model artifact still
+# fails (the strong ASR mismatches too). Tests that do not opt in keep the
+# strict behaviour.
 _PCM_SPEECH_SAMPLE_RATE_HZ = 24_000
 _MIN_PCM_SPEECH_HNR_DB = 1.0
 _PRESET_VOICE_GENDER_MAP: dict[str, str] = {
@@ -44,23 +43,6 @@ _PRESET_VOICE_GENDER_MAP: dict[str, str] = {
 }


-def _short_transcript_contains_expected(transcript: str, expected: str) -> bool:
-    """Allow minor ASR repeats/noise for very short speech clips."""
-    transcript_clean = preprocess_text(transcript)
-    expected_clean = preprocess_text(expected)
-    if not transcript_clean or not expected_clean:
-        return False
-
-    transcript_words = transcript_clean.split()
-    expected_words = expected_clean.split()
-    if not transcript_words or not expected_words:
-        return False
-
-    short_text = min(len(transcript_clean), len(expected_clean)) <= 15
-    small_word_delta = len(transcript_words) <= len(expected_words) + 2
-    return short_text and small_word_delta and expected_clean in transcript_clean
-
-
 def assert_image_diffusion_response(
    response: "DiffusionResponse",
    request_config: dict[str, Any],
@@ -638,16 +620,13 @@ def _assert_transcript_matches(
    ``transcript`` is the fast whisper-``small`` result. If it clears
    ``threshold`` the check passes immediately.

-    If the cosine check fails, very short clips get a conservative containment
-    fallback that accepts minor ASR repeats/noise only when the expected text is
-    still present and the transcript has few extra words.
-
    When ``escalation_model`` is set (opt-in via the ``transcript_escalation_model``
-    request_config key) and the fast check plus containment fallback fail, the
-    clip is re-transcribed with that stronger ASR and the assertion is decided on
-    its verdict -- so a weak whisper-``small`` mishear on a short clip does not
-    flake the gate, while a genuine model artifact still fails (the strong ASR
-    mismatches too).
+    request_config key) and the fast check fails, the clip is re-transcribed with
+    that stronger ASR and the assertion is decided on its verdict -- so a weak
+    whisper-``small`` mishear on a short clip does not flake the gate, while a
+    genuine model artifact still fails (the strong ASR mismatches too). When
+    ``escalation_model`` is ``None`` the original strict behaviour is preserved,
+    so other tests are unaffected.
    """
    expected = str(expected_text).strip().lower()
    similarity = cosine_similarity_text(transcript.strip().lower(), expected)
@@ -655,10 +634,6 @@ def _assert_transcript_matches(
    if similarity > threshold:
        return

-    if _short_transcript_contains_expected(transcript, expected):
-        print("short speech containment check passed")
-        return
-
    if escalation_model and audio_bytes:
        print(
            f"whisper-small below threshold ({similarity:.2f} <= {threshold}); "
@@ -47,7 +47,6 @@ from tests.helpers.media import (
    _merge_base64_audio_to_segment,
    decode_b64_image,
 )
-from tests.model_tests.diffusion.utils import resolve_tiny_model_path
 from vllm_omni.config.stage_config import resolve_deploy_yaml
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
 from vllm_omni.outputs import OmniRequestOutput
@@ -1065,29 +1064,6 @@ class OpenAIClientHandler:
        )
        return [resp]

-    def send_completions_http_request(
-        self,
-        request_config: dict[str, Any],
-        *,
-        err_code: int | tuple[int, ...] | list[int] | None = None,
-        err_message: str | tuple[str, ...] | list[str] | None = None,
-    ) -> list[HttpResponse]:
-        """POST ``/v1/completions`` with ``json`` or ``raw_body``."""
-        # TODO (Alex): A lot of these helpers should be consolidated as they differ only by endpoint
-        cfg = _merge_http_expectation_kwargs(
-            request_config,
-            err_code=err_code,
-            err_message=err_message,
-        )
-        r = self._post_json_endpoint("/v1/completions", cfg, default_timeout=120.0)
-        resp = self._http_response_from_requests(r)
-        assert_http_error(
-            resp,
-            err_code=cfg.get("err_code"),
-            err_message=cfg.get("err_message"),
-        )
-        return [resp]
-
    def send_omni_sleep_http_request(
        self,
        request_config: dict[str, Any],
@@ -2816,24 +2792,11 @@ def iter_omni_server(

    with omni_fixture_lock:
        params: OmniServerParams = request.param
-        # For now, when a tiny model is substituted, we preserve the original model
-        # name via --served-model-name (so that the server still accepts requests with
-        # the original name). We also do the same for server.model so that tests reading
-        # server.model send the correct name in requests.
-        #
-        # TODO: core models on this path currently do not clean up tiny models, although
-        # tiny model paths are deterministic, so it's not a huge footprint. Still, it would
-        # be ideal to cleanup consistently everywhere.
-        original_model = model_prefix + params.model
-        model = original_model
-        if run_level == "core_model":
-            model = resolve_tiny_model_path(model)
+        model = model_prefix + params.model
        port = params.port
        stage_config_path = stage_config_path_for_run_level(params.stage_config_path, run_level)

        server_args = params.server_args or []
-        if model != original_model:
-            server_args = [*server_args, "--served-model-name", original_model]
        if params.use_omni and params.stage_init_timeout is not None:
            server_args = [*server_args, "--stage-init-timeout", str(params.stage_init_timeout)]
        else:
@@ -2859,8 +2822,6 @@ def iter_omni_server(
                port=port,
                env_dict=params.env_dict,
            ) as server:
-                if model != original_model:
-                    server.model = original_model
                print("OmniServer started successfully")
                yield server
                print("OmniServer stopping...")
@@ -2884,8 +2845,6 @@ def iter_omni_server(
                    use_omni=params.use_omni,
                )
            ) as server:
-                if model != original_model:
-                    server.model = original_model
                print("OmniServer started successfully")
                yield server
                print("OmniServer stopping...")
@@ -2917,8 +2876,6 @@ def iter_omni_runner(
            extra_omni_kwargs = dict(extra) if extra is not None else {}
        stage_config_path = stage_config_path_for_run_level(stage_config_path, run_level)
        model = model_prefix + model
-        if run_level == "core_model":
-            model = resolve_tiny_model_path(model)
        with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, **extra_omni_kwargs) as runner:
            print("OmniRunner started successfully")
            yield runner
@@ -1,25 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-
-from tests.helpers.assertions import _assert_transcript_matches
-
-
-def test_short_transcript_repeat_passes_containment_fallback():
-    _assert_transcript_matches(
-        " How... how are you?",
-        audio_bytes=None,
-        expected_text="how are you",
-        threshold=0.9,
-    )
-
-
-def test_short_transcript_unrelated_text_still_fails():
-    with pytest.raises(AssertionError, match="Transcript doesn't match input"):
-        _assert_transcript_matches(
-            " I don't know, sorry.",
-            audio_bytes=None,
-            expected_text="how are you",
-            threshold=0.9,
-        )
@@ -18,11 +18,8 @@ import types

 import pytest
 import torch
-import torch.nn.functional as F
 from pytest_mock import MockerFixture

-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
 # Direct file import to avoid vllm_omni.__init__ patch dependencies.
 _MODELS = os.path.join(
    os.path.dirname(__file__),
@@ -342,91 +339,6 @@ class TestCodePredictorDtypeAlignment:
        assert not torch.equal(first[:, 1:], different[:, 1:])


-class TestCodePredictorPerRowGenerators:
-    """Seeded requests must stay deterministic inside a multi-row batch (#4883)."""
-
-    def _make_predictor(self, mocker: MockerFixture, loaded_target_classes):
-        _, _, code_predictor_wrapper, _, _ = loaded_target_classes
-        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
-        mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False)
-        cp_config, talker_config = _make_tiny_config(loaded_target_classes)
-        vllm_config = _make_vllm_config(mocker, max_num_seqs=4)
-        predictor = code_predictor_wrapper(
-            vllm_config=vllm_config,
-            config=cp_config,
-            talker_config=talker_config,
-        )
-        predictor._wrapper_config.use_cuda_graphs = False
-        return predictor, talker_config
-
-    def test_multinomial_per_row_matches_single_row_draws(self, mocker: MockerFixture, loaded_target_classes) -> None:
-        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
-        multinomial = common_mod.CodePredictorWrapper._multinomial
-        torch.manual_seed(7)
-        probs = torch.softmax(torch.randn(3, 64), dim=-1)
-
-        def seeded(seed: int) -> torch.Generator:
-            generator = torch.Generator(device=probs.device)
-            generator.manual_seed(seed)
-            return generator
-
-        batched = multinomial(probs, None, [seeded(11), None, seeded(22)])
-        assert batched.shape == (3, 1)
-        # Each seeded row must reproduce a standalone draw from the same seed,
-        # independent of what the other rows in the batch consume.
-        assert torch.equal(batched[0:1], torch.multinomial(probs[0:1], num_samples=1, generator=seeded(11)))
-        assert torch.equal(batched[2:3], torch.multinomial(probs[2:3], num_samples=1, generator=seeded(22)))
-
-    def test_forward_per_row_generators_are_row_independent(self, mocker: MockerFixture, loaded_target_classes) -> None:
-        predictor, talker_config = self._make_predictor(mocker, loaded_target_classes)
-        hidden = talker_config.hidden_size
-        torch.manual_seed(123)
-        row_embed = torch.randn(1, hidden)
-        row_hidden = torch.randn(1, hidden)
-        # Three identical rows: two share a seed, one differs.
-        layer0_code = torch.zeros(3, dtype=torch.long)
-        layer0_embed = row_embed.expand(3, hidden).contiguous()
-        last_talker_hidden = row_hidden.expand(3, hidden).contiguous()
-
-        def seeded(seed: int) -> torch.Generator:
-            generator = torch.Generator(device=layer0_code.device)
-            generator.manual_seed(seed)
-            return generator
-
-        def run(generators):
-            return predictor(
-                layer0_code=layer0_code,
-                layer0_embed=layer0_embed,
-                last_talker_hidden=last_talker_hidden,
-                do_sample=True,
-                temperature=0.9,
-                top_k=50,
-                top_p=1.0,
-                generators=generators,
-            )
-
-        first = run([seeded(1234), seeded(1234), seeded(4321)])
-        second = run([seeded(1234), seeded(1234), seeded(4321)])
-
-        # Same per-row seeds -> the whole batch reproduces across calls.
-        assert torch.equal(first, second)
-        # Identical inputs + identical seeds -> identical rows within a call.
-        assert torch.equal(first[0], first[1])
-        # A different seed on the same inputs must diverge in the residual layers.
-        assert not torch.equal(first[0, 1:], first[2, 1:])
-
-    def test_forward_rejects_mismatched_generators_length(self, mocker: MockerFixture, loaded_target_classes) -> None:
-        predictor, talker_config = self._make_predictor(mocker, loaded_target_classes)
-        hidden = talker_config.hidden_size
-        with pytest.raises(ValueError, match="one entry per row"):
-            predictor(
-                layer0_code=torch.zeros(2, dtype=torch.long),
-                layer0_embed=torch.randn(2, hidden),
-                last_talker_hidden=torch.randn(2, hidden),
-                generators=[torch.Generator()],
-            )
-
-
 class TestCodePredictorModelDtype:
    """Test the inner model forward with different dtypes."""

@@ -545,92 +457,3 @@ class TestCodePredictorWrapperConfig:
            wrapper_config=common_mod.CodePredictorWrapperConfig(use_cuda_graphs=True),
        )
        assert graph_wrapper._prefix_graphs_enabled is True
-
-
-class TestGumbelMaxSampling:
-    """Sanity tests for the shared Gumbel-max sampling helper."""
-
-    @staticmethod
-    def _sample(logits: torch.Tensor, generator=None) -> torch.Tensor:
-        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
-        return common_mod.CodePredictorWrapper._sample_codes_gumbel(logits, generator=generator)
-
-    def test_gumbel_max_picks_unique_high_logit(self, loaded_target_classes) -> None:
-        """If one logit is dominant, Gumbel-max must always pick it."""
-        torch.manual_seed(0)
-        bsz, vocab = 4, 64
-        logits = torch.full((bsz, vocab), -1e4)
-        logits[:, 7] = 0.0  # only finite entry
-        out = self._sample(logits)
-        assert out.shape == (bsz, 1)
-        assert torch.equal(out.squeeze(-1), torch.full((bsz,), 7, dtype=torch.long))
-
-    def test_gumbel_max_handles_single_finite_entry_rows(self, loaded_target_classes) -> None:
-        """Degenerate rows with one surviving finite entry should still return
-        a valid index instead of tripping multinomial-style asserts."""
-        bsz, vocab = 2, 8
-        logits = torch.full((bsz, vocab), float("-inf"))
-        logits[0, 3] = 0.0
-        logits[1, 5] = 0.0
-        out = self._sample(logits)
-        assert out[0].item() == 3
-        assert out[1].item() == 5
-
-    def test_gumbel_max_generator_makes_results_reproducible(self, loaded_target_classes) -> None:
-        """Same ``Generator.manual_seed`` should produce identical samples."""
-        torch.manual_seed(0)
-        logits = torch.randn(8, 32)
-        g1 = torch.Generator()
-        g1.manual_seed(123)
-        g2 = torch.Generator()
-        g2.manual_seed(123)
-        g3 = torch.Generator()
-        g3.manual_seed(456)
-        s1 = self._sample(logits, generator=g1)
-        s2 = self._sample(logits, generator=g2)
-        s3 = self._sample(logits, generator=g3)
-        assert torch.equal(s1, s2)
-        # Different seeds must (almost surely) yield a different sequence.
-        assert not torch.equal(s1, s3)
-
-    def test_gumbel_max_per_row_generators_match_individual_rows(self, loaded_target_classes) -> None:
-        """A per-row generator list should preserve row independence."""
-        torch.manual_seed(11)
-        logits = torch.randn(3, 16)
-        seeds = [101, 202, 303]
-        batched_generators = []
-        expected_rows = []
-        for row, seed in enumerate(seeds):
-            row_generator = torch.Generator()
-            row_generator.manual_seed(seed)
-            batched_generators.append(row_generator)
-
-            solo_generator = torch.Generator()
-            solo_generator.manual_seed(seed)
-            expected_rows.append(self._sample(logits[row : row + 1], generator=solo_generator))
-
-        batched = self._sample(logits, generator=batched_generators)
-        expected = torch.cat(expected_rows, dim=0)
-        assert torch.equal(batched, expected)
-
-    def test_gumbel_max_rejects_generator_length_mismatch(self, loaded_target_classes) -> None:
-        logits = torch.randn(2, 8)
-        generators = [torch.Generator()]
-        with pytest.raises(ValueError, match="Expected 2 per-row generators"):
-            self._sample(logits, generator=generators)
-
-    def test_gumbel_max_distribution_matches_softmax(self, loaded_target_classes) -> None:
-        """Empirical frequency under Gumbel-max should be close to the
-        softmax probability mass.  Loose tolerance because we draw a
-        finite number of samples."""
-        torch.manual_seed(7)
-        vocab = 16
-        logits = torch.randn(vocab) * 2.0
-        target = F.softmax(logits, dim=-1, dtype=torch.float32)
-
-        n_draws = 20000
-        batched_logits = logits.unsqueeze(0).expand(n_draws, vocab).contiguous()
-        samples = self._sample(batched_logits).squeeze(-1)
-        empirical = torch.bincount(samples, minlength=vocab).float() / n_draws
-        # L1 distance under 0.03 with 20k draws is comfortable for vocab=16.
-        assert (empirical - target).abs().sum().item() < 0.05
@@ -12,7 +12,6 @@ REPO_ROOT = Path(__file__).resolve().parents[4]
 RUNNER = REPO_ROOT / "vllm_omni" / "worker" / "gpu_ar_model_runner.py"
 RUNNER_ASSISTED_METADATA = REPO_ROOT / "vllm_omni" / "worker" / "runner_assisted_metadata.py"
 TALKER = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "voxcpm2_talker.py"
-VOXCPM2_RUNTIME_CONFIG = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "runtime_config.py"
 VOXCPM2_PIPELINE = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "pipeline.py"
 VOXCPM2_SCHEDULER = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "scheduler.py"
 OMNI_AR_SCHEDULER = REPO_ROOT / "vllm_omni" / "core" / "sched" / "omni_ar_scheduler.py"
@@ -75,7 +74,6 @@ def test_ar_runner_without_model_hook_stays_on_normal_path():

 def test_voxcpm2_graph_paths_fail_closed_and_preserve_deterministic_noise():
    source = TALKER.read_text()
-    runtime_config_source = VOXCPM2_RUNTIME_CONFIG.read_text()
    compact_source = "".join(source.split())

    assert "_voxcpm2_compile_without_inductor_cudagraphs" in source
@@ -88,8 +86,8 @@ def test_voxcpm2_graph_paths_fail_closed_and_preserve_deterministic_noise():
    assert '"triton.cudagraphs": False' in compile_source
    assert '"triton.cudagraph_trees": False' in compile_source

-    assert "self._enable_unified_decode_graph=self._runtime_config.unified_decode_graph_available(" in compact_source
-    assert "andnotself.deterministic_cfm_noise" in "".join(runtime_config_source.split())
+    assert "self._enable_unified_decode_graph=(use_cuda_graph" in compact_source
+    assert "andnotself._deterministic_cfm_noise" in compact_source
    assert "decode_tail_graph" not in source

    unified_source = source[source.index("def _forward_unified_decode") :]
@@ -185,9 +183,7 @@ def test_voxcpm2_scheduler_policy_stays_model_local():
    assert "def _should_defer_waiting_admission(self) -> bool:" in common_source

    assert "class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler)" in voxcpm2_scheduler_source
-    assert "from .runtime_config import _VoxCPM2RuntimeConfig" in voxcpm2_scheduler_source
-    assert "_VoxCPM2RuntimeConfig.from_vllm_config(self.vllm_config)" in voxcpm2_scheduler_source
-    assert "unified_decode_graph_available(use_cuda_graph=current_omni_platform.is_cuda())" in voxcpm2_scheduler_source
+    assert "voxcpm2_runtime_config" in voxcpm2_scheduler_source
    assert "_should_defer_waiting_for_unified_decode_graph" in voxcpm2_scheduler_source
    assert "def schedule(" not in voxcpm2_scheduler_source
    assert "create_request_queue" not in voxcpm2_scheduler_source
@@ -52,10 +52,6 @@ def _make_bare_talker():
    talker._cfm_buffers = None
    talker._last_audio_output_req_ids = []
    talker._batched_fsq_fusion_max_batch = 32
-    # Added in VoxCPM2TalkerForConditionalGeneration.__init__; set to
-    # default (False) to avoid AttributeError when test exercises forward
-    # path that checks this flag.
-    talker._enable_unified_decode_graph = False
    return talker


@@ -333,27 +329,23 @@ class TestDecodeBatchContract:

        for (expected_res, expected_meta), (actual_res, actual_meta) in zip(expected, actual, strict=True):
            torch.testing.assert_close(actual_res, expected_res)
-            torch.testing.assert_close(actual_meta.new_lm_hidden, expected_meta.new_lm_hidden)
+            torch.testing.assert_close(actual_meta["new_lm_hidden"], expected_meta["new_lm_hidden"])
        assert [state.decode_step_count for state in batch_states] == [1, 1]

    def test_batch_decode_matches_per_request_state_updates(self) -> None:
-        from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
-            _DecodeResidualMeta,
-        )
-
        seq_talker = self._make_decode_talker()
        batch_talker = self._make_decode_talker()

        seq_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
        batch_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
        metas = [
-            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[1.0, 2.0, 3.0, 4.0]])),
-            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[5.0, 6.0, 7.0, 8.0]])),
+            {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])},
+            {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])},
        ]
        batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]])

        for state, meta, res_out in zip(seq_states, metas, batch_out, strict=True):
-            seq_talker._finish_decode(state, meta, res_out.unsqueeze(0))
+            seq_talker._finish_decode(state, meta, res_out.unsqueeze(0), torch.device("cpu"))

        batch_talker._finish_decode_batch(
            [(state, False, meta) for state, meta in zip(batch_states, metas, strict=True)],
@@ -368,16 +360,12 @@ class TestDecodeBatchContract:
            torch.testing.assert_close(actual.last_audio_patch_gpu, expected.last_audio_patch_gpu)

    def test_batch_decode_preserves_per_request_cfm_rng_order(self) -> None:
-        from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
-            _DecodeResidualMeta,
-        )
-
        talker = self._make_decode_talker()
        talker._enable_batched_cfm = False
        states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
        metas = [
-            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[1.0, 2.0, 3.0, 4.0]])),
-            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[5.0, 6.0, 7.0, 8.0]])),
+            {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])},
+            {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])},
        ]
        batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]])

@@ -1,78 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import functools
-
-import pytest
-import torch
-import torch.nn as nn
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-@functools.lru_cache(maxsize=1)
-def _voxtral_tts_model_cls():
-    from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules
-
-    bootstrap_vllm_layer_custom_op_modules()
-    import vllm.model_executor.models.utils  # noqa: F401
-
-    from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts import (
-        VoxtralTTSForConditionalGeneration,
-    )
-
-    return VoxtralTTSForConditionalGeneration
-
-
-class FakeAudioGeneration(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.embed_multimodal_calls = []
-
-    def embed_multimodal(self, **kwargs):
-        self.embed_multimodal_calls.append(kwargs)
-        return [torch.full((1, 4), 7.0)]
-
-
-def _make_voxtral_tts_model():
-    model_cls = _voxtral_tts_model_cls()
-    model = model_cls.__new__(model_cls)
-    nn.Module.__init__(model)
-    model.model_stage = "audio_generation"
-    model.model = FakeAudioGeneration()
-    model._audio_token_id = 42
-    return model
-
-
-def test_tts_preprocess_consumes_nested_codes_audio_feedback():
-    model = _make_voxtral_tts_model()
-    input_ids = torch.tensor([model._audio_token_id])
-    input_embeds = torch.zeros((1, 4))
-    audio_tokens = torch.tensor([[1, 2, 3, 4]])
-
-    _, output_embeds, _ = model.tts_preprocess(
-        input_ids=input_ids,
-        input_embeds=input_embeds,
-        codes={"audio": audio_tokens},
-    )
-
-    assert len(model.model.embed_multimodal_calls) == 1
-    torch.testing.assert_close(model.model.embed_multimodal_calls[0]["audio_tokens"], audio_tokens)
-    torch.testing.assert_close(output_embeds, torch.full((1, 4), 7.0))
-
-
-def test_tts_preprocess_keeps_legacy_top_level_audio_feedback():
-    model = _make_voxtral_tts_model()
-    input_ids = torch.tensor([model._audio_token_id])
-    input_embeds = torch.zeros((1, 4))
-    audio_tokens = torch.tensor([[5, 6, 7, 8]])
-
-    _, output_embeds, _ = model.tts_preprocess(
-        input_ids=input_ids,
-        input_embeds=input_embeds,
-        audio=audio_tokens,
-    )
-
-    assert len(model.model.embed_multimodal_calls) == 1
-    torch.testing.assert_close(model.model.embed_multimodal_calls[0]["audio_tokens"], audio_tokens)
-    torch.testing.assert_close(output_embeds, torch.full((1, 4), 7.0))
@@ -8,6 +8,7 @@ import torch

 from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
    talker2code2wav_async_chunk,
+    text2flow,
    text2flow_full_payload,
    text2flow_token_only,
 )
@@ -48,13 +49,13 @@ def _transfer_manager(
    )


-def test_text2flow_token_only_supports_batched_source_outputs():
+def test_text2flow_supports_batched_source_outputs():
    source_outputs = [
        _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_token_only(source_outputs=source_outputs, prompt=None)
+    outputs = text2flow(source_outputs=source_outputs, prompt=None)

    assert len(outputs) == 2
    assert outputs[0]["prompt_token_ids"] == [1, 2, 3]
@@ -63,6 +64,16 @@ def test_text2flow_token_only_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(
@@ -7,6 +7,7 @@ import torch

 from vllm_omni.model_executor.stage_input_processors.indextts2 import (
    _strip_stop_token,
+    talker2s2mel,
    talker2s2mel_full_payload,
    talker2s2mel_token_only,
 )
@@ -107,6 +108,118 @@ def test_strip_stop_token_all_stop():
    assert lat.shape[1] == 0


+# ── talker2s2mel ──
+
+
+def test_talker2s2mel_basic():
+    mel_codes = torch.tensor([10, 20, 30, STOP_MEL_TOKEN])
+    latent = torch.randn(4, LATENT_DIM)
+    s_ref = torch.randn(1, 128)
+    ref_mel = torch.randn(1, 80, 200)
+    style = torch.randn(1, 256)
+
+    outputs = talker2s2mel(
+        source_outputs=[
+            _make_talker_output(
+                mel_codes=mel_codes,
+                latent=latent,
+                s_ref=s_ref,
+                ref_mel=ref_mel,
+                style=style,
+            )
+        ]
+    )
+
+    assert len(outputs) == 1
+    info = outputs[0]["additional_information"]
+    assert info["mel_codes"].shape == (1, 3)
+    assert info["latent"].shape == (1, 3, LATENT_DIM)
+    assert info["code_lens"].tolist() == [3]
+    assert "S_ref" in info
+    assert "ref_mel" in info
+    assert "style" in info
+    assert outputs[0]["prompt_token_ids"] == [0]
+
+
+def test_talker2s2mel_skips_unfinished():
+    mel_codes = torch.tensor([10, 20])
+    latent = torch.randn(2, LATENT_DIM)
+
+    outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent, finished=False)])
+
+    assert len(outputs) == 0
+
+
+def test_talker2s2mel_skips_empty_mel():
+    outputs = talker2s2mel(
+        source_outputs=[
+            _make_talker_output(
+                mel_codes=torch.tensor([]),
+                latent=torch.randn(0, LATENT_DIM),
+            )
+        ]
+    )
+
+    assert len(outputs) == 0
+
+
+def test_talker2s2mel_missing_optional_meta():
+    mel_codes = torch.tensor([10, 20])
+    latent = torch.randn(2, LATENT_DIM)
+
+    outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent)])
+
+    assert len(outputs) == 1
+    info = outputs[0]["additional_information"]
+    assert "S_ref" not in info
+    assert "ref_mel" not in info
+    assert "style" not in info
+
+
+def test_talker2s2mel_batch():
+    out0 = _make_talker_output(
+        mel_codes=torch.tensor([1, 2, STOP_MEL_TOKEN]),
+        latent=torch.randn(3, LATENT_DIM),
+        s_ref=torch.randn(1, 64),
+    )
+    out1 = _make_talker_output(
+        mel_codes=torch.tensor([5, 6, 7, 8]),
+        latent=torch.randn(4, LATENT_DIM),
+        ref_mel=torch.randn(1, 80, 100),
+    )
+
+    outputs = talker2s2mel(source_outputs=[out0, out1])
+
+    assert len(outputs) == 2
+    assert outputs[0]["additional_information"]["code_lens"].tolist() == [2]
+    assert outputs[1]["additional_information"]["code_lens"].tolist() == [4]
+
+
+def test_talker2s2mel_tensors_on_cpu():
+    device = torch.device("cpu")
+    mel_codes = torch.tensor([10, 20], device=device)
+    latent = torch.randn(2, LATENT_DIM, device=device)
+    s_ref = torch.randn(1, 64, device=device)
+
+    outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent, s_ref=s_ref)])
+
+    info = outputs[0]["additional_information"]
+    assert info["latent"].device.type == "cpu"
+    assert info["mel_codes"].device.type == "cpu"
+    assert info["S_ref"].device.type == "cpu"
+
+
+def test_talker2s2mel_no_multimodal_dict():
+    output = SimpleNamespace(
+        finished=True,
+        outputs=[SimpleNamespace(multimodal_output="not_a_dict")],
+    )
+
+    results = talker2s2mel(source_outputs=[output])
+
+    assert len(results) == 0
+
+
 # ── full-payload connector path ──


@@ -11,40 +11,59 @@ import torch
 from vllm_omni.model_executor.stage_input_processors import mimo_audio as sip
 from vllm_omni.model_executor.stage_input_processors.mimo_audio import (
    MAX_CODE2WAV_TOKENS,
-    llm2code2wav_full_payload,
+    llm2code2wav,
 )

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


-def test_llm2code2wav_full_payload_truncates_when_flat_exceeds_max(caplog):
-    """Flat codec sequences longer than MAX_CODE2WAV_TOKENS must be truncated."""
+def _make_source_outputs(codec_codes: torch.Tensor, request_id: str = "req-0"):
+    output = SimpleNamespace(multimodal_output={"codes": {"audio": codec_codes}})
+    """Build minimal source_outputs with one talker output carrying codec codes."""
+    talker_output = SimpleNamespace(outputs=[output], request_id=request_id)
+    return [talker_output]
+
+
+def test_llm2code2wav_truncates_when_flat_exceeds_max(caplog):
+    """Flat codec sequences longer than MAX_CODE2WAV_TOKENS must be truncated, not passed through."""
+    # prepend_and_flatten_colmajor produces 36 ids per (8, 4) codec frame:
+    # pad adds one row -> (9, 4) per frame, permuted and flattened.
+    # Pick enough frames to comfortably exceed the cap.
    frames = (MAX_CODE2WAV_TOKENS // 36) + 100
    codec_codes = torch.ones(frames, 1, 8, 4, dtype=torch.long)
-    request = SimpleNamespace(request_id="req-long")

+    source_outputs = _make_source_outputs(codec_codes, request_id="req-long")
+
+    # Attach caplog's handler directly to the module logger so the warning is
+    # captured regardless of propagation (vllm's logger configuration can
+    # interact badly with caplog.at_level's default root-handler path).
    target_logger = logging.getLogger("vllm_omni.model_executor.stage_input_processors.mimo_audio")
    target_logger.addHandler(caplog.handler)
    prev_level = target_logger.level
    target_logger.setLevel(logging.WARNING)
    try:
-        payload = llm2code2wav_full_payload(None, {"codes.audio": codec_codes}, request)
+        prompts = llm2code2wav(source_outputs)
    finally:
        target_logger.removeHandler(caplog.handler)
        target_logger.setLevel(prev_level)

-    assert payload is not None
-    assert len(payload["codes"]["audio"]) == MAX_CODE2WAV_TOKENS
+    assert len(prompts) == 1
+    assert len(prompts[0]["prompt_token_ids"]) == MAX_CODE2WAV_TOKENS
+    assert any("truncating" in rec.getMessage() for rec in caplog.records), (
+        f"Expected a 'truncating' warning; captured records: {[r.getMessage() for r in caplog.records]}"
+    )


-def test_llm2code2wav_full_payload_short_sequence_unchanged():
+def test_llm2code2wav_short_sequence_unchanged():
+    """Short codec sequences are returned without truncation."""
    codec_codes = torch.ones(4, 1, 8, 4, dtype=torch.long)
-    request = SimpleNamespace(request_id="req-short")
+    source_outputs = _make_source_outputs(codec_codes, request_id="req-short")

-    payload = llm2code2wav_full_payload(None, {"codes.audio": codec_codes}, request)
+    prompts = llm2code2wav(source_outputs)

-    assert payload is not None
-    assert 0 < len(payload["codes"]["audio"]) <= MAX_CODE2WAV_TOKENS
+    assert len(prompts) == 1
+    # 4 frames + 1 pad row, flattened col-major → well below the cap
+    assert 0 < len(prompts[0]["prompt_token_ids"]) <= MAX_CODE2WAV_TOKENS


 def test_llm2code2wav_truncation_boundary_constant_matches_yaml():
@@ -163,7 +163,7 @@ def test_talker2code2wav_full_payload_filters_by_output_token_ids() -> None:

    assert payload is not None
    assert payload["codes"]["audio"] == [10, 20, 11, 21, 12, 22]
-    assert "code_predictor_codes" not in payload
+    assert payload["code_predictor_codes"] == payload["codes"]["audio"]


 def test_talker2code2wav_full_payload_drops_count_matched_terminal_row() -> None:
@@ -204,7 +204,7 @@ def test_talker2code2wav_full_payload_drops_rows_aligned_to_non_codec_ids() -> N

    assert payload is not None
    assert payload["codes"]["audio"] == [0, 0, 0]
-    assert "code_predictor_codes" not in payload
+    assert payload["code_predictor_codes"] == payload["codes"]["audio"]


 def test_talker2code2wav_full_payload_keeps_all_zero_codec_rows() -> None:
@@ -224,7 +224,7 @@ def test_talker2code2wav_full_payload_keeps_all_zero_codec_rows() -> None:

    assert payload is not None
    assert payload["codes"]["audio"] == [0, 7, 0, 8, 0, 9]
-    assert "code_predictor_codes" not in payload
+    assert payload["code_predictor_codes"] == payload["codes"]["audio"]


 def test_thinker2talker_full_payload_packs_complete_tensors() -> None:
@@ -247,6 +247,7 @@ def test_thinker2talker_full_payload_packs_complete_tensors() -> None:
    assert payload["ids"]["all"] == [151644, 872, 3]
    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

@@ -12,9 +12,9 @@ from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import (
    max_ic_for_chunk_size,
 )
 from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
+    talker2code2wav,
    talker2code2wav_async_chunk,
    talker2code2wav_full_payload,
-    talker2code2wav_token_only,
 )

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -368,15 +368,7 @@ def test_ref_code_context_can_be_buffered_before_first_emit():
    assert rid in tm.request_payload


-def test_non_async_token_only_sizes_placeholder_for_ref_and_audio_frames():
-    """``talker2code2wav_token_only`` only allocates placeholder prompt slots.
-
-    After the connector refactor, actual codec flattening (ref prepend +
-    codebook-major layout) is performed by ``talker2code2wav_full_payload`` on
-    the worker data plane.  The orchestrator hook still derives
-    ``left_context_size`` from stage-0 multimodal_output so Code2Wav can trim
-    reference frames once the connector payload arrives.
-    """
+def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
    ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
    audio_codes = torch.tensor(
        [
@@ -395,38 +387,12 @@ def test_non_async_token_only_sizes_placeholder_for_ref_and_audio_frames():
        engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
    )

-    prompts = talker2code2wav_token_only(stage.engine_outputs)
+    prompts = talker2code2wav(stage.engine_outputs)

    assert len(prompts) == 1
    prompt = prompts[0]
    assert prompt["additional_information"] == {"meta": {"left_context_size": 2}}
-    # 2 ref frames + 2 valid audio frames (zero row filtered), 4 quantizers.
-    assert prompt["prompt_token_ids"] == [0] * (_Q * (2 + 2))
-
-
-def test_full_payload_prepends_ref_code_and_flattens_codebook_major():
-    """Worker producer is authoritative for ref prepend + codec flatten."""
-    ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
-    audio_codes = torch.tensor(
-        [
-            [0, 0, 0, 0],
-            [1, 2, 3, 4],
-            [5, 6, 7, 8],
-        ],
-        dtype=torch.long,
-    )
-    pooling_output = {
-        "codes.audio": audio_codes,
-        "codes.ref": ref_code,
-        "meta.ref_code_len": 2,
-    }
-    request = SimpleNamespace(request_id="r", output_token_ids=list(range(3)))
-
-    payload = talker2code2wav_full_payload(transfer_manager=None, pooling_output=pooling_output, request=request)
-
-    assert payload is not None
-    assert payload["meta"]["left_context_size"] == 2
-    assert payload["codes"]["audio"].tolist() == [
+    assert prompt["prompt_token_ids"] == [
        9,
        8,
        1,
@@ -467,7 +433,7 @@ def test_non_async_processor_filters_out_of_range_codec_values():
        engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
    )

-    prompts = talker2code2wav_token_only(stage.engine_outputs)
+    prompts = talker2code2wav(stage.engine_outputs)

    assert len(prompts) == 1
    prompt = prompts[0]
@@ -8,12 +8,61 @@ import pytest
 import torch

 from vllm_omni.model_executor.stage_input_processors.voxtral_tts import (
+    generator2tokenizer,
    generator2tokenizer_async_chunk,
 )

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


+# ---- Helpers for generator2tokenizer (non-streaming) ----
+
+
+def _make_stage(engine_outputs):
+    return SimpleNamespace(engine_outputs=engine_outputs)
+
+
+def _make_output(audio_tensors):
+    """Build a single generator output with multimodal_output codes.audio."""
+    return SimpleNamespace(
+        outputs=[SimpleNamespace(multimodal_output={"codes": {"audio": audio_tensors}})],
+    )
+
+
+# ---- Tests for generator2tokenizer ----
+
+
+class TestGenerator2Tokenizer:
+    def test_single_output_single_chunk(self):
+        """Single output with one audio tensor produces correct prompt_token_ids."""
+        audio = torch.tensor([10, 20, 30, 40])
+        stage = _make_stage([_make_output([audio])])
+
+        result = generator2tokenizer(stage.engine_outputs)
+
+        assert len(result) == 1
+        assert result[0]["prompt_token_ids"] == [10, 20, 30, 40]
+        assert result[0]["multi_modal_data"] is None
+
+    def test_2d_audio_tensors_flattened(self):
+        """2D audio tensors (e.g., multi-codebook frames) are flattened."""
+        audio = torch.tensor([[1, 2, 3], [4, 5, 6]])  # (2, 3)
+        stage = _make_stage([_make_output([audio])])
+
+        result = generator2tokenizer(stage.engine_outputs)
+
+        assert result[0]["prompt_token_ids"] == [1, 2, 3, 4, 5, 6]
+
+    def test_no_outputs_yet_raises(self):
+        """Stage with engine_outputs=None raises RuntimeError."""
+        stage = _make_stage(engine_outputs=None)
+        with pytest.raises(TypeError):
+            generator2tokenizer(stage.engine_outputs)
+
+
+# ---- Helpers for generator2tokenizer_async_chunk (streaming) ----
+
+
 def _req(external_req_id: str, *, finished: bool):
    return SimpleNamespace(
        external_req_id=external_req_id,
@@ -126,102 +126,6 @@ def test_magi_human_extra_registry_declares_request_and_response_params() -> Non
    assert should_init_extra_args_for_non_diffusion_stages("MagiHumanPipeline") is False


-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_ming_flash_omni_extra_registry_declares_request_and_response_params() -> None:
-    assert get_extra_body_params("MingImagePipeline") == frozenset(
-        {
-            "height",
-            "width",
-            "steps",
-            "cfg",
-            "seed",
-            "byte5_text",
-            "negative_prompt",
-        }
-    )
-    assert get_extra_output_params("MingImagePipeline") == frozenset()
-    assert should_init_extra_args_for_non_diffusion_stages("MingImagePipeline") is True
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_ming_flash_omni_declared_extra_args_route_into_sampling_params() -> None:
-    params = OmniDiffusionSamplingParams()
-    declared = get_extra_body_params("MingImagePipeline")
-    apply_declared_extra_args(
-        params,
-        declared,
-        {
-            "steps": 6,
-            "cfg": 1.5,
-            "byte5_text": ["理解与生成统一"],
-            "negative_prompt": "ugly, blurry",
-            "unknown": "ignored",
-        },
-    )
-    assert params.extra_args == {
-        "steps": 6,
-        "cfg": 1.5,
-        "byte5_text": ["理解与生成统一"],
-        "negative_prompt": "ugly, blurry",
-    }
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_ming_flash_omni_text_to_image_prompt_builder() -> None:
-    assert build_text_to_image_prompt(
-        "MingImagePipeline",
-        prompt="Please draw a cute cat.",
-        negative_prompt="ugly, blurry",
-        height=512,
-        width=768,
-    ) == {
-        "prompt": "Please draw a cute cat.",
-        "modalities": ["image"],
-        "mm_processor_kwargs": {
-            "modalities": ["image"],
-            "target_h": 512,
-            "target_w": 768,
-        },
-        "negative_prompt": "ugly, blurry",
-    }
-    # target_h/w are omitted when height/width are not supplied so the
-    # pipeline's 1024x1024 default still applies.
-    assert build_text_to_image_prompt(
-        "MingImagePipeline",
-        prompt="Draw a poster.",
-        negative_prompt=None,
-    ) == {
-        "prompt": "Draw a poster.",
-        "modalities": ["image"],
-        "mm_processor_kwargs": {"modalities": ["image"]},
-    }
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_ming_flash_omni_image_to_image_prompt_builder() -> None:
-    dummy_image = Image.new("RGB", (64, 64))
-    result = build_image_to_image_prompt(
-        "MingImagePipeline",
-        prompt="Change the background to a sandy beach at sunset.",
-        negative_prompt=None,
-        input_image=dummy_image,
-        height=256,
-        width=256,
-    )
-    assert result["modalities"] == ["img2img"]
-    assert result["multi_modal_data"]["img2img"] is dummy_image
-    assert result["mm_processor_kwargs"] == {
-        "modalities": ["img2img"],
-        "target_h": 256,
-        "target_w": 256,
-    }
-    assert "negative_prompt" not in result
-
-
@pytest.mark.core_model
@pytest.mark.cpu
 def test_cosmos3_text_to_image_prompt_builder_selects_image_modality() -> None:
@@ -1,41 +0,0 @@
-from shutil import rmtree
-
-import pytest
-
-from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
-
-
-@pytest.fixture(scope="session")
-def tiny_model_paths(request, run_level):
-    """Build or download the tiny models for the selected tests.
-
-    At core_model level, builds tiny models via the builder function.
-    At advanced_model / full_model level, uses the real HF model path.
-
-    NOTE: this is session scoped to avoid churn in tiny model creation,
-    but will ensure all the tiny models you need are created for the selected tests
-    before it starts to execute them."""
-    model_paths = {}
-    built_paths = []
-    print("Initializing tiny models...")
-    for item in request.session.items:
-        if not hasattr(item, "callspec"):
-            continue
-        model_name = item.callspec.params.get("model_name")
-        if model_name is None:
-            continue
-        if model_name not in model_paths:
-            settings = DIFFUSION_TEST_SETTINGS[model_name]
-            if run_level == "core_model":
-                print(f"Calling tiny model builder for: {model_name}")
-                path = settings.builder()
-                built_paths.append(path)
-            else:
-                print(f"Run level is {run_level}; {model_name} will use the full model")
-                path = settings.model
-            model_paths[model_name] = path
-
-    yield model_paths
-    for path in built_paths:
-        print(f"Removing tiny model: {path}")
-        rmtree(path, ignore_errors=True)
@@ -1,104 +0,0 @@
-"""
-Analogous to: https://github.com/vllm-project/vllm/blob/v0.23.0/tests/models/multimodal/generation/vlm_utils/case_filtering.py
-
-This follows a similar pattern so that we can use `get_parametrized_options` to expand our test
-configurations out, similar to the way in which vLLM expands supported input type tests. For now,
-this mostly just means parametrizing over the test groups, which may be unique per model, but
-doing so in this way lets us cleanly define test groups that may be heterogeneous + avoid having
-excessive skip marks for features that may not be supported in every pipeline.
-"""
-
-# Engine initialization is currently quite expensive. For efficiency,
-# when we run the tests, we should try to collapse compatible accelerations,
-# because even with tiny models, the overhead of spinning up a new instance
-# stacks up fast.
-import itertools
-
-import pytest
-
-from tests.helpers.mark import hardware_marks
-from tests.model_tests.diffusion.config_types import DiffusionAccs, DiffusionModelTestOpts, get_required_device_count
-from vllm_omni.platforms import current_omni_platform
-
-# These tests are intended to run on L4 GPUs since they test parallelism (i.e., need GPUs),
-# but are intended to be smoke tests that we can properly validate without real model
-# weights in the CI.
-#
-# As such, no DiffusionAccs may require more than 4 GPUs, since we would otherwise skip
-# them in the CI.
-MAX_CI_DEVICES = 4
-
-
-def get_test_group_marks(model_name: str, test_group: list[DiffusionAccs] | None, model_marks: list | None) -> list:
-    """Build the full set of pytest marks for a test group.
-
-    For now, single device groups default to core model, and multi device groups
-    default to advanced_model marks. We also validate and throw if a test group exceeds
-    the max number of GPUs for the CI to prevent accidentally adding configurations that
-    would be skipped in the CI."""
-    marks = list(model_marks) if model_marks is not None else []
-
-    required_devices = get_required_device_count(test_group)
-    if required_devices > MAX_CI_DEVICES:
-        raise ValueError(
-            f"Test group {test_group} for {model_name} requires {required_devices} devices, "
-            f"but the max CI device count is {MAX_CI_DEVICES}. "
-        )
-
-    marks.extend(hardware_marks(res={"cuda": "L4"}, num_cards=required_devices))
-    # hardware_marks only adds skipif for num_cards > 1, so we currently handle the single-device case
-    # directly here. This should probably be handled in a more common way later on.
-    assert current_omni_platform is not None and current_omni_platform.device_count is not None
-    device_count = current_omni_platform.device_count()
-    if current_omni_platform.is_cuda() and device_count < required_devices:
-        marks.append(pytest.mark.skip(reason=f"Need {required_devices} devices, got {device_count}"))
-    if required_devices > 1:
-        marks.append(pytest.mark.full_model)
-    else:
-        marks.append(pytest.mark.core_model)
-    return marks
-
-
-def get_model_parametrization(model_name: str, test_info: DiffusionModelTestOpts, online: bool):
-    """Given a model & its corresponding test options, build the list of pytest params
-    to be run for this model. The base case (no accelerations) is always included. Extra
-    test groups are always appended for offline tests, but are only added to the online
-    tests if online_base_only is overridden to False to avoid redundant testing.
-    """
-    test_groups: list[list[DiffusionAccs] | None] = [None]
-    if not (online and test_info.online_base_only) and test_info.extra_test_groups:
-        test_groups.extend(test_info.extra_test_groups)
-
-    return [
-        pytest.param(
-            model_name,
-            test_group,
-            test_info.supported_tasks,
-            test_info.check_multi_output and test_group is None,
-            test_info.check_determinism and test_group is None,
-            id=f"{model_name}[{'+'.join(test_group)}]" if test_group else model_name,
-            marks=get_test_group_marks(model_name, test_group, test_info.marks),
-        )
-        for test_group in test_groups
-    ]
-
-
-def get_parametrized_options(
-    test_settings: dict[str, DiffusionModelTestOpts],
-    online: bool = False,
-):
-    """Converts all the DiffusionModelTestOpts into an expanded list of parameters
-    based on which accelerations are available.
-
-    When online=True and a model has online_base_only=True, only the base case
-    (no accelerations) is included for that model's parametrization to minimize
-    redundancy between offline / online tests.
-    """
-    # Get a list per model type, where each entry contains a tuple of all of
-    # that model type's cases, then flatten them into the top level so that
-    # we can consume them in one mark.parametrize call.
-    parametrization = [
-        get_model_parametrization(model_name, test_info, online=online)
-        for model_name, test_info in test_settings.items()
-    ]
-    return list(itertools.chain(*parametrization))
@@ -1,187 +0,0 @@
-"""
-Common definitions for controlling what tests run where.
-"""
-
-from collections.abc import Callable
-from enum import StrEnum, auto
-from typing import NamedTuple, TypeAlias
-
-from pytest import MarkDecorator
-
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.entrypoints.omni import Omni
-
-# All builder funcs take no params and return a path
-TinyDiffusionBuilder: TypeAlias = Callable[[], str]
-
-
-class DiffusionAccs(StrEnum):
-    """Supported acceleration types / test settings for Diffusion Models."""
-
-    HSDP = auto()
-    TEA_CACHE = auto()
-    CACHE_DIT = auto()
-    SEQUENCE_PARALLEL = auto()
-    CFG_PARALLEL = auto()
-    TENSOR_PARALLEL = auto()
-    CPU_OFFLOAD = auto()
-    LAYERWISE_OFFLOAD = auto()
-    VAE_PATCH_PARALLEL = auto()
-
-
-class DiffusionTasks(StrEnum):
-    """Supported tasks for Diffusion Models."""
-
-    TEXT_TO_IMAGE = auto()
-    IMAGE_TO_IMAGE = auto()
-    # Text to video, text to audio, etc should be added here as needed
-
-
-class DiffusionModelTestOpts(NamedTuple):
-    """Configuration for one Diffusion model's tests."""
-
-    # HF model name for real-weight tests (advanced_model / full_model level).
-    # For now, whether we use the real weights vs tiny weights in the common tests
-    # depends on the run level.
-    model: str
-
-    # Creates a tiny model for the given architecture. We should always use tiny
-    # model weights for tests that do not require us to check the model quality.
-    builder: TinyDiffusionBuilder
-
-    # Actual tasks which controls the tests actually run
-    supported_tasks: list[DiffusionTasks]
-
-    # Additional acceleration groups to run beyond the base case (no acceleration).
-    # The base case is always run for every model in the test settings. None means
-    # we only run the base case.
-    extra_test_groups: list[list[DiffusionAccs]] | None = None
-
-    # Pytest Marks for this model. This may be useful for selecting which models
-    # we want to run where, similar to the way vLLM's multimodal tests mark some
-    # as core models to always run in the CI.
-    # Example: https://github.com/vllm-project/vllm/blob/v0.23.0/tests/models/multimodal/generation/test_common.py#L131
-    marks: list[MarkDecorator] | None = None
-
-    # When True (default), online tests only run the base case (no accelerations).
-    # When False, online tests run all test_groups, same as offline. This should be
-    # True unless there is a good reason for it not to be, because the execution of
-    # the acceleration should be the same on both codepaths, and CLI parsing etc should
-    # be tested by tests adding the acceleration, and not per model.
-    online_base_only: bool = True
-
-    # Additional checks to run for the base case.
-    check_multi_output: bool = True  # Runs multiple generations in one request
-    check_determinism: bool = True  # Runs 2 generations with the same seed and check determinism
-
-
-### Mappings & utils for building offline Omni() instances given a list of enabled accelerations
-ACC_OMNI_KWARGS = {
-    DiffusionAccs.VAE_PATCH_PARALLEL: {"vae_use_tiling": True},
-    DiffusionAccs.CPU_OFFLOAD: {"enable_cpu_offload": True},
-    DiffusionAccs.LAYERWISE_OFFLOAD: {"enable_layerwise_offload": True},
-    DiffusionAccs.CACHE_DIT: {"cache_backend": "cache_dit"},
-    DiffusionAccs.TEA_CACHE: {"cache_backend": "tea_cache"},
-}
-
-ACC_PARALLEL_KWARGS = {
-    DiffusionAccs.HSDP: {"use_hsdp": True, "hsdp_shard_size": 2},
-    DiffusionAccs.TENSOR_PARALLEL: {"tensor_parallel_size": 2},
-    DiffusionAccs.CFG_PARALLEL: {"cfg_parallel_size": 2},
-    DiffusionAccs.VAE_PATCH_PARALLEL: {"vae_patch_parallel_size": 2},
-    # For SP, we don't run ring here to conserve devices, since it is easy
-    # to blow up the number of needed devices fast. Compatibility for ring
-    # and ulysses together should be tested generically outside of these tests.
-    DiffusionAccs.SEQUENCE_PARALLEL: {"ulysses_degree": 2},
-}
-
-ACC_DEVICE_COUNT_KEYS = {
-    DiffusionAccs.TENSOR_PARALLEL: "tensor_parallel_size",
-    DiffusionAccs.CFG_PARALLEL: "cfg_parallel_size",
-    DiffusionAccs.SEQUENCE_PARALLEL: "ulysses_degree",
-    DiffusionAccs.VAE_PATCH_PARALLEL: "vae_patch_parallel_size",
-    DiffusionAccs.HSDP: "hsdp_shard_size",
-}
-
-### CLI args for launching an OmniServer subprocess per acceleration.
-# Follows the same pattern as e2e online serving tests, which hand-code
-# their server_args lists in OmniServerParams.
-ACC_SERVER_ARGS: dict[DiffusionAccs, list[str]] = {
-    DiffusionAccs.HSDP: ["--use-hsdp", "--hsdp-shard-size", "2"],
-    DiffusionAccs.TEA_CACHE: ["--cache-backend", "tea_cache"],
-    DiffusionAccs.CACHE_DIT: ["--cache-backend", "cache_dit"],
-    DiffusionAccs.SEQUENCE_PARALLEL: ["--usp", "2"],
-    DiffusionAccs.CFG_PARALLEL: ["--cfg-parallel-size", "2"],
-    DiffusionAccs.TENSOR_PARALLEL: ["--tensor-parallel-size", "2"],
-    DiffusionAccs.CPU_OFFLOAD: ["--enable-cpu-offload"],
-    DiffusionAccs.LAYERWISE_OFFLOAD: ["--enable-layerwise-offload"],
-    DiffusionAccs.VAE_PATCH_PARALLEL: ["--vae-use-tiling", "--vae-patch-parallel-size", "2"],
-}
-
-## TODO ^ These should be in the same object, it's getting messy.
-
-
-def get_required_device_count(accelerations: list[DiffusionAccs] | None) -> int:
-    """Compute the minimum number of devices needed for a set of accelerations.
-    The total is the product of all parallel dimensions (defaulting to 1).
-
-    If not enough devices are available for a test group's accelerations,
-    that test will be skipped."""
-    count = 1
-    if accelerations is None:
-        return count
-
-    for acc in accelerations:
-        key = ACC_DEVICE_COUNT_KEYS.get(acc)
-        if key is not None:
-            count *= ACC_PARALLEL_KWARGS[acc][key]
-    return count
-
-
-def build_parallel_config_from_diff_accelerations(accelerations: list[DiffusionAccs]) -> DiffusionParallelConfig | None:
-    """Given a list of accelerations pertaining to the current test group,
-    build the parallel config needed for the Omni() object (if any)."""
-    config_kwargs = {}
-    for acc in accelerations:
-        update_dict = ACC_PARALLEL_KWARGS.get(acc, {})
-        config_kwargs.update(update_dict)
-    if config_kwargs:
-        return DiffusionParallelConfig(**config_kwargs)
-    return None
-
-
-### Offline Omni() object builder
-def build_omni_from_diff_accelerations(accelerations: list[DiffusionAccs] | None, **kwargs) -> Omni:
-    """Given one or more acceleration types, build the corresponding Omni() object."""
-    # Coerce to a list and build the parallel config, since that depends on the accelerations
-    if accelerations is None:
-        accelerations = []
-    parallel_config = build_parallel_config_from_diff_accelerations(accelerations)
-
-    # Then add anything else that's a top-level kwarg
-    acc_kwargs = {}
-    if parallel_config is not None:
-        acc_kwargs["parallel_config"] = parallel_config
-    for acc in accelerations:
-        update_dict = ACC_OMNI_KWARGS.get(acc, {})
-        acc_kwargs.update(update_dict)
-
-    # Keys passed through should mostly be things like enforce_eager;
-    # if there's overlap, it's probably due to a misconfiguration
-    shared_keys = acc_kwargs.keys() & kwargs.keys()
-    if shared_keys:
-        raise ValueError(f"Explicit Omni kwargs and inferred Omni kwargs for accelerations overlap: {shared_keys}")
-    omni_kwargs = {**acc_kwargs, **kwargs}
-    return Omni(**omni_kwargs)
-
-
-### Online server flag builder
-def build_server_args_from_diff_accelerations(accelerations: list[DiffusionAccs] | None) -> list[str]:
-    """Given one or more acceleration types, build the corresponding CLI args
-    for launching an OmniServer subprocess."""
-    if accelerations is None:
-        return []
-    args = []
-    for acc in accelerations:
-        args.extend(ACC_SERVER_ARGS.get(acc, []))
-    return args
@@ -1,63 +0,0 @@
-import os
-import tempfile
-
-import torch
-from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2
-from diffusers.models.transformers.transformer_flux2 import Flux2Transformer2DModel
-from diffusers.pipelines.flux2.pipeline_flux2_klein import Flux2KleinPipeline
-from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
-from transformers import AutoTokenizer, Qwen3Config, Qwen3ForCausalLM
-
-TINY_MODEL_DIR = os.path.join(tempfile.gettempdir(), "vllm-omni-tiny-models")
-
-
-def _get_tiny_model_path(name: str) -> str:
-    path = os.path.join(TINY_MODEL_DIR, name)
-    os.makedirs(path, exist_ok=True)
-    return path
-
-
-def tiny_flux2_klein_builder() -> str:
-    """Build a tiny Flux2Klein model."""
-    model_dir = _get_tiny_model_path("Flux2KleinPipeline")
-
-    pipe = Flux2KleinPipeline(
-        scheduler=FlowMatchEulerDiscreteScheduler(),
-        vae=AutoencoderKLFlux2(
-            down_block_types=("DownEncoderBlock2D",),
-            up_block_types=("UpDecoderBlock2D",),
-            block_out_channels=(32,),
-            layers_per_block=1,
-            latent_channels=16,
-            norm_num_groups=16,
-        ),
-        # NOTE: For now we need 28 layers because of hardcoded stuff in the model :(
-        text_encoder=Qwen3ForCausalLM(
-            Qwen3Config(
-                hidden_size=32,
-                intermediate_size=64,
-                num_hidden_layers=28,
-                num_attention_heads=2,
-                num_key_value_heads=2,
-                head_dim=16,
-                vocab_size=151936,
-                max_position_embeddings=512,
-            )
-        ),
-        tokenizer=AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B"),
-        # NOTE: For now we need at least 2 layers for the transformer
-        # due to hardcoded hacks in CacheDiT for Flux2Klein specifically.
-        transformer=Flux2Transformer2DModel(
-            in_channels=64,
-            num_layers=2,
-            num_single_layers=2,
-            attention_head_dim=16,
-            num_attention_heads=2,
-            joint_attention_dim=96,
-            timestep_guidance_channels=32,
-            axes_dims_rope=(4, 4, 4, 4),
-        ),
-    )
-    # Need dtypes to be consistent; for now we just put it on bfloat16
-    pipe.to(torch.bfloat16).save_pretrained(model_dir)
-    return model_dir
@@ -1,42 +0,0 @@
-from tests.model_tests.diffusion import diff_model_builders
-from tests.model_tests.diffusion.config_types import (
-    DiffusionAccs,
-    DiffusionModelTestOpts,
-    DiffusionTasks,
-)
-
-# This object defines the (tiny) model configurations for common tests.
-#
-# TO ADD A NEW MODEL:
-# Map the pipeline name (i.e., the name of the pipeline in _DIFFUSION_MODELS)
-# to a DiffusionModelTestOpts and define the supported tasks. A base case with no
-# accelerations is always run. Use extra_test_groups to define additional
-# acceleration configurations to run a basic smoke test for (i.e., simple inference).
-#
-# The model will only be loaded once per test group, and will execute each of its
-# supported tasks as a pytest subtest. The base case also runs additional checks
-# that are independent of accelerations, e.g., for determinism and compatibility with
-# multi-output, unless explicitly disabled in the test settings.
-#
-# TO RUN A SUBSET OF THE TESTS:
-# These tests should be fast, but if you'd like to run a subset of the tests, you can do so
-# with `-k`; the IDs of the tests are the name of the pipeline joined by +.
-#
-# Example:
-# $ pytest test_common_offline.py -k test_pipeline_on_supported_tasks[Flux2KleinPipeline]
-#   ^ Runs only the base (no accelerations) case for Flux2KleinPipeline
-#
-# $ pytest test_common_offline.py -k test_pipeline_on_supported_tasks[Flux2KleinPipeline
-#   ^ Runs all test groups for Flux2KleinPipeline only
-DIFFUSION_TEST_SETTINGS = {
-    "Flux2KleinPipeline": DiffusionModelTestOpts(
-        model="black-forest-labs/FLUX.2-klein-4B",
-        builder=diff_model_builders.tiny_flux2_klein_builder,
-        supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE, DiffusionTasks.IMAGE_TO_IMAGE],
-        extra_test_groups=[
-            [DiffusionAccs.HSDP, DiffusionAccs.TEA_CACHE],
-            [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
-            [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
-        ],
-    )
-}
@@ -1,172 +0,0 @@
-"""
-Definitions running and validating individual tasks, e.g., text to image,
-image to image, and so on. These are called by the core test runner.
-"""
-
-import base64
-import io
-from dataclasses import replace
-
-import numpy as np
-from PIL import Image
-
-from tests.helpers.runtime import DiffusionResponse, OmniServer, OpenAIClientHandler, dummy_messages_from_mix_data
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.outputs import OmniRequestOutput
-
-PROMPT = "Dummy prompt"
-IMAGE_DIMS = (512, 512)
-HEIGHT, WIDTH = IMAGE_DIMS
-INPUT_IMAGE = Image.new("RGB", IMAGE_DIMS)
-
-# Offline sampling params
-IMAGE_GEN_SAMPLING_PARAMS = OmniDiffusionSamplingParams(
-    num_inference_steps=4,
-    height=HEIGHT,
-    width=WIDTH,
-    seed=42,
-)
-
-# Online extra_body for diffusion requests
-IMAGE_GEN_EXTRA_BODY = {
-    "height": HEIGHT,
-    "width": WIDTH,
-    "num_inference_steps": 4,
-    "seed": 42,
-}
-
-
-### Shared validation
-def _validate_images(images: list[Image.Image], expected_n: int = 1):
-    """Given a set of images, ensure we got the expected count of images,
-    and that all provided images match the expected dimensions."""
-    assert len(images) == expected_n
-    for img in images:
-        assert isinstance(img, Image.Image)
-        assert img.size == IMAGE_DIMS
-    return images
-
-
-def _validate_image_gen_determinism(images_a: list[Image.Image], images_b: list[Image.Image]):
-    """Ensure that two image sets are valid and that the results match."""
-    _validate_images(images_a)
-    _validate_images(images_b)
-    assert np.array_equal(np.array(images_a[0]), np.array(images_b[0]))
-
-
-### Output extractor utils for offline / online paths respectively
-def _get_offline_images(outputs: list[OmniRequestOutput]) -> list[Image.Image]:
-    """Extract the images from an Omni .generate() call."""
-    assert len(outputs) == 1
-    return outputs[0].images
-
-
-def _get_online_images(responses: list[DiffusionResponse]) -> list[Image.Image]:
-    """Extract the images from a server response."""
-    assert len(responses) == 1
-    images = responses[0].images
-    assert images is not None
-    return images
-
-
-### Offline helpers
-def _run_offline_t2i(omni: Omni, params: OmniDiffusionSamplingParams = IMAGE_GEN_SAMPLING_PARAMS):
-    return omni.generate({"prompt": PROMPT}, params)
-
-
-def _run_offline_i2i(omni: Omni):
-    return omni.generate(
-        {"prompt": PROMPT, "multi_modal_data": {"image": INPUT_IMAGE}},
-        IMAGE_GEN_SAMPLING_PARAMS,
-    )
-
-
-### Online helpers
-def _build_online_image_data_url() -> str:
-    """Get a valid base 64 encoded data URL corresponding to an image."""
-    buf = io.BytesIO()
-    INPUT_IMAGE.save(buf, format="PNG")
-    b64 = base64.b64encode(buf.getvalue()).decode()
-    return f"data:image/png;base64,{b64}"
-
-
-def _run_online_t2i(
-    server: OmniServer, client: OpenAIClientHandler, extra_body: dict | None = None
-) -> list[DiffusionResponse]:
-    """Run a text to image request through the server."""
-    messages = dummy_messages_from_mix_data(content_text=PROMPT)
-    request_config = {
-        "model": server.model,
-        "messages": messages,
-        "extra_body": extra_body or IMAGE_GEN_EXTRA_BODY,
-    }
-    return client.send_diffusion_request(request_config)
-
-
-def _run_online_i2i(server: OmniServer, client: OpenAIClientHandler) -> list[DiffusionResponse]:
-    """Run an image to image request through the server."""
-    image_data_url = _build_online_image_data_url()
-    messages = dummy_messages_from_mix_data(
-        content_text=PROMPT,
-        image_data_url=image_data_url,
-    )
-    request_config = {
-        "model": server.model,
-        "messages": messages,
-        "extra_body": IMAGE_GEN_EXTRA_BODY,
-    }
-    return client.send_diffusion_request(request_config)
-
-
-### Offline task runners
-def run_and_validate_text_to_image_request(omni: Omni):
-    """Run and validate a text to image request."""
-    _validate_images(_get_offline_images(_run_offline_t2i(omni)))
-
-
-def run_and_validate_image_to_image_request(omni: Omni):
-    """Run and validate an image to image request."""
-    _validate_images(_get_offline_images(_run_offline_i2i(omni)))
-
-
-def run_and_validate_text_to_image_determinism(omni: Omni):
-    """Checks for determinism; for now we just keep this for TTI."""
-    _validate_image_gen_determinism(
-        _get_offline_images(_run_offline_t2i(omni)),
-        _get_offline_images(_run_offline_t2i(omni)),
-    )
-
-
-def run_and_validate_text_to_image_multi_output(omni: Omni):
-    """Checks for multi-output; for now we just keep this for TTI."""
-    params = replace(IMAGE_GEN_SAMPLING_PARAMS, num_outputs_per_prompt=2)
-    _validate_images(_get_offline_images(_run_offline_t2i(omni, params)), expected_n=2)
-
-
-### Online task runners
-def run_and_validate_online_text_to_image_request(server: OmniServer, client: OpenAIClientHandler):
-    """Run and validate a text to image request through the server."""
-    _validate_images(_get_online_images(_run_online_t2i(server, client)))
-
-
-def run_and_validate_online_image_to_image_request(server: OmniServer, client: OpenAIClientHandler):
-    """Run and validate an image to image request through the server."""
-    _validate_images(_get_online_images(_run_online_i2i(server, client)))
-
-
-def run_and_validate_online_text_to_image_determinism(server: OmniServer, client: OpenAIClientHandler):
-    """Checks for determinism through the server; for now we just keep this for TTI."""
-    _validate_image_gen_determinism(
-        _get_online_images(_run_online_t2i(server, client)),
-        _get_online_images(_run_online_t2i(server, client)),
-    )
-
-
-def run_and_validate_online_text_to_image_multi_output(server: OmniServer, client: OpenAIClientHandler):
-    """Checks for multi-output through the server; for now we just keep this for TTI."""
-    extra_body = {**IMAGE_GEN_EXTRA_BODY, "num_outputs_per_prompt": 2}
-    _validate_images(
-        _get_online_images(_run_online_t2i(server, client, extra_body=extra_body)),
-        expected_n=2,
-    )
@@ -1,98 +0,0 @@
-"""
-Tests for alignment between _DIFFUSION_MODELS and DIFFUSION_TEST_SETTINGS; if
-tests in this file are failing, you are probably adding a new model, and need
-to account for it in DIFFUSION_TEST_SETTINGS
-"""
-
-import pytest
-
-from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
-from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
-
-pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
-
-# These diffusion models currently do not have any added (tiny) model tests.
-# In general, pipelines should only be in the list below if they have not been
-# migrated yet or there is a compelling reason that they should not be (which for
-# now may just be CI pressure etc due to the volume of models and the cost of
-# initializing the server).
-#
-# The tests here validate that all keys in the diffusion registry are accounted for;
-# if you are adding a new model and see the tests below fail, please follow the pattern
-# for adding a new tiny model builder and corresponding entry in DIFFUSION_TEST_SETTINGS.
-EXCLUDED_MODELS = [
-    "QwenImagePipeline",
-    "QwenImageEditPipeline",
-    "QwenImageEditPlusPipeline",
-    "QwenImageLayeredPipeline",
-    "GlmImagePipeline",
-    "ZImagePipeline",
-    "OvisImagePipeline",
-    "WanPipeline",
-    "WanVACEPipeline",
-    "LTX2Pipeline",
-    "LTX2ImageToVideoPipeline",
-    "LTX2TwoStagesPipeline",
-    "LTX2ImageToVideoTwoStagesPipeline",
-    "LTX2T2VDMD2Pipeline",
-    "LTX2I2VDMD2Pipeline",
-    "LTX23Pipeline",
-    "LTX23ImageToVideoPipeline",
-    "StableAudioPipeline",
-    "WanImageToVideoPipeline",
-    "WanS2VPipeline",
-    "WanT2VDMD2Pipeline",
-    "WanI2VDMD2Pipeline",
-    "LongCatImagePipeline",
-    "BagelPipeline",
-    "LancePipeline",
-    "MingImagePipeline",
-    "InternVLAA1Pipeline",
-    "LongCatImageEditPipeline",
-    "StableDiffusion3Pipeline",
-    "FluxKontextPipeline",
-    "HunyuanImage3ForCausalMM",
-    "ErnieImagePipeline",
-    "NextStep11Pipeline",
-    "FluxPipeline",
-    "FluxDMD2Pipeline",
-    "QwenImageDMD2Pipeline",
-    "OmniGen2Pipeline",
-    "HeliosPipeline",
-    "HeliosPyramidPipeline",
-    "Flux2Pipeline",
-    "DreamIDOmniPipeline",
-    "SenseNovaU1Pipeline",
-    "AudioXPipeline",
-    "HunyuanVideo15Pipeline",
-    "HunyuanVideo15ImageToVideoPipeline",
-    "MagiHumanPipeline",
-    "OmniVoicePipeline",
-    "OmniVoice",
-    "Cosmos3OmniDiffusersPipeline",
-    "DiffusersAdapterPipeline",
-    "HiDreamImagePipeline",
-    "DreamZeroPipeline",
-    "StableDiffusionXLPipeline",
-    "Gr00tN1d7Pipeline",
-    "SoulXSingerPipeline",
-    "SoulXSingerSVCPipeline",
-]
-
-
-def test_all_non_excluded_pipelines_have_tests():
-    """Ensure that every pipeline in the diffusion registry either has a test
-    configuration in DIFFUSION_TEST_SETTINGS, or is explicitly excluded for a
-    known reason."""
-    non_excluded_diff_pipes = _DIFFUSION_MODELS.keys() - set(EXCLUDED_MODELS)
-    missing_tests = non_excluded_diff_pipes - DIFFUSION_TEST_SETTINGS.keys()
-    assert len(missing_tests) == 0, f"Pipelines missing test settings: {missing_tests}"
-
-
-def test_no_excluded_models_have_test_settings():
-    """Ensure that no models in the exclude list have a test configuration in
-    DIFFUSION_TEST_SETTINGS (i.e., if they have test settings, we should not
-    exclude them)."""
-    non_excluded_diff_pipes = _DIFFUSION_MODELS.keys() - set(EXCLUDED_MODELS)
-    missing_pipes = DIFFUSION_TEST_SETTINGS.keys() - set(non_excluded_diff_pipes)
-    assert len(missing_pipes) == 0
@@ -1,69 +0,0 @@
-import pytest
-
-from tests.model_tests.diffusion.case_filtering import get_parametrized_options
-from tests.model_tests.diffusion.config_types import (
-    DiffusionAccs,
-    DiffusionTasks,
-    build_omni_from_diff_accelerations,
-)
-from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
-from tests.model_tests.diffusion.task_runners import (
-    run_and_validate_image_to_image_request,
-    run_and_validate_text_to_image_determinism,
-    run_and_validate_text_to_image_multi_output,
-    run_and_validate_text_to_image_request,
-)
-
-# NOTE : Hardware marks are added dynamically based on test requirements
-pytestmark = [pytest.mark.diffusion]
-
-
-@pytest.mark.parametrize(
-    "model_name,accelerations,supported_tasks,check_multioutput,check_determinism",
-    get_parametrized_options(DIFFUSION_TEST_SETTINGS),
-)
-def test_pipeline_on_supported_tasks(
-    model_name,
-    accelerations: list[DiffusionAccs] | None,
-    supported_tasks: list[DiffusionTasks],
-    check_multioutput: bool,
-    check_determinism: bool,
-    tiny_model_paths: dict[str, str],
-    subtests,
-):
-    """Run a smoke test on all of the pipelines supported tasks using a set of enabled accelerations."""
-    assert len(supported_tasks) > 0
-    # We initialize the Omni object before running the tasks, then run each task as a pytest subtask.
-    # This lets us init the model once, but display separate failures in pytest, and avoid halting the
-    # checks on other tasks if one fails.
-    #
-    # This allows us to have some degree of test isolation without the cost of redundant initialization,
-    # since starting the server can take 10+ seconds, even for tiny models.
-    #
-    # NOTE: Be sure to install pytest-subtests if you're running on pytest < 9
-    omni = build_omni_from_diff_accelerations(
-        accelerations=accelerations,
-        model=tiny_model_paths[model_name],
-        enforce_eager=True,
-    )
-    try:
-        for task_type in supported_tasks:
-            with subtests.test(msg=task_type):
-                if task_type == DiffusionTasks.TEXT_TO_IMAGE:
-                    run_and_validate_text_to_image_request(omni)
-                elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
-                    run_and_validate_image_to_image_request(omni)
-                else:
-                    raise ValueError(f"Task type {task_type} is not yet supported")
-
-        # NOTE: For now, we only check determinism + multi output for the base case,
-        # since checking it on every extra acceleration configuration is redundant
-        # (see case_filtering).
-        if check_determinism:
-            with subtests.test(msg="determinism"):
-                run_and_validate_text_to_image_determinism(omni)
-        if check_multioutput:
-            with subtests.test(msg="multi_output"):
-                run_and_validate_text_to_image_multi_output(omni)
-    finally:
-        omni.close()
@@ -1,75 +0,0 @@
-"""
-Analogous to test_common_offline, but for server tests. Validates the full
-online serving stack (CLI arg parsing, subprocess, API routing, response
-encoding) using tiny models.
-"""
-
-import pytest
-
-from tests.helpers.runtime import OmniServer, OpenAIClientHandler
-from tests.model_tests.diffusion.case_filtering import get_parametrized_options
-from tests.model_tests.diffusion.config_types import (
-    DiffusionAccs,
-    DiffusionTasks,
-    build_server_args_from_diff_accelerations,
-)
-from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
-from tests.model_tests.diffusion.task_runners import (
-    run_and_validate_online_image_to_image_request,
-    run_and_validate_online_text_to_image_determinism,
-    run_and_validate_online_text_to_image_multi_output,
-    run_and_validate_online_text_to_image_request,
-)
-
-# NOTE : Hardware marks are added dynamically based on test requirements
-pytestmark = [pytest.mark.diffusion]
-
-
-@pytest.mark.parametrize(
-    "model_name,accelerations,supported_tasks,check_multioutput,check_determinism",
-    get_parametrized_options(DIFFUSION_TEST_SETTINGS, online=True),
-)
-def test_online_on_supported_tasks(
-    model_name: str,
-    accelerations: list[DiffusionAccs] | None,
-    supported_tasks: list[DiffusionTasks],
-    check_multioutput: bool,
-    check_determinism: bool,
-    tiny_model_paths: dict[str, str],
-    run_level: str,
-    subtests,
-):
-    """Smoke test: start a tiny model server and run each supported task via the API."""
-    model_path = tiny_model_paths[model_name]
-    server_args = build_server_args_from_diff_accelerations(accelerations)
-    server_args.append("--enforce-eager")
-
-    with OmniServer(model_path, server_args) as server:
-        # TODO: We may want to revisit run_level validation here,
-        # because checks for things like image size etc should not
-        # depend on whether or not the weights are real or random
-        client = OpenAIClientHandler(
-            host=server.host,
-            port=server.port,
-            api_key="EMPTY",
-            run_level=run_level,
-            log_stats=server.log_stats,
-        )
-        for task_type in supported_tasks:
-            with subtests.test(msg=task_type):
-                if task_type == DiffusionTasks.TEXT_TO_IMAGE:
-                    run_and_validate_online_text_to_image_request(server, client)
-                elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
-                    run_and_validate_online_image_to_image_request(server, client)
-                else:
-                    raise ValueError(f"Task type {task_type} is not yet supported")
-
-        # NOTE: For now, we only check determinism + multi output for the base case,
-        # since checking it on every extra acceleration configuration is redundant
-        # (see case_filtering).
-        if check_determinism:
-            with subtests.test(msg="determinism"):
-                run_and_validate_online_text_to_image_determinism(server, client)
-        if check_multioutput:
-            with subtests.test(msg="multi_output"):
-                run_and_validate_online_text_to_image_multi_output(server, client)
@@ -1,33 +0,0 @@
-"""
-Utilities for resolving real models to their tiny model equivalents.
-"""
-
-import logging
-
-from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
-from vllm_omni.diffusion.data import resolve_model_class_name
-
-logger = logging.getLogger(__name__)
-
-
-def resolve_tiny_model_path(model: str) -> str:
-    """Given a real model name/path, resolve it to a tiny model path.
-
-    Raises ValueError if the pipeline class cannot be determined (invalid
-    model). Returns the original model path if no tiny builder exists yet."""
-    pipeline_class = resolve_model_class_name(model)
-    if pipeline_class is None:
-        raise ValueError(
-            f"Cannot resolve pipeline class for model: {model}. The model path may be invalid or its config unreadable."
-        )
-
-    test_opts = DIFFUSION_TEST_SETTINGS.get(pipeline_class)
-    if test_opts is None:
-        logger.warning(
-            "No tiny model builder for pipeline %s (model: %s). Using original model.",
-            pipeline_class,
-            model,
-        )
-        return model
-
-    return test_opts.builder()