ci: raise merge-tier entrypoint timeouts (H100 60->90, L4 30->45)

Both entrypoint jobs are killed mid-run by timeout_in_minutes while tests
are still passing, not by any failing test. The budgets had no headroom:
on main the H100 job timed out in 2 of the last 3 gate builds (2648, 2650)
and passed 2651 at 56/60 min; the L4 job passed at 26/30 min only with a
warm image cache. On the vLLM 0.25 wheel, model weight loading is slower
on the CI runners (like-for-like: +10 min total on H100, 1.7B TTS load
11s -> 132s on L4), which turns the coin flip into a guaranteed timeout
(build 2655).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

.buildkite/test-merge.yml

@@ -108,7 +108,10 @@ steps:
     depends_on: upload-merge-pipeline
     steps:
       - label: "Entrypoint Test with H100"
-        timeout_in_minutes: 60
+        # 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
         commands:
           - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
         agents:
@@ -146,7 +149,9 @@ steps:
                       type: DirectoryOrCreate
 
       - label: "Entrypoint Test with L4"
-        timeout_in_minutes: 30
+        # 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
         commands:
           - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4" --run-level "advanced_model"
         agents:

.buildkite/test-ready.yml

@@ -684,3 +684,24 @@ steps:
                     hostPath:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
+
+      - label: "Diffusion · GGUF Plugin Test"
+        soft_fail: true
+        source_file_dependencies:
+          - tests/plugin/test_gguf_diffusion.py
+          - vllm_omni/diffusion/model_loader
+        commands:
+          - pip install "vllm-gguf-plugin>=0.0.3"
+          - timeout 30m pytest -sv tests/plugin/test_gguf_diffusion.py
+        agents:
+          queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"

docs/configuration/composable_parallel.md

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

docs/configuration/pd_disaggregation.md

@@ -123,7 +123,7 @@ After inserting the extra thinker stage, renumber the remaining stages:
     runtime:
       devices: "2"
     engine_input_source: [1]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
+    sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only
 
   - 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]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
+    sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_token_only
 ```
 
 Compared with the default Qwen3-Omni config:

docs/contributing/model/adding_omni_model.md

@@ -475,92 +475,64 @@ 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`:
+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 |
 
 ```python
-# qwen3_omni.py
+# qwen3_omni.py (Thinker → Talker, non-async path)
 
-def thinker2talker(
-    stage_list: list[Any],
-    engine_input_source: list[int],
+def thinker2talker_token_only(
+    source_outputs: list[Any],
     prompt: OmniTokensPrompt | TextPrompt | None = None,
     requires_multimodal_data: bool = False,
+    streaming_context: Any | None = None,
 ) -> list[OmniTokensPrompt]:
-    """
-    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
+    """Allocate talker prefill slots; bulk tensors arrive via the connector."""
+    ...
 
 
-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 = []
+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."""
+    ...
 
-    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()
-        )
 
-        code2wav_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=codec_codes,
-                multi_modal_data=None,
-            )
-        )
-
-    return code2wav_inputs
+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."""
+    ...
 ```
 
+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).

docs/user_guide/quantization/gguf.md

@@ -9,6 +9,15 @@ for the transformer.
 
 GGUF is static quantization: the quantized weights are produced before serving.
 
+!!! note
+    GGUF support has migrated to OOT [vllm-gguf-plugin](https://github.com/vllm-project/vllm-gguf-plugin). Make sure you have GGUF plugin installed before serving a GGUF model.
+
+Before serving a GGUF model, make sure to install the [vllm-gguf-plugin](https://github.com/vllm-project/vllm-gguf-plugin):
+
+```bash
+uv pip install vllm-gguf-plugin>=0.0.3
+```
+
 ## Hardware Support
 
 | Device | Support |

examples/offline_inference/image_to_video/image_to_video.py

@@ -220,7 +220,7 @@ def parse_args() -> argparse.Namespace:
         "--quantization",
         type=str,
         default=None,
-        choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"],
+        choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8"],
         help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).",
     )
 

examples/offline_inference/qwen2_5_omni/README.md

@@ -60,3 +60,24 @@ 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"}}'
+```

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -553,6 +553,39 @@ 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()
 
 

examples/offline_inference/qwen2_5_omni/strategy_stage_replica.yaml

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

examples/offline_inference/qwen2_5_omni/strategy_tp2.yaml

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

examples/offline_inference/text_to_image/text_to_image.py

@@ -202,17 +202,11 @@ def parse_args() -> argparse.Namespace:
         "--quantization",
         type=str,
         default=None,
-        choices=["fp8", "int8", "gguf"],
+        choices=["fp8", "int8"],
         help="Quantization method for the transformer. "
-        "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8), 'gguf' (GGUF quantized weights). "
+        "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8). "
         "Default: None (no quantization, uses BF16).",
     )
-    parser.add_argument(
-        "--gguf-model",
-        type=str,
-        default=None,
-        help=("GGUF file path or HF reference for transformer weights. Required when --quantization gguf is set."),
-    )
     parser.add_argument(
         "--ignored-layers",
         type=str,
@@ -402,14 +396,7 @@ def main():
     # ignored_layers is specified so the list flows through OmniDiffusionConfig
     quant_kwargs: dict[str, Any] = {}
     ignored_layers = [s.strip() for s in args.ignored_layers.split(",") if s.strip()] if args.ignored_layers else None
-    if args.quantization == "gguf":
-        if not args.gguf_model:
-            raise ValueError("--gguf-model is required when --quantization gguf is set.")
-        quant_kwargs["quantization_config"] = {
-            "method": "gguf",
-            "gguf_model": args.gguf_model,
-        }
-    elif args.quantization and ignored_layers:
+    if args.quantization and ignored_layers:
         quant_kwargs["quantization_config"] = {
             "method": args.quantization,
             "ignored_layers": ignored_layers,

examples/offline_inference/text_to_video/text_to_video.md

@@ -231,7 +231,7 @@ python text_to_video.py \
 - `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
 - `--audio-sample-rate`: audio sample rate for embedded audio (when the
   pipeline returns audio; LTX-2.3 outputs 48kHz audio).
-- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
+- `--quantization`: quantization method (such as `fp8` for FP8).
 - `--flow-shift`: scheduler flow_shift parameter.
 - `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above).
 

examples/offline_inference/text_to_video/text_to_video.py

@@ -227,7 +227,7 @@ def parse_args() -> argparse.Namespace:
         "--quantization",
         type=str,
         default=None,
-        choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"],
+        choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8"],
         help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).",
     )
 

tests/config/composable_parallel/test_apply.py

@@ -0,0 +1,154 @@
+# 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")],
+            },
+        )

tests/config/composable_parallel/test_factory_integration.py

@@ -0,0 +1,141 @@
+# 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)

tests/config/composable_parallel/test_spec_types.py

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

tests/config/composable_parallel/test_strategy_loader.py

@@ -0,0 +1,97 @@
+# 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"

tests/config/composable_parallel/test_translator.py

@@ -0,0 +1,195 @@
+# 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}"

tests/config/test_omni_config.py

@@ -751,3 +751,14 @@ def test_diffusion_config_field_classification_covers_current_fields():
         "distributed_executor_backend",
     } <= omni_config_module._DIFFUSION_SHARED_CONFIG_FIELDS
     assert "prompt_file_path" in omni_config_module._DIFFUSION_RUNTIME_CONFIG_FIELDS
+
+
+def test_diffusion_config_projection_keeps_mapping_quantization_config_serializable():
+    quantization_config = {
+        "method": "example_quant",
+        "weights": "weights.bin",
+    }
+
+    cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(quantization_config=quantization_config)
+
+    assert cfg.quantization_config == quantization_config

tests/dfx/perf/tests/test_voxcpm2.json

@@ -3,7 +3,10 @@
     "test_name": "test_voxcpm2",
     "server_params": {
       "model": "openbmb/VoxCPM2",
-      "trust_remote_code": true
+      "trust_remote_code": true,
+      "extra_cli_args": [
+        "--trust-remote-code"
+      ]
     },
     "benchmark_params": [
       {

tests/diffusion/model_loader/gguf_adapters/test_flux2_klein_gguf_adapter.py

@@ -1,112 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for the Flux2-Klein GGUF adapter."""
-
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
-from vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein import (
-    Flux2KleinGGUFAdapter,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-def _make_od_config(
-    model_class_name: str = "Flux2Pipeline",
-    model_type: str = "flux",
-):
-    return SimpleNamespace(
-        model_class_name=model_class_name,
-        tf_model_config={"model_type": model_type},
-    )
-
-
-def _make_source(prefix: str = "", subfolder: str = "transformer"):
-    return SimpleNamespace(prefix=prefix, subfolder=subfolder)
-
-
-def test_flux2_klein_adapter_selected_for_flux_family():
-    adapter = get_gguf_adapter(
-        "dummy.gguf",
-        object(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    assert isinstance(adapter, Flux2KleinGGUFAdapter)
-
-
-def test_flux2_klein_adapter_matches_flux_model_type():
-    assert Flux2KleinGGUFAdapter.is_compatible(
-        _make_od_config(model_class_name="OtherPipeline", model_type="flux-dev"),
-        object(),
-        _make_source(),
-    )
-
-
-def test_flux2_klein_adapter_renames_core_projection_paths(monkeypatch: pytest.MonkeyPatch):
-    import vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein as flux_module
-
-    monkeypatch.setattr(
-        flux_module,
-        "gguf_quant_weights_iterator",
-        lambda _path: iter(
-            [
-                ("img_in.qweight", torch.ones((1, 1))),
-                ("time_in.in_layer.qweight_type", torch.tensor(1)),
-                ("double_blocks.0.img_attn.norm.query_norm.weight", torch.full((1, 1), 2.0)),
-                ("final_layer.linear.qweight", torch.full((1, 1), 3.0)),
-            ]
-        ),
-    )
-
-    adapter = Flux2KleinGGUFAdapter(
-        "dummy.gguf",
-        object(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    weights = list(adapter.weights_iterator())
-    names = [name for name, _ in weights]
-
-    assert "x_embedder.qweight" in names
-    assert "time_guidance_embed.timestep_embedder.linear_1.qweight_type" in names
-    assert "transformer_blocks.0.attn.norm_q.weight" in names
-    assert "proj_out.qweight" in names
-
-
-def test_flux2_klein_adapter_swaps_final_adaln_shift_and_scale(monkeypatch: pytest.MonkeyPatch):
-    import vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein as flux_module
-
-    monkeypatch.setattr(
-        flux_module,
-        "gguf_quant_weights_iterator",
-        lambda _path: iter(
-            [
-                (
-                    "final_layer.adaLN_modulation.1.scale",
-                    torch.tensor([1.0, 2.0, 3.0, 4.0]),
-                )
-            ]
-        ),
-    )
-
-    adapter = Flux2KleinGGUFAdapter(
-        "dummy.gguf",
-        object(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    weights = list(adapter.weights_iterator())
-
-    assert len(weights) == 1
-    assert weights[0][0] == "norm_out.linear.weight"
-    assert torch.equal(weights[0][1], torch.tensor([3.0, 4.0, 1.0, 2.0]))

tests/diffusion/model_loader/gguf_adapters/test_qwen_image_gguf_adapter.py

@@ -1,184 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import pytest
-
-from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
-from vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image import (
-    QwenImageGGUFAdapter,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-def _make_od_config(
-    model_class_name: str = "QwenImagePipeline",
-    model_type: str = "qwen_image",
-):
-    return SimpleNamespace(
-        model_class_name=model_class_name,
-        tf_model_config={"model_type": model_type},
-    )
-
-
-def _make_source(prefix: str = "", subfolder: str = "transformer"):
-    return SimpleNamespace(prefix=prefix, subfolder=subfolder)
-
-
-class _FakeTransformer:
-    def __init__(self) -> None:
-        self._params = {
-            "img_in.qweight": 1,
-            "img_in.qweight_type": 1,
-            "txt_in.qweight": 1,
-            "txt_in.qweight_type": 1,
-            "proj_out.qweight": 1,
-            "proj_out.qweight_type": 1,
-            "time_text_embed.timestep_embedder.linear_1.qweight": 1,
-            "time_text_embed.timestep_embedder.linear_1.qweight_type": 1,
-            "time_text_embed.timestep_embedder.linear_2.qweight": 1,
-            "time_text_embed.timestep_embedder.linear_2.qweight_type": 1,
-            "transformer_blocks.0.img_mod.1.qweight": 1,
-            "transformer_blocks.0.img_mod.1.qweight_type": 1,
-            "transformer_blocks.0.txt_mod.1.qweight": 1,
-            "transformer_blocks.0.txt_mod.1.qweight_type": 1,
-            "transformer_blocks.0.attn.to_qkv.qweight": 1,
-            "transformer_blocks.0.attn.to_qkv.qweight_type": 1,
-            "transformer_blocks.0.attn.add_kv_proj.qweight": 1,
-            "transformer_blocks.0.attn.add_kv_proj.qweight_type": 1,
-            "transformer_blocks.0.attn.to_out.weight": 1,
-            "transformer_blocks.0.img_mlp.net.0.proj.qweight": 1,
-            "transformer_blocks.0.img_mlp.net.0.proj.qweight_type": 1,
-        }
-
-    def named_parameters(self):
-        return list(self._params.items())
-
-    def named_buffers(self):
-        return []
-
-    def get_submodule(self, _prefix: str):
-        return self
-
-
-def test_qwen_adapter_selected_for_qwen_image_family():
-    adapter = get_gguf_adapter(
-        "dummy.gguf",
-        _FakeTransformer(),
-        _make_source(),
-        _make_od_config(),
-    )
-    assert isinstance(adapter, QwenImageGGUFAdapter)
-
-
-def test_qwen_adapter_matches_multiple_pipeline_variants():
-    for model_class_name in (
-        "QwenImagePipeline",
-        "QwenImageEditPipeline",
-        "QwenImageEditPlusPipeline",
-        "QwenImageLayeredPipeline",
-    ):
-        assert QwenImageGGUFAdapter.is_compatible(
-            _make_od_config(model_class_name=model_class_name),
-            _FakeTransformer(),
-            _make_source(),
-        )
-
-
-def test_qwen_adapter_preserves_split_projection_names(monkeypatch: pytest.MonkeyPatch):
-    import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
-
-    monkeypatch.setattr(
-        qwen_image_module,
-        "gguf_quant_weights_iterator",
-        lambda _path: iter(
-            [
-                ("transformer_blocks.0.attn.to_q.qweight_type", 1),
-                ("transformer_blocks.0.attn.to_q.qweight", 2),
-                ("transformer_blocks.0.attn.to_k.qweight_type", 3),
-                ("transformer_blocks.0.attn.to_k.qweight", 4),
-                ("transformer_blocks.0.attn.to_out.0.qweight_type", 5),
-                ("transformer_blocks.0.attn.to_out.0.qweight", 6),
-            ]
-        ),
-    )
-
-    adapter = QwenImageGGUFAdapter(
-        "dummy.gguf",
-        _FakeTransformer(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    weights = list(adapter.weights_iterator())
-
-    assert ("transformer_blocks.0.attn.to_q.qweight_type", 1) in weights
-    assert ("transformer_blocks.0.attn.to_q.qweight", 2) in weights
-    assert ("transformer_blocks.0.attn.to_k.qweight_type", 3) in weights
-    assert ("transformer_blocks.0.attn.to_k.qweight", 4) in weights
-    assert ("transformer_blocks.0.attn.to_out.0.qweight_type", 5) in weights
-    assert ("transformer_blocks.0.attn.to_out.0.qweight", 6) in weights
-
-
-def test_qwen_adapter_keeps_top_level_quantized_weights(monkeypatch: pytest.MonkeyPatch):
-    import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
-
-    monkeypatch.setattr(
-        qwen_image_module,
-        "gguf_quant_weights_iterator",
-        lambda _path: iter(
-            [
-                ("img_in.qweight_type", 1),
-                ("img_in.qweight", 2),
-                ("time_text_embed.timestep_embedder.linear_1.qweight_type", 3),
-                ("time_text_embed.timestep_embedder.linear_1.qweight", 4),
-            ]
-        ),
-    )
-
-    adapter = QwenImageGGUFAdapter(
-        "dummy.gguf",
-        _FakeTransformer(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    weights = list(adapter.weights_iterator())
-
-    assert ("img_in.qweight_type", 1) in weights
-    assert ("img_in.qweight", 2) in weights
-    assert ("time_text_embed.timestep_embedder.linear_1.qweight_type", 3) in weights
-    assert ("time_text_embed.timestep_embedder.linear_1.qweight", 4) in weights
-
-
-def test_qwen_adapter_keeps_modulation_quantized_weights(monkeypatch: pytest.MonkeyPatch):
-    import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
-
-    monkeypatch.setattr(
-        qwen_image_module,
-        "gguf_quant_weights_iterator",
-        lambda _path: iter(
-            [
-                ("transformer_blocks.0.img_mod.1.qweight_type", 1),
-                ("transformer_blocks.0.img_mod.1.qweight", 2),
-                ("transformer_blocks.0.txt_mod.1.qweight_type", 3),
-                ("transformer_blocks.0.txt_mod.1.qweight", 4),
-            ]
-        ),
-    )
-
-    adapter = QwenImageGGUFAdapter(
-        "dummy.gguf",
-        _FakeTransformer(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    weights = list(adapter.weights_iterator())
-
-    assert ("transformer_blocks.0.img_mod.1.qweight_type", 1) in weights
-    assert ("transformer_blocks.0.img_mod.1.qweight", 2) in weights
-    assert ("transformer_blocks.0.txt_mod.1.qweight_type", 3) in weights
-    assert ("transformer_blocks.0.txt_mod.1.qweight", 4) in weights

tests/diffusion/model_loader/gguf_adapters/test_z_image_gguf_adapter.py

@@ -1,83 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for the Z-Image GGUF adapter."""
-
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
-from vllm_omni.diffusion.model_loader.gguf_adapters.z_image import (
-    ZImageGGUFAdapter,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-def _make_od_config(
-    model_class_name: str = "ZImagePipeline",
-    model_type: str = "z_image",
-):
-    return SimpleNamespace(
-        model_class_name=model_class_name,
-        tf_model_config={"model_type": model_type},
-    )
-
-
-def _make_source(prefix: str = "", subfolder: str = "transformer"):
-    return SimpleNamespace(prefix=prefix, subfolder=subfolder)
-
-
-def test_z_image_adapter_selected_for_z_image_family():
-    adapter = get_gguf_adapter(
-        "dummy.gguf",
-        object(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    assert isinstance(adapter, ZImageGGUFAdapter)
-
-
-def test_z_image_adapter_matches_model_type_variants():
-    for model_type in ("z_image", "zimage", "z-image"):
-        assert ZImageGGUFAdapter.is_compatible(
-            _make_od_config(model_class_name="OtherPipeline", model_type=model_type),
-            object(),
-            _make_source(),
-        )
-
-
-def test_z_image_adapter_renames_known_gguf_tensor_paths(monkeypatch: pytest.MonkeyPatch):
-    import vllm_omni.diffusion.model_loader.gguf_adapters.z_image as z_image_module
-
-    monkeypatch.setattr(
-        z_image_module,
-        "gguf_quant_weights_iterator",
-        lambda _path: iter(
-            [
-                ("model.diffusion_model.final_layer.qweight", torch.ones((1, 1))),
-                ("model.diffusion_model.x_embedder.qweight_type", torch.tensor(1)),
-                ("transformer_blocks.0.attention.out.weight", torch.full((1, 1), 2.0)),
-                ("transformer_blocks.0.attention.qkv.qweight", torch.full((1, 1), 3.0)),
-            ]
-        ),
-    )
-
-    adapter = ZImageGGUFAdapter(
-        "dummy.gguf",
-        object(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    weights = list(adapter.weights_iterator())
-    names = [name for name, _ in weights]
-
-    assert "all_final_layer.2-1.qweight" in names
-    assert "all_x_embedder.2-1.qweight_type" in names
-    assert "transformer_blocks.0.attention.to_out.0.weight" in names
-    assert "transformer_blocks.0.attention.to_qkv.qweight" in names

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -16,7 +16,6 @@ from vllm.config.load import LoadConfig
 from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
 from vllm_omni.diffusion.models.helios import HeliosPipeline
 from vllm_omni.diffusion.registry import initialize_model
 
@@ -111,27 +110,6 @@ def test_empty_source_prefix_keeps_full_model_strict_check():
         loader.load_weights(model)
 
 
-def test_qwen_model_class_selects_qwen_gguf_adapter():
-    od_config = type(
-        "Config",
-        (),
-        {
-            "model_class_name": "QwenImagePipeline",
-            "tf_model_config": {"model_type": "qwen_image"},
-        },
-    )()
-    source = DiffusersPipelineLoader.ComponentSource(
-        model_or_path="dummy",
-        subfolder="transformer",
-        revision=None,
-        prefix="transformer.",
-    )
-
-    adapter = get_gguf_adapter("dummy.gguf", object(), source, od_config)
-
-    assert adapter.__class__.__name__ == "QwenImageGGUFAdapter"
-
-
 class _ConfigAwareModel(nn.Module):
     def __init__(self, *, od_config):
         super().__init__()
@@ -186,7 +164,6 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
     loader = DiffusersPipelineLoader(LoadConfig(), od_config)
     loader.load_weights = lambda model: None  # type: ignore[assignment]
     loader._process_weights_after_loading = lambda model, target_device: None  # type: ignore[assignment]
-    loader._is_gguf_quantization = lambda: False  # type: ignore[assignment]
 
     monkeypatch.setattr(loader_mod, "resolve_obj_by_qualname", lambda _name: _ConfigAwareModel)
     monkeypatch.setattr(loader_mod.torch, "device", lambda _name: _DeviceContext("cpu"))

tests/diffusion/model_loader/test_diffusers_loader_gguf.py

@@ -1,259 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for GGUF-specific DiffusersPipelineLoader behavior."""
-
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-import torch.nn as nn
-
-import vllm_omni.diffusion.model_loader.diffusers_loader as loader_module
-from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-class _DummyModel(nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
-        self.transformer = nn.Linear(2, 2, bias=True)
-        self.vae = nn.Linear(2, 2, bias=False)
-        self.register_buffer("transformer_buffer", torch.ones(1))
-        self.calls: list[list[str]] = []
-        self.weights_sources = [
-            DiffusersPipelineLoader.ComponentSource(
-                model_or_path="dummy",
-                subfolder="transformer",
-                revision=None,
-                prefix="transformer.",
-                fall_back_to_pt=True,
-            ),
-            DiffusersPipelineLoader.ComponentSource(
-                model_or_path="dummy",
-                subfolder="vae",
-                revision=None,
-                prefix="vae.",
-                fall_back_to_pt=True,
-            ),
-        ]
-
-    def load_weights(self, weights):
-        loadable = dict(self.named_parameters())
-        loadable.update(dict(self.named_buffers()))
-        seen: list[str] = []
-        loaded: set[str] = set()
-        for name, tensor in weights:
-            seen.append(name)
-            if name in loadable:
-                target = loadable[name]
-                target.data.copy_(tensor.to(dtype=target.dtype))
-                loaded.add(name)
-        self.calls.append(seen)
-        return loaded
-
-
-def _make_loader() -> DiffusersPipelineLoader:
-    loader = object.__new__(DiffusersPipelineLoader)
-    loader.load_config = SimpleNamespace(
-        download_dir="cache-dir",
-        ignore_patterns=["*.tmp"],
-    )
-    loader.od_config = None
-    loader.quant_config = None
-    loader.counter_before_loading_weights = 0.0
-    loader.counter_after_loading_weights = 0.0
-    return loader
-
-
-def test_is_gguf_quantization_accepts_dict_config():
-    loader = _make_loader()
-    loader.quant_config = {
-        "method": "gguf",
-        "gguf_model": "weights.gguf",
-    }
-
-    assert loader._is_gguf_quantization() is True
-
-
-def test_is_gguf_quantization_rejects_non_gguf_dict_config():
-    loader = _make_loader()
-    loader.quant_config = {
-        "method": "fp8",
-        "gguf_model": "weights.gguf",
-    }
-
-    assert loader._is_gguf_quantization() is False
-
-
-def test_is_gguf_quantization_requires_gguf_model_for_dict_config():
-    loader = _make_loader()
-    loader.quant_config = {"method": "gguf"}
-
-    with pytest.raises(ValueError, match="gguf_model"):
-        loader._is_gguf_quantization()
-
-
-def test_is_gguf_quantization_accepts_object_config():
-    loader = _make_loader()
-    loader.quant_config = SimpleNamespace(
-        get_name=lambda: "gguf",
-        gguf_model="weights.gguf",
-    )
-
-    assert loader._is_gguf_quantization() is True
-
-
-def test_is_gguf_quantization_uses_fallback_object_without_get_name():
-    loader = _make_loader()
-    loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
-
-    assert loader._is_gguf_quantization() is True
-
-
-def test_get_model_loadable_names_collects_parameters_and_buffers():
-    loader = _make_loader()
-    model = _DummyModel()
-
-    names = loader._get_model_loadable_names(model)
-
-    assert "transformer.weight" in names
-    assert "transformer.bias" in names
-    assert "transformer_buffer" in names
-
-
-def test_resolve_gguf_model_path_returns_local_file(tmp_path):
-    loader = _make_loader()
-    gguf_file = tmp_path / "model.gguf"
-    gguf_file.write_bytes(b"gguf")
-
-    resolved = loader._resolve_gguf_model_path(str(gguf_file), revision=None)
-
-    assert resolved == str(gguf_file)
-
-
-def test_resolve_gguf_model_path_downloads_explicit_gguf_filename(monkeypatch: pytest.MonkeyPatch):
-    loader = _make_loader()
-
-    monkeypatch.setattr(loader_module.os.path, "isfile", lambda _path: False)
-    monkeypatch.setattr(
-        loader_module,
-        "hf_hub_download",
-        lambda repo_id, filename, revision, cache_dir: f"{cache_dir}/{repo_id}/{filename}",
-    )
-
-    resolved = loader._resolve_gguf_model_path("owner/repo/model-Q4_0.gguf", revision="main")
-
-    assert resolved == "cache-dir/owner/repo/model-Q4_0.gguf"
-
-
-def test_resolve_gguf_model_path_downloads_by_quant_type(monkeypatch: pytest.MonkeyPatch):
-    loader = _make_loader()
-
-    monkeypatch.setattr(loader_module.os.path, "isfile", lambda _path: False)
-    monkeypatch.setattr(
-        loader_module,
-        "download_gguf",
-        lambda repo_id, quant_type, cache_dir, revision, ignore_patterns: f"{cache_dir}/{repo_id}/{quant_type}.gguf",
-    )
-
-    resolved = loader._resolve_gguf_model_path("owner/repo:Q4_0", revision="main")
-
-    assert resolved == "cache-dir/owner/repo/Q4_0.gguf"
-
-
-def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.MonkeyPatch):
-    loader = _make_loader()
-    loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
-    loader.od_config = SimpleNamespace(revision="main")
-    source = DiffusersPipelineLoader.ComponentSource(
-        model_or_path="dummy",
-        subfolder="transformer",
-        revision=None,
-        prefix="transformer.",
-    )
-
-    class _Adapter:
-        def weights_iterator(self):
-            yield "weight", torch.ones((2, 2))
-            yield "bias", torch.zeros(2)
-
-    monkeypatch.setattr(loader, "_resolve_gguf_model_path", lambda gguf_model, revision: "resolved.gguf")
-    monkeypatch.setattr(loader_module, "get_gguf_adapter", lambda gguf_file, model, source, od_config: _Adapter())
-
-    weights = list(loader._get_gguf_weights_iterator(source, object()))
-
-    assert weights[0][0] == "transformer.weight"
-    assert torch.equal(weights[0][1], torch.ones((2, 2)))
-    assert weights[1][0] == "transformer.bias"
-    assert torch.equal(weights[1][1], torch.zeros(2))
-
-
-def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(monkeypatch: pytest.MonkeyPatch):
-    loader = _make_loader()
-    model = _DummyModel()
-
-    monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
-    monkeypatch.setattr(
-        loader,
-        "_get_gguf_weights_iterator",
-        lambda source, model: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]),
-    )
-
-    def _hf_iter(source):
-        if source.subfolder == "transformer":
-            return iter([("transformer.bias", torch.zeros(2, dtype=torch.float32))])
-        return iter([("vae.weight", torch.full((2, 2), 2.0, dtype=torch.float32))])
-
-    monkeypatch.setattr(loader, "_get_weights_iterator", _hf_iter)
-    monkeypatch.setattr(
-        loader,
-        "_get_expected_parameter_names",
-        lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
-    )
-
-    loaded = loader._load_weights_with_gguf(model)
-
-    assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
-    assert model.calls[0] == ["transformer.weight"]
-    assert model.calls[1] == ["transformer.bias"]
-    assert model.calls[2] == ["vae.weight"]
-
-
-def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete(monkeypatch: pytest.MonkeyPatch):
-    loader = _make_loader()
-    model = _DummyModel()
-
-    monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
-    monkeypatch.setattr(
-        loader,
-        "_get_gguf_weights_iterator",
-        lambda source, model: iter(
-            [
-                ("transformer.weight", torch.ones((2, 2), dtype=torch.float32)),
-                ("transformer.bias", torch.zeros(2, dtype=torch.float32)),
-            ]
-        ),
-    )
-
-    hf_calls: list[str] = []
-
-    def _hf_iter(source):
-        hf_calls.append(source.subfolder or "")
-        if source.subfolder == "transformer":
-            return iter([("transformer.bias", torch.zeros(2, dtype=torch.float32))])
-        return iter([("vae.weight", torch.full((2, 2), 2.0, dtype=torch.float32))])
-
-    monkeypatch.setattr(loader, "_get_weights_iterator", _hf_iter)
-    monkeypatch.setattr(
-        loader,
-        "_get_expected_parameter_names",
-        lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
-    )
-
-    loaded = loader._load_weights_with_gguf(model)
-
-    assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
-    assert hf_calls == ["vae"]

tests/diffusion/quantization/test_fp8_config.py

@@ -154,21 +154,6 @@ def test_component_config_with_default():
     assert config.resolve("vae.encoder") is None
 
 
-def test_gguf_config():
-    from vllm_omni.quantization import build_quant_config
-    from vllm_omni.quantization.gguf_config import DiffusionGGUFConfig
-
-    config = build_quant_config(
-        {
-            "method": "gguf",
-            "gguf_model": "path/to/model.gguf",
-        }
-    )
-    assert isinstance(config, DiffusionGGUFConfig)
-    assert config.gguf_model == "path/to/model.gguf"
-    assert config.get_name() == "gguf"
-
-
 def test_integration_string():
     from vllm_omni.diffusion.data import OmniDiffusionConfig
 
@@ -236,7 +221,7 @@ def test_transformer_config_auto_detects_modelopt_fp8():
 def test_supported_methods_includes_vllm():
     from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS
 
-    for method in ["fp8", "gguf", "awq", "gptq", "bitsandbytes", "modelopt"]:
+    for method in ["fp8", "awq", "gptq", "bitsandbytes", "modelopt"]:
         assert method in SUPPORTED_QUANTIZATION_METHODS, f"{method} missing"
 
 

tests/diffusion/quantization/test_gguf_config.py

@@ -1,128 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for diffusion GGUF quantization config and linear method."""
-
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
-
-from vllm_omni.quantization import build_quant_config
-from vllm_omni.quantization.gguf_config import (
-    UNQUANTIZED_TYPES,
-    DiffusionGGUFConfig,
-    DiffusionGGUFLinearMethod,
-    dequant_gemm_gguf,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-def test_gguf_config_creation_and_delegation():
-    config = DiffusionGGUFConfig(
-        gguf_model="weights.gguf",
-        unquantized_modules=["proj_out"],
-    )
-
-    assert config.gguf_model == "weights.gguf"
-    assert config.unquantized_modules == ["proj_out"]
-    assert config.get_name() == "gguf"
-
-
-def test_build_quant_config_builds_gguf_config():
-    config = build_quant_config(
-        "gguf",
-        gguf_model="repo/model-Q4_0.gguf",
-        unquantized_modules=["norm_out"],
-    )
-
-    assert isinstance(config, DiffusionGGUFConfig)
-    assert config.gguf_model == "repo/model-Q4_0.gguf"
-    assert config.unquantized_modules == ["norm_out"]
-
-
-def test_gguf_vllm_config_returns_diffusion_linear_method_for_linear_layers():
-    linear = object.__new__(LinearBase)
-    method = DiffusionGGUFConfig(unquantized_modules=[]).get_quant_method(linear, "transformer.img_in")
-
-    assert isinstance(method, DiffusionGGUFLinearMethod)
-
-
-def test_gguf_vllm_config_respects_unquantized_modules():
-    linear = object.__new__(LinearBase)
-    method = DiffusionGGUFConfig(unquantized_modules=["proj_out"]).get_quant_method(linear, "transformer.proj_out")
-
-    assert isinstance(method, UnquantizedLinearMethod)
-
-
-def test_gguf_vllm_config_returns_none_for_non_linear_layers():
-    method = DiffusionGGUFConfig(unquantized_modules=[]).get_quant_method(torch.nn.LayerNorm(4), "norm")
-
-    assert method is None
-
-
-def test_dequant_gemm_gguf_uses_plain_matmul_for_unquantized_types():
-    qweight_type = next(iter(UNQUANTIZED_TYPES))
-    x = torch.tensor([[1.0, 2.0]], dtype=torch.float32)
-    qweight = torch.tensor([[3.0, 4.0], [5.0, 6.0]], dtype=torch.float32)
-
-    out = dequant_gemm_gguf(x, qweight, qweight_type)
-
-    assert torch.allclose(out, x @ qweight.T)
-
-
-def test_diffusion_gguf_linear_method_applies_bias_on_unquantized_weight():
-    qweight_type = next(iter(UNQUANTIZED_TYPES))
-    method = DiffusionGGUFLinearMethod(DiffusionGGUFConfig(unquantized_modules=[]))
-    layer = SimpleNamespace(
-        qweight=torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32),
-        qweight_type=SimpleNamespace(weight_type=qweight_type),
-    )
-    x = torch.tensor([[[1.0, 0.5]]], dtype=torch.float32)
-    bias = torch.tensor([0.25, -0.5], dtype=torch.float32)
-
-    out = method.apply(layer, x, bias)
-
-    expected = x @ layer.qweight.T + bias
-    assert torch.allclose(out, expected)
-
-
-def test_diffusion_gguf_linear_method_concatenates_sharded_outputs():
-    qweight_type = next(iter(UNQUANTIZED_TYPES))
-    method = DiffusionGGUFLinearMethod(DiffusionGGUFConfig(unquantized_modules=[]))
-    qweight = torch.nn.Parameter(
-        torch.tensor(
-            [
-                [1.0, 0.0],
-                [0.0, 1.0],
-                [2.0, 0.0],
-                [0.0, 2.0],
-            ],
-            dtype=torch.float32,
-        ),
-        requires_grad=False,
-    )
-    qweight.shard_id = ["left", "right"]
-    qweight.shard_offset_map = {
-        "left": (0, 2, 2),
-        "right": (2, 4, 2),
-    }
-    layer = SimpleNamespace(
-        qweight=qweight,
-        qweight_type=SimpleNamespace(
-            shard_weight_type={
-                "left": qweight_type,
-                "right": qweight_type,
-            }
-        ),
-    )
-    x = torch.tensor([[1.5, 2.0]], dtype=torch.float32)
-
-    out = method.apply(layer, x)
-
-    expected_q = x @ qweight[:2].T
-    expected_k = x @ qweight[2:4].T
-    assert torch.allclose(out, torch.cat([expected_q, expected_k], dim=-1))

tests/diffusion/quantization/test_trajectory_similarity_tool.py

@@ -99,7 +99,6 @@ def test_candidate_model_can_point_to_offline_checkpoint_without_online_quantiza
         candidate_quantization_config_json=None,
         candidate_ignored_layers=None,
         ignored_layers=None,
-        candidate_gguf_model=None,
         candidate_load_format="default",
     )
 

tests/e2e/offline_inference/test_qwen3_tts_customvoice.py

@@ -34,12 +34,10 @@ 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,
                 },

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -410,6 +410,44 @@ 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:

tests/engine/test_async_omni_engine_stage_init.py

@@ -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]),
+        lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
     )
 
     _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]),
+        lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
     )
 
     _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]),
+        lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
     )
 
     _config_path, stage_configs = engine._resolve_stage_configs(

tests/engine/test_strategy_lb_policy.py

@@ -0,0 +1,63 @@
+# 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"

tests/entrypoints/openai_api/test_serving_speech.py

@@ -39,7 +39,8 @@ from vllm_omni.entrypoints.openai.serving_speech import (
     OmniOpenAIServingSpeech,
     _create_wav_header,
 )
-from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest
+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.model_executor.models.fish_speech.prompt_utils import (
     FISH_TEXT_ONLY_SYSTEM_PROMPT,
     build_fish_voice_clone_prompt_ids,
@@ -973,6 +974,29 @@ 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")
@@ -1557,6 +1581,86 @@ 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):

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -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]),
+        return_value=("dummy.yaml", [fake_llm_stage, fake_diffusion_stage], None),
     )
 
     engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
@@ -359,3 +359,15 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo
 
     assert not hasattr(stage_configs[0].engine_args, "additional_config")
     assert stage_configs[1].engine_args.additional_config == {"torchair_graph_config": {"enabled": True}}
+
+
+def test_default_stage_config_includes_quantization_config():
+    """Ensure structured quantization_config survives default diffusion-stage creation."""
+    quantization_config = {
+        "method": "example_quant",
+        "weights": "weights.bin",
+    }
+
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"quantization_config": quantization_config})[0]
+
+    assert stage_cfg["engine_args"]["quantization_config"] == quantization_config

tests/entrypoints/test_resolve_dreamzero_config.py

@@ -32,7 +32,7 @@ def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch):
         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"

tests/entrypoints/test_serve.py

@@ -9,6 +9,7 @@ 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]
@@ -31,6 +32,20 @@ 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",
@@ -80,13 +95,97 @@ 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]),
+        return_value=("/fake/stages.yaml", [other_stage], None),
     )
 
     args = _make_headless_args(stage_id=0)
@@ -134,7 +233,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]),
+        return_value=("/fake/stages.yaml", [stage_cfg], None),
     )
     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)
@@ -210,7 +309,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]),
+        return_value=("/fake/stages.yaml", [stage_cfg], None),
     )
     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)
@@ -269,7 +368,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]),
+        return_value=("/fake/stages.yaml", [stage_cfg], None),
     )
     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)
@@ -349,7 +448,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]),
+        return_value=("/fake/stages.yaml", [stage_cfg], None),
     )
     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)

tests/entrypoints/test_utils.py

@@ -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,
+            return_value=(returned_stage_configs, None),
         )
 
-        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,6 +363,7 @@ 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

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -18,8 +18,11 @@ 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__),
@@ -542,3 +545,92 @@ 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

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py

@@ -8,7 +8,6 @@ import torch
 
 from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
     talker2code2wav_async_chunk,
-    text2flow,
     text2flow_full_payload,
     text2flow_token_only,
 )
@@ -49,13 +48,13 @@ def _transfer_manager(
     )
 
 
-def test_text2flow_supports_batched_source_outputs():
+def test_text2flow_token_only_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(source_outputs=source_outputs, prompt=None)
+    outputs = text2flow_token_only(source_outputs=source_outputs, prompt=None)
 
     assert len(outputs) == 2
     assert outputs[0]["prompt_token_ids"] == [1, 2, 3]
@@ -64,16 +63,6 @@ def test_text2flow_supports_batched_source_outputs():
     assert outputs[1]["additional_information"]["ids"]["prompt"] == [20, 21]
 
 
-def test_text2flow_strips_reference_speech_prefix_from_cumulative_ids():
-    source_outputs = [
-        _source_output("req-0", [10, 11], [8, 9, 1, 2, 3], {"speech_token": torch.tensor([[8, 9]])}),
-    ]
-
-    outputs = text2flow(source_outputs=source_outputs, prompt=None)
-
-    assert outputs[0]["prompt_token_ids"] == [1, 2, 3]
-
-
 def test_text2flow_token_only_strips_reference_speech_prefix_from_cumulative_ids():
     source_outputs = [
         _source_output(

tests/model_executor/stage_input_processors/test_indextts2.py

@@ -7,7 +7,6 @@ import torch
 
 from vllm_omni.model_executor.stage_input_processors.indextts2 import (
     _strip_stop_token,
-    talker2s2mel,
     talker2s2mel_full_payload,
     talker2s2mel_token_only,
 )
@@ -108,118 +107,6 @@ 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 ──
 
 

tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py

@@ -11,59 +11,40 @@ 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,
+    llm2code2wav_full_payload,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-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.
+def test_llm2code2wav_full_payload_truncates_when_flat_exceeds_max(caplog):
+    """Flat codec sequences longer than MAX_CODE2WAV_TOKENS must be truncated."""
     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:
-        prompts = llm2code2wav(source_outputs)
+        payload = llm2code2wav_full_payload(None, {"codes.audio": codec_codes}, request)
     finally:
         target_logger.removeHandler(caplog.handler)
         target_logger.setLevel(prev_level)
 
-    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]}"
-    )
+    assert payload is not None
+    assert len(payload["codes"]["audio"]) == MAX_CODE2WAV_TOKENS
 
 
-def test_llm2code2wav_short_sequence_unchanged():
-    """Short codec sequences are returned without truncation."""
+def test_llm2code2wav_full_payload_short_sequence_unchanged():
     codec_codes = torch.ones(4, 1, 8, 4, dtype=torch.long)
-    source_outputs = _make_source_outputs(codec_codes, request_id="req-short")
+    request = SimpleNamespace(request_id="req-short")
 
-    prompts = llm2code2wav(source_outputs)
+    payload = llm2code2wav_full_payload(None, {"codes.audio": codec_codes}, request)
 
-    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
+    assert payload is not None
+    assert 0 < len(payload["codes"]["audio"]) <= MAX_CODE2WAV_TOKENS
 
 
 def test_llm2code2wav_truncation_boundary_constant_matches_yaml():

tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py

@@ -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 payload["code_predictor_codes"] == payload["codes"]["audio"]
+    assert "code_predictor_codes" not in payload
 
 
 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 payload["code_predictor_codes"] == payload["codes"]["audio"]
+    assert "code_predictor_codes" not in payload
 
 
 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 payload["code_predictor_codes"] == payload["codes"]["audio"]
+    assert "code_predictor_codes" not in payload
 
 
 def test_thinker2talker_full_payload_packs_complete_tensors() -> None:
@@ -247,7 +247,6 @@ 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
 

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -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,7 +368,15 @@ def test_ref_code_context_can_be_buffered_before_first_emit():
     assert rid in tm.request_payload
 
 
-def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
+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.
+    """
     ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
     audio_codes = torch.tensor(
         [
@@ -387,12 +395,38 @@ def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
         engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
     )
 
-    prompts = talker2code2wav(stage.engine_outputs)
+    prompts = talker2code2wav_token_only(stage.engine_outputs)
 
     assert len(prompts) == 1
     prompt = prompts[0]
     assert prompt["additional_information"] == {"meta": {"left_context_size": 2}}
-    assert prompt["prompt_token_ids"] == [
+    # 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() == [
         9,
         8,
         1,
@@ -433,7 +467,7 @@ def test_non_async_processor_filters_out_of_range_codec_values():
         engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
     )
 
-    prompts = talker2code2wav(stage.engine_outputs)
+    prompts = talker2code2wav_token_only(stage.engine_outputs)
 
     assert len(prompts) == 1
     prompt = prompts[0]

tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py

@@ -8,61 +8,12 @@ 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,

tests/platforms/test_cuda_ir_op_priority.py

@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+from vllm.config import CompilationConfig, CompilationMode, DeviceConfig, VllmConfig
+
+from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _vllm_config(*, backend: str = "inductor", mode: CompilationMode) -> VllmConfig:
+    return VllmConfig(
+        device_config=DeviceConfig(device="cpu"),
+        compilation_config=CompilationConfig(backend=backend, mode=mode),
+    )
+
+
+@pytest.mark.parametrize(
+    "mode",
+    [CompilationMode.NONE, CompilationMode.VLLM_COMPILE, CompilationMode.STOCK_TORCH_COMPILE],
+)
+def test_cuda_default_ir_op_priority_prefers_vllm_c_when_inductor_backend(mode: CompilationMode) -> None:
+    """Regression for #4964: inductor-active configs must not switch to native-only."""
+    priority = CudaOmniPlatform.get_default_ir_op_priority(_vllm_config(mode=mode))
+    assert priority.rms_norm == ["vllm_c", "native"]
+    assert priority.fused_add_rms_norm == ["vllm_c", "native"]
+
+
+def test_cuda_default_ir_op_priority_with_oink(monkeypatch: pytest.MonkeyPatch) -> None:
+    import vllm.envs as envs
+
+    monkeypatch.setattr(envs, "VLLM_USE_OINK_OPS", True)
+    priority = CudaOmniPlatform.get_default_ir_op_priority(
+        _vllm_config(mode=CompilationMode.VLLM_COMPILE),
+    )
+    assert priority.rms_norm == ["oink", "vllm_c", "native"]
+    assert priority.fused_add_rms_norm == ["oink", "vllm_c", "native"]

tests/plugin/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

tests/plugin/test_gguf_diffusion.py

@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end tests for GGUF quantization on diffusion models.
+
+Validates that GGUF-quantized diffusion models generate valid images and
+use less peak GPU memory than BF16 baseline.
+
+Requires vllm-omni to be installed alongside the plugin.
+
+Usage:
+    pytest tests/diffusion/test_gguf_memory.py -v
+"""
+
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import numpy as np
+import pytest
+import torch
+
+from tests.helpers.env import DeviceMemoryMonitor
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+
+class DiffusionGGUFTestConfig(NamedTuple):
+    artifact_prefix: str
+    gguf_model: str
+    hf_model: str
+    min_cosine_similarity: float = 0.95
+
+
+Z_IMAGE_CONFIG = DiffusionGGUFTestConfig(
+    artifact_prefix="zimage",
+    hf_model="Tongyi-MAI/Z-Image-Turbo",
+    gguf_model="unsloth/Z-Image-Turbo-GGUF:Q4_0",
+)
+
+
+FLUX_CONFIG = DiffusionGGUFTestConfig(
+    artifact_prefix="flux2_klein",
+    hf_model="black-forest-labs/FLUX.2-klein-4B",
+    gguf_model="unsloth/FLUX.2-klein-4B-GGUF:Q8_0",
+)
+
+
+def _image_cosine_similarity(hf_image, gguf_image) -> float:
+    hf_tensor = torch.as_tensor(np.array(hf_image.convert("RGB"), copy=True), dtype=torch.float32).flatten()
+    gguf_tensor = torch.as_tensor(np.array(gguf_image.convert("RGB"), copy=True), dtype=torch.float32).flatten()
+    return torch.nn.functional.cosine_similarity(hf_tensor, gguf_tensor, dim=0).item()
+
+
+def _generate_single_stage_image(
+    model: str,
+    height: int = 256,
+    width: int = 256,
+    num_inference_steps: int = 20,
+    seed: int = 42,
+    **extra_kwargs,
+) -> tuple[list, float]:
+    """Generate an image with a single-stage diffusion model.
+
+    Returns (images, peak_memory_gib).
+    """
+    omni_kwargs = dict(extra_kwargs)
+
+    memory_monitor = DeviceMemoryMonitor(device_index=0)
+    memory_monitor.start()
+    try:
+        with OmniRunner(model, **omni_kwargs) as runner:
+            generator = torch.Generator(
+                device=current_omni_platform.device_type,
+            ).manual_seed(seed)
+            outputs = runner.omni.generate(
+                "a photo of a cat sitting on a laptop keyboard",
+                OmniDiffusionSamplingParams(
+                    height=height,
+                    width=width,
+                    num_inference_steps=num_inference_steps,
+                    guidance_scale=0.0,
+                    generator=generator,
+                ),
+            )
+    finally:
+        memory_monitor.stop()
+    peak_mem = memory_monitor.peak_used_mb / 1024
+
+    first_output = outputs[0]
+    assert first_output.final_output_type == "image"
+    if hasattr(first_output, "images") and first_output.images:
+        images = first_output.images
+    else:
+        assert hasattr(first_output, "request_output") and first_output.request_output
+        request_output = first_output.request_output
+        if isinstance(request_output, list):
+            req_out = request_output[0]
+        else:
+            req_out = request_output
+        assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
+        images = req_out.images
+    assert len(images) >= 1
+    assert images[0].width == width
+    assert images[0].height == height
+
+    return images, peak_mem
+
+
+@hardware_test(res={"cuda": "L4"})
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.slow
+@pytest.mark.parametrize("model", [Z_IMAGE_CONFIG, FLUX_CONFIG], ids=["Z-Image-Turbo", "FLUX.2-klein"])
+def test_single_stage_diffusion_gguf(model: DiffusionGGUFTestConfig, monkeypatch: pytest.MonkeyPatch) -> None:
+    """Z-Image-Turbo GGUF generates valid images and uses less memory than BF16."""
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+
+    # BF16 baseline
+    hf_images, mem_bf16 = _generate_single_stage_image(
+        model=model.hf_model,
+    )
+
+    # GGUF
+    images, mem_gguf = _generate_single_stage_image(
+        model=model.hf_model,
+        diffusion_quantization_config={
+            "method": "gguf",
+            "gguf_model": model.gguf_model,
+        },
+    )
+
+    hf_image_path = f"test_{model.artifact_prefix}_hf.png"
+    gguf_image_path = f"test_{model.artifact_prefix}_gguf.png"
+    assert len(hf_images) >= 1
+    hf_images[0].save(hf_image_path)
+    assert len(images) >= 1
+    images[0].save(gguf_image_path)
+    print(f"Saved HF image: {hf_image_path}")
+    print(f"Saved GGUF image: {gguf_image_path}")
+
+    image_cosine_similarity = _image_cosine_similarity(hf_images[0], images[0])
+    print(f"{model.artifact_prefix} image cosine similarity: {image_cosine_similarity:.4f}")
+    assert image_cosine_similarity >= model.min_cosine_similarity, (
+        f"GGUF image cosine similarity ({image_cosine_similarity:.4f}) should be >= {model.min_cosine_similarity:.4f}"
+    )
+
+    print(f"{model.artifact_prefix} BF16 peak VRAM delta: {mem_bf16:.2f} GiB")
+    print(f"{model.artifact_prefix} GGUF peak VRAM delta: {mem_gguf:.2f} GiB")
+    reduction = (mem_bf16 - mem_gguf) / mem_bf16 * 100
+    print(f"VRAM reduction: {reduction:.1f}%")
+    assert mem_gguf < mem_bf16, f"GGUF ({mem_gguf:.2f} GiB) should use less VRAM than BF16 ({mem_bf16:.2f} GiB)"

tests/test_config_factory.py

@@ -1166,7 +1166,10 @@ class TestQwen3OmniPipeline:
         assert isinstance(s, StagePipelineConfig)
         assert s.input_sources == (0,)
         assert s.sampling_constraints["stop_token_ids"] == [2150]
-        assert s.custom_process_input_func is not None
+        # thinker2talker was removed: sync_process_input_func always wins.
+        assert s.custom_process_input_func is None
+        assert s.sync_process_input_func is not None
+        assert s.sync_process_input_func.endswith("thinker2talker_token_only")
         assert s.custom_process_next_stage_input_func is not None
 
     def test_code2wav(self):
@@ -1180,7 +1183,7 @@ class TestQwen3OmniPipeline:
         assert isinstance(s, StagePipelineConfig)
         assert s.execution_type == StageExecutionType.LLM_GENERATION
         assert s.final_output_type == "audio"
-        assert s.custom_process_input_func is not None
+        assert s.engine_output_type == "audio"
 
 
 class TestQwen2_5OmniPipeline:
@@ -1718,7 +1721,7 @@ class TestAuraOmniDeploy:
     def test_aura_omni_deploy_resolves_four_native_stages(self):
         pipeline_cfg = StageConfigFactory.resolve_pipeline_config("aura_omni")
 
-        stages = StageConfigFactory._create_from_registry(
+        stages, _ = StageConfigFactory._create_from_registry(
             "qwen3_tts",
             pipeline_cfg,
             cli_overrides={},
@@ -1940,11 +1943,12 @@ class TestSentinelDefaultPrecedence:
             model_type,
             Q3_OMNI_ALL_STAGES_HF_CONFIG,
         )
-        return StageConfigFactory._create_from_registry(
+        stages, _ = StageConfigFactory._create_from_registry(
             "qwen3_omni_moe",
             pipeline_cfg,
             cli_overrides=cli_overrides,
         )
+        return stages
 
     def test_typed_kwarg_overrides_yaml(self):
         stages = self._stages({"max_num_seqs": 999})

tests/test_version.py

@@ -56,3 +56,26 @@ def test_warning_contains_version_strings():
     msg = str(record[0].message)
     assert "0.19.0rc2.dev21" in msg
     assert "0.18.0" in msg
+
+
+def test_omni_plugin_loader_invokes_vllm_general_plugins_first(monkeypatch):
+    import vllm.plugins as vllm_plugins
+
+    import vllm_omni.plugins as omni_plugins
+
+    calls = []
+    monkeypatch.setattr(omni_plugins, "omni_plugins_loaded", False)
+    monkeypatch.setattr(vllm_plugins, "load_general_plugins", lambda: calls.append("vllm"))
+    monkeypatch.setattr(
+        omni_plugins,
+        "load_omni_plugins_by_group",
+        lambda group: {"omni": lambda: calls.append(group)},
+    )
+
+    omni_plugins.load_omni_general_plugins()
+
+    assert calls == ["vllm", omni_plugins.OMNI_DEFAULT_PLUGINS_GROUP]
+
+    omni_plugins.load_omni_general_plugins()
+
+    assert calls == ["vllm", omni_plugins.OMNI_DEFAULT_PLUGINS_GROUP]

tests/utils/test_tracking_parser.py

@@ -591,7 +591,7 @@ def test_explicit_cli_arg_reaches_runtime_overrides(mock_stages):
     ns = p.parse_args(["--max-num-seqs", "999"])
 
     explicit_kwargs = ns.get_explicit_kwargs_dict()
-    stages = StageConfigFactory._create_from_registry(
+    stages, _ = StageConfigFactory._create_from_registry(
         _TEST_MODEL,
         _TEST_PIPELINE,
         explicit_kwargs,
@@ -608,7 +608,7 @@ def test_omitted_default_not_in_runtime_overrides(mock_stages):
     ns = p.parse_args([])
 
     explicit_kwargs = ns.get_explicit_kwargs_dict()
-    stages = StageConfigFactory._create_from_registry(
+    stages, _ = StageConfigFactory._create_from_registry(
         _TEST_MODEL,
         _TEST_PIPELINE,
         explicit_kwargs,
@@ -630,7 +630,7 @@ def test_config_file_args_reach_runtime_overrides(mock_stages):
         ns = p.parse_args(["--config", "fake.yaml"])
 
     explicit_kwargs = ns.get_explicit_kwargs_dict()
-    stages = StageConfigFactory._create_from_registry(
+    stages, _ = StageConfigFactory._create_from_registry(
         _TEST_MODEL,
         _TEST_PIPELINE,
         explicit_kwargs,
@@ -648,7 +648,7 @@ def test_per_stage_override_routes_correctly(mock_stages):
     ns = p.parse_args(["--stage-0-gpu-memory-utilization", "0.42"])
 
     explicit_kwargs = ns.get_explicit_kwargs_dict()
-    stages = StageConfigFactory._create_from_registry(
+    stages, _ = StageConfigFactory._create_from_registry(
         _TEST_MODEL,
         _TEST_PIPELINE,
         explicit_kwargs,
@@ -670,7 +670,7 @@ def test_explicit_args_omitted_from_yaml(mock_stages):
     ns = p.parse_args(["--enforce-eager"])
 
     explicit_kwargs = ns.get_explicit_kwargs_dict()
-    stages = StageConfigFactory._create_from_registry(
+    stages, _ = StageConfigFactory._create_from_registry(
         _TEST_MODEL,
         _TEST_PIPELINE,
         explicit_kwargs,

vllm_omni/config/composable_parallel/__init__.py

@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Composable parallel strategies for vLLM-Omni.
+
+A small, runtime-agnostic vocabulary for declaring a stage's parallel layout
+(:class:`StrategySpec` + mesh/routing/aggregation patterns), plus the logic to
+translate a per-role stack of specs into concrete engine sizing and overlay it
+onto merged stage configs.
+
+The spec/pattern *types* are deliberately pure data with no engine imports; the
+translation and apply logic lives alongside them and consumes vLLM-Omni's own
+config objects.
+"""
+
+from vllm_omni.config.composable_parallel.aggregation import (
+    AGGREGATION_PATTERN_VARIANTS,
+    AggregationConflictError,
+    AggregationPattern,
+    AllGather,
+    Combine,
+    FanInByStage,
+    GatherDim,
+    StitchPipeline,
+    StitchSpatial,
+    TakeRank,
+    Union,
+)
+from vllm_omni.config.composable_parallel.apply import (
+    StrategyApplyError,
+    StrategyApplyResult,
+    apply_strategy_specs,
+    check_device_layout,
+)
+from vllm_omni.config.composable_parallel.routing import (
+    ROUTING_PATTERN_VARIANTS,
+    Broadcast,
+    DuplicateWithCondTag,
+    PartitionByHash,
+    PipelineMicrobatch,
+    RouteByExpert,
+    RouteByStage,
+    RoutingKeyError,
+    RoutingPattern,
+    ShardSequence,
+    ShardSpatial,
+)
+from vllm_omni.config.composable_parallel.spec import (
+    MESH_AXIS_KINDS,
+    KernelSpec,
+    LayerHookSpec,
+    MeshAxisKind,
+    MeshAxisSpec,
+    SpecMergeConflictError,
+    StrategySpec,
+)
+from vllm_omni.config.composable_parallel.translator import (
+    AxisTranslationError,
+    OmniParallelConfig,
+    translate_strategy_stack,
+)
+
+__all__ = [
+    # spec types
+    "MeshAxisKind",
+    "MESH_AXIS_KINDS",
+    "MeshAxisSpec",
+    "LayerHookSpec",
+    "KernelSpec",
+    "StrategySpec",
+    "SpecMergeConflictError",
+    # routing patterns
+    "RoutingPattern",
+    "RoutingKeyError",
+    "Broadcast",
+    "PartitionByHash",
+    "ShardSequence",
+    "ShardSpatial",
+    "RouteByExpert",
+    "PipelineMicrobatch",
+    "DuplicateWithCondTag",
+    "RouteByStage",
+    "ROUTING_PATTERN_VARIANTS",
+    # aggregation patterns
+    "AggregationPattern",
+    "AggregationConflictError",
+    "TakeRank",
+    "Union",
+    "GatherDim",
+    "AllGather",
+    "StitchSpatial",
+    "StitchPipeline",
+    "Combine",
+    "FanInByStage",
+    "AGGREGATION_PATTERN_VARIANTS",
+    # translation
+    "OmniParallelConfig",
+    "translate_strategy_stack",
+    "AxisTranslationError",
+    # apply
+    "apply_strategy_specs",
+    "StrategyApplyResult",
+    "check_device_layout",
+    "StrategyApplyError",
+]

vllm_omni/config/composable_parallel/aggregation.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Closed aggregation-pattern hierarchy for declarative parallel strategies.
+
+An aggregation pattern describes how the per-worker results of one mesh axis
+are recombined into a single logical result (e.g. take one tensor-parallel
+rank's output, union disjoint data-parallel shards). The hierarchy is *closed*:
+subclasses may only be defined in this module.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class AggregationPattern:
+    """Base class for closed aggregation patterns. Subclass only within this module."""
+
+    def __init_subclass__(cls, **kwargs: object) -> None:
+        super().__init_subclass__(**kwargs)
+        if cls.__module__ != __name__:
+            raise TypeError("AggregationPattern is closed to external subclassing")
+
+    def pattern_kind(self) -> str:
+        raise NotImplementedError
+
+
+class AggregationConflictError(ValueError):
+    """Raised when two aggregation inputs conflict for the same request id."""
+
+
+@dataclass(frozen=True)
+class TakeRank(AggregationPattern):
+    """Select one rank's result (TP — all ranks agree)."""
+
+    rank: int = 0
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.rank, int) or isinstance(self.rank, bool):
+            raise TypeError(f"TakeRank.rank must be int, got {type(self.rank).__name__}")
+
+    def pattern_kind(self) -> str:
+        return "take_rank"
+
+
+@dataclass(frozen=True)
+class Union(AggregationPattern):
+    """Union of disjoint replica results (DP, HSDP)."""
+
+    def pattern_kind(self) -> str:
+        return "union"
+
+
+@dataclass(frozen=True)
+class GatherDim(AggregationPattern):
+    """Gather along a dimension and materialize on one designated rank.
+
+    Default designated rank is axis-rank 0 unless adapter overrides.
+    """
+
+    dim: int = 1
+
+    def pattern_kind(self) -> str:
+        return "gather_dim"
+
+
+@dataclass(frozen=True)
+class AllGather(AggregationPattern):
+    """All ranks materialize the same aggregate result across the axis."""
+
+    dim: int | None = None
+
+    def pattern_kind(self) -> str:
+        return "all_gather"
+
+
+@dataclass(frozen=True)
+class StitchSpatial(AggregationPattern):
+    """Stitch spatial patches (VAE-PP)."""
+
+    grid: tuple[int, int] = (1, 1)
+
+    def pattern_kind(self) -> str:
+        return "stitch_spatial"
+
+
+@dataclass(frozen=True)
+class StitchPipeline(AggregationPattern):
+    """Stitch pipeline stage outputs (PP)."""
+
+    def pattern_kind(self) -> str:
+        return "stitch_pipeline"
+
+
+@dataclass(frozen=True)
+class Combine(AggregationPattern):
+    """Custom reducer combination (CFG). Adapter resolves reducer_id."""
+
+    reducer_id: str = "default"
+
+    def pattern_kind(self) -> str:
+        return "combine"
+
+
+@dataclass(frozen=True)
+class FanInByStage(AggregationPattern):
+    """Fan-in results routed by stage (stage replicas)."""
+
+    def pattern_kind(self) -> str:
+        return "fan_in_by_stage"
+
+
+AGGREGATION_PATTERN_VARIANTS: tuple[type[AggregationPattern], ...] = (
+    TakeRank,
+    Union,
+    GatherDim,
+    AllGather,
+    StitchSpatial,
+    StitchPipeline,
+    Combine,
+    FanInByStage,
+)
+
+ReducerFn = Callable[..., object]

vllm_omni/config/composable_parallel/apply.py

@@ -0,0 +1,264 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Apply declarative parallel strategies onto merged stage configs.
+
+This is the *override-after-merge* seam. ``merge_pipeline_deploy`` first fuses
+the pipeline topology + deploy YAML into a list of ``StageConfig`` objects; this
+module then overlays the engine sizing derived from a per-role
+``StrategySpec`` stack onto those stages.
+
+Design rules (one writer per axis):
+
+* A strategy only writes the axes it actually declares. If a stack has no ``tp``
+  axis, this module never touches ``tensor_parallel_size`` — it does not force a
+  default of 1.
+* **Conflict-on-explicit.** If a knob was already set explicitly by the deploy
+  YAML and the strategy derives a different value, we raise rather than silently
+  override. Equal values are a no-op; unset values are filled.
+* **Pre-spawn device check.** When a stage declares ``devices``, the device
+  count must be consistent with the engine world size (``tp * dp * pp``) and the
+  replica count, so misconfigurations fail at config time instead of at spawn.
+
+``omni_lb_policy`` is *not* a per-stage config knob: omni reads it once at
+orchestrator construction (``AsyncOmniEngine``), not from stage configs. So a
+stage_replica axis's derived policy is surfaced on :class:`StrategyApplyResult`
+for the caller to pass to the orchestrator; it is never silently written into a
+stage where it would be ignored.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+from typing import Any, NoReturn
+
+from vllm.logger import init_logger
+
+from vllm_omni.config.composable_parallel.spec import StrategySpec
+from vllm_omni.config.composable_parallel.translator import (
+    OmniParallelConfig,
+    translate_strategy_stack,
+)
+
+logger = init_logger(__name__)
+
+# Role key -> the StageConfig fields each declared axis writes.
+_ENGINE_FIELD_BY_KIND: dict[str, str] = {
+    "tp": "tensor_parallel_size",
+    "dp": "data_parallel_size",
+    "pp": "pipeline_parallel_size",
+}
+
+RoleKey = str
+
+
+class StrategyApplyError(ValueError):
+    """Error for applying a strategy onto stage configs.
+
+    A single error type (rather than per-cause subclasses) keeps the public
+    surface small and consistent with the rest of the codebase; the specific
+    cause — an unmatched role, a conflict with an explicit deploy value, or a
+    device-count mismatch — is in the message and is logged before the raise so
+    it stays visible even when the type is unavailable to a caller (e.g. across a
+    server boundary).
+    """
+
+
+def _fail(msg: str) -> NoReturn:
+    """Log and raise :class:`StrategyApplyError`."""
+    logger.error("[composable_parallel] %s", msg)
+    raise StrategyApplyError(msg)
+
+
+@dataclass
+class StrategyApplyResult:
+    """Outcome of applying strategies to a stage list.
+
+    ``stages`` is the same list passed in (mutated in place and returned for
+    convenience). ``omni_lb_policy`` is the pipeline-wide load-balancer policy
+    derived from any stage_replica axes (``None`` if none declared one); the
+    caller wires it into the orchestrator, since omni does not read it from
+    stage configs.
+
+    ``per_role_config`` is keyed by the role key the caller supplied — a
+    ``model_stage`` name (str). ``per_stage_config`` is the same information
+    keyed by the resolved integer ``stage_id`` — handy for re-validating a
+    stage after later layers (e.g. CLI overrides) have run.
+    """
+
+    stages: list[Any]
+    omni_lb_policy: str | None = None
+    per_role_config: dict[RoleKey, OmniParallelConfig] = field(default_factory=dict)
+    per_stage_config: dict[int, OmniParallelConfig] = field(default_factory=dict)
+
+
+def _resolve_stage(stages: Sequence[Any], key: RoleKey) -> Any:
+    """Find the single stage matching ``key`` by ``model_stage`` name.
+
+    Role keys are ``model_stage`` names (str) — the same human-readable names
+    that strategy.yaml authors use to refer to a stage (e.g. ``"thinker"``).
+    """
+    matches = [s for s in stages if getattr(s, "model_stage", None) == key]
+    descriptor = f"model_stage={key!r}"
+
+    if not matches:
+        available = ", ".join(f"{getattr(s, 'model_stage', '?')!r}(id={getattr(s, 'stage_id', '?')})" for s in stages)
+        _fail(f"strategy role {descriptor} did not match any stage; available stages: {available}")
+    if len(matches) > 1:
+        _fail(f"strategy role {descriptor} is ambiguous; it matched {len(matches)} stages")
+    return matches[0]
+
+
+def _set_explicit(engine_args: dict[str, Any], field_name: str, value: Any, *, role: RoleKey) -> None:
+    """Write ``field_name`` with conflict-on-explicit semantics.
+
+    ``engine_args`` is the raw mapping loaded straight from the deploy YAML
+    (before the engine-args dataclass is constructed), so a key is *present* only
+    when the YAML set it explicitly. We therefore key off membership rather than
+    ``.get``: an explicit ``None`` in the YAML is a real, present value and must
+    not be conflated with a key miss (which ``.get`` would report as ``None``).
+    """
+    if field_name in engine_args and engine_args[field_name] != value:
+        existing = engine_args[field_name]
+        _fail(
+            f"role {role!r}: strategy derives {field_name}={value} but the deploy config already "
+            f"set {field_name}={existing}. Remove one of them (the strategy is the single writer "
+            "for the axes it declares)."
+        )
+    engine_args[field_name] = value
+
+
+def _set_num_replicas(runtime: dict[str, Any], value: int, *, role: RoleKey) -> None:
+    """Write ``num_replicas`` with conflict-on-explicit (treating 1 as unset)."""
+    existing = runtime.get("num_replicas", 1)
+    if existing not in (1, value):
+        _fail(
+            f"role {role!r}: strategy derives num_replicas={value} but the deploy config already "
+            f"set num_replicas={existing}. Remove one of them."
+        )
+    runtime["num_replicas"] = value
+
+
+def _parse_device_count(devices: Any) -> int | None:
+    """Return the number of device ids in a ``devices`` value, or None."""
+    if devices is None:
+        return None
+    if isinstance(devices, (list, tuple)):
+        return len([d for d in devices])
+    text = str(devices).strip()
+    if not text:
+        return None
+    return len([d for d in text.split(",") if d.strip()])
+
+
+def check_device_layout(
+    devices: Any,
+    *,
+    tensor_parallel_size: int,
+    data_parallel_size: int,
+    pipeline_parallel_size: int,
+    num_replicas: int,
+    role: RoleKey,
+) -> None:
+    """Validate a stage's device count against its world size (× replicas in pool mode).
+
+    Raises :class:`StrategyApplyError` when the declared ``devices`` count is
+    neither the per-replica world size nor the full ``num_replicas`` pool.
+    Exposed publicly so the resolved (post-CLI) layout can be re-checked after
+    later override layers, not just the strategy-derived snapshot.
+    """
+    count = _parse_device_count(devices)
+    if count is None:
+        return
+    world = int(tensor_parallel_size) * int(data_parallel_size) * int(pipeline_parallel_size)
+    replicas = int(num_replicas or 1)
+    # Two valid shapes: a single per-replica template (== world) or the full
+    # pool across replicas (== replicas * world). This mirrors omni's own
+    # per-replica device splitter.
+    valid = {world, replicas * world}
+    if count not in valid:
+        _fail(
+            f"role {role!r}: declared {count} device id(s) but the strategy world size is {world} "
+            f"(tp={tensor_parallel_size} * dp={data_parallel_size} * pp={pipeline_parallel_size}). "
+            f"Provide either {world} (per-replica template) or {replicas * world} "
+            f"(num_replicas={replicas} pool)."
+        )
+
+
+def _check_devices(runtime: dict[str, Any], cfg: OmniParallelConfig, *, role: RoleKey) -> None:
+    """Pre-spawn check on the strategy-derived snapshot (see :func:`check_device_layout`)."""
+    check_device_layout(
+        runtime.get("devices"),
+        tensor_parallel_size=cfg.tensor_parallel_size,
+        data_parallel_size=cfg.data_parallel_size,
+        pipeline_parallel_size=cfg.pipeline_parallel_size,
+        num_replicas=int(runtime.get("num_replicas", 1) or 1),
+        role=role,
+    )
+
+
+def _apply_to_stage(stage: Any, cfg: OmniParallelConfig, *, role: RoleKey) -> None:
+    engine_args = stage.yaml_engine_args
+    runtime = stage.yaml_runtime
+
+    declared = set(cfg.l1_owners.keys())
+    for kind, field_name in _ENGINE_FIELD_BY_KIND.items():
+        if kind in declared:
+            _set_explicit(engine_args, field_name, getattr(cfg, field_name), role=role)
+    if "ep" in declared and cfg.enable_expert_parallel:
+        _set_explicit(engine_args, "enable_expert_parallel", True, role=role)
+    if "stage_replica" in declared:
+        _set_num_replicas(runtime, cfg.stage_replica_size, role=role)
+
+    _check_devices(runtime, cfg, role=role)
+
+
+def apply_strategy_specs(
+    stages: list[Any],
+    strategy_specs: Mapping[RoleKey, Sequence[StrategySpec]],
+) -> StrategyApplyResult:
+    """Overlay per-role strategy specs onto a merged stage list.
+
+    Args:
+        stages: the ``list[StageConfig]`` returned by ``merge_pipeline_deploy``.
+        strategy_specs: maps a role (a ``model_stage`` name, e.g. ``"thinker"``)
+            to that stage's stack of ``StrategySpec`` (one per declared mesh
+            axis). Role keys are ``model_stage`` names (str).
+
+    Returns:
+        A :class:`StrategyApplyResult` with the mutated stages and the derived
+        pipeline-wide ``omni_lb_policy`` (if any stage_replica axis set one).
+
+    Raises:
+        StrategyApplyError: a role matched zero or multiple stages, a derived
+            value conflicts with an explicit deploy value (including two roles
+            deriving different ``omni_lb_policy`` values), or a stage's device
+            count is inconsistent with its world size.
+        AxisTranslationError: the spec stack is invalid/unsupported.
+        NotImplementedError: the spec stack requests routing not built yet.
+    """
+    result = StrategyApplyResult(stages=stages)
+    lb_policy: str | None = None
+    lb_owner: RoleKey | None = None
+
+    for key, specs in strategy_specs.items():
+        stage = _resolve_stage(stages, key)
+        cfg = translate_strategy_stack(specs)
+        _apply_to_stage(stage, cfg, role=key)
+        result.per_role_config[key] = cfg
+        stage_id = getattr(stage, "stage_id", None)
+        if stage_id is not None:
+            result.per_stage_config[int(stage_id)] = cfg
+
+        if cfg.stage_replica_size > 1 and cfg.omni_lb_policy is not None:
+            if lb_policy is not None and lb_policy != cfg.omni_lb_policy:
+                _fail(
+                    f"role {key!r} derives omni_lb_policy={cfg.omni_lb_policy!r} but role "
+                    f"{lb_owner!r} already derived {lb_policy!r}. The load-balancer policy is "
+                    "pipeline-wide; only one value is allowed."
+                )
+            lb_policy = cfg.omni_lb_policy
+            lb_owner = key
+
+    result.omni_lb_policy = lb_policy
+    return result

vllm_omni/config/composable_parallel/routing.py

@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Closed routing-pattern hierarchy for declarative parallel strategies.
+
+A routing pattern describes how a request batch is distributed across the
+workers of one mesh axis (e.g. broadcast to every tensor-parallel rank, hash
+partition across data-parallel ranks). The hierarchy is *closed*: subclasses
+may only be defined in this module, so the set of patterns is a fixed,
+exhaustively matchable contract.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Literal
+
+
+@dataclass(frozen=True)
+class RoutingPattern:
+    """Base class for closed routing patterns. Subclass only within this module."""
+
+    def __init_subclass__(cls, **kwargs: object) -> None:
+        super().__init_subclass__(**kwargs)
+        if cls.__module__ != __name__:
+            raise TypeError("RoutingPattern is closed to external subclassing")
+
+    def pattern_kind(self) -> str:
+        raise NotImplementedError
+
+
+class RoutingKeyError(KeyError):
+    """Raised when a routing key is missing from a request."""
+
+    def __init__(self, request_id: object, missing_key: str) -> None:
+        self.request_id = request_id
+        self.missing_key = missing_key
+        super().__init__(f"Missing routing key {missing_key!r} for request_id={request_id!r}")
+
+
+@dataclass(frozen=True)
+class Broadcast(RoutingPattern):
+    """Every worker in the group sees the whole batch (TP, dense EP)."""
+
+    def pattern_kind(self) -> str:
+        return "broadcast"
+
+
+@dataclass(frozen=True)
+class PartitionByHash(RoutingPattern):
+    """Stable hash partition by request key (DP, HSDP)."""
+
+    key: str = "request_id"
+
+    def pattern_kind(self) -> str:
+        return "partition_by_hash"
+
+
+@dataclass(frozen=True)
+class ShardSequence(RoutingPattern):
+    """Shard along a sequence dimension (SP-Ulysses, SP-Ring, CP)."""
+
+    dim: int = 1
+
+    def pattern_kind(self) -> str:
+        return "shard_sequence"
+
+
+@dataclass(frozen=True)
+class ShardSpatial(RoutingPattern):
+    """Shard along a spatial grid (VAE-PP)."""
+
+    grid: tuple[int, int] = (1, 1)
+
+    def pattern_kind(self) -> str:
+        return "shard_spatial"
+
+
+@dataclass(frozen=True)
+class RouteByExpert(RoutingPattern):
+    """Sparse MoE expert routing."""
+
+    router_id: str = "default"
+
+    def pattern_kind(self) -> str:
+        return "route_by_expert"
+
+
+@dataclass(frozen=True)
+class PipelineMicrobatch(RoutingPattern):
+    """Pipeline parallelism microbatch routing (PP)."""
+
+    microbatch_size: int = 1
+
+    def pattern_kind(self) -> str:
+        return "pipeline_microbatch"
+
+
+@dataclass(frozen=True)
+class DuplicateWithCondTag(RoutingPattern):
+    """Duplicate batch with a conditional tag (CFG)."""
+
+    flag_name: str = "guidance_branch"
+
+    def pattern_kind(self) -> str:
+        return "duplicate_with_cond_tag"
+
+
+@dataclass(frozen=True)
+class RouteByStage(RoutingPattern):
+    """Route requests by pipeline stage policy (stage replicas).
+
+    The ``random``/``round_robin``/``least_queue`` policies describe stateless
+    load balancing that a runtime (e.g. an engine's built-in DP load balancer,
+    or the omni StagePool balancer) can realize. ``hash`` describes
+    deterministic, key-stable routing that a load balancer cannot honor on its
+    own (see ``PartitionByHash``).
+    """
+
+    routing_policy: Literal["random", "round_robin", "least_queue", "hash"] = "round_robin"
+
+    def pattern_kind(self) -> str:
+        return "route_by_stage"
+
+
+ROUTING_PATTERN_VARIANTS: tuple[type[RoutingPattern], ...] = (
+    Broadcast,
+    PartitionByHash,
+    ShardSequence,
+    ShardSpatial,
+    RouteByExpert,
+    PipelineMicrobatch,
+    DuplicateWithCondTag,
+    RouteByStage,
+)
+
+# Optional callable slot for Combine-style custom routing reducers in adapters.
+RouterFn = Callable[..., object]

vllm_omni/config/composable_parallel/spec.py

@@ -0,0 +1,136 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Core declarative types for parallel-strategy specification.
+
+A :class:`StrategySpec` declares one parallelism scheme: the mesh axis it sizes
+(kind + degree), how a batch is routed across that axis, and how per-worker
+results are aggregated back. A *stack* of specs (one per axis) fully describes a
+stage's parallel layout in a runtime-agnostic, data-only form that can be
+translated into concrete engine sizing.
+
+The mesh-axis kinds below enumerate every parallelism dimension the contract can
+describe. Only a subset is wired end-to-end today (see ``translator.py``):
+
+* **Wired** (translatable now): ``tp``, ``dp``, ``pp``, ``ep``,
+  ``stage_replica``.
+* **Reserved** (declared in the type system but not yet translatable — the
+  translator raises ``AxisTranslationError`` for them): ``sp_ulysses``,
+  ``sp_ring``, ``cfg``, ``vae_pp``, ``hsdp``, ``stage_pp``, ``cp``.
+
+Reserved kinds fail fast at translation time rather than silently doing nothing,
+so declaring one is an explicit "not yet" rather than a no-op.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Literal
+
+from vllm_omni.config.composable_parallel.aggregation import AggregationPattern
+from vllm_omni.config.composable_parallel.routing import RoutingPattern
+
+
+class MeshAxisKind(str, Enum):
+    """Enumerates every parallelism dimension the strategy contract can describe.
+
+    Subclassing ``str`` keeps existing string comparisons and dict keys working
+    (``MeshAxisKind.TP == "tp"`` and both hash the same), while making the type
+    refactor-safe: a typo'd kind is an ``AttributeError`` on the enum rather than
+    a silently-wrong bare string.
+    """
+
+    TP = "tp"
+    DP = "dp"
+    PP = "pp"
+    EP = "ep"
+    SP_ULYSSES = "sp_ulysses"
+    SP_RING = "sp_ring"
+    CFG = "cfg"
+    VAE_PP = "vae_pp"
+    HSDP = "hsdp"
+    STAGE_PP = "stage_pp"
+    STAGE_REPLICA = "stage_replica"
+    CP = "cp"
+
+
+# Bare string values, kept as a plain tuple so membership checks and error
+# messages read as the kind names (``"tp"``) rather than enum reprs. ``str``-enum
+# equality means both ``"tp"`` and ``MeshAxisKind.TP`` satisfy ``in`` checks.
+MESH_AXIS_KINDS: tuple[str, ...] = tuple(k.value for k in MeshAxisKind)
+
+HookCategory = Literal[
+    "linear",
+    "attention_pre",
+    "attention_post",
+    "ffn_pre",
+    "ffn_post",
+    "moe_dispatch",
+    "other",
+]
+
+HOOK_CATEGORY_ORDER: dict[HookCategory, int] = {
+    "linear": 0,
+    "attention_pre": 1,
+    "attention_post": 2,
+    "ffn_pre": 3,
+    "ffn_post": 4,
+    "moe_dispatch": 5,
+    "other": 6,
+}
+
+
+class SpecMergeConflictError(ValueError):
+    """Raised when merged hook/kernel specs have conflicting declarations."""
+
+
+@dataclass(frozen=True)
+class MeshAxisSpec:
+    """Declares one axis of a process mesh (kind + size)."""
+
+    kind: MeshAxisKind
+    size: int
+
+    def __post_init__(self) -> None:
+        if self.kind not in MESH_AXIS_KINDS:
+            raise ValueError(f"MeshAxisSpec.kind must be one of {MESH_AXIS_KINDS}, got {self.kind!r}")
+        if self.size <= 0:
+            raise ValueError(f"MeshAxisSpec.size must be > 0, got {self.size}")
+
+
+@dataclass(frozen=True)
+class LayerHookSpec:
+    """L2 hook slot referenced by StrategySpec; resolved by the model walker."""
+
+    hook_id: str
+    target: str | None = None
+    category: HookCategory = "other"
+    priority: int = 0
+    axis_index: int = 0
+
+
+@dataclass(frozen=True)
+class KernelSpec:
+    """L3 kernel slot referenced by StrategySpec; resolved by the kernel registry."""
+
+    kernel_id: str
+    target: str | None = None
+    category: HookCategory = "other"
+    priority: int = 0
+    axis_index: int = 0
+    group_axis_kind: MeshAxisKind | None = None
+    requires_collective: bool = False
+
+
+@dataclass(frozen=True)
+class StrategySpec:
+    """Declarative contract for one parallelism scheme (data + hook/kernel slots)."""
+
+    name: str
+    mesh_axis: MeshAxisSpec
+    routing: RoutingPattern
+    aggregation: AggregationPattern
+    layer_hook_specs: tuple[LayerHookSpec, ...] = ()
+    kernel_specs: tuple[KernelSpec, ...] = ()
+    shard_extension: Mapping[str, Any] = field(default_factory=dict)

vllm_omni/config/composable_parallel/strategy_loader.py

@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Load per-role parallel strategies from a strategy file.
+
+The strategy file is a small, optional companion to the deploy config. It maps a
+logical role (``model_stage``, e.g. ``thinker``) to a list of axis declarations,
+each of which becomes one :class:`StrategySpec`. YAML is the only format wired up
+today, but the module/symbol names stay format-agnostic so a different source can
+be added without renaming the public surface::
+
+    strategies:
+      thinker:
+        - axis: tp
+          size: 2
+      talker:
+        - axis: stage_replica
+          size: 2
+          routing: round_robin   # random | round_robin | least_queue
+      code2wav:
+        - axis: stage_replica
+          size: 2
+          routing: round_robin
+
+Per-axis fields:
+
+* ``axis`` (required): mesh-axis kind (``tp``, ``dp``, ``pp``, ``ep``,
+  ``stage_replica`` are translatable today).
+* ``size`` (required): axis degree (> 0).
+* ``routing`` (optional): stateless policy for ``dp`` / ``stage_replica`` axes
+  (``random`` / ``round_robin`` / ``least_queue``). Ignored for axes whose
+  routing is fixed (``tp`` broadcast, ``pp`` microbatch).
+* ``l1_owner`` (optional): ``engine`` or ``delegated``; overrides the default
+  owner for the axis (recorded in ``shard_extension``).
+* ``name`` (optional): a label for the spec; defaults to the axis kind.
+
+Routing/aggregation patterns are filled in from sensible per-kind defaults so
+the file stays compact; the translator validates the combination.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+from vllm_omni.config.composable_parallel.aggregation import (
+    AggregationPattern,
+    FanInByStage,
+    StitchPipeline,
+    TakeRank,
+    Union,
+)
+from vllm_omni.config.composable_parallel.routing import (
+    Broadcast,
+    PipelineMicrobatch,
+    RouteByStage,
+    RoutingPattern,
+)
+from vllm_omni.config.composable_parallel.spec import MeshAxisSpec, StrategySpec
+from vllm_omni.config.yaml_util import load_yaml_config, to_dict
+
+# Per-kind defaults for routing + aggregation, so the file only needs axis+size.
+_ROUTING_POLICY_KINDS = frozenset({"dp", "stage_replica"})
+
+
+class StrategyLoadError(ValueError):
+    """Raised when a strategy file is malformed."""
+
+
+def _default_routing(kind: str, routing_policy: str | None) -> RoutingPattern:
+    if kind in _ROUTING_POLICY_KINDS:
+        return RouteByStage(routing_policy=routing_policy or "round_robin")  # type: ignore[arg-type]
+    if kind == "pp":
+        return PipelineMicrobatch()
+    # tp, ep, and any other kind default to broadcast (the translator rejects
+    # kinds it cannot handle before it ever inspects routing).
+    return Broadcast()
+
+
+def _default_aggregation(kind: str) -> AggregationPattern:
+    if kind == "stage_replica":
+        return FanInByStage()
+    if kind == "pp":
+        return StitchPipeline()
+    if kind in ("dp", "ep"):
+        return Union()
+    return TakeRank()
+
+
+def _build_spec(role: str, entry: Mapping[str, Any]) -> StrategySpec:
+    if "axis" not in entry:
+        raise StrategyLoadError(f"role {role!r}: each strategy entry needs an 'axis' field")
+    if "size" not in entry:
+        raise StrategyLoadError(f"role {role!r}: strategy entry for axis {entry['axis']!r} needs a 'size'")
+
+    kind = str(entry["axis"])
+    try:
+        size = int(entry["size"])
+    except (TypeError, ValueError) as exc:
+        raise StrategyLoadError(f"role {role!r}: axis {kind!r} size must be an integer, got {entry['size']!r}") from exc
+
+    routing_policy = entry.get("routing")
+    if routing_policy is not None and kind not in _ROUTING_POLICY_KINDS:
+        raise StrategyLoadError(
+            f"role {role!r}: axis {kind!r} does not accept a 'routing' policy (only {sorted(_ROUTING_POLICY_KINDS)} do)"
+        )
+
+    shard_extension: dict[str, Any] = {}
+    l1_owner = entry.get("l1_owner")
+    if l1_owner is not None:
+        shard_extension["l1_owner"] = str(l1_owner)
+
+    return StrategySpec(
+        name=str(entry.get("name", kind)),
+        mesh_axis=MeshAxisSpec(kind=kind, size=size),  # type: ignore[arg-type]
+        routing=_default_routing(kind, routing_policy),
+        aggregation=_default_aggregation(kind),
+        shard_extension=shard_extension,
+    )
+
+
+def parse_strategy_specs(data: Mapping[str, Any]) -> dict[str, list[StrategySpec]]:
+    """Parse an already-loaded strategy mapping into per-role spec stacks."""
+    strategies = data.get("strategies", data)
+    if not isinstance(strategies, Mapping):
+        raise StrategyLoadError("strategy file must contain a 'strategies' mapping of role -> list of axes")
+
+    result: dict[str, list[StrategySpec]] = {}
+    for role, entries in strategies.items():
+        if not isinstance(entries, (list, tuple)):
+            raise StrategyLoadError(
+                f"role {role!r}: expected a list of axis declarations, got {type(entries).__name__}"
+            )
+        specs: list[StrategySpec] = []
+        for entry in entries:
+            if not isinstance(entry, Mapping):
+                raise StrategyLoadError(
+                    f"role {role!r}: each axis declaration must be a mapping, got {type(entry).__name__}"
+                )
+            specs.append(_build_spec(str(role), dict(entry)))
+        result[str(role)] = specs
+    return result
+
+
+def load_strategy_specs(path: str) -> dict[str, list[StrategySpec]]:
+    """Load and parse a strategy file into per-role spec stacks."""
+    data = to_dict(load_yaml_config(path))
+    if not isinstance(data, Mapping):
+        raise StrategyLoadError(f"strategy file {path!r} did not parse to a mapping")
+    return parse_strategy_specs(data)

vllm_omni/config/composable_parallel/translator.py

@@ -0,0 +1,378 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Translate a ``StrategySpec`` stack into vLLM-Omni parallel sizing.
+
+Read a stack of strategy specs (one per mesh axis) and work out the
+tensor/data/pipeline parallel sizes, the dense-EP flag, and the number of stage
+replicas, plus who owns each axis's request routing. The result is a plain,
+CPU-computable sizing struct (:class:`OmniParallelConfig`) keyed by the real
+``OmniEngineArgs`` / ``EngineArgs`` field names, so the deploy layer can splat it
+onto a stage's engine args.
+
+The distinction this module enforces — engine data parallelism (a true vLLM
+intra-engine world dimension) vs. omni stage replicas (independent engines fanned
+out by omni's coordinator) — is documented on the axis validators that depend on
+it (see :func:`_validate_dp` and :func:`_stage_replica_lb_policy`). Routing we do
+not support yet (key-stable / affinity routing) raises ``NotImplementedError``;
+any other invalid spec raises :class:`AxisTranslationError`.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+from typing import Literal, NoReturn, cast, get_args
+
+from vllm.logger import init_logger
+
+from vllm_omni.config.composable_parallel.routing import (
+    Broadcast,
+    PartitionByHash,
+    PipelineMicrobatch,
+    RouteByStage,
+    RoutingPattern,
+)
+from vllm_omni.config.composable_parallel.spec import MeshAxisKind, StrategySpec
+
+logger = init_logger(__name__)
+
+# Who owns an axis's request routing. A closed string type so an unexpected raw
+# value is caught statically and at runtime (``_VALID_L1_OWNERS`` is derived from
+# it, keeping one source of truth) rather than silently flowing through.
+L1Owner = Literal["delegated", "engine"]
+
+# DP, TP, PP, (dense) EP and stage_replica translate today. DP/TP/PP are true
+# world dimensions vLLM realizes intra-engine; dense EP is a flag over the
+# existing TP*DP ranks; stage_replica is an omni-coordinator-level fan-out of
+# independent engine replicas (NOT a vLLM world dimension). The remaining kinds
+# (sequence/context parallel, CFG, VAE pipelines, ...) need per-layer hooks or
+# custom collectives and arrive in later stages.
+_SUPPORTED_KINDS: tuple[MeshAxisKind, ...] = ("dp", "tp", "pp", "ep", "stage_replica")
+
+# Which EngineArgs field each axis kind *sizes*. EP and stage_replica are
+# deliberately absent: EP is not an independent world dimension but an
+# ``enable_expert_parallel`` flag whose degree equals tp*dp; stage_replica is a
+# per-stage deploy ``num_replicas`` count, not an intra-engine world dimension.
+_AXIS_TO_ENGINE_FIELD: dict[MeshAxisKind, str] = {
+    "tp": "tensor_parallel_size",
+    "dp": "data_parallel_size",
+    "pp": "pipeline_parallel_size",
+}
+
+# Default L1 owner per axis kind. ``stage_replica`` is the only delegated axis:
+# omni's coordinator/StagePool owns routing across replicas. ``dp`` is engine
+# data parallelism, realized by vLLM's own intra-engine DP load balancer.
+_DEFAULT_L1_OWNER: dict[MeshAxisKind, L1Owner] = {
+    "dp": "engine",
+    "tp": "engine",
+    "pp": "engine",
+    "ep": "engine",
+    "stage_replica": "delegated",
+}
+
+# How a RouteByStage policy maps onto omni's load balancer policy string. These
+# are the stateless options omni can actually do; note there's no "hash" here,
+# because omni has no key-stable balancer.
+_STAGE_POLICY_TO_OMNI_LB: dict[str, str] = {
+    "random": "random",
+    "round_robin": "round-robin",
+    "least_queue": "least-queue-length",
+}
+
+_VALID_L1_OWNERS: frozenset[str] = frozenset(get_args(L1Owner))
+
+
+class AxisTranslationError(ValueError):
+    """Error for an invalid or unsupported strategy spec.
+
+    A single error type (rather than a tree of subclasses) keeps the public
+    surface small and consistent with the rest of the codebase; the specific
+    cause is in the message and is logged before the raise so it is visible even
+    when the type is unavailable to a caller (e.g. across a server boundary).
+
+    Strategies that are *valid but not built yet* — key-stable / affinity routing
+    today — raise ``NotImplementedError`` instead, to distinguish "we haven't
+    implemented this" from "your config is wrong".
+    """
+
+
+def _fail(msg: str) -> NoReturn:
+    """Log and raise :class:`AxisTranslationError` for an invalid/unsupported spec."""
+    logger.error("[composable_parallel] %s", msg)
+    raise AxisTranslationError(msg)
+
+
+def _not_implemented(msg: str) -> NoReturn:
+    """Log and raise ``NotImplementedError`` for a valid-but-unbuilt strategy."""
+    logger.error("[composable_parallel] %s", msg)
+    raise NotImplementedError(msg)
+
+
+@dataclass(frozen=True)
+class OmniParallelConfig:
+    """Result of translating a spec stack into omni parallel sizing."""
+
+    tensor_parallel_size: int = 1
+    data_parallel_size: int = 1
+    pipeline_parallel_size: int = 1
+    # Dense expert parallelism: a flag over the TP*DP ranks, not a world axis.
+    enable_expert_parallel: bool = False
+    # Number of independent omni stage replicas (the ``stage_replica`` axis).
+    # This is NOT a vLLM world dimension; it is the per-stage deploy
+    # ``num_replicas`` count. 1 means a single (un-replicated) engine.
+    stage_replica_size: int = 1
+    # The omni StagePool LB policy string for the (delegated) stage_replica
+    # axis, if any. Only meaningful when ``stage_replica_size > 1``.
+    omni_lb_policy: str | None = None
+    # axis kind -> resolved L1 owner ("delegated" | "engine").
+    l1_owners: Mapping[MeshAxisKind, L1Owner] = field(default_factory=dict)
+
+    @property
+    def world_size(self) -> int:
+        # EP and stage_replica are intentionally excluded. EP reuses the TP*DP
+        # ranks rather than adding a dimension; stage_replica spins up separate
+        # engines (each its own world), not extra ranks in this engine's group.
+        # This product is what vLLM calls ``world_size_across_dp`` (tp * pp * dp).
+        return self.tensor_parallel_size * self.data_parallel_size * self.pipeline_parallel_size
+
+    @property
+    def delegated_axes(self) -> tuple[MeshAxisKind, ...]:
+        return tuple(kind for kind, owner in self.l1_owners.items() if owner == "delegated")
+
+    def as_engine_kwargs(self) -> dict[str, object]:
+        """Return per-stage kwargs keyed by real OmniEngineArgs/EngineArgs field names.
+
+        Two derived values are intentionally *not* emitted here because they are
+        not per-stage engine args:
+
+        * ``stage_replica_size`` — a per-stage deploy ``num_replicas`` knob
+          (``StageDeployConfig``); the deploy layer consumes it separately.
+        * ``omni_lb_policy`` — a pipeline-wide load-balancer policy the engine
+          reads once at construction (see ``StrategyApplyResult.omni_lb_policy``);
+          it is applied at the orchestrator level, not folded into per-stage args.
+        """
+        kwargs: dict[str, object] = {
+            "tensor_parallel_size": self.tensor_parallel_size,
+            "data_parallel_size": self.data_parallel_size,
+            "pipeline_parallel_size": self.pipeline_parallel_size,
+        }
+        if self.enable_expert_parallel:
+            kwargs["enable_expert_parallel"] = True
+        return kwargs
+
+
+def _is_affinity_dp_routing(routing: RoutingPattern) -> bool:
+    """True when DP routing demands key-stable (hash) placement."""
+    if isinstance(routing, PartitionByHash):
+        return True
+    if isinstance(routing, RouteByStage) and routing.routing_policy == "hash":
+        return True
+    return False
+
+
+def _resolve_l1_owner(spec: StrategySpec) -> L1Owner:
+    kind = spec.mesh_axis.kind
+    declared_owner = spec.shard_extension.get("l1_owner")
+    owner = declared_owner
+    if owner is None:
+        owner = _DEFAULT_L1_OWNER.get(kind, "engine")
+        if kind not in _DEFAULT_L1_OWNER:
+            logger.debug(
+                "[composable_parallel] axis %r has unknown kind %r with no default l1_owner; falling back to %r",
+                spec.name,
+                kind,
+                owner,
+            )
+        else:
+            logger.debug(
+                "[composable_parallel] axis %r (kind %r) declared no l1_owner; using default %r",
+                spec.name,
+                kind,
+                owner,
+            )
+    if owner not in _VALID_L1_OWNERS:
+        # Distinguish a bad value supplied via shard_extension from a bad default
+        # indexed out of _DEFAULT_L1_OWNER, so the source of the invalid owner is
+        # debuggable rather than swallowed.
+        source = "shard_extension" if declared_owner is not None else "_DEFAULT_L1_OWNER"
+        logger.debug(
+            "[composable_parallel] axis %r (kind %r) resolved invalid l1_owner %r from %s; valid owners are %s",
+            spec.name,
+            kind,
+            owner,
+            source,
+            sorted(_VALID_L1_OWNERS),
+        )
+        _fail(f"axis {kind!r} has invalid l1_owner {owner!r}; expected one of {sorted(_VALID_L1_OWNERS)}")
+    return cast(L1Owner, owner)
+
+
+def _validate_dp(spec: StrategySpec, owner: L1Owner) -> None:
+    """Validate an engine data-parallel axis.
+
+    DP is realized intra-engine by vLLM's own DP load balancer, so it is
+    engine-owned and emits no ``omni_lb_policy`` (that string configures omni's
+    *replica* balancer, a different layer). Key-stable (hash) request affinity is
+    not something vLLM's DP LB guarantees, so it is rejected rather than silently
+    dropped.
+    """
+    if owner != "engine":
+        _fail(
+            f"dp axis {spec.name!r} is engine data parallelism realized intra-engine by "
+            f"vLLM's DP load balancer; l1_owner must be 'engine', got {owner!r}. For "
+            "omni-coordinator-level request fan-out across independent replicas, use a "
+            "'stage_replica' axis."
+        )
+    if _is_affinity_dp_routing(spec.routing):
+        _not_implemented(
+            f"dp axis {spec.name!r} requests key-stable (hash) routing, which vLLM's "
+            "intra-engine DP load balancer does not guarantee — not supported yet. Use "
+            "RouteByStage(random|round_robin|least_queue) for stateless DP balancing."
+        )
+    if not isinstance(spec.routing, RouteByStage):
+        _fail(
+            f"dp axis {spec.name!r} expects RouteByStage(random|round_robin|least_queue) routing, "
+            f"got {type(spec.routing).__name__}"
+        )
+    if spec.routing.routing_policy not in _STAGE_POLICY_TO_OMNI_LB:
+        # Recognized-but-unimplemented routing (key-stable/hash) is already
+        # handled above via _is_affinity_dp_routing -> _not_implemented. Anything
+        # left here is an unknown/invalid policy value (e.g. a typo from YAML),
+        # i.e. invalid input -> AxisTranslationError, not NotImplementedError.
+        _fail(
+            f"dp axis {spec.name!r} has invalid routing_policy "
+            f"{spec.routing.routing_policy!r}; expected one of "
+            f"{sorted(_STAGE_POLICY_TO_OMNI_LB)}."
+        )
+
+
+def _stage_replica_lb_policy(spec: StrategySpec, owner: L1Owner) -> str:
+    """Return the omni StagePool LB policy for a delegated stage_replica axis.
+
+    A ``stage_replica`` axis is *not* a vLLM world dimension: it stands up N
+    independent engine replicas of one pipeline stage, coordinated by the omni
+    coordinator and balanced over a StagePool with stateless policies only
+    (random / round-robin / least-queue-length). It maps to the per-stage
+    ``num_replicas`` count plus the pipeline-level ``omni_lb_policy`` string, so
+    its ``l1_owner`` must be ``"delegated"`` (omni owns the routing). Key-stable
+    (hash) routing is rejected because omni has no key-stable balancer yet.
+    """
+    if owner != "delegated":
+        _fail(
+            f"stage_replica axis {spec.name!r} must be 'delegated' to omni's StagePool load "
+            f"balancer; got owner {owner!r}. Replica routing is owned by omni's coordinator."
+        )
+
+    routing = spec.routing
+    if _is_affinity_dp_routing(routing):
+        _not_implemented(
+            f"stage_replica axis {spec.name!r} requests key-stable (hash) routing, which needs a "
+            "dedicated load balancer — not implemented yet. Use "
+            "RouteByStage(random|round_robin|least_queue) to delegate to omni's load balancer."
+        )
+    if not isinstance(routing, RouteByStage):
+        _fail(
+            f"stage_replica axis {spec.name!r} expects RouteByStage(random|round_robin|least_queue) "
+            f"routing, got {type(routing).__name__}"
+        )
+    policy = _STAGE_POLICY_TO_OMNI_LB.get(routing.routing_policy)
+    if policy is None:
+        _not_implemented(
+            f"stage_replica axis {spec.name!r} routing_policy {routing.routing_policy!r} has no omni LB policy"
+        )
+    return policy
+
+
+def _validate_tp(spec: StrategySpec, owner: L1Owner) -> None:
+    if not isinstance(spec.routing, Broadcast):
+        _fail(f"tp axis {spec.name!r} expects Broadcast routing, got {type(spec.routing).__name__}")
+    if owner != "engine":
+        _fail(f"tp axis {spec.name!r} is realized intra-engine; l1_owner must be 'engine', got {owner!r}")
+
+
+def _validate_pp(spec: StrategySpec, owner: L1Owner) -> None:
+    if not isinstance(spec.routing, PipelineMicrobatch):
+        _fail(f"pp axis {spec.name!r} expects PipelineMicrobatch routing, got {type(spec.routing).__name__}")
+    if owner != "engine":
+        _fail(f"pp axis {spec.name!r} is realized intra-engine; l1_owner must be 'engine', got {owner!r}")
+
+
+def _validate_ep(spec: StrategySpec, owner: L1Owner) -> None:
+    # Dense EP: every rank still sees the whole batch, experts are sharded
+    # across ranks inside the engine (sparse MoE all-to-all is a later stage).
+    if not isinstance(spec.routing, Broadcast):
+        _fail(
+            f"ep axis {spec.name!r} expects Broadcast routing (dense expert parallel), "
+            f"got {type(spec.routing).__name__}"
+        )
+    if owner != "engine":
+        _fail(f"ep axis {spec.name!r} is realized intra-engine; l1_owner must be 'engine', got {owner!r}")
+
+
+def translate_strategy_stack(specs: Sequence[StrategySpec]) -> OmniParallelConfig:
+    """Translate a spec stack into an ``OmniParallelConfig``.
+
+    Supported kinds: dp (engine data parallel), tp, pp, (dense) ep, and
+    stage_replica (omni replicas). Raises ``NotImplementedError`` for deferred
+    (affinity / key-stable) routing, and :class:`AxisTranslationError` for any
+    other invalid spec (a kind not yet translatable, a repeated kind, an owner
+    incompatible with the axis kind, or unsupported routing). The EP degree must
+    equal tensor_parallel_size * data_parallel_size.
+    """
+    sizes: dict[str, int] = {"tensor_parallel_size": 1, "data_parallel_size": 1, "pipeline_parallel_size": 1}
+    owners: dict[MeshAxisKind, L1Owner] = {}
+    omni_lb_policy: str | None = None
+    stage_replica_size = 1
+    enable_expert_parallel = False
+    ep_size: int | None = None
+
+    for spec in specs:
+        kind = spec.mesh_axis.kind
+        if kind not in _SUPPORTED_KINDS:
+            _fail(
+                f"axis kind {kind!r} is not translatable yet (supported: {list(_SUPPORTED_KINDS)}); "
+                "it is designed-for and lands in a later stage"
+            )
+        if kind in owners:
+            _fail(f"axis kind {kind!r} appears more than once in the spec stack")
+
+        owner = _resolve_l1_owner(spec)
+        if kind == "dp":
+            _validate_dp(spec, owner)
+        elif kind == "tp":
+            _validate_tp(spec, owner)
+        elif kind == "pp":
+            _validate_pp(spec, owner)
+        elif kind == "ep":
+            _validate_ep(spec, owner)
+            enable_expert_parallel = True
+            ep_size = spec.mesh_axis.size
+        elif kind == "stage_replica":
+            omni_lb_policy = _stage_replica_lb_policy(spec, owner)
+            stage_replica_size = spec.mesh_axis.size
+
+        if kind in _AXIS_TO_ENGINE_FIELD:
+            sizes[_AXIS_TO_ENGINE_FIELD[kind]] = spec.mesh_axis.size
+        owners[kind] = owner
+
+    if enable_expert_parallel:
+        # Dense EP shards experts across exactly the TP*DP ranks, so the declared
+        # EP degree must match that product (it is not its own world dimension,
+        # and PP is excluded). An EP axis of size 1 is a degenerate no-op that
+        # still sets the flag; it is only valid when TP*DP == 1.
+        ep_ranks = sizes["tensor_parallel_size"] * sizes["data_parallel_size"]
+        if ep_size != ep_ranks:
+            _fail(
+                f"ep axis size {ep_size} must equal tensor_parallel_size*data_parallel_size "
+                f"(={ep_ranks}); dense expert parallelism shards experts across exactly those ranks"
+            )
+
+    return OmniParallelConfig(
+        tensor_parallel_size=sizes["tensor_parallel_size"],
+        data_parallel_size=sizes["data_parallel_size"],
+        pipeline_parallel_size=sizes["pipeline_parallel_size"],
+        enable_expert_parallel=enable_expert_parallel,
+        stage_replica_size=stage_replica_size,
+        omni_lb_policy=omni_lb_policy,
+        l1_owners=owners,
+    )

vllm_omni/config/config_factory.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 
 import dataclasses
 import functools
+from collections.abc import Mapping
 from dataclasses import asdict
 from pathlib import Path
 from typing import Any
@@ -33,6 +34,16 @@ from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
 logger = init_logger(__name__)
 
 
+# Default degree for any parallel axis / replica count that isn't set anywhere
+# (CLI, deploy YAML, or pipeline default): a single, un-parallelized rank.
+# TODO(composable_parallel): this "1" default and the per-axis device-layout
+# fallbacks are currently re-derived in the merge, apply, and reconcile layers.
+# Centralize them in one schema so the pre-spawn device guard and the engine-args
+# defaults can't drift apart. This is the light slice; the full device-layout
+# centralization is tracked as a follow-up.
+_DEFAULT_PARALLEL_DEGREE = 1
+
+
 class StageConfigFactory:
     """Factory that loads pipeline YAML and merges CLI overrides.
 
@@ -267,14 +278,27 @@ class StageConfigFactory:
         trust_remote_code: bool = False,
         cli_overrides: dict[str, Any] | None = None,
         deploy_config_path: str | None = None,
+        strategy_specs: Mapping[Any, Any] | None = None,
         **deprecated_kwargs: Any,
-    ) -> list[StageConfig] | None:
+    ) -> tuple[list[StageConfig] | None, str | None]:
         """Load pipeline + deploy config, merge with CLI overrides.
 
         Checks OMNI_PIPELINES first, since supported models should be explicitly
         registered. If a model is not registered in OMNI_PIPELINES, tries to fall
         back to using the Transformers config & finding pipelines that have overlapping
         supported architectures.
+
+        When ``strategy_specs`` is provided (a mapping of role -> list of
+        ``StrategySpec``), the derived parallel sizing is overlaid onto the
+        merged stages (see ``vllm_omni.config.composable_parallel``). This is
+        opt-in: omitting it leaves the existing merge path untouched.
+
+        Returns ``(stages, omni_lb_policy)``: the merged stages (``None`` when the
+        model is not in the pipeline registry and the caller should fall back to
+        the legacy YAML path) and the strategy-derived, pipeline-wide
+        ``omni_lb_policy`` (``None`` when no stage_replica axis set one). The
+        policy is returned rather than threaded through a mutable out-param so the
+        engine can apply it without every intermediate call carrying the dict.
         """
         if cli_overrides is None:
             cli_overrides = {}
@@ -290,8 +314,9 @@ class StageConfigFactory:
                 pipeline_cfg,
                 cli_overrides,
                 deploy_config_path,
+                strategy_specs=strategy_specs,
             )
-        return None
+        return None, None
 
     @classmethod
     def _create_from_registry(
@@ -300,12 +325,17 @@ class StageConfigFactory:
         pipeline_cfg: PipelineConfig,
         cli_overrides: dict[str, Any],
         deploy_config_path: str | None = None,
+        strategy_specs: Mapping[Any, Any] | None = None,
         **deprecated_kwargs: Any,
-    ) -> list[StageConfig]:
+    ) -> tuple[list[StageConfig], str | None]:
         """Create StageConfigs from pipeline registry + deploy YAML.
 
         Precedence: caller-typed (non-None) value > deploy YAML >
         StageDeployConfig dataclass default.
+
+        Returns ``(stages, omni_lb_policy)`` — the strategy-derived pipeline-wide
+        load-balancer policy (``None`` when no strategy set one) travels with the
+        stages instead of through a mutable out-param.
         """
         # Resolve deploy config path
         if deploy_config_path is None:
@@ -338,12 +368,127 @@ class StageConfigFactory:
 
         stages = merge_pipeline_deploy(pipeline_cfg, deploy_cfg, cli_overrides)
 
+        # Overlay declarative parallel strategies (opt-in) before CLI overrides.
+        applied = cls._apply_strategy_specs(stages, strategy_specs)
+
         explicit_overrides = {k: v for k, v in cli_overrides.items() if v is not None}
 
         for stage in stages:
             stage.runtime_overrides = cls._merge_cli_overrides(stage, explicit_overrides)
 
-        return stages
+        # Re-validate the resolved layout now that CLI overrides are on top.
+        cls._reconcile_strategy_with_cli(stages, applied)
+
+        omni_lb_policy = applied.omni_lb_policy if applied is not None else None
+        return stages, omni_lb_policy
+
+    @staticmethod
+    def _apply_strategy_specs(
+        stages: list[StageConfig],
+        strategy_specs: Mapping[Any, Any] | None,
+    ) -> Any:
+        """Overlay derived parallel sizing onto merged stages (opt-in).
+
+        ``omni_lb_policy`` cannot be set from stage configs (omni reads it once
+        at orchestrator construction), so a derived non-default policy is logged
+        here and carried on the returned ``StrategyApplyResult`` for the caller
+        to hand to the orchestrator.
+
+        Returns the ``StrategyApplyResult`` (or ``None`` when no strategy was
+        supplied) so the caller can re-validate the resolved layout once CLI
+        overrides have been merged on top, and read its ``omni_lb_policy``.
+        """
+        if not strategy_specs:
+            return None
+        from vllm_omni.config.composable_parallel import apply_strategy_specs
+
+        applied = apply_strategy_specs(stages, strategy_specs)
+        if applied.omni_lb_policy is not None:
+            logger.info(
+                "[composable_parallel] strategy derived omni_lb_policy=%r; it will be applied "
+                "to the orchestrator unless an explicit --omni-lb-policy was given.",
+                applied.omni_lb_policy,
+            )
+        return applied
+
+    @staticmethod
+    def _reconcile_strategy_with_cli(
+        stages: list[StageConfig],
+        applied: Any,
+    ) -> None:
+        """Reconcile CLI overrides applied *after* a strategy overlay.
+
+        CLI overrides are applied last (in ``to_omegaconf``), so for any stage a
+        strategy touched we must (a) warn loudly when a CLI arg overrides a
+        strategy-declared axis — the strategy is meant to be the single writer
+        for the axes it declares, so a silent CLI win is surprising — and
+        (b) re-run the device-count check against the *effective* (post-CLI)
+        ``tp``/``dp``/``pp``/``num_replicas``/``devices`` so a CLI
+        ``--devices``/``--tensor-parallel-size``/``--num-replicas`` cannot slip
+        past the pre-spawn guard that exists to prevent silent OOMs.
+        """
+        if applied is None:
+            return
+        from vllm_omni.config.composable_parallel import check_device_layout
+
+        # Axis kind -> (engine-arg field, strategy-derived attribute on the cfg).
+        axis_fields = {
+            "tp": ("tensor_parallel_size", "tensor_parallel_size"),
+            "dp": ("data_parallel_size", "data_parallel_size"),
+            "pp": ("pipeline_parallel_size", "pipeline_parallel_size"),
+        }
+
+        for stage in stages:
+            cfg = applied.per_stage_config.get(stage.stage_id)
+            if cfg is None:
+                continue
+            overrides = stage.runtime_overrides or {}
+            declared = set(cfg.l1_owners.keys())
+
+            for kind, (field_name, attr) in axis_fields.items():
+                if kind in declared and overrides.get(field_name) is not None:
+                    cli_val = overrides[field_name]
+                    derived = getattr(cfg, attr)
+                    if cli_val != derived:
+                        logger.warning(
+                            "[composable_parallel] stage %s: CLI %s=%s overrides the "
+                            "strategy-derived %s=%s. The CLI value wins; remove one to avoid ambiguity.",
+                            stage.stage_id,
+                            field_name,
+                            cli_val,
+                            field_name,
+                            derived,
+                        )
+            if "stage_replica" in declared and overrides.get("num_replicas") is not None:
+                cli_val = overrides["num_replicas"]
+                if cli_val != cfg.stage_replica_size:
+                    logger.warning(
+                        "[composable_parallel] stage %s: CLI num_replicas=%s overrides the "
+                        "strategy-derived num_replicas=%s. The CLI value wins; remove one to avoid ambiguity.",
+                        stage.stage_id,
+                        cli_val,
+                        cfg.stage_replica_size,
+                    )
+
+            def _eff(field_name: str, fallback: Any) -> Any:
+                val = overrides.get(field_name)
+                return val if val is not None else fallback
+
+            def _eff_degree(field_name: str, source: dict[str, Any]) -> int:
+                # Single place the per-axis "default to 1" fallback is applied,
+                # for both the override and the YAML-default sides (see the
+                # _DEFAULT_PARALLEL_DEGREE TODO).
+                value = _eff(field_name, source.get(field_name, _DEFAULT_PARALLEL_DEGREE))
+                return int(value or _DEFAULT_PARALLEL_DEGREE)
+
+            check_device_layout(
+                _eff("devices", stage.yaml_runtime.get("devices")),
+                tensor_parallel_size=_eff_degree("tensor_parallel_size", stage.yaml_engine_args),
+                data_parallel_size=_eff_degree("data_parallel_size", stage.yaml_engine_args),
+                pipeline_parallel_size=_eff_degree("pipeline_parallel_size", stage.yaml_engine_args),
+                num_replicas=_eff_degree("num_replicas", stage.yaml_runtime),
+                role=stage.model_stage,
+            )
 
     @classmethod
     def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]:

vllm_omni/config/omni_config.py

@@ -551,7 +551,7 @@ class _DiffusionConfigProjection:
             elif isinstance(self.quantization_config, str):
                 self.quantization_config = build_quant_config(self.quantization_config)
             elif isinstance(self.quantization_config, Mapping):
-                self.quantization_config = build_quant_config(dict(self.quantization_config))
+                self.quantization_config = dict(self.quantization_config)
             else:
                 raise TypeError(
                     "quantization_config must be str, dict, QuantizationConfig, or None, "

vllm_omni/deploy/qwen3_omni_moe.yaml

@@ -47,7 +47,6 @@ stages:
       temperature: 0.9
       top_k: 50
       max_tokens: 4096
-      seed: 42
       repetition_penalty: 1.05
 
   - stage_id: 2
@@ -67,7 +66,6 @@ stages:
       top_p: 1.0
       top_k: -1
       max_tokens: 65536
-      seed: 42
       repetition_penalty: 1.1
 
 platforms:

vllm_omni/deploy/qwen3_tts.yaml

@@ -35,7 +35,7 @@ connectors:
 
 stages:
   - stage_id: 0
-    max_num_seqs: 10
+    max_num_seqs: 64
     gpu_memory_utilization: 0.3
     trust_remote_code: true
     # Prefix caching on by default. The HS opt-out (talker only consumes the
@@ -43,9 +43,9 @@ stages:
     # PR make per-step bookkeeping negligible, and any text repetition across
     # requests (same voice or shared ref_audio) is a net throughput win.
     # Disable explicitly if memory headroom is tight on a single-GPU deploy.
-    enable_prefix_caching: true
+    enable_prefix_caching: false
     async_scheduling: true
-    max_num_batched_tokens: 512
+    max_num_batched_tokens: 32768
     max_model_len: 4096
     devices: "0"
     output_connectors:
@@ -67,7 +67,7 @@ stages:
 
   - stage_id: 1
     # Tuned for Base voice clone; CustomVoice / VoiceDesign are TTFA-optimal at 1.
-    max_num_seqs: 10
+    max_num_seqs: 64
     gpu_memory_utilization: 0.3
     enforce_eager: false
     trust_remote_code: true
@@ -95,3 +95,8 @@ platforms:
     stages:
       - stage_id: 0
         enforce_eager: false
+        compilation_config:
+          cudagraph_mode: PIECEWISE
+      - stage_id: 1
+        compilation_config:
+          cudagraph_mode: PIECEWISE

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -11,7 +11,6 @@ from pathlib import Path
 from typing import cast
 
 import torch
-from huggingface_hub import hf_hub_download
 from torch import nn
 from vllm.config import ModelConfig
 from vllm.config.load import LoadConfig
@@ -36,7 +35,6 @@ from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp
 from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
     get_checkpoint_adapter,
 )
-from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
 from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
 from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 from vllm_omni.diffusion.registry import initialize_model
@@ -381,10 +379,7 @@ class DiffusersPipelineLoader:
                     # DiffusionPipeline.from_pretrained() internally; it does
                     # not use our native customized pipeline classes.
                     cast(DiffusersAdapterPipeline, model).load_weights()
-                elif self._is_gguf_quantization():
-                    self._load_weights_with_gguf(model)
                 else:
-                    # Quantization does not happen in `load_weights` but after it
                     self.load_weights(model)
 
             # Process weights after loading for quantization (e.g., FP8 online quantization)
@@ -454,7 +449,7 @@ class DiffusersPipelineLoader:
     def _is_expected_quantized_weight(name: str) -> bool:
         """Return True if *name* is a quantization-specific parameter.
 
-        Quantization methods (GPTQ, AWQ, FP8, GGUF, Autoround, etc.) create extra
+        Quantization methods (GPTQ, AWQ, FP8, Autoround, etc.) create extra
         parameters that have no counterpart in an unquantized checkpoint.
         These are expected to be absent and should not trigger a load error.
         """
@@ -467,8 +462,6 @@ class DiffusersPipelineLoader:
             ".weight_scale",
             ".weight_scale_inv",
             ".input_scale",
-            # GGUF
-            ".qweight_type",
             # INT8  (weight_scale already covered above)
         )
         return name.endswith(_QUANTIZED_WEIGHT_SUFFIXES)
@@ -500,116 +493,6 @@ class DiffusersPipelineLoader:
         if unexpected_missing:
             raise ValueError(f"Following weights were not initialized from checkpoint: {unexpected_missing}")
 
-    def _get_gguf_model_from_config(self) -> str | None:
-        """Extract gguf_model from quant_config (handles both dict and object)."""
-        if self.quant_config is None:
-            return None
-        if isinstance(self.quant_config, dict):
-            return self.quant_config.get("gguf_model")
-        return getattr(self.quant_config, "gguf_model", None)
-
-    def _is_gguf_quantization(self) -> bool:
-        """Check whether or not this pipeline loader is pointing at a GGUF config."""
-        from vllm_omni.quantization.gguf_config import DiffusionGGUFConfig
-
-        maybe_gguf_model = self._get_gguf_model_from_config()
-        is_gguf = (
-            isinstance(self.quant_config, DiffusionGGUFConfig)
-            or (isinstance(self.quant_config, dict) and self.quant_config.get("method") == "gguf")
-            or (hasattr(self.quant_config, "get_name") and self.quant_config.get_name() == "gguf")
-        )
-        if is_gguf:
-            if not maybe_gguf_model:
-                raise ValueError("GGUF quantization requires gguf_model")
-            return True
-
-        if isinstance(self.quant_config, dict):
-            return False
-
-        return bool(maybe_gguf_model)
-
-    def _is_transformer_source(self, source: "ComponentSource") -> bool:
-        if source.subfolder == "transformer":
-            return True
-        return source.prefix.startswith("transformer.")
-
-    def _get_model_loadable_names(self, model: nn.Module) -> set[str]:
-        # Avoid model.state_dict() here because GGUF uses UninitializedParameter
-        # which raises during detach(). Collect names directly.
-        return {name for name, _ in model.named_parameters()} | {name for name, _ in model.named_buffers()}
-
-    def _resolve_gguf_model_path(self, gguf_model: str, revision: str | None) -> str:
-        if os.path.isfile(gguf_model):
-            return gguf_model
-        # repo_id/filename.gguf
-        if "/" in gguf_model and gguf_model.endswith(".gguf"):
-            repo_id, filename = gguf_model.rsplit("/", 1)
-            return hf_hub_download(
-                repo_id=repo_id,
-                filename=filename,
-                revision=revision,
-                cache_dir=self.load_config.download_dir,
-            )
-        # repo_id:quant_type
-        if "/" in gguf_model and ":" in gguf_model:
-            repo_id, quant_type = gguf_model.rsplit(":", 1)
-            return download_gguf(
-                repo_id,
-                quant_type,
-                cache_dir=self.load_config.download_dir,
-                revision=revision,
-                ignore_patterns=self.load_config.ignore_patterns,
-            )
-        raise ValueError(
-            f"Unrecognized GGUF reference: {gguf_model!r} (expected local file, "
-            "<repo_id>/<filename>.gguf, or <repo_id>:<quant_type>)"
-        )
-
-    def _get_gguf_weights_iterator(
-        self,
-        source: "ComponentSource",
-        model: nn.Module,
-    ) -> Generator[tuple[str, torch.Tensor], None, None]:
-        gguf_model = self._get_gguf_model_from_config()
-        if gguf_model is None:
-            raise ValueError("GGUF quantization requires quantization_config.gguf_model")
-        gguf_file = self._resolve_gguf_model_path(gguf_model, self.od_config.revision)
-        adapter = get_gguf_adapter(gguf_file, model, source, self.od_config)
-        weights_iter = adapter.weights_iterator()
-        return ((source.prefix + name, tensor) for (name, tensor) in weights_iter)
-
-    def _load_weights_with_gguf(self, model: nn.Module) -> set[str]:
-        sources = self._get_weight_sources(model)
-        loaded: set[str] = set()
-        loadable_names: set[str] | None = None
-
-        for source in sources:
-            if self._is_transformer_source(source):
-                loaded |= model.load_weights(self._get_gguf_weights_iterator(source, model))
-
-                # GGUF checkpoints can be transformer-only or partially quantized.
-                # Only fall back to HF if this source still has missing loadable weights.
-                loadable_names = loadable_names or self._get_model_loadable_names(model)
-                has_missing_for_source = any(
-                    name.startswith(source.prefix) and name not in loaded for name in loadable_names
-                )
-                if not has_missing_for_source:
-                    continue
-
-                hf_iter = self._get_weights_iterator(source)
-                hf_iter = (
-                    (name, tensor) for (name, tensor) in hf_iter if name in loadable_names and name not in loaded
-                )
-                loaded |= model.load_weights(hf_iter)
-            else:
-                loaded |= model.load_weights(self._get_weights_iterator(source))
-
-        weights_to_load = self._get_expected_parameter_names(model)
-        weights_not_loaded = weights_to_load - loaded
-        if weights_not_loaded:
-            raise ValueError(f"Following weights were not initialized from checkpoint: {weights_not_loaded}")
-        return loaded
-
     def _init_from_load_format(
         self,
         load_format: str,

vllm_omni/diffusion/model_loader/gguf_adapters/__init__.py

@@ -1,47 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-from typing import TYPE_CHECKING
-
-import torch
-
-from .base import GGUFAdapter
-from .flux2_klein import Flux2KleinGGUFAdapter
-from .qwen_image import QwenImageGGUFAdapter
-from .z_image import ZImageGGUFAdapter
-
-if TYPE_CHECKING:
-    from vllm_omni.diffusion.data import OmniDiffusionConfig
-    from vllm_omni.diffusion.model_loader.diffusers_loader import (
-        DiffusersPipelineLoader,
-    )
-
-
-def get_gguf_adapter(
-    gguf_file: str,
-    model: torch.nn.Module,
-    source: DiffusersPipelineLoader.ComponentSource,
-    od_config: OmniDiffusionConfig,
-) -> GGUFAdapter:
-    adapter_classes = (QwenImageGGUFAdapter, ZImageGGUFAdapter, Flux2KleinGGUFAdapter)
-    for adapter_cls in adapter_classes:
-        if adapter_cls.is_compatible(od_config, model, source):
-            return adapter_cls(gguf_file, model, source, od_config)
-    model_type = None
-    if od_config.tf_model_config is not None:
-        model_type = od_config.tf_model_config.get("model_type")
-    supported = ", ".join(cls.__name__ for cls in adapter_classes)
-    raise ValueError(
-        "No GGUF adapter matched diffusion model "
-        f"(model_class_name={od_config.model_class_name!r}, model_type={model_type!r}). "
-        f"Supported adapters: {supported}."
-    )
-
-
-__all__ = [
-    "GGUFAdapter",
-    "Flux2KleinGGUFAdapter",
-    "QwenImageGGUFAdapter",
-    "ZImageGGUFAdapter",
-    "get_gguf_adapter",
-]

vllm_omni/diffusion/model_loader/gguf_adapters/base.py

@@ -1,109 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-from abc import ABC, abstractmethod
-from collections.abc import Generator
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
-
-import numpy as np
-import torch
-
-try:
-    import gguf
-except ImportError:
-    # Optional dependency: only required for GGUF-quantized diffusion models.
-    # Keep this import lazy so a missing ``gguf`` does not break importing the
-    # diffusion model loader (and thus inspection/loading of every other
-    # diffusion model that imports through this package).
-    gguf = None
-
-if TYPE_CHECKING:
-    from vllm_omni.diffusion.data import OmniDiffusionConfig
-    from vllm_omni.diffusion.model_loader.diffusers_loader import (
-        DiffusersPipelineLoader,
-    )
-
-
-@dataclass
-class MappedTensor:
-    name: str
-    tensor: Any
-    tensor_type: Any
-    row_slice: slice | None = None
-    swap_scale_shift: bool = False
-
-
-class GGUFAdapter(ABC):
-    """Base class for model-specific GGUF adapters."""
-
-    def __init__(
-        self,
-        gguf_file: str,
-        model: torch.nn.Module,
-        source: DiffusersPipelineLoader.ComponentSource,
-        od_config: OmniDiffusionConfig,
-    ) -> None:
-        self.gguf_file = gguf_file
-        self.model = model
-        self.source = source
-        self.od_config = od_config
-
-    @staticmethod
-    def is_compatible(
-        od_config: OmniDiffusionConfig,
-        model: torch.nn.Module,
-        source: DiffusersPipelineLoader.ComponentSource,
-    ) -> bool:
-        return False
-
-    @abstractmethod
-    def weights_iterator(self) -> Generator[tuple[str, torch.Tensor], None, None]:
-        raise NotImplementedError
-
-
-# FIXME(Isotr0py): Sync implemnentation with upstream vLLM?
-def gguf_quant_weights_iterator(gguf_file: str) -> Generator[tuple[str, torch.Tensor]]:
-    """
-    Iterate over the quant weights in the model gguf files and convert
-    them to torch tensors.
-    Be careful of the order of yielding weight types and weights data,
-    we have to yield all weight types first before yielding any weights.
-    Otherwise it would cause issue when loading weights with for packed
-    layer with different quant types.
-    """
-
-    if gguf is None:
-        raise ImportError(
-            "The 'gguf' package is required to load GGUF-quantized diffusion models. "
-            "Install it with `pip install gguf`."
-        )
-    reader = gguf.GGUFReader(gguf_file)
-
-    for tensor in reader.tensors:
-        weight_type = tensor.tensor_type
-        name = tensor.name
-
-        if weight_type.name not in ("F32", "F16"):
-            weight_type_name = name.replace("weight", "qweight_type")
-            weight_type = torch.tensor(weight_type)
-            yield weight_type_name, weight_type
-
-    for tensor in reader.tensors:
-        weight = tensor.data
-        weight_type = tensor.tensor_type
-        name = tensor.name
-        if weight_type.name not in ("F32", "F16"):
-            name = name.replace("weight", "qweight")
-        if weight_type.name == "BF16" and tensor.data.dtype == np.uint8:
-            # BF16 is currently the only "quantization" type that isn't
-            # actually quantized but is read as a raw byte tensor.
-            # Reinterpret as `torch.bfloat16` tensor.
-            weight = weight.view(np.uint16)
-            if reader.byte_order == "S":
-                # GGUF endianness != system endianness
-                weight = weight.byteswap()
-            param = torch.tensor(weight).view(torch.bfloat16)
-        else:
-            param = torch.tensor(weight)
-        yield name, param

vllm_omni/diffusion/model_loader/gguf_adapters/flux2_klein.py

@@ -1,98 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-from collections.abc import Iterable
-
-import torch
-from vllm.model_executor.models.utils import WeightsMapper
-
-from .base import GGUFAdapter, gguf_quant_weights_iterator
-
-FLUX2_TRANSFORMER_KEYS_RENAME_DICT = {
-    "single_blocks.": "single_transformer_blocks.",
-    # Image and text input projections
-    "img_in": "x_embedder",
-    "txt_in": "context_embedder",
-    # Timestep and guidance embeddings
-    "time_in.in_layer": "time_guidance_embed.timestep_embedder.linear_1",
-    "time_in.out_layer": "time_guidance_embed.timestep_embedder.linear_2",
-    "guidance_in.in_layer": "time_guidance_embed.guidance_embedder.linear_1",
-    "guidance_in.out_layer": "time_guidance_embed.guidance_embedder.linear_2",
-    # Modulation parameters
-    "double_stream_modulation_img.lin": "double_stream_modulation_img.linear",
-    "double_stream_modulation_txt.lin": "double_stream_modulation_txt.linear",
-    "single_stream_modulation.lin": "single_stream_modulation.linear",
-    # Final output layer
-    # "final_layer.adaLN_modulation.1": "norm_out.linear",  # Handle separately since we need to swap mod params
-    "final_layer.linear": "proj_out",
-}
-
-FLUX2_TRANSFORMER_ADA_LAYER_NORM_KEY_MAP = {
-    "final_layer.adaLN_modulation.1": "norm_out.linear",
-}
-
-FLUX2_TRANSFORMER_DOUBLE_BLOCK_KEY_MAP = {
-    "double_blocks.": "transformer_blocks.",
-    # Handle fused QKV projections separately as we need to break into Q, K, V projections
-    "img_attn.norm.query_norm": "attn.norm_q",
-    "img_attn.norm.key_norm": "attn.norm_k",
-    "img_attn.proj": "attn.to_out.0",
-    "img_mlp.0": "ff.linear_in",
-    "img_mlp.2": "ff.linear_out",
-    "txt_attn.norm.query_norm": "attn.norm_added_q",
-    "txt_attn.norm.key_norm": "attn.norm_added_k",
-    "txt_attn.proj": "attn.to_add_out",
-    "txt_mlp.0": "ff_context.linear_in",
-    "txt_mlp.2": "ff_context.linear_out",
-    # Additional for fuse qkv
-    "img_attn.qkv": "attn.to_qkv",
-    "txt_attn.qkv": "attn.add_kv_proj",
-}
-
-FLUX2_TRANSFORMER_SINGLE_BLOCK_KEY_MAP = {
-    "linear1": "attn.to_qkv_mlp_proj",
-    "linear2": "attn.to_out",
-    "norm.query_norm": "attn.norm_q",
-    "norm.key_norm": "attn.norm_k",
-}
-
-
-class Flux2KleinGGUFAdapter(GGUFAdapter):
-    """GGUF adapter for Flux2-Klein models with qkv splitting and adaLN swap."""
-
-    @staticmethod
-    def is_compatible(od_config, model: torch.nn.Module, source) -> bool:
-        model_class = od_config.model_class_name or ""
-        if model_class.startswith("Flux2"):
-            return True
-        cfg = od_config.tf_model_config
-        if cfg is not None:
-            model_type = str(cfg.get("model_type", "")).lower()
-            if model_type.startswith("flux"):
-                return True
-        return False
-
-    gguf_to_hf_mapper = WeightsMapper(
-        # double_stream_modulation
-        orig_to_new_prefix=FLUX2_TRANSFORMER_KEYS_RENAME_DICT | FLUX2_TRANSFORMER_ADA_LAYER_NORM_KEY_MAP,
-        orig_to_new_substr=FLUX2_TRANSFORMER_DOUBLE_BLOCK_KEY_MAP | FLUX2_TRANSFORMER_SINGLE_BLOCK_KEY_MAP,
-    )
-
-    def weights_iterator(self) -> Iterable[tuple[str, torch.Tensor]]:
-        def custom_weights_adapter(
-            weights: Iterable[tuple[str, torch.Tensor]],
-        ) -> Iterable[tuple[str, torch.Tensor]]:
-            for name, weight in weights:
-                # Handle the special case for adaLN modulation parameters that require swapping shift and scale
-                if name.endswith(".scale"):
-                    name = name.replace(".scale", ".weight")
-                if name == "norm_out.linear.weight":
-                    shift, scale = weight.chunk(2, dim=0)
-                    weight = torch.cat([scale, shift], dim=0)
-                    yield name, weight
-                else:
-                    yield name, weight
-
-        weights = gguf_quant_weights_iterator(self.gguf_file)
-        weights = self.gguf_to_hf_mapper.apply(weights)
-        yield from custom_weights_adapter(weights)

vllm_omni/diffusion/model_loader/gguf_adapters/qwen_image.py

@@ -1,27 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-from collections.abc import Iterable
-
-import torch
-
-from .base import GGUFAdapter, gguf_quant_weights_iterator
-
-
-class QwenImageGGUFAdapter(GGUFAdapter):
-    """GGUF adapter for the Qwen-Image transformer family."""
-
-    @staticmethod
-    def is_compatible(od_config, model: torch.nn.Module, source) -> bool:
-        model_class = od_config.model_class_name or ""
-        if model_class.startswith("QwenImage"):
-            return True
-        cfg = od_config.tf_model_config
-        if cfg is not None:
-            model_type = str(cfg.get("model_type", "")).lower()
-            if model_type.startswith("qwen_image"):
-                return True
-        return False
-
-    def weights_iterator(self) -> Iterable[tuple[str, torch.Tensor]]:
-        yield from gguf_quant_weights_iterator(self.gguf_file)

vllm_omni/diffusion/model_loader/gguf_adapters/z_image.py

@@ -1,43 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-from collections.abc import Iterable
-
-import torch
-from vllm.model_executor.models.utils import WeightsMapper
-
-from .base import GGUFAdapter, gguf_quant_weights_iterator
-
-Z_IMAGE_KEYS_RENAME_DICT = {
-    "final_layer.": "all_final_layer.2-1.",
-    "x_embedder.": "all_x_embedder.2-1.",
-    ".attention.qkv": ".attention.to_qkv",
-    ".attention.k_norm": ".attention.norm_k",
-    ".attention.q_norm": ".attention.norm_q",
-    ".attention.out": ".attention.to_out.0",
-    "model.diffusion_model.": "",
-}
-
-
-class ZImageGGUFAdapter(GGUFAdapter):
-    """GGUF adapter for Z-Image models with QKV/FFN shard support."""
-
-    @staticmethod
-    def is_compatible(od_config, model: torch.nn.Module, source) -> bool:
-        model_class = od_config.model_class_name or ""
-        if model_class.startswith("ZImage"):
-            return True
-        cfg = od_config.tf_model_config
-        if cfg is not None:
-            model_type = str(cfg.get("model_type", "")).lower()
-            if model_type in {"z_image", "zimage", "z-image"}:
-                return True
-        return False
-
-    gguf_to_hf_mapper = WeightsMapper(
-        orig_to_new_substr=Z_IMAGE_KEYS_RENAME_DICT,
-    )
-
-    def weights_iterator(self) -> Iterable[tuple[str, torch.Tensor]]:
-        weights = gguf_quant_weights_iterator(self.gguf_file)
-        yield from self.gguf_to_hf_mapper.apply(weights)

vllm_omni/diffusion/models/flux2/flux2_transformer.py

@@ -1034,7 +1034,7 @@ class Flux2Transformer2DModel(nn.Module):
             name = original_name
             if name not in params_dict and ".to_out.0." in name:
                 name = name.replace(".to_out.0.", ".to_out.")
-            # Some GGUF checkpoints include quantized tensors for modules that
+            # Some quantized checkpoints include tensors for modules that
             # are intentionally left unquantized in this model.
             param = params_dict.get(name)
             if param is None:

vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py

@@ -1044,7 +1044,7 @@ class Flux2Transformer2DModel(nn.Module):
             name = original_name
             if name not in params_dict and ".to_out.0." in name:
                 name = name.replace(".to_out.0.", ".to_out.")
-            # Some GGUF checkpoints include quantized tensors for modules that
+            # Some quantized checkpoints include tensors for modules that
             # are intentionally left unquantized in this model.
             param = params_dict.get(name)
             if param is None:

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -17,7 +17,6 @@ from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous
 from vllm.logger import init_logger
-from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     QKVParallelLinear,
@@ -555,8 +554,8 @@ class QwenImageCrossAttention(nn.Module):
         self.query_num_heads = self.to_qkv.num_heads
         self.kv_num_heads = self.to_qkv.num_kv_heads
 
-        self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_q = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
 
         self.inner_dim = out_dim if out_dim is not None else head_dim * self.total_num_heads
 
@@ -593,8 +592,8 @@ class QwenImageCrossAttention(nn.Module):
             prefix=f"{prefix}.to_out.0",
         )
 
-        self.norm_added_q = RMSNorm(head_dim, eps=eps)
-        self.norm_added_k = RMSNorm(head_dim, eps=eps)
+        self.norm_added_q = nn.RMSNorm(head_dim, eps=eps)
+        self.norm_added_k = nn.RMSNorm(head_dim, eps=eps)
 
         self.attn = Attention(
             num_heads=self.query_num_heads,
@@ -1007,7 +1006,7 @@ class QwenImageTransformer2DModel(CachedTransformer):
             quant_config=quant_config,
         )
 
-        self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6)
+        self.txt_norm = nn.RMSNorm(joint_attention_dim, eps=1e-6)
 
         # Entry projections (image/text) are kept full precision —
         # small sensitive layers at the network boundary (see #2728).

vllm_omni/engine/arg_utils.py

@@ -434,6 +434,9 @@ class OrchestratorArgs:
     stage_configs_path: str | None = None
     deploy_config: str | None = None
     stage_overrides: str | None = None  # raw JSON string; parsed downstream
+    # Optional composable-parallel strategy.yaml; orchestrator reads it, overlays
+    # derived sizing onto merged stages, then drops it before per-stage engine args.
+    strategy_config: str | None = None
 
     # === Mode Switches (orchestrator reads, DeployConfig redistributes) ===
     async_chunk: bool | None = None

vllm_omni/engine/async_omni_engine.py

@@ -58,7 +58,10 @@ from vllm_omni.engine.stage_runtime import (
     create_stage_runtime,
 )
 from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
-from vllm_omni.entrypoints.utils import load_and_resolve_stage_configs
+from vllm_omni.entrypoints.utils import (
+    load_and_resolve_stage_configs,
+    parse_stage_overrides,
+)
 from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.metrics.prometheus import OmniRequestCounter
 
@@ -1040,6 +1043,7 @@ class AsyncOmniEngine:
             "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True),
             "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4),
             "quantization": kwargs.get("quantization", None),
+            "quantization_config": kwargs.get("quantization_config", None),
             "diffusion_kv_cache_dtype": kwargs.get("diffusion_kv_cache_dtype", None),
             "diffusion_kv_cache_skip_steps": kwargs.get("diffusion_kv_cache_skip_steps", None),
             "diffusion_kv_cache_skip_layers": kwargs.get("diffusion_kv_cache_skip_layers", None),
@@ -1123,11 +1127,41 @@ class AsyncOmniEngine:
 
         return result
 
+    def _apply_strategy_lb_policy(self, derived: str | None, kwargs: dict[str, Any]) -> None:
+        """Apply a strategy-derived ``omni_lb_policy`` to the engine.
+
+        Precedence: an explicit ``--omni-lb-policy`` always wins. ``"random"`` is
+        the engine default and is treated as "unset" (indistinguishable from no
+        flag), so a strategy value overrides it. If the user explicitly passed a
+        non-default policy that conflicts with the strategy-derived one, raise so
+        the mismatch is not silently ignored.
+        """
+        if not derived:
+            return
+        explicit = kwargs.get("omni_lb_policy")
+        user_set = explicit is not None and str(explicit) != "random"
+        if user_set:
+            if str(explicit) != str(derived):
+                raise ValueError(
+                    f"Conflicting load-balancer policy: --omni-lb-policy={explicit!r} was given "
+                    f"but the composable-parallel strategy derived omni_lb_policy={derived!r}. "
+                    "Drop --omni-lb-policy to use the strategy value, or make them match."
+                )
+            return
+        if self._omni_lb_policy != str(derived):
+            logger.info(
+                "[composable_parallel] applying strategy-derived omni_lb_policy=%r (was %r).",
+                derived,
+                self._omni_lb_policy,
+            )
+            self._omni_lb_policy = str(derived)
+
     def _resolve_stage_configs(self, model: str, kwargs: dict[str, Any]) -> tuple[str, list[Any]]:
         """Resolve stage configs and inject defaults shared by orchestrator/headless."""
 
         stage_configs_path = kwargs.get("stage_configs_path", None)
         deploy_config_path = kwargs.pop("deploy_config", None)
+        strategy_config_path = kwargs.pop("strategy_config", None)
         stage_overrides_json = kwargs.pop("stage_overrides", None)
         explicit_stage_configs = kwargs.pop("stage_configs", None)
         if explicit_stage_configs is not None:
@@ -1142,27 +1176,23 @@ class AsyncOmniEngine:
             base_kwargs = kwargs
 
         # Parse --stage-overrides JSON string if provided
-        stage_overrides = None
-        if stage_overrides_json:
-            if isinstance(stage_overrides_json, str):
-                try:
-                    stage_overrides = json.loads(stage_overrides_json)
-                except json.JSONDecodeError as exc:
-                    raise ValueError(
-                        f"--stage-overrides is not valid JSON: {exc}. Got: {stage_overrides_json!r}"
-                    ) from exc
-            else:
-                stage_overrides = stage_overrides_json
+        stage_overrides = parse_stage_overrides(stage_overrides_json)
 
-        config_path, stage_configs = load_and_resolve_stage_configs(
+        config_path, stage_configs, strategy_lb_policy = load_and_resolve_stage_configs(
             model,
             stage_configs_path,
             base_kwargs,
             default_stage_cfg_factory=lambda: self._create_default_diffusion_stage_cfg(kwargs),
             deploy_config_path=deploy_config_path,
             stage_overrides=stage_overrides,
+            strategy_config_path=strategy_config_path,
         )
 
+        # A strategy.yaml may derive a pipeline-wide load-balancer policy. It is
+        # an orchestrator-level knob (read once at construction), so apply it here
+        # rather than as a per-stage config field.
+        self._apply_strategy_lb_policy(strategy_lb_policy, kwargs)
+
         # Inject diffusion LoRA-related knobs from kwargs if not present in the stage config.
         for cfg in stage_configs:
             try:
@@ -1204,7 +1234,7 @@ class AsyncOmniEngine:
                             kwargs.get("diffusion_attention_config"),
                             attention_backend=kwargs.get("diffusion_attention_backend"),
                         )
-                quantization_config = kwargs.get("diffusion_quantization_config")
+                quantization_config = kwargs.get("diffusion_quantization_config") or kwargs.get("quantization_config")
                 if quantization_config is not None:
                     if (
                         not hasattr(cfg.engine_args, "quantization_config")

vllm_omni/engine/stage_engine_core_client.py

@@ -230,7 +230,7 @@ class StageEngineCoreClientBase(StageClientBase):
 
     async def add_request_async(self, request: EngineCoreRequest) -> None:
         """Add request to the stage engine core."""
-        logger.info(
+        logger.debug(
             "[%s] stage-%s [rep-%s] add request: %s",
             self.__class__.__name__,
             self.stage_id,

vllm_omni/entrypoints/cli/serve.py

@@ -254,6 +254,15 @@ class OmniServeCommand(CLISubcommand):
             help="Path to a deploy config YAML (new format with stages/engine_args). "
             "Mutually exclusive with --stage-configs-path.",
         )
+        omni_config_group.add_argument(
+            "--strategy-config",
+            type=str,
+            default=None,
+            help="Path to a composable-parallel strategy.yaml. Only applies to "
+            "registry-based models (e.g. qwen2_5_omni): its derived parallel sizing "
+            "is overlaid onto the registry-merged stages before per-stage engine "
+            "args are built.",
+        )
         omni_config_group.add_argument(
             "--stage-overrides",
             type=str,
@@ -469,7 +478,7 @@ class OmniServeCommand(CLISubcommand):
             default=None,
             help=(
                 "JSON string for diffusion quantization_config. "
-                'Example: \'{"method":"gguf","gguf_model":"/path/to/model.gguf"}\'.'
+                'Example: \'{"method":"fp8","activation_scheme":"dynamic"}\'.'
             ),
         )
         omni_config_group.add_argument(
@@ -766,7 +775,10 @@ def run_headless(args: TrackingNamespace) -> None:
         load_omni_transfer_config_for_model,
         prepare_engine_environment,
     )
-    from vllm_omni.entrypoints.utils import load_and_resolve_stage_configs
+    from vllm_omni.entrypoints.utils import (
+        load_and_resolve_stage_configs,
+        parse_stage_overrides,
+    )
 
     model = args.model
     stage_id: int | None = args.stage_id
@@ -802,11 +814,18 @@ def run_headless(args: TrackingNamespace) -> None:
             args.replica_id,
         )
 
-    config_path, stage_configs = load_and_resolve_stage_configs(
+    # Parse --stage-overrides (raw JSON string) exactly like the standard
+    # engine path (AsyncOmniEngine._resolve_stage_configs) so headless and
+    # standard launches resolve to the same per-stage device layout.
+    stage_overrides = parse_stage_overrides(args_dict.get("stage_overrides"))
+
+    config_path, stage_configs, _ = load_and_resolve_stage_configs(
         model,
         stage_configs_path,
         args_dict,
         deploy_config_path=args_dict.get("deploy_config"),
+        stage_overrides=stage_overrides,
+        strategy_config_path=args_dict.get("strategy_config"),
     )
 
     # Locate the stage config that matches stage_id.

vllm_omni/entrypoints/openai/serving_speech.py

@@ -1395,7 +1395,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         emb_dim = len(embedding)
         if self._tts_model_type == "ming_tts":
             if emb_dim != 192:
-                logger.warning("speaker_embedding has %d dimensions; Ming dense expects 192", emb_dim)
+                raise ValueError(f"Ming speaker embedding must have 192 dims, got {emb_dim}")
         else:
             dim_err = self._validate_qwen_tts_speaker_embedding_dim(emb_dim)
             if dim_err is not None:

vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py

@@ -25,9 +25,18 @@ class MingTTSAdapter(ARTTSAdapter):
         ref_audio_source = request.ref_audio
         voice_lower = request.voice.lower() if isinstance(request.voice, str) else None
         if ref_audio_source is None and voice_lower in server.uploaded_speakers:
-            ref_audio_source = server._get_uploaded_audio_data(request.voice)
-            if request.ref_text is None:
-                request.ref_text = server.uploaded_speakers[voice_lower].get("ref_text")
+            speaker_info = server.uploaded_speakers[voice_lower]
+            if speaker_info.get("embedding_source") == "direct":
+                if request.speaker_embedding is None:
+                    request.speaker_embedding = server._get_uploaded_speaker_embedding(request.voice)
+                if request.speaker_embedding is None:
+                    raise ValueError(f"Speaker embedding for uploaded voice '{request.voice}' is missing")
+            else:
+                ref_audio_source = server._get_uploaded_audio_data(request.voice)
+                if not ref_audio_source:
+                    raise ValueError(f"Audio file for uploaded voice '{request.voice}' is missing")
+                if request.ref_text is None:
+                    request.ref_text = speaker_info.get("ref_text")
         ref_audio_data = None
         if isinstance(ref_audio_source, list):
             ref_audio_data = await server._resolve_ref_audio_many(ref_audio_source)

vllm_omni/entrypoints/utils.py

@@ -1,3 +1,4 @@
+import json
 import os
 import types
 from collections import Counter
@@ -316,7 +317,8 @@ def load_stage_configs_from_model(
     base_engine_args: dict | None = None,
     deploy_config_path: str | None = None,
     stage_overrides: dict[str, dict[str, Any]] | None = None,
-) -> list:
+    strategy_config_path: str | None = None,
+) -> tuple[list, str | None]:
     """Load stage configurations from model's default config file.
 
     For models registered in the pipeline registry (new path), uses
@@ -330,9 +332,15 @@ def load_stage_configs_from_model(
         base_engine_args: Base engine args to merge as CLI overrides.
         deploy_config_path: Optional explicit deploy config path.
         stage_overrides: Per-stage overrides from --stage-overrides.
+        strategy_config_path: Optional path to a composable-parallel
+            ``strategy.yaml`` whose derived sizing is overlaid onto the
+            registry-merged stages (opt-in; ignored on the legacy YAML path).
 
     Returns:
-        List of stage configuration dictionaries
+        ``(stage_configs, omni_lb_policy)``: the list of stage configuration
+        dictionaries plus the strategy-derived pipeline-wide ``omni_lb_policy``
+        (``None`` when no strategy set one). The policy is returned rather than
+        written into a caller-provided mutable dict.
     """
     if base_engine_args is None:
         base_engine_args = {}
@@ -343,26 +351,44 @@ def load_stage_configs_from_model(
             for key, val in overrides.items():
                 cli_overrides[f"stage_{stage_id_str}_{key}"] = val
 
-    stages = StageConfigFactory.create_from_model(
+    strategy_specs = None
+    if strategy_config_path is not None:
+        from vllm_omni.config.composable_parallel.strategy_loader import load_strategy_specs
+
+        strategy_specs = load_strategy_specs(strategy_config_path)
+
+    stages, omni_lb_policy = StageConfigFactory.create_from_model(
         model,
         trust_remote_code=cli_overrides.get("trust_remote_code", False),
         cli_overrides=cli_overrides,
         deploy_config_path=deploy_config_path,
+        strategy_specs=strategy_specs,
     )
     if stages is not None:
         # Convert StageConfig objects to OmegaConf for backward compat
-        return [stage.to_omegaconf() for stage in stages]
+        return [stage.to_omegaconf() for stage in stages], omni_lb_policy
 
-    # Legacy fallback: load from YAML
+    # Legacy fallback: load from YAML. A composable-parallel strategy cannot be
+    # applied here (it overlays onto registry-merged stages), so warn rather than
+    # silently dropping the operator's --strategy-config.
+    if strategy_config_path is not None:
+        logger.warning(
+            "--strategy-config (%s) was provided but model %r 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.",
+            strategy_config_path,
+            model,
+        )
     stage_config_path = resolve_model_config_path(model)
     if stage_config_path is None:
-        return []
+        return [], None
     stage_configs = load_stage_configs_from_yaml(
         config_path=stage_config_path,
         base_engine_args=base_engine_args,
         prefer_stage_engine_args=True,
     )
-    return stage_configs
+    return stage_configs, None
 
 
 def load_stage_configs_from_yaml(
@@ -501,6 +527,25 @@ def filter_stages(
         return stage_configs
 
 
+def parse_stage_overrides(value: Any) -> dict[str, dict[str, Any]] | None:
+    """Parse the ``--stage-overrides`` value into a per-stage override dict.
+
+    ``value`` may be a raw JSON string (as supplied on the CLI) or an
+    already-parsed mapping. Returns ``None`` when no overrides are given.
+
+    Raises:
+        ValueError: when ``value`` is a string that is not valid JSON.
+    """
+    if not value:
+        return None
+    if isinstance(value, str):
+        try:
+            return json.loads(value)
+        except json.JSONDecodeError as exc:
+            raise ValueError(f"--stage-overrides is not valid JSON: {exc}. Got: {value!r}") from exc
+    return value
+
+
 def load_and_resolve_stage_configs(
     model: str,
     stage_configs_path: str | None,
@@ -508,7 +553,8 @@ def load_and_resolve_stage_configs(
     default_stage_cfg_factory: Any = None,
     deploy_config_path: str | None = None,
     stage_overrides: dict[str, dict[str, Any]] | None = None,
-) -> tuple[str, list]:
+    strategy_config_path: str | None = None,
+) -> tuple[str, list, str | None]:
     """Load stage configurations from model or YAML file with fallback to defaults.
 
     Args:
@@ -521,9 +567,13 @@ def load_and_resolve_stage_configs(
             Mutually exclusive with ``stage_configs_path``.
         stage_overrides: Per-stage overrides from ``--stage-overrides`` JSON.
             Keys are stage_id strings, values are dicts of overrides.
+        strategy_config_path: Optional path to a composable-parallel
+            ``strategy.yaml`` overlaid onto the registry-merged stages.
 
     Returns:
-        Tuple of (config_path, stage_configs)
+        Tuple of ``(config_path, stage_configs, omni_lb_policy)`` — the last is
+        the strategy-derived pipeline-wide load-balancer policy (``None`` when no
+        strategy set one), returned for the engine to apply.
     """
     if stage_configs_path is not None and deploy_config_path is not None:
         raise ValueError(
@@ -551,13 +601,15 @@ def load_and_resolve_stage_configs(
                 stage_configs_path,
             )
 
+    omni_lb_policy: str | None = None
     if deploy_config_path is not None:
         config_path = deploy_config_path
-        stage_configs = load_stage_configs_from_model(
+        stage_configs, omni_lb_policy = load_stage_configs_from_model(
             model,
             base_engine_args=kwargs,
             deploy_config_path=deploy_config_path,
             stage_overrides=stage_overrides,
+            strategy_config_path=strategy_config_path,
         )
         if not stage_configs:
             if default_stage_cfg_factory is not None:
@@ -567,10 +619,11 @@ def load_and_resolve_stage_configs(
                 stage_configs = []
     elif stage_configs_path is None:
         config_path = resolve_model_config_path(model)
-        stage_configs = load_stage_configs_from_model(
+        stage_configs, omni_lb_policy = load_stage_configs_from_model(
             model,
             base_engine_args=kwargs,
             stage_overrides=stage_overrides,
+            strategy_config_path=strategy_config_path,
         )
         if not stage_configs:
             if default_stage_cfg_factory is not None:
@@ -585,7 +638,7 @@ def load_and_resolve_stage_configs(
     stage_configs = filter_stages(config_path, stage_configs, kwargs)
     logger.debug(f"stage_configs: {stage_configs}")
 
-    return config_path, stage_configs
+    return config_path, stage_configs, omni_lb_policy
 
 
 def get_final_stage_id_for_e2e(

vllm_omni/model_executor/models/common/qwen3_code_predictor.py

@@ -28,6 +28,9 @@ from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
+_GeneratorLike = torch.Generator | Sequence[torch.Generator | None] | None
+_UNIFORM_EPS = 1e-20
+
 
 # ===================================================================
 # HF-numerics-compatible layers for code predictor
@@ -605,6 +608,30 @@ class CodePredictorWrapper(nn.Module):
                 values.add(parsed)
         return values
 
+    @staticmethod
+    def _normalize_generators(
+        generator: _GeneratorLike, batch_size: int
+    ) -> torch.Generator | list[torch.Generator | None] | None:
+        if generator is None or isinstance(generator, torch.Generator):
+            return generator
+
+        row_generators = list(generator)
+        if len(row_generators) != batch_size:
+            raise ValueError(f"Expected {batch_size} per-row generators, but got {len(row_generators)}.")
+        return row_generators
+
+    @classmethod
+    def _sample_codes_gumbel(cls, logits: torch.Tensor, generator: _GeneratorLike = None) -> torch.Tensor:
+        """Sample ``logits`` via Gumbel-max with optional per-row generators."""
+        row_generators = cls._normalize_generators(generator, int(logits.shape[0]))
+        u = torch.empty_like(logits, dtype=torch.float32)
+        if isinstance(row_generators, list):
+            for row, row_generator in enumerate(row_generators):
+                u[row : row + 1].uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generator)
+        else:
+            u.uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generators)
+        return (logits.float() - torch.log(-torch.log(u))).argmax(dim=-1, keepdim=True)
+
     def _prefix_seq_lens(self, max_seq: int) -> list[int]:
         all_seq_lens = list(range(2, max_seq))
         if not self._prefix_graph_seq_lens:
@@ -780,6 +807,7 @@ class CodePredictorWrapper(nn.Module):
         bsz = int(layer0_code.shape[0])
         if generators is not None and len(generators) != bsz:
             raise ValueError(f"generators must have one entry per row: got {len(generators)} for batch {bsz}")
+        sample_generator: _GeneratorLike = generators if generators is not None else generator
         num_groups = self._num_groups
         device = layer0_code.device
 
@@ -856,9 +884,20 @@ class CodePredictorWrapper(nn.Module):
 
             logits = lm_heads[step - 1](hidden_out[:bsz, step, :])
 
-            # Sample next code
+            # Sample next code via Gumbel-max.
+            #
+            # ``argmax_i(logits_i + Gumbel_i)`` with
+            # ``Gumbel_i = -log(-log(u_i)), u_i ~ Uniform(0, 1)`` is
+            # distributionally identical to sampling from ``softmax(logits)``.
+            # In this file the motivations are practical rather than graph
+            # related: it is measurably cheaper than ``softmax + multinomial``
+            # on the B x 2048 shapes used here, it stays well-defined for
+            # degenerate masked rows with a surviving finite entry (and is more
+            # defensive than ``multinomial`` around fully-masked/NaN inputs),
+            # and the helper below can honor either one batch generator or one
+            # generator per seeded row.
             if stored_mode:
-                # "stored" mode: top-k -> top-p -> softmax -> multinomial
+                # "stored" mode: top-k -> top-p -> Gumbel-max
                 if s_top_k > 0:
                     topk_vals, _ = logits.topk(s_top_k, dim=-1)
                     logits = logits.masked_fill(logits < topk_vals[:, -1:], float("-inf"))
@@ -869,17 +908,15 @@ class CodePredictorWrapper(nn.Module):
                     remove_mask = (cumulative_probs - sorted_probs) >= s_top_p
                     sorted_logits[remove_mask] = float("-inf")
                     logits = sorted_logits.scatter(1, sorted_idx, sorted_logits)
-                probs = F.softmax(logits, dim=-1, dtype=torch.float32)
-                code = self._multinomial(probs, generator, generators)
+                code = self._sample_codes_gumbel(logits, generator=sample_generator)
             else:
-                # "per_call" mode: temperature-scaled + top-k
+                # "per_call" mode: temperature-scaled + top-k -> Gumbel-max
                 if use_sampling:
                     scaled = logits * inv_temperature
                     if top_k > 0:
                         topk_vals, _ = scaled.topk(top_k, dim=-1)
                         scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf"))
-                    probs = F.softmax(scaled, dim=-1, dtype=torch.float32)
-                    code = self._multinomial(probs, generator, generators)
+                    code = self._sample_codes_gumbel(scaled, generator=sample_generator)
                 else:
                     code = logits.argmax(dim=-1, keepdim=True)
 

vllm_omni/model_executor/models/cosyvoice3/pipeline.py

@@ -5,9 +5,8 @@
 Stage 0: Talker   — text prompt → speech tokens (LLM autoregressive).
 Stage 1: Code2Wav — flow-matching decoder → acoustic features → waveform.
   * ``sync_process_input_func`` (``text2flow_token_only``) runs when
-    ``deploy.async_chunk=false``: stage 1 builds full-sequence flow input
-    from the emitted speech tokens. ``text2flow`` is the
-    ``custom_process_input_func`` variant.
+    ``deploy.async_chunk=false``: stage 1 allocates placeholder slots while
+    bulk tensors arrive via ``text2flow_full_payload`` on the connector.
   * ``async_chunk_process_next_stage_input_func`` runs when
     ``deploy.async_chunk=true``: stage 0 streams codec chunks to stage 1
     through the shared-memory connector.
@@ -47,7 +46,6 @@ COSYVOICE3_PIPELINE = PipelineConfig(
             final_output=True,
             final_output_type="audio",
             engine_output_type="latent",
-            custom_process_input_func=f"{_PROC}.text2flow",
             sync_process_input_func=f"{_PROC}.text2flow_token_only",
         ),
     ),

vllm_omni/model_executor/models/mimo_audio/pipeline.py

@@ -56,7 +56,6 @@ MIMO_AUDIO_PIPELINE = PipelineConfig(
             final_output=True,
             final_output_type="audio",
             engine_output_type="audio",
-            custom_process_input_func=f"{_PROC}.llm2code2wav",
             sync_process_input_func=f"{_PROC}.llm2code2wav_token_only",
             sampling_constraints={"detokenize": False},
         ),

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py

@@ -17,7 +17,6 @@ from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin,
-    Qwen2_5OmniThinkerProcessingInfo,
 )
 from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
 from vllm.model_executor.models.vision import (
@@ -39,6 +38,7 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
     Qwen2_5OmniThinkerDummyInputsBuilder,
     Qwen2_5OmniThinkerMultiModalProcessor,
+    Qwen2_5OmniThinkerProcessingInfo,
 )
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges
 from vllm_omni.platforms import current_omni_platform

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py

@@ -4,6 +4,7 @@ from collections.abc import Iterable, Iterator, Mapping
 from typing import Any
 
 import numpy as np
+import PIL.Image as PILImage
 import torch
 from torch import nn
 from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
@@ -26,16 +27,21 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniAudioFeatureInputs,
     Qwen2_5OmniThinkerDummyInputsBuilder,
-    Qwen2_5OmniThinkerProcessingInfo,
     check_interleaved_audio_video,
     merge_interleaved_embeddings,
 )
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin as Qwen2_5OmniConditionalGenerationMixinBase,
 )
+from vllm.model_executor.models.qwen2_5_omni_thinker import (
+    Qwen2_5OmniThinkerMultiModalDataParser as _Qwen2_5OmniThinkerMultiModalDataParserBase,
+)
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniThinkerMultiModalProcessor as _Qwen2_5OmniThinkerMultiModalProcessorBase,
 )
+from vllm.model_executor.models.qwen2_5_omni_thinker import (
+    Qwen2_5OmniThinkerProcessingInfo as _Qwen2_5OmniThinkerProcessingInfoBase,
+)
 from vllm.model_executor.models.qwen2_5_vl import (
     Qwen2_5_VisionTransformer,
     Qwen2_5_VLImageEmbeddingInputs,
@@ -56,12 +62,13 @@ from vllm.multimodal.inputs import (
     MultiModalFeatureSpec,
     MultiModalKwargsItems,
 )
-from vllm.multimodal.parse import MultiModalDataItems
+from vllm.multimodal.parse import MultiModalDataItems, VideoProcessorItems
 from vllm.multimodal.processing.processor import (
     MultiModalPromptUpdates,
     PlaceholderFeaturesInfo,
 )
 from vllm.sequence import IntermediateTensors
+from vllm.utils.collection_utils import is_list_of
 
 from vllm_omni.quantization.component_config import (
     resolve_encoder_quant_config,
@@ -74,11 +81,143 @@ except (ImportError, ModuleNotFoundError):
 logger = init_logger(__name__)
 
 
+def _presampled_videos_hf_kwargs(
+    mm_data: Mapping[str, object],
+    mm_kwargs: Mapping[str, object],
+) -> Mapping[str, object]:
+    """Adjust HF video kwargs for videos pre-sampled by vLLM's video loader.
+
+    When ``video_metadata`` is present (emitted by
+    ``Qwen2_5OmniVideoProcessorItems``), the frames were already sampled
+    according to ``media_io_kwargs``, so the HF processor must not re-sample
+    them, and it needs the sampled fps (instead of its default) to compute
+    ``video_second_per_grid`` correctly. Otherwise the audio/video temporal
+    alignment breaks under ``use_audio_in_video=True``.
+    """
+    video_metadata = mm_data.get("video_metadata")
+    if not video_metadata:
+        return mm_kwargs
+
+    mm_kwargs = dict(mm_kwargs)
+    videos_kwargs = dict(mm_kwargs.get("videos_kwargs") or {})
+    videos_kwargs["do_sample_frames"] = False
+
+    def _compute_sampled_video_fps(metadata) -> float | None:
+        duration = getattr(metadata, "duration", None)
+        indices = getattr(metadata, "frames_indices", None)
+        if not duration or not indices or duration <= 0:
+            return None
+        return len(indices) / float(duration)
+
+    # An explicit user-provided fps takes precedence.
+    if "fps" not in videos_kwargs:
+        fps_values = [_compute_sampled_video_fps(m) for m in video_metadata]
+        # HF accepts a single fps per call and uses it for every video's
+        # video_second_per_grid.
+        known_fps = [fps for fps in fps_values if fps is not None]
+        unique_fps = set(known_fps)
+        if len(unique_fps) == 1:
+            videos_kwargs["fps"] = known_fps[0]
+        elif len(unique_fps) > 1:
+            logger.warning(
+                f"Mixed sampled FPS {sorted(unique_fps)} in one request; HF accepts a single fps, using {known_fps[0]}."
+            )
+            videos_kwargs["fps"] = known_fps[0]
+
+    mm_kwargs["videos_kwargs"] = videos_kwargs
+    return mm_kwargs
+
+
+class Qwen2_5OmniVideoProcessorItems(VideoProcessorItems):
+    """Video items that carry the loader's ``(frames, metadata)`` tuples.
+
+    The tuples are kept in ``data`` so that mm hashing and cache-miss
+    re-parsing retain the metadata; ``get_processor_data`` unpacks them into
+    the ``videos`` + ``video_metadata`` arguments of the HF processor.
+    """
+
+    def get_processor_data(self) -> Mapping[str, object]:
+        from transformers.video_utils import VideoMetadata
+
+        videos = [frames for frames, _ in self.data]
+        video_metadata = [
+            VideoMetadata(**{k: v for k, v in metadata.items() if k != "do_sample_frames"}) for _, metadata in self.data
+        ]
+        return {"videos": videos, "video_metadata": video_metadata}
+
+
+class Qwen2_5OmniThinkerMultiModalDataParser(_Qwen2_5OmniThinkerMultiModalDataParserBase):
+    def _parse_video_data(self, data):
+        if data is None or isinstance(data, dict) or self.is_embeddings(data):
+            return super()._parse_video_data(data)
+
+        # Normalize to a list of per-video items (same as the base parser).
+        if (is_list_of(data, PILImage.Image) and len(data) > 0) or (
+            isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 4
+        ):
+            data_items = [data]
+        elif isinstance(data, (np.ndarray, torch.Tensor)):
+            data_items = [elem for elem in data]
+        elif isinstance(data, tuple) and len(data) == 2:
+            data_items = [data]
+        else:
+            data_items = data
+
+        videos_with_metadata = [self._get_video_with_metadata(item) for item in data_items]
+
+        # Metadata is optional: videos fetched by vLLM's media connector
+        # arrive as (frames, metadata) tuples, while plain arrays (e.g.
+        # offline `multi_modal_data` or dummy profiling inputs) carry no
+        # metadata and parse as before.
+        if any(metadata is None for _, metadata in videos_with_metadata):
+            return super()._parse_video_data(data)
+
+        return Qwen2_5OmniVideoProcessorItems(
+            videos_with_metadata,
+            metadata=[metadata for _, metadata in videos_with_metadata],
+        )
+
+
+class Qwen2_5OmniThinkerProcessingInfo(_Qwen2_5OmniThinkerProcessingInfoBase):
+    def get_data_parser(self):
+        feature_extractor = self.get_feature_extractor()
+
+        return Qwen2_5OmniThinkerMultiModalDataParser(
+            spatial_merge_size=self.get_hf_config().vision_config.spatial_merge_size,
+            target_sr=feature_extractor.sampling_rate,
+            target_channels=self.get_target_channels(),
+            expected_hidden_size=self._get_expected_hidden_size(),
+        )
+
+
 class Qwen2_5OmniThinkerMultiModalProcessor(
     _Qwen2_5OmniThinkerMultiModalProcessorBase,
 ):
     """Override to fix use_audio_in_video detection when mm cache returns None."""
 
+    def _cached_apply_hf_processor(self, inputs, timing_ctx):
+        # If use_audio_in_video, process video and audio together; otherwise,
+        # skip cache to avoid errors.
+        # Prevents partial cache hits from causing processing failures.
+        if inputs.hf_processor_mm_kwargs.get("use_audio_in_video"):
+            return self._apply_hf_processor(inputs, timing_ctx)
+        return super()._cached_apply_hf_processor(inputs, timing_ctx)
+
+    def _call_hf_processor(
+        self,
+        prompt: str,
+        mm_data: Mapping[str, object],
+        mm_kwargs: Mapping[str, object],
+        tok_kwargs: Mapping[str, object],
+    ):
+        mm_kwargs = _presampled_videos_hf_kwargs(mm_data, mm_kwargs)
+        return super()._call_hf_processor(
+            prompt=prompt,
+            mm_data=mm_data,
+            mm_kwargs=mm_kwargs,
+            tok_kwargs=tok_kwargs,
+        )
+
     def _maybe_apply_prompt_updates(
         self,
         mm_items: MultiModalDataItems,

vllm_omni/model_executor/models/qwen3_omni/pipeline.py

@@ -51,7 +51,6 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
             input_sources=(0,),
             hf_config_name="talker_config",
             engine_output_type="latent",
-            custom_process_input_func=f"{_PROC}.thinker2talker",
             sync_process_input_func=f"{_PROC}.thinker2talker_token_only",
             custom_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_full_payload"),
             async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"),
@@ -69,7 +68,6 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
             final_output_type="audio",
             hf_config_name="thinker_config",
             engine_output_type="audio",
-            custom_process_input_func=f"{_PROC}.talker2code2wav",
             sampling_constraints={"detokenize": True},
         ),
     ),

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -74,7 +74,6 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniAudioFeatureInputs,
     Qwen2_5OmniThinkerDummyInputsBuilder,
-    Qwen2_5OmniThinkerMultiModalProcessor,
     check_interleaved_audio_video,
     merge_interleaved_embeddings,
 )
@@ -120,6 +119,9 @@ from vllm.transformers_utils.processor import cached_processor_from_config
 
 from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin,
+    Qwen2_5OmniThinkerMultiModalDataParser,
+    Qwen2_5OmniThinkerMultiModalProcessor,
+    _presampled_videos_hf_kwargs,
 )
 from vllm_omni.quantization.component_config import (
     PRE_QUANTIZED_METHODS,
@@ -621,6 +623,19 @@ class Qwen3OmniMoeThinkerProcessingInfo(Qwen2AudioProcessingInfo, Qwen2_5_VLProc
         assert isinstance(feature_extractor, WhisperFeatureExtractor)
         return feature_extractor
 
+    def get_data_parser(self):
+        feature_extractor = self.get_feature_extractor()
+
+        # Install the Omni parser, which keeps video metadata from vLLM's
+        # video loader (see Qwen2_5OmniVideoProcessorItems); the parser
+        # inherited from Qwen2AudioProcessingInfo would drop it.
+        return Qwen2_5OmniThinkerMultiModalDataParser(
+            spatial_merge_size=self.get_hf_config().vision_config.spatial_merge_size,
+            target_sr=feature_extractor.sampling_rate,
+            target_channels=1,
+            expected_hidden_size=self._get_expected_hidden_size(),
+        )
+
     def get_supported_mm_limits(self) -> Mapping[str, int | None]:
         return {"audio": None, "image": None, "video": None}
 
@@ -671,8 +686,11 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
         tok_kwargs: Mapping[str, object],
     ) -> BatchFeature:
         mm_data = dict(mm_data)
+        mm_kwargs = dict(mm_kwargs)
         audios = mm_data.pop("audios", [])
 
+        mm_kwargs = dict(_presampled_videos_hf_kwargs(mm_data, mm_kwargs))
+
         def pad_to_hop_length(x: np.ndarray, hop_length: int) -> np.ndarray:
             length = x.shape[-1]
             if length % hop_length != 0:
@@ -927,9 +945,13 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
 
             audio_in_video_item_idx += 1
 
-            second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts", None)
-            if second_per_grid_ts:
-                video_second_per_grid_t = second_per_grid_ts[item_idx]
+            # Prefer the HF-computed value (temporal_patch_size / sampled fps)
+            # over request kwargs;
+            second_per_grid_ts = out_mm_data.get("second_per_grid_ts")
+            if second_per_grid_ts is None:
+                second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts", None)
+            if second_per_grid_ts is not None:
+                video_second_per_grid_t = float(second_per_grid_ts[item_idx])
             else:
                 video_second_per_grid_t = 2.0
 

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py

@@ -13,6 +13,7 @@ from vllm_omni.model_executor.models.common.qwen3_code_predictor import (
     CodePredictorWrapper,
     CodePredictorWrapperConfig,
 )
+from vllm_omni.platforms import current_omni_platform
 
 from .configuration_qwen3_tts import Qwen3TTSTalkerCodePredictorConfig, Qwen3TTSTalkerConfig
 
@@ -21,7 +22,7 @@ Qwen3TTSTalkerCodePredictorModelVLLM = CodePredictorBaseModel
 
 
 class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapper):
-    """Qwen3-TTS code predictor (CUDA graphs, per-call sampling, projection)."""
+    """Qwen3-TTS code predictor (per-call sampling, projection)."""
 
     def __init__(
         self,
@@ -35,7 +36,7 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapp
             vllm_config=vllm_config,
             cp_config=config,
             wrapper_config=CodePredictorWrapperConfig(
-                use_cuda_graphs=True,
+                use_cuda_graphs=current_omni_platform.is_npu(),
                 use_parallel_embedding=False,
                 use_projection=(config.hidden_size != talker_config.hidden_size),
                 return_proj_buf=False,

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -324,13 +324,33 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # OmniGPUModelRunner will store talker_mtp output under this key in
         # per-request additional_information.
         self.talker_mtp_output_key = ("codes", "audio")
+        self.talker_mtp_graph_safe = True
         # talker_mtp samples with per-row generators, so explicitly-seeded
         # requests stay batched instead of one scalar forward per row (#4883).
-        # Only valid while talker_mtp receives the unpadded active batch (this
-        # talker is not graph-wrapped); a padded batch would need the runner to
-        # pad the generators list as well.
-        self.talker_mtp_accepts_per_row_generators = True
-
+        # This is only valid while talker_mtp receives the *unpadded* active
+        # batch, i.e. when it is NOT graph-wrapped. The runner wraps talker_mtp
+        # (padded batches + a single captured RNG stream) whenever full
+        # cudagraphs are enabled, so disable per-row generators in that case.
+        # Consequence: under full cudagraphs, per-request ``tts_local_seed`` is
+        # not reproducible (matching Qwen3-Omni, which has no per-row seeding);
+        # eager mode keeps the per-row generator fast-path.
+        # TODO(#4923): the model should not read ``cudagraph_mode`` — it is an
+        # upstream (runner) concern, and the runner already re-derives the same
+        # wrap decision in ``_init_talker_mtp``. Once all TTS models declare
+        # ``talker_mtp_graph_safe`` we should move this gating into the model
+        # runner (models only expose the static capability flag, and users can
+        # then configure ``cudagraph_mode`` freely). This also means mtp seeding
+        # is silently dropped when the decode batch > 1 under full cudagraphs;
+        # the follow-up should refactor the per-row generator in the runner to
+        # be cudagraph-safe so seeded batches stay reproducible.
+        cudagraph_mode = getattr(vllm_config.compilation_config, "cudagraph_mode", None)
+        talker_mtp_graph_wrapped = (
+            self.talker_mtp_graph_safe and cudagraph_mode is not None and cudagraph_mode.has_full_cudagraphs()
+        )
+        self.talker_mtp_accepts_per_row_generators = not talker_mtp_graph_wrapped
+        self.use_async_omni_output = True
+        self.eager_omni_postprocess_before_async_output = True
+        self.omni_pooler_payload_include_hidden = False
         self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
 
         if get_pp_group().is_last_rank:

vllm_omni/model_executor/stage_input_processors/cosyvoice3.py

@@ -119,30 +119,6 @@ def _decode_additional_information(raw_info: Any) -> dict[str, Any]:
     return decoded
 
 
-def text2flow(
-    source_outputs: list[Any],
-    prompt: OmniTokensPrompt | TextPrompt = None,
-    _requires_multimodal_data: bool = True,
-):
-    """Build stage-1 inputs by prefixing stage-0 prompt ids to its outputs."""
-    engine_inputs: list[OmniTokensPrompt] = []
-    for source_output in source_outputs:
-        output = source_output.outputs[0]
-        multi_modal_data = output.multimodal_output
-        if multi_modal_data is None:
-            raise RuntimeError(f"Missing multimodal_output for request {source_output.request_id}")
-
-        prefix_ids = _ensure_list(source_output.prompt_token_ids)
-        raw_output_ids = _ensure_list(output.cumulative_token_ids)
-        prompt_speech_ids = _prompt_speech_token_ids(multi_modal_data)
-        output_ids = _strip_prompt_prefix(raw_output_ids, prefix_ids)
-        output_ids = _strip_prompt_prefix(output_ids, prompt_speech_ids)
-        additional_info = dict(multi_modal_data)
-        additional_info.setdefault("ids", {})["prompt"] = prefix_ids
-        engine_inputs.append(OmniTokensPrompt(prompt_token_ids=output_ids, additional_information=additional_info))
-    return engine_inputs
-
-
 def talker2code2wav_async_chunk(
     transfer_manager: Any,
     multimodal_output: dict[str, Any] | None,

vllm_omni/model_executor/stage_input_processors/fish_speech.py

@@ -114,37 +114,6 @@ def _extract_last_frame(multimodal_output: dict[str, Any]) -> torch.Tensor | Non
     raise ValueError(f"Invalid audio_codes shape for Fish Speech async_chunk: {tuple(audio_codes.shape)}")
 
 
-def slow_ar_to_dac_decoder(
-    source_outputs: list[Any],
-    _prompt: Any = None,
-    _requires_multimodal_data: bool = False,
-) -> list[Any]:
-    """Non-async processor: wait for Slow AR to finish, then pass all codes to DAC decoder."""
-    from vllm_omni.inputs.data import OmniTokensPrompt
-
-    slow_ar_outputs = source_outputs
-    dac_inputs: list[OmniTokensPrompt] = []
-
-    for output in slow_ar_outputs:
-        out = output.outputs[0]
-        # audio_codes shape: [num_frames, num_codebooks]
-        audio_codes = out.multimodal_output["audio_codes"].to(torch.long)
-        # Filter zero-padded frames.
-        valid_mask = audio_codes.any(dim=1)
-        audio_codes = audio_codes[valid_mask]
-        # Codebook-major flat: [num_codebooks * num_frames]
-        codec_codes = audio_codes.transpose(0, 1).cpu().reshape(-1).tolist()
-        dac_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=codec_codes,
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-                additional_information=None,
-            )
-        )
-    return dac_inputs
-
-
 def slow_ar_to_dac_decoder_async_chunk(
     transfer_manager: Any,
     multimodal_output: dict[str, Any] | None,

vllm_omni/model_executor/stage_input_processors/indextts2.py

@@ -143,65 +143,6 @@ def _request_id(request: Any) -> str:
     return str(getattr(request, "external_req_id", None) or getattr(request, "request_id", "?"))
 
 
-def talker2s2mel(
-    source_outputs: list[Any],
-    prompt: Any = None,
-    _requires_multimodal_data: bool = False,
-) -> list[Any]:
-    """Legacy orchestrator path: collect all Stage 0 output, format Stage 1 input."""
-    _old_nt = torch.get_num_threads()
-    torch.set_num_threads(1)
-    try:
-        return _talker2s2mel_impl(source_outputs)
-    finally:
-        torch.set_num_threads(_old_nt)
-
-
-def _talker2s2mel_impl(source_outputs: list[Any]) -> list[Any]:
-    from vllm_omni.inputs.data import OmniTokensPrompt
-
-    s2mel_inputs: list[OmniTokensPrompt] = []
-
-    for i, talker_output in enumerate(source_outputs):
-        if not talker_output.finished:
-            continue
-
-        output = talker_output.outputs[0]
-        mm = output.multimodal_output
-        if not isinstance(mm, dict):
-            logger.warning("Talker output %d has no multimodal_output dict", i)
-            continue
-
-        codes_dict = mm.get("codes", {})
-        mel_codes = codes_dict.get("mel")
-        if not isinstance(mel_codes, torch.Tensor) or mel_codes.numel() == 0:
-            logger.warning("Talker output %d has empty mel_codes", i)
-            continue
-        mel_codes = mel_codes.to(torch.long)
-
-        hs_dict = mm.get("hidden_states", {})
-        latent = hs_dict.get("latent")
-        if not isinstance(latent, torch.Tensor) or latent.numel() == 0:
-            logger.warning("Talker output %d has empty latent", i)
-            continue
-
-        meta = mm.get("meta", {})
-        if not isinstance(meta, dict):
-            meta = {}
-
-        additional_information = _build_s2mel_additional_information(mel_codes, latent, meta, context="talker2s2mel")
-        s2mel_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=[0],
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-                additional_information=additional_information,
-            )
-        )
-
-    return s2mel_inputs
-
-
 def talker2s2mel_token_only(
     source_outputs: list[Any],
     prompt: Any = None,

vllm_omni/model_executor/stage_input_processors/mimo_audio.py

@@ -1,7 +1,6 @@
 from typing import Any
 
 import torch
-from vllm.inputs import TextPrompt
 from vllm.logger import init_logger
 
 from vllm_omni.data_entry_keys import (
@@ -10,7 +9,6 @@ from vllm_omni.data_entry_keys import (
     OmniPayload,
     OmniPayloadStruct,
 )
-from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import TALKER_CODEC_PAD_TOKEN_ID
 
 logger = init_logger(__name__)
@@ -230,88 +228,6 @@ def llm2code2wav_async_chunk(
     )
 
 
-def llm2code2wav(
-    source_outputs: list[Any],
-    prompt: OmniTokensPrompt | TextPrompt | None = None,
-    requires_multimodal_data: bool = False,
-) -> list[OmniTokensPrompt]:
-    """
-    Process talker outputs to create code2wav inputs.
-
-    Workflow:
-    1. Extract talker's codec code outputs (8-layer RVQ codes)
-    2. Flatten codes for code2wav input
-    3. Package for code2wav stage
-
-    Args:
-        prompt: Original prompt data
-        requires_multimodal_data: Whether multimodal data is required
-
-    Returns:
-        List of OmniTokensPrompt for code2wav stage
-    """
-    talker_outputs = source_outputs
-    code2wav_inputs = []
-
-    # Process each talker output
-    for i, talker_output in enumerate(talker_outputs):
-        output = talker_output.outputs[0]
-
-        # Extract codec codes from talker output
-        # Expected shape: [8, seq_len] (8-layer RVQ codes)
-        mm = output.multimodal_output
-        mm_codes = mm.get("codes", {})
-        mm_hs = mm.get("hidden_states", {})
-        if "audio" in mm_codes:
-            codec_codes = mm_codes["audio"].to(torch.long)  # [seq_batch_size, 1, 8, 4]
-            is_all_zero = (codec_codes == 0).all(dim=(1, 2, 3))
-            non_zero_indices = (~is_all_zero).nonzero(as_tuple=True)[0]
-            if len(non_zero_indices) == 0:
-                # All codec codes are zero - skip this request with a warning
-                request_id = getattr(talker_output, "request_id", f"unknown_{i}")
-                logger.warning(
-                    "Skipping request %s: all codec codes are zero (empty output from Stage-0). "
-                    "This may indicate the model failed to generate valid audio codes.",
-                    request_id,
-                )
-            else:
-                if len(non_zero_indices) < codec_codes.shape[0]:
-                    codec_codes = codec_codes[non_zero_indices]
-        elif "output" in mm_hs and "audio" not in mm_codes:
-            codec_codes = torch.zeros(1, 1, 8, 4, dtype=torch.long)
-        else:
-            raise ValueError(f"Invalid multimodal_output: {output.multimodal_output}")
-
-        pad_vec = torch.tensor([TALKER_CODEC_PAD_TOKEN_ID] * 4)
-
-        code_final = prepend_and_flatten_colmajor(codec_codes, pad_vec)
-        code_final = code_final.tolist()
-
-        # Guard against flattened sequences longer than code2wav's max_model_len.
-        # Without this, add_request raises ``could not broadcast input array
-        # from shape (N,) into shape (max_model_len,)`` and kills the engine
-        # core (see issue #2683). Mirrors the offline end2end.py safeguard.
-        if len(code_final) > MAX_CODE2WAV_TOKENS:
-            request_id = getattr(talker_output, "request_id", f"unknown_{i}")
-            logger.warning(
-                "Request %s: code_final len=%d > MAX_CODE2WAV_TOKENS=%d, truncating.",
-                request_id,
-                len(code_final),
-                MAX_CODE2WAV_TOKENS,
-            )
-            code_final = code_final[:MAX_CODE2WAV_TOKENS]
-
-        code2wav_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=code_final,
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-            )
-        )
-
-    return code2wav_inputs
-
-
 # ============================================================================
 # Worker-connector data plane (non-async-chunk path).
 # AR runner's `flatten_payload` converts the model emit
@@ -329,8 +245,9 @@ _FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset()
 def _filter_zero_codec_rows(codec_codes: torch.Tensor) -> torch.Tensor:
     """Drop zero-padded codec rows from a 4-D `[N, 1, 8, 4]` tensor.
 
-    Mirrors the zero-row filter in the orchestrator-path `llm2code2wav`
-    body (see this file's ``llm2code2wav`` around line 224).
+    Shared by the sync placeholder builder (``llm2code2wav_token_only``) and
+    the full-payload producer (``llm2code2wav_full_payload``) so both size the
+    downstream codec sequence off the same non-zero frames.
     """
     if codec_codes.ndim != 4 or codec_codes.numel() == 0:
         return codec_codes

vllm_omni/model_executor/stage_input_processors/qwen3_omni.py

@@ -10,7 +10,6 @@ from typing import Any
 
 import torch
 from vllm.inputs import TextPrompt
-from vllm.platforms import current_platform
 
 from vllm_omni.data_entry_keys import (
     CodesStruct,
@@ -25,9 +24,7 @@ from vllm_omni.data_entry_keys import (
 from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.stage_input_processors.tts_utils import (
-    extract_language_from_prompt,
     extract_language_from_request,
-    extract_speaker_from_prompt,
     extract_speaker_from_request,
 )
 
@@ -569,25 +566,31 @@ def thinker2talker_full_payload(
         output_token_ids = _ensure_list(getattr(request, "output_token_ids", []) or [])
         all_token_ids = list(prompt_token_ids) + list(output_token_ids)
 
-    # Trim the trailing stop-token row from the accumulated thinker output.
-    # The accumulator captures one hidden-state row per executed thinker
-    # forward (prefill + every decode step including the one that emitted
-    # the stop_token), so for a finished request thinker_emb has exactly one
-    # row more than the rows the talker should consume. async_chunk's
-    # chunk-0 path naturally captures only the prefill / non-stop portion,
-    # which is why the [async_chunk] parametrization passes while [default]
-    # over-generates one codec frame on short outputs (e.g.
-    # test_one_word_prompt_001[default]: audio extends "London" with
-    # spurious phonemes).
-    if isinstance(thinker_emb, torch.Tensor) and thinker_emb.shape[0] > 0:
+    # Drop the terminal stop-token row only when more than one row was
+    # accumulated; trimming a single row would ship 0 conditioning tensors
+    # while ids still has tokens and break talker prefill alignment.
+    if isinstance(thinker_emb, torch.Tensor) and thinker_emb.shape[0] > 1:
         thinker_emb_prefill = thinker_emb[:-1]
     else:
         thinker_emb_prefill = thinker_emb
-    if isinstance(thinker_hid, torch.Tensor) and thinker_hid.shape[0] > 0:
+    if isinstance(thinker_hid, torch.Tensor) and thinker_hid.shape[0] > 1:
         thinker_hid_prefill = thinker_hid[:-1]
     else:
         thinker_hid_prefill = thinker_hid
 
+    emb_rows = int(thinker_emb_prefill.shape[0]) if isinstance(thinker_emb_prefill, torch.Tensor) else 0
+    hid_rows = int(thinker_hid_prefill.shape[0]) if isinstance(thinker_hid_prefill, torch.Tensor) else 0
+    if len(all_token_ids) > 0 and (emb_rows == 0 or hid_rows == 0):
+        logger.warning(
+            "thinker2talker_full_payload: empty thinker conditioning for req=%s "
+            "(ids_len=%s embed_rows=%s hidden_rows=%s); withholding payload.",
+            rid,
+            len(all_token_ids),
+            emb_rows,
+            hid_rows,
+        )
+        return None
+
     payload: OmniPayload = {
         "embed": {
             "prefill": thinker_emb_prefill.detach().cpu(),
@@ -599,7 +602,6 @@ def thinker2talker_full_payload(
         "ids": {"all": list(all_token_ids), "prompt": list(prompt_token_ids)},
         "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
     }
-    payload["next_stage_prompt_len"] = _compute_talker_prompt_ids_length(payload, device="cpu")
     speaker = extract_speaker_from_request(request)
     if speaker is not None:
         payload["speaker"] = speaker
@@ -609,154 +611,34 @@ def thinker2talker_full_payload(
     return payload
 
 
-def thinker2talker(
-    source_outputs: list[Any],
-    prompt: OmniTokensPrompt | TextPrompt | None = None,
-    requires_multimodal_data: bool = False,
-    streaming_context: Any | None = None,
-) -> list[OmniTokensPrompt]:
-    """
-    Process thinker outputs to create talker inputs.
-
-    Workflow:
-    1. Extract thinker's text generation outputs (token IDs + hidden states)
-    2. Split hidden states into: prompt embeddings + generated embeddings
-    3. Package for talker with additional information
-
-    In PD disaggregation mode, merges prefill-stage prompt embeddings with
-    decode-stage generated embeddings before handing off to the talker.
-
-    Args:
-        prompt: Original prompt data
-        requires_multimodal_data: Whether multimodal data is required
-
-    Returns:
-        List of OmniTokensPrompt for talker stage
-    """
-    thinker_outputs = source_outputs
-    talker_inputs: list[OmniTokensPrompt] = []
-
-    device = torch.device(current_platform.device_type)
-
-    # Process each thinker output
-    for i, thinker_output in enumerate(thinker_outputs):
-        output = thinker_output.outputs[0]
-        req_id = str(getattr(thinker_output, "request_id", f"idx-{i}"))
-        prompt_token_ids = _ensure_list(thinker_output.prompt_token_ids)
-        output_ids = _ensure_list(output.cumulative_token_ids)
-        is_streaming_session = bool(getattr(streaming_context, "enabled", False))
-        if is_streaming_session:
-            prompt_token_ids, output_ids = _get_streaming_talker_tokens(
-                req_id,
-                prompt_token_ids,
-                output_ids,
-                getattr(streaming_context, "new_prompt_len_snapshot", None),
-                streaming_context,
-                clear_state=bool(getattr(thinker_output, "finished", False)),
-            )
-        thinker_sequences = prompt_token_ids + output_ids
-        thinker_input_ids = prompt_token_ids
-        new_seq_length = len(prompt_token_ids + output_ids) - 1
-        thinker_mm_raw = getattr(output, "multimodal_output", None)
-        if not isinstance(thinker_mm_raw, Mapping):
-            logger.debug("thinker2talker: skip req=%s due to empty multimodal_output", req_id)
-            continue
-        thinker_mm: OmniPayload = thinker_mm_raw
-        mm_hs = thinker_mm.get("hidden_states", {})
-        mm_layers = mm_hs.get("layers", {}) if isinstance(mm_hs, Mapping) else {}
-        emb_layer = _layer_tensor(mm_layers, _EMBED_LAYER_KEY)
-        hid_layer = _layer_tensor(mm_layers, _HIDDEN_LAYER_KEY)
-        if emb_layer is None or hid_layer is None:
-            logger.debug("thinker2talker: skip req=%s due to missing hidden-state layers", req_id)
-            continue
-        thinker_emb = emb_layer.detach().to(device=device, dtype=torch.float)[-new_seq_length:]
-        thinker_hid = hid_layer.detach().to(device=device, dtype=torch.float)[-new_seq_length:]
-
-        prefill_mm: dict[str, Any] | None = None
-        prefill_mm = _get_prefill_multimodal_output(req_id, streaming_context)
-
-        if prefill_mm is not None:
-            expected_total = len(prompt_token_ids) + len(output_ids)
-            try:
-                thinker_emb, thinker_hid = _merge_pd_embeddings(
-                    thinker_emb, thinker_hid, prefill_mm, device, expected_total=expected_total
-                )
-            except Exception as exc:
-                logger.warning("[PD] Could not merge prefill embeddings: %s", exc)
-
-        payload = OmniPayloadStruct(
-            embed=EmbeddingsStruct(
-                prefill=thinker_emb,
-                tts_bos=_resolve_tts_token_embedding(
-                    "tts_bos", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
-                ),
-                tts_eos=_resolve_tts_token_embedding(
-                    "tts_eos", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
-                ),
-                tts_pad=_resolve_tts_token_embedding(
-                    "tts_pad", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
-                ),
-            ),
-            hidden_states=HiddenStatesStruct(output=thinker_hid),
-            ids=IdsStruct(all=thinker_sequences, prompt=thinker_input_ids),
-            speaker=extract_speaker_from_prompt(prompt, index=i),
-            language=extract_language_from_prompt(prompt, index=i),
-        )
-        info = to_dict(payload)
-        prompt_len = _compute_talker_prompt_ids_length(info, device=device)
-
-        talker_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=[0] * prompt_len,
-                additional_information=info,
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-            )
-        )
-
-    return talker_inputs
-
-
 def thinker2talker_token_only(
     source_outputs: list[Any],
     prompt: OmniTokensPrompt | TextPrompt | None = None,
     requires_multimodal_data: bool = False,
     streaming_context: Any | None = None,
 ) -> list[OmniTokensPrompt]:
-    """Non-async-chunk Stage-1 input builder for the connector data plane.
+    """Orchestrator-side placeholder builder for Stage-1 (Talker) when
+    ``async_chunk=False``.
 
-    The worker connector (Stage-0 ``thinker2talker_full_payload`` →
-    ``_sync_local_stage_payloads``) supplies the bulk talker conditioning
-    tensors (embed / hidden_states / ids) via ``model_intermediate_buffer``.
-    The orchestrator only needs to ship a placeholder prefill prompt of the
-    correct length so the scheduler can allocate KV-cache slots.
+    After the communication-layer refactor, this function only allocates a
+    placeholder ``prompt_token_ids`` of the correct length so the scheduler can
+    reserve KV-cache slots. It does **not** forward bulk tensors or voice
+    metadata.
 
-    Small per-request voice metadata (``speaker`` / ``language``) is forwarded
-    here from the user prompt so the worker's line-408 buffer seed picks it
-    up. The connector-side ``extract_speaker_from_request`` reads the
-    strongly-typed ``request.additional_information.entries["speaker"]`` which
-    currently does not always round-trip the user-supplied voice; until that
-    plumbing is normalized, providing the small fields directly preserves
-    voice selection (regression discovered on Buildkite 9668:
-    ``test_speaker_002[default]`` lost the preset voice).
+    Bulk talker conditioning (embed / hidden_states / ids) and per-request
+    voice metadata (``speaker`` / ``language``) are packed by Stage-0
+    ``thinker2talker_full_payload`` via ``extract_speaker_from_request`` /
+    ``extract_language_from_request``, sent over the connector, and merged into
+    the Talker worker's ``model_intermediate_buffer`` on recv. The orchestrator
+    path therefore sets ``additional_information=None``.
+
+    ``prompt`` / ``requires_multimodal_data`` are kept for call-site signature
+    compatibility with other orchestrator input processors; they are unused.
     """
     talker_inputs: list[OmniTokensPrompt] = []
     for i, thinker_output in enumerate(source_outputs):
         output = thinker_output.outputs[0]
         req_id = str(getattr(thinker_output, "request_id", f"idx-{i}"))
-        # Skip-on-missing parity with thinker2talker_full_payload: if the
-        # connector builder would drop this request (no MM dict or missing
-        # hidden-state layers), do the same here so the worker buffer
-        # presence agrees with the orchestrator's scheduling decision.
-        thinker_mm_raw = getattr(output, "multimodal_output", None)
-        if not isinstance(thinker_mm_raw, Mapping):
-            logger.debug("thinker2talker_token_only: skip req=%s due to empty multimodal_output", req_id)
-            continue
-        mm_hs = thinker_mm_raw.get("hidden_states", {})
-        mm_layers = mm_hs.get("layers", {}) if isinstance(mm_hs, Mapping) else {}
-        if _layer_tensor(mm_layers, _EMBED_LAYER_KEY) is None or _layer_tensor(mm_layers, _HIDDEN_LAYER_KEY) is None:
-            logger.debug("thinker2talker_token_only: skip req=%s due to missing hidden-state layers", req_id)
-            continue
         prompt_token_ids = _ensure_list(thinker_output.prompt_token_ids)
         output_ids = _ensure_list(output.cumulative_token_ids)
         is_streaming_session = bool(getattr(streaming_context, "enabled", False))
@@ -773,21 +655,10 @@ def thinker2talker_token_only(
         thinker_input_ids = prompt_token_ids
         info_for_len = {"ids": {"all": thinker_sequences, "prompt": thinker_input_ids}}
         prompt_len = _compute_talker_prompt_ids_length(info_for_len, device="cpu")
-
-        # Forward only small voice metadata; bulk tensors come from the
-        # connector path via _sync_local_stage_payloads.
-        small_info: dict[str, Any] = {}
-        speaker = extract_speaker_from_prompt(prompt, index=i)
-        if speaker is not None:
-            small_info["speaker"] = speaker
-        language = extract_language_from_prompt(prompt, index=i)
-        if language is not None:
-            small_info["language"] = language
-
         talker_inputs.append(
             OmniTokensPrompt(
                 prompt_token_ids=[0] * prompt_len,
-                additional_information=(small_info if small_info else None),
+                additional_information=None,
                 multi_modal_data=None,
                 mm_processor_kwargs=None,
             )
@@ -949,59 +820,5 @@ def talker2code2wav_full_payload(
     )
     return {
         "codes": {"audio": codec_codes},
-        "code_predictor_codes": codec_codes,
         "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
     }
-
-
-def talker2code2wav(
-    source_outputs: list[Any],
-    _prompt: OmniTokensPrompt | TextPrompt | None = None,
-    _requires_multimodal_data: bool = False,
-    streaming_context: Any | None = None,
-) -> list[OmniTokensPrompt]:
-    """
-    Process talker outputs to create code2wav inputs.
-
-    Workflow:
-    1. Extract talker's codec code outputs (8-layer RVQ codes)
-    2. Flatten codes for code2wav input
-    3. Package for code2wav stage
-
-    Args:
-    Returns:
-        List of OmniTokensPrompt for code2wav stage
-    """
-    talker_outputs = source_outputs
-    code2wav_inputs: list[OmniTokensPrompt] = []
-    # Process each talker output
-    for i, talker_output in enumerate(talker_outputs):
-        output = talker_output.outputs[0]
-        req_id = str(getattr(talker_output, "request_id", f"idx-{i}"))
-        cur_seq_len = len(output.cumulative_token_ids) - 1
-        seq_len = cur_seq_len
-        is_streaming_session = bool(getattr(streaming_context, "enabled", False))
-        if is_streaming_session:
-            seq_len = _get_streaming_codec_delta_len(cur_seq_len, req_id, talker_output, streaming_context)
-        mm_raw = getattr(output, "multimodal_output", None)
-        if not isinstance(mm_raw, Mapping):
-            logger.debug("talker2code2wav: skip req=%s due to empty multimodal_output", req_id)
-            continue
-        mm: OmniPayload = mm_raw
-        if "codes" not in mm or not isinstance(mm.get("codes"), dict) or "audio" not in mm["codes"]:
-            logger.debug("talker2code2wav: skip req=%s due to missing codes.audio", req_id)
-            continue
-        # Extract codec codes from talker output
-        # Expected shape: [8, seq_len] (8-layer RVQ codes)
-        codec_codes = (
-            mm["codes"]["audio"][-seq_len:].to(torch.long).transpose(0, 1).cpu().to(torch.long).reshape(-1).tolist()
-        )  # 16, seq_len
-        code2wav_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=codec_codes,
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-            )
-        )
-
-    return code2wav_inputs

vllm_omni/model_executor/stage_input_processors/qwen3_tts.py

@@ -38,103 +38,6 @@ def _qwen3_tts_empty_finished_payload():
     }
 
 
-def talker2code2wav(
-    source_outputs: list[Any],
-    prompt: Any = None,
-    _requires_multimodal_data: bool = False,
-) -> list[Any]:
-    """Non-async: collect all talker codes, then pass to code2wav at once."""
-    from vllm_omni.inputs.data import OmniTokensPrompt
-
-    talker_outputs = source_outputs
-    code2wav_inputs: list[OmniTokensPrompt] = []
-    for i, talker_output in enumerate(talker_outputs):
-        if not talker_output.finished:
-            # Non-async decode should only run once, after talker has
-            # accumulated the final code sequence.
-            continue
-        output = talker_output.outputs[0]
-        mm = output.multimodal_output
-        mm_codes = mm.get("codes", {})
-
-        # audio_codes shape: [num_frames, Q] where Q=num_quantizers (16)
-        audio_codes = mm_codes["audio"].to(torch.long)
-        token_ids = output.cumulative_token_ids
-
-        # token_ids provides an upper bound on the newly generated codec span.
-        # audio_codes may still contain zero-padded / invalid rows, so trim only
-        # after filtering valid frames instead of trying to align EOS indices.
-        seq_len = max(len(token_ids) - 1, 0)
-        # Filter invalid frames: zero-padded (EOS), out-of-range values (e.g.
-        # stop_token_id=2150 exceeds codebook_size=2048), and negative
-        # sentinels (e.g. -1 padding).
-        _CODEBOOK_SIZE = 2048
-        valid_mask = (
-            (audio_codes >= 0).all(dim=1) & audio_codes.any(dim=1) & (audio_codes.max(dim=1).values < _CODEBOOK_SIZE)
-        )
-        audio_codes = audio_codes[valid_mask]
-        if seq_len > 0 and audio_codes.ndim == 2 and int(audio_codes.shape[0]) > seq_len:
-            audio_codes = audio_codes[-seq_len:]
-        ref_code = mm_codes.get("ref")
-        ref_code_len = mm.get("meta", {}).get("ref_code_len")
-        if isinstance(ref_code_len, torch.Tensor):
-            ref_code_len = int(ref_code_len.reshape(-1)[-1].item()) if ref_code_len.numel() > 0 else 0
-        elif ref_code_len is None:
-            ref_code_len = 0
-        else:
-            ref_code_len = int(ref_code_len)
-        if isinstance(ref_code, list):
-            ref_code = ref_code[0] if ref_code else None
-        if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
-            ref_code = ref_code.to(torch.long).cpu().contiguous()
-            if ref_code.ndim == 1:
-                num_quantizers = int(audio_codes.shape[1]) if audio_codes.ndim == 2 and audio_codes.shape[1] > 0 else 16
-                if ref_code.numel() % num_quantizers != 0:
-                    logger.warning(
-                        "Ignoring malformed ref_code with %d elements not divisible by num_quantizers=%d",
-                        ref_code.numel(),
-                        num_quantizers,
-                    )
-                    ref_code = None
-                else:
-                    ref_code = ref_code.reshape(-1, num_quantizers)
-            elif ref_code.ndim != 2:
-                logger.warning("Ignoring malformed ref_code shape %s", tuple(ref_code.shape))
-                ref_code = None
-            if isinstance(ref_code, torch.Tensor) and ref_code_len > 0 and int(ref_code.shape[0]) > ref_code_len:
-                logger.warning(
-                    "Trimming ref_code from %d frames to ref_code_len=%d before Code2Wav.",
-                    int(ref_code.shape[0]),
-                    ref_code_len,
-                )
-                ref_code = ref_code[:ref_code_len]
-            if not isinstance(ref_code, torch.Tensor):
-                ref_code_len = 0
-            else:
-                ref_code_len = int(ref_code.shape[0])
-                audio_codes = torch.cat([ref_code.to(audio_codes.device), audio_codes], dim=0)
-        else:
-            ref_code_len = 0
-        # Code2Wav expects codebook-major flat: [Q*num_frames]
-        codec_codes = audio_codes.transpose(0, 1).cpu().reshape(-1).tolist()
-        additional_information = to_dict(
-            OmniPayloadStruct(
-                meta=MetaStruct(left_context_size=ref_code_len) if ref_code_len > 0 else None,
-                speaker=extract_speaker_from_prompt(prompt, index=i),
-                language=extract_language_from_prompt(prompt, index=i),
-            )
-        )
-        code2wav_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=codec_codes,
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-                additional_information=additional_information if additional_information else None,
-            )
-        )
-    return code2wav_inputs
-
-
 def _extract_last_frame(multimodal_output: OmniPayload | dict[str, Any]) -> torch.Tensor | None:
     audio_codes = multimodal_output.get("codes", {}).get("audio")
     if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
@@ -351,10 +254,7 @@ _NUM_QUANTIZERS_DEFAULT = 16
 
 
 def _filter_audio_codes_qwen3_tts(audio_codes: torch.Tensor) -> torch.Tensor:
-    """Filter zero-padded, out-of-range, and negative-padded codec frames.
-
-    Mirrors the orchestrator-path body in `talker2code2wav` above.
-    """
+    """Filter zero-padded, out-of-range, and negative-padded codec frames."""
     if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
         return audio_codes
     if audio_codes.ndim != 2:
@@ -367,10 +267,8 @@ def _filter_audio_codes_qwen3_tts(audio_codes: torch.Tensor) -> torch.Tensor:
 
 def _coerce_ref_code_len(raw) -> int:
     """Coerce mm["meta"]["ref_code_len"] / pooling_output["meta.ref_code_len"]
-    raw value (Tensor | int | None) into a non-negative int.  Mirrors the
-    extraction inlined in the legacy ``talker2code2wav`` path; clamps any
-    negative input to 0 since downstream code treats this as a non-negative
-    frame count."""
+    raw value (Tensor | int | None) into a non-negative int; clamps any
+    negative input to 0 since downstream code treats this as a frame count."""
     if isinstance(raw, torch.Tensor):
         value = int(raw.reshape(-1)[-1].item()) if raw.numel() > 0 else 0
     elif raw is None:
@@ -381,7 +279,7 @@ def _coerce_ref_code_len(raw) -> int:
 
 
 def _normalize_ref_code(ref_code, num_quantizers: int, ref_code_len: int):
-    """Coerce ref_code into a [ref_len, Q] tensor or None.  Mirrors orchestrator path."""
+    """Coerce ref_code into a [ref_len, Q] tensor or None."""
     if isinstance(ref_code, list):
         ref_code = ref_code[0] if ref_code else None
     if not isinstance(ref_code, torch.Tensor) or ref_code.numel() == 0:
@@ -407,8 +305,7 @@ def talker2code2wav_token_only(
 
     Sized to the expected codec token count (codebook-major flat:
     Q * (ref_frames + audio_frames)).  Speaker / language metadata are
-    extracted from `prompt` and threaded via `additional_information`
-    (orchestrator-style; same as the legacy `talker2code2wav` builder).
+    extracted from `prompt` and threaded via `additional_information`.
     Actual codec ids are delivered via the worker connector payload built
     by `talker2code2wav_full_payload`.
     """
@@ -473,8 +370,8 @@ def talker2code2wav_full_payload(
     Reads accumulated codec from `pooling_output["codes.audio"]` (CONCAT
     across steps via flatten_payload), latest `pooling_output["codes.ref"]`
     (prefill-emitted), and latest `pooling_output["meta.ref_code_len"]`.
-    Replicates the orchestrator-path body of `talker2code2wav` (filter,
-    crop to seq_len, prepend ref, codebook-major flatten).
+    Filters invalid frames, crops to seq_len, prepends ref, and flattens
+    codebook-major for code2wav consumption.
     """
     del transfer_manager
     rid = getattr(request, "request_id", "?")

vllm_omni/model_executor/stage_input_processors/voxtral_tts.py

@@ -2,7 +2,6 @@ from collections.abc import Mapping
 from typing import Any
 
 import torch
-from vllm.inputs import TextPrompt
 from vllm.logger import init_logger
 
 from vllm_omni.data_entry_keys import (
@@ -11,7 +10,6 @@ from vllm_omni.data_entry_keys import (
     OmniPayload,
     OmniPayloadStruct,
 )
-from vllm_omni.inputs.data import OmniTokensPrompt
 
 logger = init_logger(__name__)
 
@@ -31,29 +29,6 @@ def _codec_audio_tensors(multimodal_output: OmniPayload | dict[str, Any]) -> lis
     return None
 
 
-def generator2tokenizer(
-    source_outputs: list[Any],
-    _prompt: OmniTokensPrompt | TextPrompt = None,
-    _requires_multimodal_data: bool = False,
-):
-    tokenizer_inputs = []
-    for generator_output in source_outputs:
-        output = generator_output.outputs[0]
-        audio_tensors = _codec_audio_tensors(output.multimodal_output)
-        if not audio_tensors:
-            raise ValueError("Missing codes.audio in generator multimodal_output")
-        audio_tokens = torch.cat(audio_tensors, dim=-1).flatten().detach().cpu().tolist()
-        tokenizer_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=audio_tokens,
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-            )
-        )
-
-    return tokenizer_inputs
-
-
 def _extract_last_frame(multimodal_output: OmniPayload | dict[str, Any]) -> torch.Tensor | None:
     audio_tensors = _codec_audio_tensors(multimodal_output)
     if not audio_tensors:

vllm_omni/platforms/cuda/platform.py

@@ -249,16 +249,8 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
 
     @classmethod
     def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig:
-        """Copied from upstream CudaPlatformBase with inductor-aware logic.
-
-        When inductor is active (compiling) use native as the default;
-        otherwise prefer vllm_c kernels where available.
-        """
-        from vllm.config.compilation import CompilationMode
-
-        cc = vllm_config.compilation_config
-        using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
-        default = ["native"] if using_inductor else ["vllm_c", "native"]
+        """Prefer ``vllm_c`` CUDA kernels over ``native`` for diffusion IR ops."""
+        default = ["vllm_c", "native"]
 
         # Use oink if enabled for rms_norm
         # TODO(Laurawly/luka): remove this env var,

vllm_omni/plugins/__init__.py

@@ -68,6 +68,10 @@ def load_omni_general_plugins() -> None:
         return
     omni_plugins_loaded = True
 
+    from vllm.plugins import load_general_plugins
+
+    load_general_plugins()
+
     plugins = load_omni_plugins_by_group(group=OMNI_DEFAULT_PLUGINS_GROUP)
     # General plugins: we only need to execute the loaded functions.
     for func in plugins.values():

vllm_omni/quantization/__init__.py

@@ -12,10 +12,10 @@ Adds per-component quantization for multi-stage models.
 """
 
 from .component_config import ComponentQuantizationConfig
-from .factory import SUPPORTED_QUANTIZATION_METHODS, build_quant_config
+from .factory import SUPPORTED_QUANTIZATION_METHODS, build_quant_config, register_quantization_override
 from .inc_config import OmniINCConfig
 
-# Heavy configs (GGUF, MXFP8) are NOT imported here to avoid pulling in
+# Heavy configs are NOT imported here to avoid pulling in
 # optional dependencies (pynvml, torch_npu) at module load time.
 # Import them directly when needed:
 #   from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
@@ -25,4 +25,5 @@ __all__ = [
     "ComponentQuantizationConfig",
     "OmniINCConfig",
     "SUPPORTED_QUANTIZATION_METHODS",
+    "register_quantization_override",
 ]

vllm_omni/quantization/factory.py

@@ -3,8 +3,8 @@
 """Factory for building quantization configs.
 
 build_quant_config() delegates to vLLM's quantization registry.
-The only extension point is _OVERRIDES for methods that need a
-different QuantizationConfig subclass in the OMNI context (e.g. GGUF).
+Out-of-tree integrations can register Omni-specific builders with
+register_quantization_override().
 """
 
 from __future__ import annotations
@@ -81,13 +81,6 @@ from .component_config import ComponentQuantizationConfig  # noqa: E402
 logger = init_logger(__name__)
 
 
-def _build_gguf(**kw: Any) -> QuantizationConfig:
-    """Lazy import to avoid pulling in CUDA/pynvml at module load time."""
-    from .gguf_config import DiffusionGGUFConfig
-
-    return DiffusionGGUFConfig(**kw)
-
-
 def _build_int8(**kw: Any) -> QuantizationConfig:
     """Lazy import for Int8 diffusion config (supports CUDA + NPU)."""
     from .int8_config import DiffusionInt8Config
@@ -140,7 +133,6 @@ def _build_inc(**kw: Any) -> QuantizationConfig:
 
 
 _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = {
-    "gguf": _build_gguf,
     "int8": _build_int8,
     "mxfp8": _build_mxfp8,
     "mxfp4": _build_mxfp4,
@@ -150,7 +142,24 @@ _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = {
     "auto_round": _build_inc,
 }
 
-SUPPORTED_QUANTIZATION_METHODS: list[str] = list(dict.fromkeys(QUANTIZATION_METHODS + list(_OVERRIDES.keys())))
+
+def _compute_supported_quantization_methods() -> list[str]:
+    return list(dict.fromkeys(QUANTIZATION_METHODS + list(_OVERRIDES.keys())))
+
+
+SUPPORTED_QUANTIZATION_METHODS: list[str] = _compute_supported_quantization_methods()
+
+
+def _refresh_registered_methods() -> None:
+    SUPPORTED_QUANTIZATION_METHODS[:] = _compute_supported_quantization_methods()
+    global _CACHED_ALIAS_MAP
+    _CACHED_ALIAS_MAP = None
+
+
+def register_quantization_override(method: str, builder: Callable[..., QuantizationConfig]) -> None:
+    """Register an Omni-specific quantization config builder."""
+    _OVERRIDES[_normalize_method_name(method)] = builder
+    _refresh_registered_methods()
 
 
 def _build_reverse_alias_map() -> dict[str, str]:

vllm_omni/quantization/gguf_config.py

@@ -1,323 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""GGUF quantization config for diffusion transformers.
-
-Uses dequant+GEMM instead of the fused kernel path (which expects 2D inputs).
-
-The GGUF quantization module was migrated to an external plugin in upstream
-vLLM (commit 6635279d8).  This file inlines the base classes that
-DiffusionGGUFConfig / DiffusionGGUFLinearMethod depend on, reproduced from
-the last upstream version before the migration.
-"""
-
-from __future__ import annotations
-
-from collections.abc import Mapping
-from types import MappingProxyType
-from typing import TYPE_CHECKING, Any
-
-import gguf
-import torch
-from gguf import GGMLQuantizationType as WeightType
-from torch.nn.parameter import Parameter, UninitializedParameter
-from vllm import _custom_ops as ops
-from vllm.logger import init_logger
-from vllm.model_executor.layers.linear import (
-    LinearBase,
-    LinearMethodBase,
-    UnquantizedLinearMethod,
-)
-from vllm.model_executor.layers.quantization.base_config import (
-    QuantizationConfig,
-    QuantizeMethodBase,
-)
-from vllm.model_executor.utils import set_weight_attrs
-
-if TYPE_CHECKING:
-    from vllm.model_executor.layers.quantization import QuantizationMethods
-
-logger = init_logger(__name__)
-
-# ---------------------------------------------------------------------------
-# GGUF base types (inlined from upstream vllm before migration to plugin)
-# ---------------------------------------------------------------------------
-
-UNQUANTIZED_TYPES: set[WeightType] = {WeightType.F32, WeightType.F16, WeightType.BF16}
-
-STANDARD_QUANT_TYPES: set[WeightType] = {
-    WeightType.Q4_0,
-    WeightType.Q4_1,
-    WeightType.Q5_0,
-    WeightType.Q5_1,
-    WeightType.Q8_0,
-    WeightType.Q8_1,
-}
-KQUANT_TYPES: set[WeightType] = {
-    WeightType.Q2_K,
-    WeightType.Q3_K,
-    WeightType.Q4_K,
-    WeightType.Q5_K,
-    WeightType.Q6_K,
-}
-IMATRIX_QUANT_TYPES: set[WeightType] = {
-    WeightType.IQ1_M,
-    WeightType.IQ1_S,
-    WeightType.IQ2_XXS,
-    WeightType.IQ2_XS,
-    WeightType.IQ2_S,
-    WeightType.IQ3_XXS,
-    WeightType.IQ3_S,
-    WeightType.IQ4_XS,
-    WeightType.IQ4_NL,
-}
-DEQUANT_TYPES: set[WeightType] = STANDARD_QUANT_TYPES | KQUANT_TYPES | IMATRIX_QUANT_TYPES
-
-
-def is_layer_skipped_gguf(
-    prefix: str,
-    unquantized_modules: list[str],
-    fused_mapping: Mapping[str, list[str]] = MappingProxyType({}),
-) -> bool:
-    # Fused layers like gate_up_proj or qkv_proj will not be fused
-    # in the safetensors checkpoint. So, we convert the name
-    # from the fused version to unfused + check to make sure that
-    # each shard of the fused layer has the same scheme.
-    proj_name = prefix.split(".")[-1]
-    if proj_name in fused_mapping:
-        shard_prefixes = [prefix.replace(proj_name, shard_proj_name) for shard_proj_name in fused_mapping[proj_name]]
-
-        is_skipped = None
-        for shard_prefix in shard_prefixes:
-            is_shard_skipped = any(shard_prefix in module_name for module_name in unquantized_modules)
-
-            if is_skipped is None:
-                is_skipped = is_shard_skipped
-            elif is_shard_skipped != is_skipped:
-                raise ValueError(
-                    f"Detected some but not all shards of {prefix} "
-                    "are quantized. All shards of fused layers "
-                    "to have the same precision."
-                )
-    else:
-        is_skipped = any(module_name in prefix for module_name in unquantized_modules)
-
-    assert is_skipped is not None
-    return is_skipped
-
-
-class _GGUFUninitializedParameter(UninitializedParameter):
-    """UninitializedParameter subclass that stores GGUF weight data in a list.
-
-    The actual tensor is materialized later via _create_padded_weight_param.
-    """
-
-    cls_to_become = Parameter
-    data_container: list[torch.Tensor]
-
-
-class GGUFConfig(QuantizationConfig):
-    """Config class for GGUF (inlined from upstream vllm)."""
-
-    def __init__(self, unquantized_modules: list[str] | None = None) -> None:
-        super().__init__()
-        self.unquantized_modules = unquantized_modules or []
-
-    def __repr__(self) -> str:
-        return "GGUFConfig()"
-
-    def get_name(self) -> QuantizationMethods:
-        return "gguf"
-
-    def get_supported_act_dtypes(self) -> list[torch.dtype]:
-        return [torch.half, torch.bfloat16, torch.float32]
-
-    @classmethod
-    def get_min_capability(cls) -> int:
-        return 60
-
-    @classmethod
-    def get_config_filenames(cls) -> list[str]:
-        return []  # no extra configs.
-
-    @classmethod
-    def from_config(cls, config: dict[str, Any]) -> GGUFConfig:
-        return cls()
-
-    @classmethod
-    def override_quantization_method(
-        cls, hf_quant_cfg: dict[str, Any], user_quant: str | None, hf_config=None
-    ) -> QuantizationMethods | None:
-        if user_quant == "gguf":
-            return "gguf"
-        return None
-
-    def get_quant_method(self, layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None:
-        if isinstance(layer, LinearBase):
-            if is_layer_skipped_gguf(prefix, self.unquantized_modules, self.packed_modules_mapping):
-                return UnquantizedLinearMethod()
-            return GGUFLinearMethod(self)
-        return None
-
-
-class GGUFLinearMethod(LinearMethodBase):
-    """Linear method for GGUF (inlined from upstream vllm)."""
-
-    def __init__(self, quant_config: GGUFConfig):
-        self.quant_config = quant_config
-
-    def create_weights(
-        self,
-        layer: torch.nn.Module,
-        input_size_per_partition: int,
-        output_partition_sizes: list[int],
-        input_size: int,
-        output_size: int,
-        params_dtype: torch.dtype,
-        **extra_weight_attrs,
-    ):
-        self.params_dtype = params_dtype
-        output_size_per_partition = sum(output_partition_sizes)
-
-        tensor_shape = (output_size_per_partition, input_size_per_partition)
-        qweight = _GGUFUninitializedParameter(requires_grad=False)
-        set_weight_attrs(
-            qweight,
-            {
-                "input_dim": 1,
-                "output_dim": 0,
-                "tensor_shape": tensor_shape,
-                "is_gguf_weight": True,
-                "data_container": [],
-                "shard_id": [],
-                "shard_id_map": {},
-            },
-        )
-        set_weight_attrs(qweight, extra_weight_attrs)
-        layer.register_parameter("qweight", qweight)
-
-        qweight_type = Parameter(
-            torch.empty(len(output_partition_sizes), dtype=torch.uint8),
-            requires_grad=False,
-        )
-        set_weight_attrs(
-            qweight_type,
-            {
-                "is_gguf_weight_type": True,
-                "weight_type": 0,
-                "shard_weight_type": {},
-                "ignore_warning": True,
-            },
-        )
-        set_weight_attrs(qweight_type, extra_weight_attrs)
-        layer.register_parameter("qweight_type", qweight_type)
-
-    def process_weights_after_loading(self, layer: torch.nn.Module):
-        qweight_type = layer.qweight_type.weight_type
-        if not (qweight_type in UNQUANTIZED_TYPES or qweight_type in DEQUANT_TYPES):
-            qweight_type = WeightType(qweight_type)
-            raise ValueError(f"Unsupported GGUF quantization type {qweight_type} in layer {layer}.")
-        self._create_padded_weight_param(layer)
-
-    def _create_padded_weight_param(self, layer: torch.nn.Module):
-        """Create padded weight parameter for GGUF MergedLinear layer."""
-        qweight = layer.qweight
-        shard_id_map = qweight.shard_id_map
-        shard_id = qweight.shard_id
-        if len(data_container := qweight.data_container) > 1:
-            dtype = {data.dtype for data in data_container}
-            assert len(dtype) == 1, ValueError(f"Data container has mixed dtypes: {dtype}")
-            dtype = next(iter(dtype))
-            padded_side = max(x.size(1) for x in data_container)
-            concat_side = sum(x.size(0) for x in data_container)
-            padded_data = torch.zeros((concat_side, padded_side), dtype=dtype, device=qweight.device)
-            shard_offset_map = dict[str, tuple[int, int, int]]()
-            for idx in shard_id:
-                id_in_container = shard_id_map[idx]
-                start = sum(x.size(0) for x in data_container[:id_in_container])
-                end = start + data_container[id_in_container].size(0)
-                size = data_container[id_in_container].size(1)
-                padded_data[start:end, :size] = data_container[id_in_container]
-                shard_offset_map[idx] = (start, end, size)
-            qweight.data_container.clear()
-            padded_param = Parameter(padded_data, requires_grad=False)
-            set_weight_attrs(padded_param, vars(qweight))
-            set_weight_attrs(padded_param, {"shard_offset_map": shard_offset_map})
-            layer.register_parameter("qweight", padded_param)
-
-    def apply(
-        self,
-        layer: torch.nn.Module,
-        x: torch.Tensor,
-        bias: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        # Base apply: overridden by DiffusionGGUFLinearMethod.
-        raise NotImplementedError("Use DiffusionGGUFLinearMethod for diffusion GGUF")
-
-
-# ---------------------------------------------------------------------------
-# Diffusion-specific GGUF (dequant+GEMM, replaces fused kernel path)
-# ---------------------------------------------------------------------------
-
-
-def dequant_gemm_gguf(x: torch.Tensor, qweight: torch.Tensor, qweight_type: int) -> torch.Tensor:
-    if qweight_type in UNQUANTIZED_TYPES:
-        return x @ qweight.T
-    block_size, type_size = gguf.GGML_QUANT_SIZES[qweight_type]
-    shape = (qweight.shape[0], qweight.shape[1] // type_size * block_size)
-    weight = ops.ggml_dequantize(qweight, qweight_type, *shape, x.dtype)
-    return x @ weight.T
-
-
-class DiffusionGGUFLinearMethod(GGUFLinearMethod):
-    """GGUF linear method using dequant+GEMM for N-D diffusion tensors."""
-
-    def apply(
-        self,
-        layer: torch.nn.Module,
-        x: torch.Tensor,
-        bias: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        shard_id = getattr(layer.qweight, "shard_id", None)
-
-        if shard_id:
-            shard_id = ["q", "k", "v"] if "q" in shard_id else shard_id
-            qweight = layer.qweight
-            result = []
-            for idx in shard_id:
-                start, end, offset = layer.qweight.shard_offset_map[idx]
-                qweight_type = layer.qweight_type.shard_weight_type[idx]
-                result.append(
-                    dequant_gemm_gguf(
-                        x,
-                        qweight[start:end, :offset].contiguous(),
-                        qweight_type,
-                    )
-                )
-            out = torch.cat(result, axis=-1)
-        else:
-            qweight = layer.qweight
-            qweight_type = layer.qweight_type.weight_type
-            out = dequant_gemm_gguf(x, qweight, qweight_type)
-
-        if bias is not None:
-            out.add_(bias)
-        return out
-
-
-class DiffusionGGUFConfig(GGUFConfig):
-    """GGUF config that carries gguf_model path and uses dequant+GEMM."""
-
-    def __init__(
-        self,
-        gguf_model: str | None = None,
-        unquantized_modules: list[str] | None = None,
-    ) -> None:
-        super().__init__(unquantized_modules=unquantized_modules or [])
-        self.gguf_model = gguf_model
-
-    def get_quant_method(self, layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None:
-        if isinstance(layer, LinearBase):
-            if is_layer_skipped_gguf(prefix, self.unquantized_modules, self.packed_modules_mapping):
-                return UnquantizedLinearMethod()
-            return DiffusionGGUFLinearMethod(self)
-        return None

vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py

@@ -280,7 +280,6 @@ def _build_quantization_config(
     quantization: str | None,
     quantization_config_json: str | None,
     ignored_layers: str | None,
-    gguf_model: str | None,
 ) -> dict[str, Any] | None:
     explicit = _load_json_arg(quantization_config_json)
     if explicit is not None:
@@ -292,8 +291,6 @@ def _build_quantization_config(
     ignored = _split_csv(ignored_layers)
     if ignored:
         config["ignored_layers"] = ignored
-    if gguf_model is not None:
-        config["gguf_model"] = gguf_model
     return config
 
 
@@ -305,7 +302,6 @@ def _build_variant_config(args: argparse.Namespace, variant: str) -> VariantConf
         quantization=quantization,
         quantization_config_json=getattr(args, f"{variant}_quantization_config_json"),
         ignored_layers=ignored_layers,
-        gguf_model=getattr(args, f"{variant}_gguf_model"),
     )
     return VariantConfig(
         label=variant,
@@ -554,8 +550,6 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--reference-ignored-layers", help="Comma-separated reference ignored layer patterns.")
     parser.add_argument("--candidate-ignored-layers", help="Comma-separated candidate ignored layer patterns.")
     parser.add_argument("--ignored-layers", help="Alias for --candidate-ignored-layers.")
-    parser.add_argument("--reference-gguf-model", help="Reference GGUF file or HF reference when using gguf.")
-    parser.add_argument("--candidate-gguf-model", help="Candidate GGUF file or HF reference when using gguf.")
     parser.add_argument("--reference-load-format", default="default", help="Reference diffusion load format.")
     parser.add_argument("--candidate-load-format", default="default", help="Candidate diffusion load format.")
     parser.add_argument("--prompt", default="a cup of coffee on the table")