diff --git a/docs/cli/serve.md b/docs/cli/serve.md
index cf3d068a6..f77503091 100644
--- a/docs/cli/serve.md
+++ b/docs/cli/serve.md
@@ -1,5 +1,67 @@
# vllm-omni serve
+## Multiple API frontends
+
+For local EngineCore pipelines with one or more stages, `--api-server-count` starts
+multiple API frontend processes that share one parent-owned set of stage
+engines:
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+ --api-server-count 2 \
+ --port 8091
+```
+
+This mode is intended for workloads whose request parsing, multimodal input
+processing, or response encoding can bottleneck a single frontend. It
+currently requires local multi-process EngineCore stages with one local
+process group per replica. A stage may still use multiple `num_replicas`.
+It cannot be combined with diffusion, headless or
+remote stages, intra-stage data parallelism, Ray, fault tolerance, elastic
+expert parallelism, sleep mode, or runtime LoRA updating. Runtime voice upload
+and deletion are also disabled because those mutations are process-local;
+built-in voices, inline reference audio, and voices restored at startup remain
+available. The `/v1/omni/sleep` and `/v1/omni/wakeup` control routes return HTTP
+409 for the same reason: their bookkeeping is process-local while stage engines
+are shared.
+
+A single-stage model pipeline is supported through the normal launch command;
+this is distinct from the distributed `--stage-id` / `--headless` launch mode.
+Every stage in the pipeline must be an EngineCore stage. Mixed AR-to-DiT
+pipelines, including MiniMax-H3, remain unsupported even when their first
+stage is an EngineCore stage. Diffusion support requires per-frontend channels
+and shared-engine lifecycle handling in the diffusion backend.
+
+### Workloads and performance validation
+
+Each API process owns its frontend processor, Orchestrator, request state,
+and per-replica EngineCore channels. Requests return through the channel of
+the originating API process; model weights and GPU stage engines are shared.
+This distributes frontend and orchestration work across CPU processes.
+
+| Workload | Potential benefit | Limitation |
+| ---------- | ------------------- | ------------ |
+| Concurrent cold image/video inputs on supported EngineCore pipelines | Parallel frontend processing when one API process is CPU-limited | A single cold request and GPU computation may be unchanged |
+| Many replicas producing frequent text/audio outputs | Less API/Orchestrator work and GIL contention per process | Routing remains serial within each Orchestrator |
+| Warm inputs or low-concurrency, GPU-limited generation | Additional frontend capacity | More API processes may add overhead without improving throughput |
+
+Here, a cold input means media or preprocessing caches miss, not that model
+weights must be loaded or GPU kernels compiled. Each frontend has its own
+processor/cache state, so cache duplication, CPU thread contention, and memory
+usage also matter. Multi-API serving does not itself change the output queue,
+move response encoding to an executor, or change the model's numerical scheduler.
+
+Compare `--api-server-count 1`, `2`, and higher counts with the same hardware,
+model/stage configuration, media, sampling parameters, output lengths, and
+offered load. Keep cold-cache and warmed-cache runs separate and report repeated
+measurements. Record CPU and request counts for every API worker, input and
+output queue delay, GPU utilization, first-output latency, and end-to-end latency.
+Confirm requests reach every worker: connection reuse can concentrate traffic
+on one process. Queue-only or single-API experiments do not establish a
+multi-API throughput gain. The multi-replica workload in
+[#4680](https://github.com/vllm-project/vllm-omni/issues/4680) remains a relevant
+validation target.
+
## Stage-based CLI quickstart
The stage-based CLI is designed for deployments that require launching each pipeline stage in an isolated process
@@ -61,3 +123,10 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
## Arguments
--8<-- "docs/generated/argparse_omni/omni_serve.inc.md"
+
+Async video jobs use process-local metadata and task handles. With multiple API
+workers, creating, listing, retrieving, downloading, and deleting these jobs
+returns HTTP 409. This restriction is independent of the current diffusion
+stage launch restriction: supporting multi-worker diffusion also requires a
+shared job lifecycle. Missing frontend topology state returns HTTP 503 for
+operations that require process-local state protection.
diff --git a/docs/design/architecture_overview.md b/docs/design/architecture_overview.md
index 862516831..8b7168824 100644
--- a/docs/design/architecture_overview.md
+++ b/docs/design/architecture_overview.md
@@ -323,7 +323,7 @@ inputs, fields, and ownership rules are described below.
flowchart TB
layer1["Layer 1 · Authoring inputs
PipelineConfig + DeployConfig"]
layer2["Layer 2 · Production resolution boundary
resolve_omni_config()
(StageConfigFactory + VllmOmniConfig internally)"]
- layer3["Layer 3 · Startup hand-off
OmniConfigResolution
PipelineConfig + temporary OmegaConf stage bridge"]
+ layer3["Layer 3 · Startup hand-off
OmniConfigResolution
PipelineConfig + typed stage configs"]
layer4["Layer 4 · Runtime launch planning
StageRuntime"]
layer5["Layer 5 · Engine materialization
VllmConfig / OmniDiffusionConfig"]
@@ -335,10 +335,10 @@ The production resolution boundary is `resolve_omni_config()` in
It delegates typed construction to `StageConfigFactory.create_from_model()` and
`VllmOmniConfig.from_pipeline_config()`, then returns an
`OmniConfigResolution` consumed by both `AsyncOmniEngine` and headless startup.
-Until `StageRuntime` consumes typed stage configs directly, this envelope carries
-the effective `PipelineConfig` alongside OmegaConf-compatible `stage_configs` as
-a temporary runtime bridge; both views describe the same resolved topology,
-including injected stages.
+This envelope carries the effective `PipelineConfig` alongside typed
+`stage_configs`, which `StageRuntime` consumes directly. Both views describe the
+same resolved topology, including injected stages. Backend arguments are
+projected from these typed stages when each engine is initialized.
The legacy `stage_args` YAML path has been removed. Model topology now resolves
through `PipelineConfig`, with runtime overrides supplied by `DeployConfig`.
@@ -360,7 +360,7 @@ The important ownership rules are:
3. CLI and Python overrides are applied at the resolution boundary, with
per-stage overrides taking precedence over global values where supported.
4. `OmniConfigResolution` is the sole startup hand-off. Its `pipeline_config`
- and temporary `stage_configs` compatibility view must describe the same
+ and typed `stage_configs` must describe the same
topology.
5. `StageRuntime` owns launch planning and replica lifecycle. `ReplicaInitPlan`
is runtime-private state, not a user configuration object.
diff --git a/docs/design/fullduplex-personaplex.md b/docs/design/fullduplex-personaplex.md
index 33d125c97..6d32cd8e3 100644
--- a/docs/design/fullduplex-personaplex.md
+++ b/docs/design/fullduplex-personaplex.md
@@ -271,3 +271,111 @@ The required evidence is:
Audio that is empty, all zero, non-finite, or only a protocol `listen` event is
not a successful end-to-end result.
+
+### Paced multi-session load driver
+
+The existing driver keeps its two-session admission and slot-reuse validation
+when `--sessions` is omitted. Passing `--sessions N` selects a separate,
+bounded load mode; it does not replace the lifecycle test or change server
+admission capacity. Start an existing PersonaPlex server configured to admit
+the requested number of sessions, then run from a source checkout with the
+vLLM-Omni client dependencies installed:
+
+```bash
+python -m tests.e2e.online_serving.personaplex_realtime_duplex \
+ --model nvidia/personaplex-7b-v1 \
+ --input-wav /path/to/input.wav \
+ --sessions 2 --load-frames 1000 --drain-s 2 \
+ --server-revision SERVER_COMMIT --server-hardware A100-80GB \
+ --output-dir /tmp/personaplex-load-n2
+```
+
+Use a new output directory for every run. The driver refuses to overwrite an
+existing `load-result.json`. The server labels are supplied by the operator,
+not detected or verified by the client. Record the server command, package
+versions, checkpoint revision and hardware alongside a published measurement.
+
+All attempted sessions finish admission before admitted sessions share a
+common input start. A rejected session remains a failed row in the report;
+it is not removed from the denominator. Each session receives at most
+`--load-frames` input frames, without looping the WAV, and uses 80 ms pacing.
+A stalled send cannot cause a catch-up burst. The timeline records send
+schedule lateness so an overloaded client is visible rather than mistaken
+for a faster server. This mode measures steady concurrency, not join/leave
+interference; the original lifecycle mode still checks slot reuse.
+Load mode does not probe an additional N+1 session for `resource_exhausted`;
+the lifecycle mode owns that admission-limit check. A per-session
+`client_pacing_warning` and the summary's `client_pacing_warning_sessions`
+identify send-lateness p99 above one nominal frame (80 ms). This is a warning,
+not an extra acceptance gate: read it alongside server-facing measurements
+so a slow client is not mistaken for a slow server.
+
+The fixed drain window is part of the measurement configuration. Audio
+arriving during cleanup is excluded, so closing cannot rescue a session that
+missed its observation deadline. Cleanup is attempted for admitted sessions
+even when audio validation fails. Protocol errors, missing/invalid audio,
+shared response IDs and cleanup failures make the report fail and the CLI
+exit nonzero. Cancellation closes transports and propagates to the caller.
+`--cleanup-timeout-s` (default 5 s) bounds session-close send and acknowledgement
+together, independently of `--timeout-s`. The subsequent WebSocket close
+handshake has the same separate bound; an unresponsive close acknowledgement
+cannot hold cancellation for the full inference timeout.
+
+`load-result.json` retains per-session send and audio-packet timelines using
+the client's monotonic clock. It does not export handshake credentials or
+raw audio payloads. Metric definitions are explicit:
+
+Both modes validate each nonempty audio packet's strict base64 encoding, PCM16
+sample alignment, 24 kHz sample rate and response identity. Valid aggregate
+audio cannot hide a malformed packet. In load mode, rejected packets make the session fail
+and appear in `invalid_audio_packets` with their event index, receive time and
+fixed diagnostic code; they do not contribute samples or timing intervals.
+Empty audio deltas remain legal and contribute no samples. This check does
+not change the public client's decoding. The lifecycle scenario and dispatch
+remain the same, but malformed packets that aggregate checks once tolerated
+now fail that mode as well.
+
+The probe accepts both the current OpenAI Realtime audio event name
+`response.output_audio.delta` and the legacy `response.audio.delta`. Validation
+uses the wire payload directly rather than relying on a particular client
+library's alias table, so a server-side event-name migration cannot silently
+turn real audio into a zero-output measurement.
+
+| Field | Meaning |
+| --- | --- |
+| `client_audio_packet_interval_ms` | Consecutive nonempty packet arrivals: median, linear p99, max, and zero-based index of the first maximal interval. Interval index 0 ends at packet 1. One packet can contain multiple codec frames; this is not server tick latency. |
+| `client_first_audio_after_stream_start_ms` | First audio receipt minus first input send start; admission is excluded. |
+| `client_stream_rtf` | First input send start to last audio receipt, divided by received audio duration. This includes real-time input pacing and is not an inference-only RTF. |
+| `client_send_lateness_ms` | Send start minus its original shared schedule, clipped at zero. |
+| `frame_deficit` | Integer input-frame minus output-frame count from the shared session validation; missing when valid whole-frame accounting is unavailable. |
+| `output_frames`, `voiced_frames`, `silent_frames`, `frame_coverage_ratio`, `audio_rms`, `audio_chunks` | Session-validation statistics retained on success and on acceptance failure once computed. |
+| `acceptance_check` | Failed local check, with numeric diagnostics in `error`; other exceptions retain phase/type only, without remote handshake or metadata text. |
+| `audio_underrun_s`, `audio_underrun_event_count`, `audio_continuity_ok` | Shared `compute_continuity_stats` on valid PCM16 packet arrivals, with its default 100 ms underrun threshold; absent audio gives missing values rather than a continuity pass. |
+
+No output gives missing latency/RTF values, not zero. By default RTF is
+diagnostic; `--max-client-rtf` adds an explicit ceiling. The existing frame
+coverage, audible-frame and scheduler-data-plane checks still apply. A longer
+`--drain-s` can demonstrate eventual completion under backlog, but it must not
+be reported as real-time capacity when the fixed-window run failed or client
+RTF exceeded the stated target. Likewise, the first live session can pay
+one-time model/codec/JIT costs; either use the server's configured duplex
+warmup path or label that run as cold-start evidence before comparing steady
+capacity. Passing client-only metrics does not certify model quality or a
+real-time capacity target on another GPU. Begin with N=1 and N=2; higher N is a
+capacity experiment, not a promised property of an A100 or of this patch.
+Continuity is a client-arrival diagnostic, not proof of speaker playback and
+not an additional pass/fail gate. Packet-interval and send-lateness summaries
+retain their documented linear p99; the older client interval helper reports
+rounded nearest-rank p95 instead and is not substituted for that schema.
+
+Deterministic driver and localhost WebSocket checks need no model or GPU:
+
+```bash
+python -m pytest tests/e2e/online_serving/test_personaplex_load_driver.py \
+ -m 'core_model and cpu' --run-level core_model -q
+```
+
+These checks test the driver against controlled protocol responses. Run the
+load-driver command against real weights separately before reporting model-serving
+performance. Server tick timestamps and generic TTFT/TPOT aggregation remain
+the responsibility of the existing duplex metrics paths, not this driver.
diff --git a/docs/design/module/entrypoints.md b/docs/design/module/entrypoints.md
index db62f9cb6..f37ddc6ba 100644
--- a/docs/design/module/entrypoints.md
+++ b/docs/design/module/entrypoints.md
@@ -94,6 +94,13 @@ This document owns offline API semantics, CLI and serve composition,
supported OpenAI-compatible routes, request validation and normalization,
response conversion, streaming/session behavior, and engine handoff.
+When `--api-server-count` is greater than one, the serve composition root
+starts the frontend process manager while a parent-owned `StageRuntime`
+launches the shared local stage engines. Each frontend receives only its own
+stage-client channel configuration and does not launch or retire backend
+processes. CLI validation rejects unsupported distributed, diffusion,
+fault-tolerant, elastic-EP, Ray, and runtime-LoRA combinations before startup.
+
It does not own configuration precedence, cross-stage routing, stage
placement, payload implementation, or semantic error classification.
`entrypoints/openai/errors.py` is an explicit primary-path exception owned by
@@ -141,6 +148,9 @@ stage so `wake_up(stage_ids=[0])` does not skip a later `wake_up(stage_ids=[1])`
Streaming input pumps take an admission slot immediately before each EngineCore
ADD or update, not while waiting for the next client chunk. Frontend abort
keeps `request_states` until `generate()` consumes the terminal output.
+Multi-API serving must preserve one client rank and one channel set per
+frontend, include every frontend in readiness/failure observation, and keep
+stage lifecycle ownership in the parent composition root.
## Promotion gate
diff --git a/docs/design/module/stage_runtime.md b/docs/design/module/stage_runtime.md
index aefba5a40..38255880f 100644
--- a/docs/design/module/stage_runtime.md
+++ b/docs/design/module/stage_runtime.md
@@ -77,6 +77,22 @@ readiness, one pool per logical stage, replica identity and membership,
selection and affinity, stage-client/process lifecycle, liveness, draining,
and shutdown.
+Local EngineCore initialization and multi-API serving use
+`StageRuntime.launch_stage_engines(num_api_servers=1)` as their common launch
+entry. It resolves plans when not supplied by the runtime's device-group
+scheduler, coordinates device locks and environment overlays, starts engines,
+and rolls back resources on startup or attachment failure. Single-API clients
+attach locally after readiness and take ownership of the engine resources;
+there is no additional API subprocess. Diffusion and remote attachment retain
+their backend-specific initialization.
+
+For local multi-API serving, the parent runtime owns and launches each stage
+engine once. It supplies a distinct input/output channel for every API
+frontend and stage replica; frontend workers attach clients to those engines
+without acquiring stage-process ownership. Readiness observation and shutdown
+remain parent-owned. This topology currently excludes diffusion and
+remote/headless stages, intra-stage data parallelism, and Ray backends.
+
It does not own cross-stage request policy, model scheduler policy, connector
transfer semantics, configuration precedence, or public error rendering.
@@ -114,6 +130,10 @@ affinity, membership changes, abort routing, process failure, draining, and
repeated shutdown. Cross-stage policy changes belong in
`engine_orchestration.md`.
+For multi-API changes, additionally exercise per-frontend channel assignment,
+shared-engine startup handshakes, frontend failure observation, and parent
+cleanup without duplicate stage shutdown.
+
## Promotion gate
- Reconcile terminology and ownership after #5441 reaches a final state.
diff --git a/docs/user_guide/examples/online_serving/qwen3_omni.md b/docs/user_guide/examples/online_serving/qwen3_omni.md
index 2c85c33fe..654802c5c 100644
--- a/docs/user_guide/examples/online_serving/qwen3_omni.md
+++ b/docs/user_guide/examples/online_serving/qwen3_omni.md
@@ -2,7 +2,6 @@
Source .
-
## 🛠️ Installation
Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
@@ -15,6 +14,19 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
+For a local deployment with multiple API frontend processes sharing one set
+of stage engines, add `--api-server-count`:
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+ --api-server-count 2 --port 8091
+```
+
+This mode supports local EngineCore stages with one local process group per
+replica, including stages configured with multiple `num_replicas`. Headless,
+remote, diffusion, intra-replica data-parallel, and Ray stage deployments
+remain single-frontend.
+
The default deployment configuration situated at `vllm_omni/deploy/qwen3_omni_moe.yaml` is resolved and loaded
automatically via the model registry, obviating the necessity for the `--deploy-config` flag in standard deployment topologies.
The bundled Qwen3-Omni setup explicitly sets `moe_backend: triton` for the Thinker and Talker stages. This keeps
@@ -27,6 +39,7 @@ on the Thinker stage. Keep Talker and Code2Wav at `pipeline_parallel_size: 1`. T
default asynchronous scheduling and audio chunk streaming can remain enabled.
To explicitly utilize a custom deployment YAML, specify the configuration path:
+
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config /path/to/deploy_config_file
@@ -50,7 +63,7 @@ Adopt the stage-based CLI architecture to independently instantiate execution pr
The example below pins Stage 0 to GPU 0 and Stage 1/2 to GPU 1 via
`CUDA_VISIBLE_DEVICES`.
-**1. Stage 0 (Thinker + API server)**
+#### 1. Stage 0 (Thinker + API server)
```bash
CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
@@ -60,7 +73,7 @@ CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
--omni-master-port 26000
```
-**2. Stage 1 (Talker)**
+#### 2. Stage 1 (Talker)
```bash
CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
@@ -70,7 +83,7 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
--omni-master-port 26000
```
-**3. Stage 2 (Code2Wav)**
+#### 3. Stage 2 (Code2Wav)
```bash
CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
@@ -95,6 +108,9 @@ To experiment with another supported MoE backend, set `moe_backend` in the
deployment YAML for both the Thinker and Talker stages. For the stage-based
CLI, pass `--moe-backend` directly on each stage command.
+For other stage settings, pass the option on that stage command. For example,
+to change its GPU memory utilization:
+
```bash
CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
--stage-id 1 \
@@ -107,11 +123,12 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
### Send Multi-modal Request
Get into the example folder
+
```bash
cd examples/online_serving/qwen3_omni
```
-#### Send request via python
+#### Send request via python
```bash
python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image --port 8091 --host "localhost"
@@ -126,7 +143,6 @@ The Python client supports the following command-line arguments:
- `--audio-path` (or `-a`): Path to local audio file or URL. If not provided and query-type is `use_audio`, uses default audio URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs and common audio formats: MP3, WAV, OGG, FLAC, M4A. Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3`
- `--prompt` (or `-p`): Custom text prompt/question. If not provided, uses default prompt for the selected query type. Example: `--prompt "What are the main activities shown in this video?"`
-
For example, to use a local video file with custom prompt:
```bash
@@ -136,19 +152,20 @@ python openai_chat_completion_client_for_multimodal_generation.py \
--prompt "What are the main activities shown in this video?"
```
-#### Send request via curl
+#### Send request via curl
```bash
bash run_curl_multimodal_generation.sh use_image
```
## Modality control
+
You can control output modalities to specify which types of output the model should generate. This is useful when you only need text output and want to skip audio generation stages for better performance.
### Supported modalities
| Modalities | Output |
-|------------|--------|
+| ------------ | -------- |
| `["text"]` | Text only |
| `["audio"]` | Audio only |
| `["text", "audio"]` | Text + Audio |
@@ -226,7 +243,9 @@ print(response.choices[1].message.audio) # Audio response
```
## Streaming Output
+
If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. We support both text streaming output and audio streaming output. Other modalities can output normally.
+
```bash
python openai_chat_completion_client_for_multimodal_generation.py \
--query-type use_image \
@@ -259,12 +278,14 @@ The convenience script launches both the vLLM server and Gradio demo together:
```
This script will:
+
1. Start the vLLM server in the background
2. Wait for the server to be ready
3. Launch the Gradio demo
4. Handle cleanup when you press Ctrl+C
The script supports the following arguments:
+
- `--model`: Model name/path (default: Qwen/Qwen3-Omni-30B-A3B-Instruct)
- `--server-port`: Port for vLLM server (default: 8091)
- `--gradio-port`: Port for Gradio demo (default: 7861)
@@ -275,18 +296,19 @@ The script supports the following arguments:
#### Option 2: Manual Launch (Two-Step Process)
-**Step 1: Launch the vLLM API server**
+##### Step 1: Launch the vLLM API server
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
If you have a custom deploy config file:
+
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --deploy-config /path/to/deploy_config_file
```
-**Step 2: Run the Gradio demo**
+##### Step 2: Run the Gradio demo
In a separate terminal:
@@ -299,7 +321,7 @@ Then open `http://localhost:7861/` on your local browser to interact with the we
The gradio script supports the following arguments:
- `--model`: Model name/path (should match the server model)
-- `--api-base`: Base URL for the vLLM API server (default: http://localhost:8091/v1)
+- `--api-base`: Base URL for the vLLM API server (default: )
- `--ip`: Host/IP for Gradio server (default: 127.0.0.1)
- `--port`: Port for Gradio server (default: 7861)
- `--share`: Share the Gradio demo publicly (creates a public link)
diff --git a/docs/user_guide/examples/online_serving/text_to_speech.md b/docs/user_guide/examples/online_serving/text_to_speech.md
index e9499551f..c1107ad60 100644
--- a/docs/user_guide/examples/online_serving/text_to_speech.md
+++ b/docs/user_guide/examples/online_serving/text_to_speech.md
@@ -413,6 +413,22 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
./qwen3_tts/run_server.sh Base
```
+For a local deployment with multiple API frontend processes sharing one set
+of TTS stage engines, add `--api-server-count`:
+
+```bash
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+ --omni --api-server-count 2 --port 8091
+```
+
+This mode supports local EngineCore stages with one local process group per
+replica, including stages configured with multiple `num_replicas`. Headless,
+remote, intra-replica data-parallel, and Ray stage deployments are not
+supported. Runtime voice upload and deletion are disabled with multiple
+API frontends because their registries are process-local. Built-in voices,
+inline `ref_audio`, and voices restored from `custom_voice_dir` at startup are
+supported.
+
### Executor backend
Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config ` to override.
@@ -478,6 +494,10 @@ checkpoint. When a request names an uploaded voice, the server infers
`task_type="Base"`. Built-in presets such as `vivian` and `ryan` remain
CustomVoice speakers and require a CustomVoice checkpoint.
+The runtime upload and delete routes require a single API frontend; with
+`--api-server-count > 1`, use inline `ref_audio` or restore precomputed voices
+from `custom_voice_dir` at startup.
+
### Precomputed custom voices
For reused Base voice-cloning speakers, precompute the reference artifacts once
diff --git a/docs/user_guide/quantization/mxfp4.md b/docs/user_guide/quantization/mxfp4.md
index b18b3e7b3..00ff97c01 100644
--- a/docs/user_guide/quantization/mxfp4.md
+++ b/docs/user_guide/quantization/mxfp4.md
@@ -6,22 +6,25 @@ W4A4 MXFP4 (Microscaling FP4) quantizes both weights and activations to FP4
(`float4_e2m1fn_x2`, packed 2 values per byte) using the OCP MX format: groups
of 32 K-dimension elements share a single `float8_e8m0fnu` exponent scale.
+On Ascend, single-scale MXFP4 also supports Universal Optimal Scaling
+(UOS, `Qmax=7.25`), described in the single-scale scaling section below.
+
vLLM-Omni provides two quantization methods with different scale structures:
| Method | Scale structure | Mode | Use case |
-|--------|----------------|------|----------|
-| `mxfp4` | Single-scale (per-32 fine only) | Online only | Quick accuracy baseline; no checkpoint prep needed |
-| `mxfp4_dualscale` | Dual-scale (fine per-32 + coarse per-512 + per-channel `mul_scale`) | Online + Offline | Production; better accuracy; offline recommended |
+| -------- | ---------------- | ------ | ---------- |
+| `mxfp4` | Single-scale (per-32 fine only) | Online + pre-quantized single-scale weights | Online baseline, or an existing single-scale checkpoint in the layout below |
+| `mxfp4_dualscale` | Dual-scale (fine per-32 + coarse per-512 + per-channel `mul_scale`) | Online + Offline | W4A4 with dual-level scales; offline mode loads calibrated Smooth tensors |
-!!! tip "Recommended: `mxfp4_dualscale` offline"
- For production deployments, use the `mxfp4_dualscale` offline mode with a
+!!! tip "Existing W4A4-only DualScale deployments"
+ For W4A4-only deployments, the existing `mxfp4_dualscale` offline mode uses a
pre-quantized checkpoint produced by msModelSlim. Offline checkpoints load
- calibrated `mul_scale` tensors from disk, providing measurably better accuracy
- than any online method. The one-time preprocessing cost amortises across all
- subsequent inference runs.
+ calibrated `mul_scale` tensors from disk. Quality depends on the model,
+ calibration recipe and BF16 fallback policy; evaluate the checkpoint against
+ your quality target before deployment.
- Use `mxfp4` online only for quick experimentation where preprocessing time
- is not acceptable and accuracy loss is tolerable.
+ W4A8 layer/step fallback requires single-scale `mxfp4`, online or offline.
+ DualScale supports W4A4 and BF16 routing only.
!!! warning "Online single-scale ≠ Offline dual-scale"
`mxfp4_dualscale` offline mode uses `NPUMxfp4DualScaleLinearMethod`:
@@ -34,61 +37,244 @@ vLLM-Omni provides two quantization methods with different scale structures:
## Hardware Support
-| Device | Support |
-|--------|---------|
-| NVIDIA Blackwell GPU (SM 100+) | ⭕ |
-| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ |
-| NVIDIA Ampere GPU (SM 80+) | ⭕ |
-| AMD ROCm (gfx950 / MI355X) | ✅ |
-| Intel XPU | ⭕ |
-| Ascend NPU (Atlas 950 A5) | ✅ |
+The following matrix describes these vLLM-Omni implementations, rather than
+hardware support for the MXFP4 format in general.
-Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
+| Device | Single-scale online OCP | Single-scale offline | UOS / W4A8 fallback | DualScale |
+| -------- | -------- | -------- | -------- | -------- |
+| Ascend NPU (Atlas 950 A5) | ✅ | ✅ | ✅ | ✅ |
+| AMD ROCm (gfx950 / MI355X) | ✅ | ❌ | ❌ | ❌ |
+| Other AMD architectures | ❌ | ❌ | ❌ | ❌ |
+| NVIDIA GPU | ❌ | ❌ | ❌ | ❌ |
+| Intel XPU | ❌ | ❌ | ❌ | ❌ |
+
+Legend: `✅` supported by the implementation, `❌` unsupported.
+ROCm uses the existing AITER online path and requires `mxfp4_scale_alg=0`.
## Model Type Support
### Diffusion Model (Wan2.2)
| Model | Online | Offline | Notes |
-|-------|--------|---------|-------|
-| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade (`transformer` + `transformer_2`); both transformers quantized with the same config |
-| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade; same scheme as T2V-A14B |
-| Wan2.2-TI2V-5B | ❌ | ❌ | Parameter count too small; W4A4 causes unacceptable accuracy loss |
+| ------- | -------- | --------- | ------- |
+| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4` / `mxfp4_dualscale` | Cascade experts resolve checkpoint metadata independently; explicit global runtime policy applies to both |
+| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4` / `mxfp4_dualscale` | Per-expert configuration uses the same interface as T2V; validate I2V checkpoints separately |
+| Wan2.2-TI2V-5B | Requires validation | Requires validation | This guide does not establish checkpoint compatibility or generation quality for TI2V-5B |
-The choice between `mxfp4` and `mxfp4_dualscale` in **online mode** is about
-quantization quality, not model compatibility — both work on cascade (A14B) and
-single-transformer models alike, the same as `mxfp8` online:
+Both methods can be integrated with single-transformer and cascade models.
+This alone does not establish support or quality for every model variant:
-- `mxfp4`: single-scale, lower overhead, simpler compute, online only
-- `mxfp4_dualscale`: dual-scale + optional BF16 fallback, better accuracy, online **and** offline
+- `mxfp4`: single-scale; online and pre-quantized weights, with optional W4A8 fallback on NPU
+- `mxfp4_dualscale`: dual-scale; online and offline, with BF16 fallback
-**Offline** checkpoints for A14B are always in `mxfp4_dualscale` format (produced
-by the merge script); there is no offline `mxfp4` single-scale format.
+The existing `merge_mxfp4_dualscale_checkpoint.py` produces **offline
+`mxfp4_dualscale`** checkpoints. Single-scale weights must satisfy the separate
+layout described below; changing a DualScale checkpoint's method name to
+`mxfp4` does not convert its weights. Prepare single-scale checkpoints with
+external tooling. Real exported checkpoint compatibility and model quality
+require validation of that checkpoint; the synthetic tests do not establish
+UOS export or video acceptance.
!!! note "Per-layer BF16 fallback in offline cascade models"
- The A14B offline checkpoint uses `quant_method: mxfp4_dualscale`. Most
+ An A14B offline DualScale checkpoint uses `quant_method: mxfp4_dualscale`. Most
linear layers are stored as W4A4 MXFP4 DualScale; precision-sensitive layers
retain their original BF16 weights and are listed in `ignored_layers` inside
each transformer's `config.json`. The two transformers may have different
`ignored_layers` sets — the pipeline reads each transformer's own `config.json`
- and rebuilds the config locally when they differ, so routing is always
- per-transformer-accurate.
-
-!!! warning "TI2V-5B not supported"
- Wan2.2-TI2V-5B is excluded from W4A4 quantization. Its smaller parameter
- count makes it significantly more sensitive to 4-bit quantization noise,
- resulting in unacceptable accuracy loss. Use [MXFP8](mxfp8.md) for TI2V-5B.
+ and resolves the storage format and BF16 layer policy independently.
## Configuration
+### W4A8 fallback at selected layers and denoising steps (Ascend NPU)
+
+For single-scale `mxfp4`, set `w4a8_fallback_steps` to a list of zero-based
+request denoising loop indices (not scheduler timestep values).
+Set `w4a8_fallback_layers` to exact runtime Linear paths, for example
+`["blocks.10.attn1.to_qkv", "blocks.12.ffn.net_2"]`. A selected layer uses W4A8
+at every step; at a selected step, all other quantized layers also use W4A8.
+The rules are combined with **OR**. W4A8 uses MXFP8 activations with MXFP4
+weights. Omitting both lists or passing `[]` keeps the existing W4A4 behavior. Existing
+`ignored_layers` take priority: **BF16 > (layer OR step) W4A8 > W4A4**.
+DualScale's existing leading BF16 block rule is unchanged; it does not enable
+W4A8.
+
+```python
+from vllm_omni import Omni
+
+omni = Omni(
+ model="",
+ quantization_config={
+ "method": "mxfp4",
+ "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"],
+ "w4a8_fallback_steps": [0, 1, 38, 39],
+ },
+)
+```
+
+Both policies also work with offline single-scale `mxfp4`. Any nonempty
+W4A8 list with `mxfp4_dualscale`, online or offline, raises an explicit error
+at configuration construction. Empty lists preserve its existing W4A4 path.
+Each expert's checkpoint still determines its serialized
+format and BF16 layer routing. A serialized checkpoint config is the complete
+storage contract: `ignored_layers`, or its `modules_to_not_convert` alias,
+selects the BF16 layers, and omitting both means no explicit BF16 layers. A
+method-only checkpoint value such as `"mxfp4"` or
+`{"quant_method": "mxfp4"}` carries no storage policy and leaves a caller's
+online `ignored_layers` unchanged. An explicitly supplied quantization config
+retains its W4A8 layer and step lists when the loader rebuilds either expert's
+config from disk, including empty lists that disable checkpoint policies.
+Within an explicit runtime config, an omitted W4A8 list defaults to `[]`;
+individual omitted fields do not inherit saved policies. Without an explicit
+runtime config, auto-detection uses the saved checkpoint policies.
+
+Layer names are relative to each expert (`blocks.10...`, without a
+`transformer.` or `transformer_2.` prefix). A global config applies the same
+list to both experts; per-component `transformer`/`transformer_2` configs can
+select different lists. Wan selects these expert names exactly: a `transformer`
+entry does not match `transformer_2`; an omitted expert uses `default` if set.
+A component mapped to `null` (or an unmatched component
+without a default) stays unquantized and requires a BF16 checkpoint; loading a
+quantized checkpoint for that disabled component raises an error. A top-level
+`quantization_config=None` retains checkpoint auto-detection.
+Names are exact, with no regex, wildcard, block expansion,
+or checkpoint-name remapping. Wan validates them against its runtime Linear
+modules and logs the selected paths/count and BF16 overrides per PP rank.
+Use `attn1.to_qkv` for fused self-attention; `attn1.to_q`, `to_k`, or `to_v`
+checkpoint names raise an error. Cross-attention `attn2.to_q/to_k/to_v` remains
+separate. Layer-only fallback does not require a denoise ForwardContext.
+
+The text-to-video example accepts the JSON directly:
+
+```bash
+ASCEND_RT_VISIBLE_DEVICES=0,1,2,3 \
+python examples/offline_inference/text_to_video/text_to_video.py \
+ --model /path/to/Wan2.2-T2V-A14B-Diffusers \
+ --quantization-config '{"method":"mxfp4","w4a8_fallback_steps":[0,2]}' \
+ --tensor-parallel-size 1 --ulysses-degree 2 --cfg-parallel-size 2 \
+ --enforce-eager --height 480 --width 832 --num-frames 81 \
+ --num-inference-steps 40 --seed 42 --output wan22-w4a8.mp4
+```
+
+Use either `--quantization` or `--quantization-config` in this example.
+Serving accepts the same JSON through `--diffusion-quantization-config`.
+
+The list uses one request's loop index across both high- and low-noise experts,
+and both CFG branches see the same index. A new request starts at step zero.
+Indices outside the request's step range are never selected. A pipeline must
+publish `ForwardContext.denoise_step_idx`; requesting step fallback without
+that context raises an error instead of silently ignoring the list.
+
+Single-scale W4A4 and W4A8 reuse the same prepared `weight` and `weight_scale`.
+Online loading quantizes the original BF16/FP16 weight once; offline loading
+packs the numeric FP4 values and reshapes the saved E8M0 scale without
+requantizing the weight. Only activation precision changes during forward.
+There is no DualScale decoding, second W4 representation or `w4a8_weight` /
+`w4a8_weight_scale` cache. No weight quantization runs inside the denoising loop.
+
+The mixed FP8/FP4 GEMM uses BF16 output and a matching BF16 row bias for
+compatibility with the target A5 runtime. An FP16 caller receives that result
+converted back to FP16; its output shape and dtype are preserved.
+
+The single-scale fallback validation target is eager Wan2.2 T2V A14B on A5.
+Quality scores, optimal fallback lists, real UOS exports and additional
+compile/offload combinations require separate evaluation. Existing W4A4-only
+DualScale still needs complete 512-channel input groups; use TP=1 with
+sequence/CFG parallelism for its initial multi-card run.
+
+### `mxfp4` — Pre-Quantized Single-Scale Weights (Ascend NPU)
+
+Set the following fields in each transformer's `config.json` when a
+single-scale checkpoint uses calibrated Smooth scaling:
+
+```json
+{
+ "quantization_config": {
+ "quant_method": "mxfp4",
+ "is_checkpoint_mxfp4_serialized": true,
+ "require_smooth_scale": true,
+ "ignored_layers": []
+ }
+}
+```
+
+Each quantized Linear must provide tensors under its model parameter name:
+
+| Tensor | Checkpoint layout | Meaning |
+| -------- | ------------------- | --------- |
+| `weight` | `(N, K)` numeric FP4 values | Packed to FP4 after loading |
+| `weight_scale` | `(N, ceil(K / 32))`, `uint8` | Raw E8M0 exponent bytes, required |
+| `mul_scale` | `(K,)`, FP32 | Smooth pre-scale, required when declared above |
+
+Both W4A4 and W4A8 quantize `x * mul_scale`. The Smooth tensor is converted
+to the layer's activation dtype after loading. Row-parallel layers load the
+corresponding input-channel shard. Pre-fused QKV weights must use one shared
+Smooth tensor for the fused projection.
+
+For backward compatibility, `require_smooth_scale` defaults to `false`:
+an absent tensor uses identity scaling, while a present tensor is still loaded.
+Without a declaration, an absent Smooth tensor cannot be distinguished from a
+legacy unsmoothed checkpoint. Calibrated checkpoints should declare
+`require_smooth_scale: true` so missing or misnamed Smooth tensors fail loading;
+incorrect tensor shapes, non-floating tensors, NaN/Inf and nonpositive values
+also fail. Values that overflow or underflow to zero in the activation dtype
+are rejected during weight processing. The flag requires serialized single-scale
+weights and is rejected for online quantization. It does not make DualScale
+scales optional or change missing-parameter policies for other backends.
+Either the checkpoint or an explicit caller can require Smooth: a `true`
+declaration is retained when per-expert metadata rebuilds the configuration,
+including when the other source omits the flag or sets it to `false`.
+
+Prepare the canonical numeric FP4 checkpoint before inference using external
+tooling. Runtime loading does not parse producer-specific export formats.
+Retain the producer recipe to identify the weight-generation algorithm and
+Smooth calibration; tensor storage alone does not establish UOS provenance.
+
+### `mxfp4` — Single-Scale Block Scaling (OCP / UOS)
+
+`mxfp4_scale_alg` selects the MXFP4 block-scaling method on NPU.
+UOS is a scaling rule, not a separate activation dtype or an activation-only
+algorithm. Both choices use FP4 E2M1 values with shared E8M0 scales.
+
+| Mode | Weight preparation | A4 activation quantization | W4A8 fallback |
+| --- | --- | --- | --- |
+| Online | Selected scaling method, once at loading | Selected scaling method | Same W4; OCP MXFP8 A8 |
+| Offline | Load existing W4 and scales unchanged | Selected scaling method | Same W4; OCP MXFP8 A8 |
+
+The available settings are:
+
+| Setting | MXFP4 scaling method | NPU parameters |
+| --- | --- | --- |
+| `0` (default) | OCP MX | `scale_alg=0` |
+| `2` | UOS (`Qmax=7.25`) | `scale_alg=2, dst_type_max=7.25` |
+
+Universal Optimal Scaling (UOS, `Qmax=7.25`) is derived in
+[MXAttention: Data-Free Optimal Scaling and Pre-Normalization Quantization for MXFP4 Attention](https://arxiv.org/abs/2607.24377),
+Section 4.1. The value `7.25` is the boundary used to select the shared
+power-of-two scale; the maximum finite E2M1 value remains `6`.
+This option applies UOS to Linear quantization. It does not implement the
+paper's full MXAttention pipeline or Pre-Normalization Quantization (PNQ).
+
+Both use FP4 E2M1 with E8M0 scales, `axis=-1`, `block_size=32`,
+`round_mode="rint"`. W4A8 always quantizes the activation to
+`float8_e4m3fn` with `scale_alg=0` and the same grouping/rounding, without
+UOS parameters. A4 and A8 share the prepared W4 and scale. UOS is currently
+NPU-only; ROCm retains its existing AITER online algorithm.
+
+An explicit UOS online configuration is
+`{"method": "mxfp4", "mxfp4_scale_alg": 2}`. For offline checkpoints the
+setting affects **only A4 activations**. Existing weights are never requantized:
+a UOS checkpoint can run UOS W4 + UOS A4 or UOS W4 + OCP A4, and both fall back to
+UOS W4 + OCP A8. Record the weight-generation algorithm separately from the
+runtime activation setting. Smooth remains active with either setting. The
+UOS configuration explicitly selects `2`; attention computation is unchanged.
+
### `mxfp4` — Single-Scale Online Mode
-Online mode requires no pre-processing. vLLM-Omni quantizes BF16 weights to
-MXFP4 at load time using `npu_dynamic_mx_quant`. A single block scale
-(`float8_e8m0fnu`, one per 32 K elements) is computed on the fly; no
-calibration `mul_scale` is available. Applies equally to single-transformer
-and cascade (A14B) models — both transformers in a cascade receive the same
-quantization config automatically.
+Online mode quantizes BF16/FP16 weights once at loading and computes activation
+scales during forward. Each block of 32 K elements shares one E8M0 scale;
+no calibrated `mul_scale` is available. An explicit global quantization config
+applies to both cascade experts. Per-component configs can select different
+policies; checkpoint auto-detection resolves each expert independently.
```python
from vllm_omni import Omni
@@ -103,7 +289,7 @@ outputs = omni.generate(
```bash
# Single-transformer or cascade model — same command
-python text_to_video.py --model --quantization mxfp4
+python examples/offline_inference/text_to_video/text_to_video.py --model --quantization mxfp4
# Online serving
vllm serve --omni --quantization mxfp4
@@ -113,19 +299,19 @@ vllm serve --omni --quantization mxfp4
Online DualScale mode computes both fine and coarse scales on the fly from BF16
weights using `npu_dynamic_dual_level_mx_quant`. Applies equally to
-single-transformer and cascade (A14B) models. Compared to `mxfp4` online,
-DualScale provides better quantization accuracy at higher compute cost.
+single-transformer and cascade (A14B) models with compatible Linear layers.
+Compare quality and performance using the intended model and workload.
The default configuration keeps the leading 5 transformer blocks in BF16
-(`num_bf16_fallback_layers=5`). Accuracy evaluation on Wan2.2-A14B shows this
-is sufficient to meet quality requirements and is the recommended setting.
+(`num_bf16_fallback_layers=5`). This is an implementation default; select and
+validate the BF16 policy against your own quality target.
```python
omni = Omni(model="", quantization="mxfp4_dualscale")
```
```bash
-python text_to_video.py --model --quantization mxfp4_dualscale
+python examples/offline_inference/text_to_video/text_to_video.py --model --quantization mxfp4_dualscale
```
If accuracy debugging identifies additional precision-sensitive layers, they
@@ -136,21 +322,22 @@ omni = Omni(
model="",
quantization_config={
"method": "mxfp4_dualscale",
- "ignored_layers": ["blocks.10.attn1.to_q"], # explicit per-layer override
+ "ignored_layers": ["blocks.10.attn1.to_qkv"], # explicit per-layer override
},
)
```
BF16 fallback routing in online mode applies two rules in priority order:
-1. **`ignored_layers`** (explicit per-layer override): any layer whose prefix
- matches is kept in BF16 regardless of block index.
+1. **`ignored_layers`** (explicit per-layer override): a listed runtime Linear
+ module is kept in BF16 regardless of block index. Use exact module paths,
+ not arbitrary block prefixes or names ending in `.weight`.
2. **`num_bf16_fallback_layers`** (coarse leading-block rule): the first N
transformer blocks (`blocks.0` … `blocks.N-1`) fall back to BF16. Defaults
- to `5` (recommended). Layers outside `blocks.N.*`
- (e.g. `condition_embedder`) are always quantized.
+ to `5`. Layers outside `blocks.N.*` are quantized unless explicitly listed
+ in `ignored_layers`.
-### `mxfp4_dualscale` — DualScale Offline Mode (Recommended)
+### `mxfp4_dualscale` — DualScale Offline Mode
Offline mode loads a pre-quantized DualScale checkpoint from msModelSlim. A
preprocessing step converts the raw quantized output to the diffusers format
@@ -169,8 +356,8 @@ inside `config.json`. At runtime, each layer's prefix is matched against
Each quantized linear layer stores four tensors:
| Tensor | Shape | dtype | Description |
-|--------|-------|-------|-------------|
-| `weight` | `(N, K)` | float8_e4m3fn | FP4 packed (2 values per byte) |
+| -------- | ------- | ------- | ------------- |
+| `weight` | `(N, K)` | float8_e4m3fn | Numeric FP4 values; packed into two values per byte after loading |
| `weight_scale` | `(N, K//32)` | uint8 | Fine block scale (`float8_e8m0fnu` bit pattern) |
| `weight_dual_scale` | `(N, K//512, 1)` | float32 | Coarse block scale |
| `mul_scale` | `(K,)` | float32 | Per-input-channel smooth pre-scale (from calibration) |
@@ -223,7 +410,7 @@ python vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py \
```
| Argument | Description |
-|----------|-------------|
+| ---------- | ------------- |
| `--model-type` | Model variant: `Wan2.2-T2V-A14B` or `Wan2.2-I2V-A14B` |
| `--original-model` | Root directory of the original BF16 diffusers model |
| `--quant-path` | Root directory of the msModelSlim quantized output |
@@ -251,14 +438,14 @@ The `quantization_config` injected into `config.json` for each transformer:
```
`ignored_layers` lists every linear layer that retains its original BF16 weight,
-using vllm-omni model parameter names (QKV-fused, FFN underscored, `to_out`
+using vllm-omni runtime module paths (QKV-fused, FFN underscored, `to_out`
unindexed). The exact entries are determined by the quantization tool (msModelSlim)
and may differ between `transformer` and `transformer_2` in a cascade model.
#### Step 3 — Serve
```bash
-python text_to_video.py --model /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale
+python examples/offline_inference/text_to_video/text_to_video.py --model /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale
# Online serving
vllm serve /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale --omni
@@ -275,25 +462,32 @@ omni = Omni(model="/path/to/Wan2.2-T2V-A14B-MXFP4-DualScale")
## Parameters
-### `mxfp4` (single-scale, online only)
+### `mxfp4` (single-scale, online + offline)
| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
+| ----------- | ------ | --------- | ------------- |
| `method` | str | — | `"mxfp4"` |
-| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16 |
+| `mxfp4_scale_alg` | int | `0` | MXFP4 block scaling: `0` OCP MX, `2` UOS (`Qmax=7.25`, NPU only). Online: W4 and A4; offline: A4 only. A8 remains OCP MXFP8 |
+| `is_checkpoint_mxfp4_serialized` | bool | `false` | Load pre-quantized single-scale weights and scales; declared in checkpoint `config.json` |
+| `require_smooth_scale` | bool | `false` | Require calibrated Smooth tensors for offline single-scale weights |
+| `ignored_layers` | list[str] | `[]` | Exact runtime Linear module paths to keep in BF16 |
+| `w4a8_fallback_layers` | list[str] | `[]` | Exact runtime Linear paths that always use A8; NPU only |
+| `w4a8_fallback_steps` | list[int] | `[]` | Zero-based denoise steps that use A8; NPU only |
### `mxfp4_dualscale` (dual-scale, online + offline)
| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
+| ----------- | ------ | --------- | ------------- |
| `method` | str | — | `"mxfp4_dualscale"` |
| `is_checkpoint_serialized` | bool | `False` | `True` for offline DualScale checkpoints; auto-set from `config.json` when using the preprocessing script |
-| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16. **Works in both modes**: offline — populated by the merge script for interleaved sensitive layers; online — user-supplied for explicit per-layer precision override |
-| `num_bf16_fallback_layers` | int | `5` | **Online mode only**: leading N transformer blocks (`blocks.0` … `blocks.N-1`) kept in BF16. Applied after `ignored_layers`; ignored in offline mode. Default of `5` is the evaluated recommended value for Wan2.2-A14B |
+| `ignored_layers` | list[str] | `[]` | Exact runtime Linear module paths to keep in BF16. **Works in both modes**: offline — populated by the merge script for interleaved sensitive layers; online — user-supplied for explicit per-layer precision override |
+| `num_bf16_fallback_layers` | int | `5` | **Online mode only**: leading N transformer blocks (`blocks.0` … `blocks.N-1`) kept in BF16. Applied after `ignored_layers`; ignored in offline mode. Validate the default of `5` against the intended model and quality target |
+| `w4a8_fallback_layers` | list[str] | `[]` | Must remain empty; DualScale does not support W4A8 fallback |
+| `w4a8_fallback_steps` | list[int] | `[]` | Must remain empty; DualScale does not support W4A8 fallback |
#### BF16 fallback priority (online mode)
-```
+```text
for each linear layer:
if prefix in ignored_layers → BF16 (explicit override, highest priority)
elif block_idx < num_bf16_fallback_layers → BF16 (coarse leading-block rule)
@@ -307,8 +501,9 @@ unless they appear in `ignored_layers`.
1. **Online single-scale (`mxfp4`)** quantizes BF16 weights at load time using
`npu_dynamic_mx_quant` (single-scale). No calibration `mul_scale` is
- available — all output partitions receive an identity pre-scale. No offline
- checkpoint format exists for this method.
+ available — all output partitions receive an identity pre-scale. Offline
+ single-scale checkpoints use the canonical numeric FP4 layout
+ described above.
2. **Online dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=False`)**
quantizes BF16 weights using `npu_dynamic_dual_level_mx_quant` (fine + coarse
@@ -316,7 +511,7 @@ unless they appear in `ignored_layers`.
explicit `ignored_layers` stay in BF16 for accuracy.
3. **Offline dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=True`)** —
- **recommended for production** — loads four tensors per quantized layer: FP4
+ loads four tensors per quantized layer: FP4
weight, fine scale (`uint8` reinterpreted as `float8_e8m0fnu`), coarse scale
(`float32`), and per-input-channel `mul_scale` (`float32`). BF16 fallback
layers have no quantization tensors and are routed via `ignored_layers`.
@@ -325,21 +520,21 @@ unless they appear in `ignored_layers`.
layout as `float8_e8m0fnu`) and reinterpreted at load time without a lossy
float32 round-trip.
-5. **Cascade model config propagation**: in a cascade model (transformer +
- transformer_2), vLLM-Omni reads each transformer's own `config.json` and
- rebuilds the quant config locally when `ignored_layers` differs between
- transformers, ensuring per-layer routing is accurate for each. The first
- transformer's config is propagated to `od_config` so the second transformer
- can reuse it as a starting point.
+5. **Cascade model configuration**: checkpoint auto-detection resolves each
+ expert's own `config.json` independently; the first expert's detected config
+ is not reused as the second expert's policy. An explicit global runtime
+ config applies to both experts, while per-component configs select policies
+ by the exact `transformer` / `transformer_2` names. Serialized checkpoint
+ metadata still determines each expert's storage format and BF16 layers.
6. **Self-attention QKV fusion**: Q, K, V projection weights are fused into a
single `QKVParallelLinear` layer at runtime. `ignored_layers` entries use the
fused name (`attn1.to_qkv`), written automatically by the merge script.
-7. W4A4 carries higher quantization noise than W8A8 (16 vs 256 levels). The
- DualScale offline method mitigates this with calibrated `mul_scale` smooth
- quantization. Use `ignored_layers` and `num_bf16_fallback_layers` to trade
- off compression vs. accuracy for precision-sensitive layers.
+7. FP4 has fewer representable values than FP8, so evaluate quantization
+ error and generation quality for the chosen recipe. Offline Smooth scaling
+ and BF16 fallback can help with sensitive layers.
+ `num_bf16_fallback_layers` applies only to online DualScale.
## Adapting MXFP4 for a New Model
@@ -350,27 +545,26 @@ a merge script for offline checkpoints.
### Step 1 — Discover runtime layer names
-`ignored_layers` entries must match the **runtime parameter names** used inside
-vllm-omni, which may differ from the names stored in the diffusers checkpoint.
-The canonical source of truth is the model's own `named_parameters()`.
+`ignored_layers` entries use exact **runtime Linear module paths**, without
+`.weight`. Inspect the actual transformer instance in model initialization or
+a developer harness; `Omni` does not expose an `omni.pipeline.transformer` API.
+For a transformer instance named `transformer`:
```python
-from vllm_omni import Omni
+from vllm.model_executor.layers.linear import LinearBase
-# Load the model without quantization to inspect parameter names.
-omni = Omni(model="/path/to/your-model") # no --quantization flag
-for name, _ in omni.pipeline.transformer.named_parameters():
- if "weight" in name and "scale" not in name:
+for name, module in transformer.named_modules():
+ if isinstance(module, LinearBase):
print(name)
```
-Compare the printed names against the diffusers checkpoint keys
+Compare these module paths against the corresponding diffusers checkpoint keys
(`safetensors.safe_open` or `torch.load`) to identify any renames your model
applies. Common patterns that differ in Wan2.2 (and may appear in other
models):
| Diffusers checkpoint name | vllm-omni runtime name | Reason |
-|---------------------------|------------------------|--------|
+| --------------------------- | ------------------------ | -------- |
| `attn1.to_q`, `attn1.to_k`, `attn1.to_v` | `attn1.to_qkv` | Self-attention Q/K/V fused into `QKVParallelLinear` |
| `ffn.net.0.proj` | `ffn.net_0.proj` | Dots in sub-module names replaced with underscores |
| `ffn.net.2` | `ffn.net_2` | Same underscore rule |
@@ -382,8 +576,9 @@ fused runtime parameters.
!!! warning "Partial QKV fallback is not allowed"
If your model fuses Q, K, V into a single layer, `ignored_layers` must
- include **all three or none**. A partial fallback (e.g. `to_q` in BF16 but
- `to_k`, `to_v` quantized) cannot be expressed at runtime because they share
+ name the fused runtime module (for Wan, `attn1.to_qkv`). The source
+ checkpoint must keep **all three or none** in BF16. A partial fallback
+ (e.g. `to_q` in BF16 but `to_k`, `to_v` quantized) cannot be expressed at runtime because they share
one `QKVParallelLinear`. The merge script enforces this and raises an error
if only some of the trio appear as non-quantized.
@@ -392,12 +587,13 @@ fused runtime parameters.
#### Online mode
Pass `ignored_layers` directly in the quantization config using the **runtime
-names** discovered in Step 1. No code changes to the model are required.
+names** discovered in Step 1. This assumes the model already passes its
+quantization config and module prefixes to compatible vLLM Linear layers.
```python
omni = Omni(
model="/path/to/your-model",
- quantization={
+ quantization_config={
"method": "mxfp4_dualscale",
"ignored_layers": [
"blocks.0.attn1.to_qkv", # runtime name, not diffusers name
@@ -409,9 +605,14 @@ omni = Omni(
```
```bash
-# CLI does not support list-typed ignored_layers directly.
-# Use the Python API or set ignored_layers in config.json (offline).
-python your_script.py --model /path/to/your-model --quantization mxfp4_dualscale
+# The text-to-video example accepts list-valued fields through JSON.
+python examples/offline_inference/text_to_video/text_to_video.py \
+ --model /path/to/your-model \
+ --quantization-config '{"method":"mxfp4_dualscale","ignored_layers":["blocks.0.attn1.to_qkv"]}'
+
+# Serving accepts the same configuration.
+vllm serve /path/to/your-model --omni \
+ --diffusion-quantization-config '{"method":"mxfp4_dualscale","ignored_layers":["blocks.0.attn1.to_qkv"]}'
```
The `num_bf16_fallback_layers` coarse rule is an alternative to listing layers
@@ -436,9 +637,13 @@ the merge script is correct. The injected block:
}
```
-To add a layer manually (e.g. to pin an additional layer to BF16 without
-re-running the merge script), edit `config.json` inside the transformer
-subfolder. Use runtime names, not diffusers checkpoint names.
+Adding an offline BF16 fallback layer requires both the original BF16 weight
+from the base model and a matching `ignored_layers` entry. Regenerate the
+checkpoint with the revised storage policy, restoring the original weight
+(and bias, if present) and removing that layer's quantization tensors.
+Editing `config.json` alone does not reconstruct BF16 weights from stored FP4
+values. For fused QKV, restore all three projections together. Use runtime
+module paths in `ignored_layers`.
### Step 3 — Write a merge script for offline mode
@@ -474,6 +679,7 @@ renaming rules and collect the results.
!!! tip "Validate before serving"
After producing the offline checkpoint, load it without a `--quantization`
flag and verify that vLLM-Omni auto-detects the correct method. Check that
- the layer count reported in the startup log matches expectations: quantized
- layer count + `ignored_layers` count should equal total linear layer count.
- Any mismatch indicates a name-mapping bug in the merge script.
+ each intended runtime Linear module resolves to the expected quantized or
+ BF16 method and that every BF16 entry has original BF16 weights. Account
+ for fused modules and, in online DualScale, the leading-block fallback rule;
+ the raw length of `ignored_layers` alone is not a reliable layer count.
diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py
index 06179ef9e..7b7bcb890 100644
--- a/examples/offline_inference/text_to_video/text_to_video.py
+++ b/examples/offline_inference/text_to_video/text_to_video.py
@@ -384,13 +384,20 @@ def parse_args() -> argparse.Namespace:
default=None,
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
)
- parser.add_argument(
+ quantization_group = parser.add_mutually_exclusive_group()
+ quantization_group.add_argument(
"--quantization",
type=str,
default=None,
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).",
)
+ quantization_group.add_argument(
+ "--quantization-config",
+ type=json.loads,
+ default=None,
+ help='Quantization JSON, e.g. \'{"method":"mxfp4","w4a8_fallback_steps":[0,2]}\'.',
+ )
# Distributed and parallel execution
parser.add_argument(
@@ -572,6 +579,8 @@ def main():
omni_kwargs["flow_shift"] = args.flow_shift
if args.quantization is not None:
omni_kwargs["quantization"] = args.quantization
+ if args.quantization_config is not None:
+ omni_kwargs["quantization_config"] = args.quantization_config
if args.cache_backend is not None:
omni_kwargs["cache_backend"] = args.cache_backend
omni_kwargs["cache_config"] = cache_config
diff --git a/examples/online_serving/qwen3_omni/README.md b/examples/online_serving/qwen3_omni/README.md
index 107e1f098..8d6014db2 100644
--- a/examples/online_serving/qwen3_omni/README.md
+++ b/examples/online_serving/qwen3_omni/README.md
@@ -14,6 +14,17 @@ Please refer to [README.md](../../../README.md)
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
+For a local deployment with multiple API frontend processes sharing one set
+of stage engines, add `--api-server-count`:
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --api-server-count 2 --port 8091
+```
+
+This mode currently supports local EngineCore stages. It is not available
+with headless or remote stage deployments, and pure diffusion stages remain
+single-frontend.
+
The default deployment configuration, situated at `vllm_omni/deploy/qwen3_omni_moe.yaml`, is resolved and loaded
automatically via the model registry, obviating the `--deploy-config` flag in standard deployment topologies.
Asynchronous chunk streaming operates as **enabled by default** within this bundled configuration.
@@ -156,7 +167,7 @@ parser defaults). If you don't pass a flag, the YAML value wins.
> chunked vs end-to-end modes (e.g. qwen3_tts code2wav) dispatch
> automatically based on that bool — no extra flag or variant yaml is
> needed.
-
+>
> ⚠️ **For multi-stage models that share GPUs (qwen3_omni_moe by default
> shares cuda:1 between stages 1 and 2), avoid using global memory flags.**
> A global `--gpu-memory-utilization 0.85` would apply to every stage and
diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md
index b749731e0..77e9ec4f1 100644
--- a/examples/online_serving/text_to_speech/README.md
+++ b/examples/online_serving/text_to_speech/README.md
@@ -13,7 +13,7 @@ For the full list of supported architectures across all modalities, see
## Supported Models
| Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo |
-|---|---|---|---|---|---|
+| --- | --- | --- | --- | --- | --- |
| Breeze-TTS-2 | `BreezeBlue/Breeze-TTS-2` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | speaker tags (`S0`..`S9`, default `S0`) | — |
| Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | uploaded audio voice only; no presets | ✓ |
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
@@ -781,6 +781,20 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
./qwen3_tts/run_server.sh Base
```
+For a local deployment with multiple API frontend processes sharing one set
+of TTS stage engines, add `--api-server-count`:
+
+```bash
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+ --omni --api-server-count 2 --port 8091
+```
+
+This mode supports local EngineCore stages only; headless or remote stage
+deployments are not supported. Runtime voice upload and deletion are disabled
+with multiple API frontends because their registries are process-local.
+Built-in voices, inline `ref_audio`, and voices restored from `custom_voice_dir`
+at startup are supported.
+
### Executor backend
Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config ` to override.
@@ -835,6 +849,10 @@ checkpoint. When a request names an uploaded voice, the server infers
`task_type="Base"`. Built-in presets such as `vivian` and `ryan` remain
CustomVoice speakers and require a CustomVoice checkpoint.
+The runtime upload and delete routes require a single API frontend; with
+`--api-server-count > 1`, use inline `ref_audio` or restore precomputed voices
+from `custom_voice_dir` at startup.
+
### Precomputed custom voices
For reused Base voice-cloning speakers, precompute the reference artifacts once
diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md
index b48a1f472..934e882f7 100644
--- a/recipes/MiniMaxAI/MiniMax-H3.md
+++ b/recipes/MiniMaxAI/MiniMax-H3.md
@@ -1106,6 +1106,68 @@ generations move the end-to-end figure toward the denoising one. Fusing the
adapter does not measurably change startup: weight loading took 77.3 s with it
against 85.8 s without.
+### FastH3 8-Step V2 full checkpoint
+
+[FastH3 8-Step V2](https://huggingface.co/FastVideo/FastVideo-FastH3-8-Step-V2)
+ships a full Diffusers-format transformer, including learned VSA compression
+gates. Serve the Hugging Face model ID or a downloaded snapshot directly:
+
+```bash
+export FASTH3_V2_MODEL=FastVideo/FastVideo-FastH3-8-Step-V2
+```
+
+Use a Hugging Face account with access to the release, and pin `--revision`
+when reproducing a result. The existing component loader reads its safetensors
+shards; the H3 weight loader maps names and loads Q/K/V and MLP projections into
+native parameters in memory. No conversion command or second transformer
+checkpoint is needed.
+
+The transformer and text components come from the V2 release. To retain Omni's
+native tiled and parallel VAE execution, the loader fetches only the video/audio
+VAEs from the `MiniMaxAI/MiniMax-H3` revision pinned in V2's `provenance.json`.
+Those components use the normal Hugging Face cache. No separate base-model
+argument or full base-transformer download is required. Download time is a
+preparation cost, separate from warmup and request latency.
+
+Install the optional `fastvideo-kernel` build required by the
+`FASTVIDEO_VSA` backend, including its H3 block-map entry point. A four-GPU
+configuration is:
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2,3 vllm serve "${FASTH3_V2_MODEL}" --omni \
+ --host 127.0.0.1 --port 8095 --trust-remote-code --task-type fl2va \
+ --served-model-name FastH3-V2 \
+ --num-gpus 4 --usp 4 --ring 1 \
+ --vae-patch-parallel-size 4 --vae-parallel-mode tile --vae-use-tiling \
+ --diffusion-attention-backend FASTVIDEO_VSA
+
+curl --fail-with-body -sS http://127.0.0.1:8095/v1/videos/sync \
+ -F 'model=FastH3-V2' \
+ -F 'prompt=A red fox runs through fresh snow at dawn, with fast pawsteps, winter wind and distant birds.' \
+ -F 'width=1344' -F 'height=768' -F 'aspect_ratio=16:9' -F 'fps=24' \
+ -F 'num_inference_steps=8' -F 'guidance_scale=1' -F 'flow_shift=10' \
+ -F 'seed=1101' \
+ -F 'extra_params={"task":"t2va","duration":4.4,"audio_flow_shift":3.0}' \
+ -o fasth3_v2.mp4
+```
+
+V2 pins video/audio shifts **10/3**, guidance **1**, and the unshifted ladder
+`[0.999, 0.874, 0.749, 0.624, 0.5, 0.375, 0.25, 0.125, 0.0]`.
+Use **`num_inference_steps=8`** in Omni: these nine nodes bound eight model
+evaluations. FastVideo's `--steps 9` counts nodes instead. The existing H3 ODE
+update is reused without stochastic re-noising.
+
+The trained attention policy is **80% sparsity with 64-token video tiles**.
+Omni derives the number of selected video tiles from each request's geometry;
+do not supply a fixed `--fastvideo-vsa-topk`. Text/condition/audio prefix keys
+remain available to every query, and prefix queries remain dense. Missing VSA
+geometry or a failed H3 kernel raises an error instead of serving the student
+with dense attention.
+
+This release supports T2VA only, with local or pure Ulysses attention. Additional
+LoRA adapters, Ref2VA and FL2VA conditioning are unsupported. The legacy
+four-step adapter keeps its original sampling and fixed-top-k behavior.
+
## Key parameters
| Parameter | Recommended value | Notes |
diff --git a/tests/config/composable_parallel/test_factory_integration.py b/tests/config/composable_parallel/test_factory_integration.py
index 468b3db2c..d49d32042 100644
--- a/tests/config/composable_parallel/test_factory_integration.py
+++ b/tests/config/composable_parallel/test_factory_integration.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Integration tests: strategy overlay on the real merge_pipeline_deploy seam."""
from __future__ import annotations
@@ -20,6 +20,7 @@ from vllm_omni.config.composable_parallel import (
apply_strategy_specs,
)
from vllm_omni.config.config_factory import StageConfigFactory
+from vllm_omni.config.omni_config import VllmOmniConfig
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
@@ -137,3 +138,31 @@ def test_cli_overrides_strategy_with_warning():
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)
+
+
+def test_typed_cli_overrides_strategy():
+ pipeline = OMNI_PIPELINES["qwen2_5_omni"]
+ config = VllmOmniConfig.from_pipeline_config(
+ pipeline,
+ user_deploy_config=load_deploy_config(_DEPLOY),
+ cli_overrides={"stage_1_num_replicas": 3},
+ strategy_specs={"talker": [_stage_replica(2, "round_robin")]},
+ )
+
+ talker = next(stage for stage in config.stage_configs if stage.model_stage == "talker")
+ assert talker.runtime_config.num_replicas == 3
+ assert config.orchestrator_config.omni_lb_policy == "round-robin"
+ assert config.strategy_omni_lb_policy == "round-robin"
+
+
+def test_typed_tp_only_strategy_has_no_derived_lb_policy():
+ pipeline = OMNI_PIPELINES["qwen2_5_omni"]
+ config = VllmOmniConfig.from_pipeline_config(
+ pipeline,
+ user_deploy_config=load_deploy_config(_DEPLOY),
+ cli_overrides={"omni_lb_policy": "round-robin"},
+ strategy_specs={"thinker": [_tp(1)]},
+ )
+
+ assert config.orchestrator_config.omni_lb_policy == "round-robin"
+ assert config.strategy_omni_lb_policy is None
diff --git a/tests/config/test_forced_aligner_injection.py b/tests/config/test_forced_aligner_injection.py
index 65c315229..2037b22f1 100644
--- a/tests/config/test_forced_aligner_injection.py
+++ b/tests/config/test_forced_aligner_injection.py
@@ -2,13 +2,19 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import pytest
+from vllm.pooling_params import PoolingParams
import vllm_omni.config.pipeline_registry # noqa: F401 (populate registry)
+from vllm_omni.config.omni_config import VllmOmniConfig
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES as _PIPELINE_REGISTRY
from vllm_omni.config.stage_config import (
DeployConfig,
StageExecutionType,
)
+from vllm_omni.engine.stage_init_utils import (
+ build_engine_args_dict_from_omni_stage_config,
+ extract_stage_metadata_from_omni_stage_config,
+)
from vllm_omni.model_executor.stage_input_processors.forced_aligner import (
POOLING_OUTPUT_DECODER_PATH,
TIMESTAMPS_MODALITY,
@@ -94,3 +100,33 @@ def test_aligner_uses_completed_audio_without_disabling_upstream_chunks(async_ch
assert stages[2].engine_args.runner == "pooling"
assert len(pipeline.stages) == 2
assert deploy.stages == []
+
+
+def test_injected_aligner_survives_typed_config_and_runtime_projections():
+ pipeline, deploy = inject_forced_aligner_stage(
+ _PIPELINE_REGISTRY["qwen3_tts"],
+ DeployConfig(),
+ {"forced_aligner": "/models/Qwen3-ForcedAligner-0.6B"},
+ )
+ config = VllmOmniConfig.from_pipeline_config(
+ pipeline,
+ user_deploy_config=deploy,
+ cli_overrides={"model": "/models/Qwen3-TTS"},
+ )
+ stage = config.stage_configs[-1]
+
+ assert stage.pooling_config.runner == "pooling"
+ assert stage.pooling_config.pooling_output_decoder == POOLING_OUTPUT_DECODER_PATH
+ assert isinstance(stage.pooling_config.default_pooling_params, PoolingParams)
+ assert stage.pooling_config.default_pooling_params.task == deploy.stages[-1].default_pooling_params["task"]
+
+ metadata = extract_stage_metadata_from_omni_stage_config(stage)
+ assert isinstance(metadata.default_sampling_params, PoolingParams)
+ assert metadata.default_sampling_params.task == deploy.stages[-1].default_pooling_params["task"]
+
+ engine_args = build_engine_args_dict_from_omni_stage_config(
+ stage,
+ stage.model_config.model or "/models/Qwen3-ForcedAligner-0.6B",
+ )
+ assert engine_args["runner"] == "pooling"
+ assert engine_args["pooling_output_decoder"] == POOLING_OUTPUT_DECODER_PATH
diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py
index e0a5ad9b0..97f9a5370 100644
--- a/tests/config/test_omni_config.py
+++ b/tests/config/test_omni_config.py
@@ -11,6 +11,7 @@ from pathlib import Path
import msgspec
import pytest
+import torch
from pydantic import ValidationError
from pydantic.fields import FieldInfo
from transformers import Qwen3OmniMoeConfig
@@ -123,6 +124,18 @@ def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) -
assert [stage.model_config.session_mode for stage in omni_config.stage_configs] == ["turn", "turn"]
+def test_nested_stage_override_deep_merges_structured_model_config() -> None:
+ config = _from_pipeline_key(
+ "cosmos3_policy",
+ deploy_config_path=get_deploy_config_path("cosmos3_policy_droid.yaml"),
+ cli_overrides={"stage_0_model_config": {"guardrails": False}},
+ )
+
+ model_config = config.stage_by_id(0).diffusion_config.model_config
+ assert model_config["guardrails"] is False
+ assert model_config["policy_server_config"]["action_space"] == "joint_position"
+
+
@pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(model_type: str):
pipeline = _resolve_pipeline_or_skip(model_type)
@@ -654,6 +667,7 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations():
"load_config",
"cache_config",
"scheduler_config",
+ "pooling_config",
"connector_config",
"runtime_config",
"parallel_config",
@@ -889,6 +903,17 @@ def test_structured_llm_stage_registration_payloads_remain_msgpack_transport_saf
msgspec.msgpack.encode(payload)
+def test_structured_diffusion_torch_dtype_is_msgpack_transport_safe():
+ stage_config = _from_pipeline_key(
+ "hunyuan_image3_dit",
+ cli_overrides={"dtype": torch.bfloat16},
+ ).stage_by_id(0)
+
+ serialized = _serialize_stage_config(stage_config)
+ assert serialized["diffusion_config"]["dtype"] == "bfloat16"
+ msgspec.msgpack.encode(serialized)
+
+
def test_diffusion_parallel_config_fields_cover_legacy_surface():
from vllm_omni.diffusion.data import DiffusionParallelConfig
@@ -1480,7 +1505,7 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
- engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model")
+ engine_args = build_engine_args_dict_from_omni_stage_config(stage, model=str(tmp_path))
od_config = OmniDiffusionConfig.from_kwargs(**engine_args)
assert od_config.kv_transfer_config.engine_id == "dit-engine-1"
@@ -1739,3 +1764,31 @@ def test_async_chunk_rejects_mismatched_connector_edge(disabled_stage, builder):
builder(pipeline, deploy)
else:
builder(pipeline, user_deploy_config=deploy)
+
+
+@pytest.mark.parametrize("explicit", [False, True])
+def test_diffusion_quantization_origin_survives_projection_and_transport(monkeypatch, explicit):
+ from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
+ from vllm_omni.quantization import build_quant_config
+
+ checkpoint = TransformerConfig.from_dict(
+ {
+ "quantization_config": {
+ "quant_method": "mxfp4",
+ "is_checkpoint_mxfp4_serialized": True,
+ "w4a8_fallback_steps": [37],
+ }
+ }
+ )
+ requested = build_quant_config("mxfp4", w4a8_fallback_steps=[]) if explicit else None
+ cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
+ tf_model_config=checkpoint,
+ quantization_config=requested,
+ )
+ assert cfg.quantization_config_is_auto_detected is not explicit
+ monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
+ monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda self: self.set_tf_model_config(checkpoint))
+ cfg.enrich_config()
+ restored = ForkingPickler.loads(ForkingPickler.dumps(cfg))
+ assert restored.quantization_config_is_auto_detected is not explicit
+ assert restored.quantization_config.w4a8_fallback_steps == ([] if explicit else [37])
diff --git a/tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py b/tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py
new file mode 100644
index 000000000..9a530bf90
--- /dev/null
+++ b/tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py
@@ -0,0 +1,679 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import json
+from collections.abc import Callable, Iterable
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, DiffusionParallelConfig, OmniDiffusionConfig
+from vllm_omni.diffusion.models.minimax_h3.fasth3_checkpoint import (
+ FASTH3_BASE_MODEL,
+ FASTH3_V2_BASE_SCHEDULE,
+ FASTH3_V2_MODEL_ID,
+ FastH3CheckpointSpec,
+)
+from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+ MiniMaxH3DiTArchConfig,
+ MiniMaxH3DiTModel,
+)
+from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+ resolve_minimax_h3_diffusion_model_path,
+)
+from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule
+from vllm_omni.errors import OmniClientError
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.lora.request import LoRARequest
+from vllm_omni.model_executor.models.minimax_h3.checkpoint import (
+ is_minimax_h3_modular,
+ resolve_minimax_h3_model_root,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _fastvideo_metadata() -> dict[str, object]:
+ """Return the release-owned metadata read from fastvideo_inference.json."""
+ return {
+ "model_id": FASTH3_V2_MODEL_ID,
+ "schema_version": "fasth3-inference-contract-v1",
+ "dmd_denoising_steps": [999, 874, 749, 624, 500, 375, 250, 125],
+ "transformer_forwards": 8,
+ "video_scheduler_shift": 10.0,
+ "audio_scheduler_shift": 3.0,
+ "guidance_scale": 1.0,
+ "attention_backend": "VIDEO_SPARSE_ATTN_H3",
+ "vsa_sparsity": 0.8,
+ "vsa_tile_size": 64,
+ "task": "t2av",
+ }
+
+
+def _sampling(
+ *,
+ num_inference_steps: int | str | None = 8,
+ guidance_scale: float | None = 1.0,
+ extra_args: dict[str, object] | None = None,
+ lora_request: LoRARequest | None = None,
+ timesteps: torch.Tensor | None = None,
+ sigmas: list[float] | None = None,
+) -> OmniDiffusionSamplingParams:
+ return OmniDiffusionSamplingParams(
+ # Keep the string form to exercise malformed client input validation.
+ num_inference_steps=num_inference_steps, # type: ignore[arg-type]
+ guidance_scale=guidance_scale,
+ extra_args=extra_args or {},
+ lora_request=lora_request,
+ timesteps=timesteps,
+ sigmas=sigmas,
+ )
+
+
+def _od_config(
+ *,
+ backend: str = "FASTVIDEO_VSA",
+ topk: int | None = None,
+ per_role: dict[str, AttentionSpec] | None = None,
+ lora_path: str | None = None,
+ ring_degree: int = 1,
+ allgather_degree: int = 1,
+) -> OmniDiffusionConfig:
+ attention_config = AttentionConfig(
+ default=AttentionSpec(backend=backend, fastvideo_vsa_topk=topk),
+ per_role=per_role or {},
+ )
+ return OmniDiffusionConfig(
+ diffusion_attention_config=attention_config,
+ lora_path=lora_path,
+ parallel_config=DiffusionParallelConfig(ring_degree=ring_degree, allgather_degree=allgather_degree),
+ )
+
+
+def test_v2_release_has_nine_sigma_nodes_and_eight_intervals():
+ assert FASTH3_V2_BASE_SCHEDULE.base_schedule == (
+ 0.999,
+ 0.874,
+ 0.749,
+ 0.624,
+ 0.5,
+ 0.375,
+ 0.25,
+ 0.125,
+ 0.0,
+ )
+ assert len(FASTH3_V2_BASE_SCHEDULE.base_schedule) == 9
+ assert FASTH3_V2_BASE_SCHEDULE.num_inference_steps == 8
+ assert FASTH3_V2_BASE_SCHEDULE.shifted_sigmas(10.0) == pytest.approx(
+ [
+ 0.9998999099189271,
+ 0.985788405143244,
+ 0.9675752486758817,
+ 0.94316807738815,
+ 0.9090909090909091,
+ 0.8571428571428571,
+ 0.7692307692307693,
+ 0.5882352941176471,
+ 0.0,
+ ],
+ abs=1e-12,
+ )
+ assert FASTH3_V2_BASE_SCHEDULE.shifted_sigmas(3.0) == pytest.approx(
+ [
+ 0.9996664442961973,
+ 0.9541484716157204,
+ 0.8995196156925539,
+ 0.8327402135231315,
+ 0.75,
+ 0.6428571428571429,
+ 0.5,
+ 0.3,
+ 0.0,
+ ],
+ abs=1e-12,
+ )
+
+
+def test_fastvideo_metadata_builds_release_sampling_metadata():
+ spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
+
+ assert spec.vsa_sparsity == 0.8
+ assert spec.release_metadata() == {
+ "partition": "fl2va",
+ "tasks": ["t2va"],
+ "sigma_shift_scales": {"video": 10.0, "audio": 3.0},
+ "base_schedule": list(FASTH3_V2_BASE_SCHEDULE.base_schedule),
+ }
+
+
+@pytest.mark.parametrize(
+ ("field", "value"),
+ [
+ ("model_id", "FastVideo/FastH3-other"),
+ ("schema_version", "fasth3-inference-contract-v0"),
+ ("dmd_denoising_steps", [999, 749, 500, 250]),
+ ("transformer_forwards", 4),
+ ("video_scheduler_shift", 12.0),
+ ("audio_scheduler_shift", 4.0),
+ ("guidance_scale", 2.0),
+ ("attention_backend", "FLASH_ATTN"),
+ ("vsa_sparsity", 0.9),
+ ("vsa_tile_size", 32),
+ ("task", "t2va"),
+ ],
+)
+def test_fastvideo_metadata_rejects_contract_drift(field, value):
+ metadata = _fastvideo_metadata()
+ metadata[field] = value
+
+ with pytest.raises(ValueError, match=field):
+ FastH3CheckpointSpec.from_metadata(metadata)
+
+
+def test_fastvideo_metadata_rejects_missing_contract_field():
+ metadata = _fastvideo_metadata()
+ metadata.pop("task")
+
+ with pytest.raises(ValueError, match="task"):
+ FastH3CheckpointSpec.from_metadata(metadata)
+
+
+def test_legacy_h3_metadata_is_not_treated_as_fastvideo_release():
+ with pytest.raises(ValueError, match="model_id"):
+ FastH3CheckpointSpec.from_metadata(
+ {
+ "schema_version": 1,
+ "partition": "fl2va",
+ "tasks": ["t2va", "fl2va"],
+ "sigma_shift_scales": {"video": 12.0, "audio": 3.0},
+ }
+ )
+
+
+def test_v2_serving_contract_accepts_backend_selected_for_h3_self_role():
+ config = _od_config(
+ backend="TORCH_SDPA",
+ per_role={"self": AttentionSpec(backend="FASTVIDEO_VSA")},
+ )
+ FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).check_serving_contract(
+ partition="fl2va", od_config=config
+ )
+
+
+@pytest.mark.parametrize(
+ ("kwargs", "partition", "match"),
+ [
+ ({}, "ref2va", "task-type fl2va"),
+ ({"lora_path": "/tmp/adapter"}, "fl2va", "LoRA"),
+ ({"backend": "TORCH_SDPA"}, "fl2va", "FASTVIDEO_VSA"),
+ ({"topk": 2}, "fl2va", "fixed fastvideo_vsa_topk"),
+ ({"ring_degree": 2}, "fl2va", "local attention"),
+ ({"allgather_degree": 2}, "fl2va", "local attention"),
+ ],
+)
+def test_v2_serving_contract_rejects_incompatible_runtime(kwargs, partition, match):
+ spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
+ with pytest.raises(ValueError, match=match):
+ spec.check_serving_contract(partition=partition, od_config=_od_config(**kwargs))
+
+
+def test_v2_serving_contract_reads_a_per_role_fixed_topk_override():
+ spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
+ config = _od_config(
+ per_role={"self": AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=1)},
+ )
+ with pytest.raises(ValueError, match="fixed fastvideo_vsa_topk"):
+ spec.check_serving_contract(partition="fl2va", od_config=config)
+
+
+def test_v2_request_contract_accepts_omitted_steps_and_default_shifts():
+ spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
+ spec.check_request(_sampling(num_inference_steps=None))
+ spec.check_request(_sampling(extra_args={"flow_shift": 10, "audio_flow_shift": 3}, guidance_scale=1))
+
+
+@pytest.mark.parametrize(
+ "sampling",
+ [
+ _sampling(num_inference_steps=None),
+ _sampling(timesteps=torch.tensor([0.9, 0.5, 0.0])),
+ _sampling(sigmas=[1.0, 0.5, 0.0]),
+ ],
+)
+def test_v2_step_execution_requires_pinned_count_and_schedule(sampling):
+ spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
+ with pytest.raises(OmniClientError, match="step execution"):
+ spec.check_request(sampling, step_execution=True)
+
+
+def test_v2_step_execution_accepts_only_explicit_pinned_count():
+ spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
+ spec.check_request(_sampling(num_inference_steps=8), step_execution=True)
+
+
+@pytest.mark.parametrize(
+ ("kwargs", "match"),
+ [
+ ({"num_inference_steps": 7}, "num_inference_steps=8"),
+ ({"num_inference_steps": 9}, "num_inference_steps=8"),
+ ({"num_inference_steps": "8"}, "num_inference_steps=8"),
+ ({"extra_args": {"flow_shift": 9}}, "flow_shift=10"),
+ ({"extra_args": {"audio_flow_shift": 4}}, "audio_flow_shift=3"),
+ ({"extra_args": {"flow_shift": "bad"}}, "flow_shift=10"),
+ ({"guidance_scale": 2.0}, "guidance_scale=1"),
+ (
+ {"lora_request": LoRARequest(lora_name="test", lora_int_id=1, lora_path="/tmp/test.safetensors")},
+ "per-request LoRA",
+ ),
+ ],
+)
+def test_v2_request_contract_rejects_sampling_drift(kwargs, match):
+ spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
+ with pytest.raises(OmniClientError, match=match):
+ spec.check_request(_sampling(**kwargs))
+
+
+def test_resolve_native_vaes_uses_pinned_revision_and_only_vae_patterns(tmp_path, monkeypatch):
+ model_root = tmp_path / "FastH3"
+ model_root.mkdir()
+ revision = "base-revision-123"
+ (model_root / "provenance.json").write_text(
+ json.dumps({"base_model": f"hf://{FASTH3_BASE_MODEL}@{revision}"}),
+ encoding="utf-8",
+ )
+ snapshot = tmp_path / "hf-snapshot"
+ calls: dict[str, object] = {}
+
+ def fake_download(**kwargs: object) -> str:
+ calls.update(kwargs)
+ snapshot.mkdir()
+ return str(snapshot)
+
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.models.minimax_h3.fasth3_checkpoint.download_weights_from_hf_specific",
+ fake_download,
+ )
+
+ result = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).resolve_native_vaes(model_root)
+
+ assert result == snapshot / "FL2VA"
+ assert calls == {
+ "model_name_or_path": FASTH3_BASE_MODEL,
+ "cache_dir": None,
+ "allow_patterns": ["FL2VA/video_vae/**", "FL2VA/audio_vae/**"],
+ "revision": revision,
+ "require_all": True,
+ }
+ assert sorted(path.name for path in model_root.iterdir()) == ["provenance.json"]
+ assert not (model_root / "model_index.json").exists()
+ assert not (model_root / "conversion.json").exists()
+
+
+@pytest.mark.parametrize(
+ "base_model",
+ [
+ "MiniMaxAI/MiniMax-H3@main",
+ "hf://MiniMaxAI/MiniMax-H3@",
+ "hf://Other/MiniMax-H3@revision",
+ "",
+ None,
+ ],
+)
+def test_resolve_native_vaes_rejects_unpinned_provenance(tmp_path, base_model):
+ model_root = tmp_path / "FastH3"
+ model_root.mkdir()
+ (model_root / "provenance.json").write_text(json.dumps({"base_model": base_model}), encoding="utf-8")
+
+ with pytest.raises(ValueError, match="provenance"):
+ FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).resolve_native_vaes(model_root)
+
+
+def test_modular_local_model_root_uses_text_encoder_without_partition_files(tmp_path):
+ root = tmp_path / "FastH3"
+ (root / "text_encoder").mkdir(parents=True)
+ (root / "text_encoder" / "config.json").write_text("{}", encoding="utf-8")
+ (root / "modular_model_index.json").write_text(
+ json.dumps({"_class_name": "MiniMaxH3ModularPipeline"}),
+ encoding="utf-8",
+ )
+
+ assert is_minimax_h3_modular(str(root)) is True
+ assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "text_encoder")
+ assert not (root / "FL2VA").exists()
+
+
+def test_modular_hub_model_root_downloads_only_text_encoder(monkeypatch, tmp_path):
+ from vllm_omni.model_executor.models.minimax_h3 import checkpoint as checkpoint_module
+
+ snapshot = tmp_path / "modular-snapshot"
+ calls: dict[str, object] = {}
+
+ def fake_index(model: str, *, revision: str | None = None) -> dict[str, str]:
+ assert model == "FastVideo/FastVideo-FastH3-8-Step-V2"
+ assert revision == "release-revision"
+ return {"_class_name": "MiniMaxH3ModularPipeline"}
+
+ def fake_download(**kwargs: object) -> str:
+ calls.update(kwargs)
+ snapshot.mkdir()
+ return str(snapshot)
+
+ monkeypatch.setattr(checkpoint_module, "get_diffusion_model_index", fake_index)
+ monkeypatch.setattr(checkpoint_module, "download_weights_from_hf_specific", fake_download)
+
+ result = resolve_minimax_h3_model_root(
+ "FastVideo/FastVideo-FastH3-8-Step-V2",
+ "release-revision",
+ "fl2va",
+ )
+
+ assert result == str(snapshot / "text_encoder")
+ assert calls == {
+ "model_name_or_path": "FastVideo/FastVideo-FastH3-8-Step-V2",
+ "cache_dir": None,
+ "allow_patterns": ["text_encoder/**"],
+ "revision": "release-revision",
+ "require_all": True,
+ }
+
+
+def test_native_hub_model_with_modular_class_keeps_partition_layout(monkeypatch):
+ from vllm_omni.model_executor.models.minimax_h3 import checkpoint as checkpoint_module
+
+ monkeypatch.setattr(
+ checkpoint_module,
+ "get_diffusion_model_index",
+ lambda _model, *, revision=None: {"_class_name": "MiniMaxH3ModularPipeline"},
+ )
+
+ assert checkpoint_module.is_minimax_h3_modular("MiniMaxAI/MiniMax-H3") is False
+
+
+def test_native_local_model_root_keeps_partition_layout(tmp_path):
+ root = tmp_path / "MiniMax-H3"
+ (root / "FL2VA" / "text_encoder").mkdir(parents=True)
+ (root / "FL2VA" / "text_encoder" / "config.json").write_text("{}", encoding="utf-8")
+
+ assert is_minimax_h3_modular(str(root)) is False
+ assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "FL2VA" / "text_encoder")
+
+
+def test_diffusion_resolver_returns_local_modular_root_without_partition_files(tmp_path):
+ root = tmp_path / "FastH3"
+ root.mkdir()
+ (root / "modular_model_index.json").write_text(
+ json.dumps({"_class_name": "MiniMaxH3ModularPipeline"}),
+ encoding="utf-8",
+ )
+
+ assert resolve_minimax_h3_diffusion_model_path(str(root), None, "fl2va") == str(root)
+ assert not (root / "FL2VA").exists()
+ assert not (root / "Ref2VA").exists()
+
+
+def test_diffusion_resolver_hub_modular_downloads_transformer_contract_only(monkeypatch, tmp_path):
+ from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
+
+ snapshot = tmp_path / "modular-snapshot"
+ calls: dict[str, object] = {}
+
+ def fake_modular(_model: str, _revision: str | None = None) -> bool:
+ return True
+
+ def fake_download(**kwargs: object) -> str:
+ calls.update(kwargs)
+ snapshot.mkdir()
+ return str(snapshot)
+
+ monkeypatch.setattr(pipeline_module, "is_minimax_h3_modular", fake_modular)
+ monkeypatch.setattr(pipeline_module, "download_weights_from_hf_specific", fake_download)
+
+ result = resolve_minimax_h3_diffusion_model_path(
+ "FastVideo/FastVideo-FastH3-8-Step-V2", "release-revision", "fl2va"
+ )
+
+ assert result == str(snapshot)
+ assert calls == {
+ "model_name_or_path": "FastVideo/FastVideo-FastH3-8-Step-V2",
+ "cache_dir": None,
+ "allow_patterns": [
+ "modular_model_index.json",
+ "fastvideo_inference.json",
+ "provenance.json",
+ "transformer/**",
+ ],
+ "revision": "release-revision",
+ "require_all": True,
+ }
+ patterns = calls["allow_patterns"]
+ assert isinstance(patterns, list)
+ assert all("video_vae" not in pattern and "audio_vae" not in pattern for pattern in patterns)
+
+
+def test_diffusers_arch_config_normalizes_all_aliases():
+ arch = MiniMaxH3DiTArchConfig.from_mapping(
+ {
+ "num_refiner_layers": 3,
+ "ffn_dim": 12,
+ "in_channels": 4,
+ "audio_in_channels": 6,
+ "freq_dim": 8,
+ "time_embed_hidden_dim": 10,
+ "rope_freq_dim": 5,
+ "patch_size": [1, 2, 2],
+ }
+ )
+
+ assert arch.token_refiner_num_layers == 3
+ assert arch.ffn_hidden_size == 12
+ assert arch.latents_dim == 4
+ assert arch.audio_latents_dim == 6
+ assert arch.timestep_input_dim == 8
+ assert arch.time_embed_hidden_size == 10
+ assert arch.rope_inv_freq_len == 5
+ assert arch.patch_size == (1, 2, 2)
+
+
+class _WeightTarget(nn.Module):
+ def __init__(self, shape: tuple[int, ...], loader: Callable[..., None]) -> None:
+ super().__init__()
+ self.weight = nn.Parameter(torch.zeros(shape))
+ self.weight.weight_loader = loader
+
+
+class _PipelineLoadTransformer(nn.Module):
+ def __init__(self, loaded: set[str]) -> None:
+ super().__init__()
+ self.arch = MiniMaxH3DiTArchConfig(num_layers=1)
+ self._loaded = loaded
+
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+ list(weights)
+ return self._loaded
+
+ def post_load_weights(self) -> None:
+ pass
+
+
+def test_pipeline_load_weights_rejects_missing_fast_h3_compression_gate():
+ from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ nn.Module.__init__(pipeline)
+ pipeline.transformer = _PipelineLoadTransformer({"blocks.0.attn.qkv_proj.weight"})
+ pipeline.video_vae = None
+ pipeline.audio_vae = None
+ pipeline._fasth3 = None
+ pipeline._fasth3_checkpoint = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
+
+ with pytest.raises(ValueError, match="missing compression gates"):
+ pipeline.load_weights([("transformer.blocks.0.attn.qkv_proj.weight", torch.zeros(1))])
+
+
+def _load_fixture(
+ *, diffusers_weights: bool
+) -> tuple[MiniMaxH3DiTModel, list[tuple[str | None, torch.Tensor]], list[tuple[int, torch.Tensor]]]:
+ model = object.__new__(MiniMaxH3DiTModel)
+ nn.Module.__init__(model)
+ model.arch = MiniMaxH3DiTArchConfig(
+ num_layers=1,
+ hidden_size=3,
+ num_attention_heads=2,
+ attention_head_dim=2,
+ ffn_hidden_size=6,
+ rope_inv_freq_len=2,
+ )
+ model._diffusers_weights = diffusers_weights
+ model._rope_theta = 10000.0
+ qkv_calls: list[tuple[str | None, torch.Tensor]] = []
+ fc1_calls: list[tuple[int, torch.Tensor]] = []
+
+ def qkv_loader(param: nn.Parameter, loaded_weight: torch.Tensor, shard_id: str | None = None) -> None:
+ if shard_id is None:
+ assert loaded_weight.shape == (12, 3)
+ param.data.copy_(loaded_weight)
+ qkv_calls.append((None, loaded_weight.clone()))
+ return
+ assert loaded_weight.shape == (4, 3)
+ assert shard_id in {"q", "k", "v"}
+ offset = {"q": 0, "k": 4, "v": 8}[shard_id]
+ param.data[offset : offset + 4].copy_(loaded_weight)
+ qkv_calls.append((shard_id, loaded_weight.clone()))
+
+ def fc1_loader(param: nn.Parameter, loaded_weight: torch.Tensor, shard_id: int) -> None:
+ assert loaded_weight.shape == (3, 3)
+ param.data[shard_id * 3 : (shard_id + 1) * 3].copy_(loaded_weight)
+ fc1_calls.append((shard_id, loaded_weight.clone()))
+
+ model.blocks = nn.ModuleList([nn.Module()])
+ model.blocks[0].attn = nn.Module()
+ model.blocks[0].attn.qkv_proj = _WeightTarget((12, 3), qkv_loader)
+ model.blocks[0].mlp = nn.Module()
+ model.blocks[0].mlp.fc1 = _WeightTarget((6, 3), fc1_loader)
+ model.rope = nn.Module()
+ model.rope.register_buffer("inv_freq", torch.zeros(2))
+ return model, qkv_calls, fc1_calls
+
+
+def test_diffusers_load_weights_uses_shape_aware_qkv_and_mlp_slices():
+ model, qkv_calls, fc1_calls = _load_fixture(diffusers_weights=True)
+ q = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 10
+ k = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 100
+ v = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 200
+ mlp = torch.arange(18, dtype=torch.float32).reshape(6, 3) + 300
+
+ loaded = model.load_weights(
+ [
+ ("transformer_blocks.0.attn.to_q.weight", q),
+ ("transformer_blocks.0.attn.to_k.weight", k),
+ ("transformer_blocks.0.attn.to_v.weight", v),
+ ("transformer_blocks.0.ff.net.0.proj.weight", mlp),
+ ]
+ )
+
+ assert loaded == {
+ "blocks.0.attn.qkv_proj.weight",
+ "blocks.0.mlp.fc1.weight",
+ "rope.inv_freq",
+ }
+ assert [shard_id for shard_id, _ in qkv_calls] == ["q", "k", "v"]
+ assert [shard_id for shard_id, _ in fc1_calls] == [0, 1]
+ torch.testing.assert_close(model.blocks[0].attn.qkv_proj.weight, torch.cat((q, k, v)))
+ torch.testing.assert_close(model.blocks[0].mlp.fc1.weight, torch.cat((mlp[3:], mlp[:3])))
+ torch.testing.assert_close(
+ model.rope.inv_freq,
+ 1.0 / (10000.0 ** (torch.arange(0, 4, 2, dtype=torch.float32) / 4)),
+ rtol=0,
+ atol=0,
+ )
+
+
+def test_native_load_weights_keeps_grouped_qkv_and_gate_first_layouts():
+ model, qkv_calls, fc1_calls = _load_fixture(diffusers_weights=False)
+ grouped_qkv = torch.arange(36, dtype=torch.float32).reshape(12, 3)
+ fc1 = torch.arange(18, dtype=torch.float32).reshape(6, 3)
+
+ loaded = model.load_weights(
+ [
+ ("blocks.0.attn.qkv_proj.weight", grouped_qkv),
+ ("blocks.0.mlp.fc1.weight", fc1),
+ ]
+ )
+
+ expected_qkv = torch.cat(
+ (
+ grouped_qkv[:2],
+ grouped_qkv[6:8],
+ grouped_qkv[2:4],
+ grouped_qkv[8:10],
+ grouped_qkv[4:6],
+ grouped_qkv[10:],
+ )
+ )
+ assert loaded == {"blocks.0.attn.qkv_proj.weight", "blocks.0.mlp.fc1.weight"}
+ assert [shard_id for shard_id, _ in qkv_calls] == [None]
+ assert [shard_id for shard_id, _ in fc1_calls] == [0, 1]
+ torch.testing.assert_close(model.blocks[0].attn.qkv_proj.weight, expected_qkv)
+ torch.testing.assert_close(model.blocks[0].mlp.fc1.weight, fc1)
+ assert torch.equal(model.rope.inv_freq, torch.zeros(2))
+
+
+def test_diffusers_load_weights_rejects_unknown_name():
+ model, _, _ = _load_fixture(diffusers_weights=True)
+
+ with pytest.raises(ValueError, match="unsupported Diffusers H3 weight"):
+ model.load_weights([("transformer_blocks.0.attn.not_a_projection.weight", torch.zeros(4, 3))])
+
+
+def test_diffusers_load_weights_rejects_duplicate_source_name():
+ model, _, _ = _load_fixture(diffusers_weights=True)
+ weight = torch.zeros(4, 3)
+
+ with pytest.raises(ValueError, match="duplicate Diffusers H3 weight"):
+ model.load_weights(
+ [
+ ("transformer_blocks.0.attn.to_q.weight", weight),
+ ("transformer_blocks.0.attn.to_q.weight", weight),
+ ]
+ )
+
+
+def test_diffusers_load_weights_rejects_incomplete_qkv_group():
+ model, _, _ = _load_fixture(diffusers_weights=True)
+
+ with pytest.raises(ValueError, match="incomplete Diffusers H3 QKV"):
+ model.load_weights(
+ [
+ ("transformer_blocks.0.attn.to_q.weight", torch.zeros(4, 3)),
+ ("transformer_blocks.0.attn.to_k.weight", torch.zeros(4, 3)),
+ ]
+ )
+
+
+def test_diffusers_load_weights_rejects_missing_compression_gate_parameter():
+ model, _, _ = _load_fixture(diffusers_weights=True)
+
+ with pytest.raises(ValueError, match="has no model parameter"):
+ model.load_weights([("transformer_blocks.0.attn.to_gate_compress.weight", torch.zeros(4, 3))])
+
+
+def test_resolve_sigma_positions_uses_release_metadata_schedule():
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ nn.Module.__init__(pipeline)
+ pipeline._fasth3 = None
+ pipeline._lora_sigma_schedules = {}
+ spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
+ pipeline._base_schedule_by_partition = {
+ "fl2va": DMD2SigmaSchedule.from_metadata(spec.release_metadata()),
+ }
+
+ positions, steps = pipeline._resolve_sigma_positions("t2va", _sampling())
+ assert positions == FASTH3_V2_BASE_SCHEDULE.base_schedule
+ assert steps == 8 == len(positions) - 1
+ with pytest.raises(OmniClientError, match="must be 8"):
+ pipeline._resolve_sigma_positions("t2va", _sampling(num_inference_steps=9))
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
index 4b6fc0d88..cea4862d1 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
@@ -423,6 +423,7 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership(
assert pipeline.partition == expected_partition
assert pipeline.supported_tasks == expected_tasks
assert len(created["dit"]) == expected_dits
+ assert all(kwargs["diffusers_weights"] is False for _, kwargs in created["dit"])
component_path = tmp_path / component_partition
assert created["video_vae"] == [str(component_path / "video_vae")]
assert created["audio_vae"] == [str(component_path / "audio_vae")]
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py
index 5c874ff02..c89c37db9 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py
@@ -262,7 +262,10 @@ def test_a_vsa_variant_is_recognised_by_its_compression_gates(tmp_path):
def test_a_tensor_naming_no_h3_parameter_is_an_error(tmp_path):
path = tmp_path / "fasth3" / "adapter_model.safetensors"
- _write_adapter(path, tensors={"transformer_blocks.0.attn.norm_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))})
+ _write_adapter(
+ path,
+ tensors={"transformer_blocks.0.attn.not_a_projection.lora_A.weight": torch.ones((_RANK, _HIDDEN))},
+ )
# Dropping it silently would load a model that is not the distilled student.
with pytest.raises(FastH3AdapterError, match="name no known"):
_load(path.parent)
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py
index 92bd09ac2..e1f83fd01 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py
@@ -1,4 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
from unittest.mock import Mock
@@ -41,7 +42,7 @@ class _FakeAttention(nn.Module):
super().__init__()
-def _small_od_config():
+def _small_od_config(*, class_name=None):
arch = {
"num_layers": 1,
"token_refiner_num_layers": 1,
@@ -60,6 +61,8 @@ def _small_od_config():
"final_adaln_out_features": 2 * 8,
"rope_inv_freq_len": 2,
}
+ if class_name is not None:
+ arch["_class_name"] = class_name
return SimpleNamespace(
tf_model_config=arch,
parallel_config=SimpleNamespace(ulysses_degree=1),
@@ -133,6 +136,24 @@ def test_fp8_scope_and_prefix_propagation(monkeypatch):
)
+def test_explicit_native_weight_format_overrides_diffusers_class_name(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+
+ monkeypatch.setattr(h3, "ColumnParallelLinear", _FakeLinear)
+ monkeypatch.setattr(h3, "MergedColumnParallelLinear", _FakeLinear)
+ monkeypatch.setattr(h3, "QKVParallelLinear", _FakeLinear)
+ monkeypatch.setattr(h3, "RowParallelLinear", _FakeLinear)
+ monkeypatch.setattr(h3, "Attention", _FakeAttention)
+ monkeypatch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: 1)
+
+ model = h3.MiniMaxH3DiTModel(
+ _small_od_config(class_name="MiniMaxH3Transformer3DModel"),
+ diffusers_weights=False,
+ )
+
+ assert model._diffusers_weights is False
+
+
class _WeightTarget(nn.Module):
def __init__(self, loader):
super().__init__()
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py
index ba0e0ddf5..3beda5acc 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py
@@ -187,8 +187,14 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig):
"model": model_ref,
"enforce_eager": True,
# The fused BF16 baseline only fits on H100-80GB with encoder TP.
- "stage_0_tensor_parallel_size": 2,
+ # This test loads the direct FL2VA diffusion pipeline (not the
+ # disaggregated two-stage MiniMax-H3 pipeline), so configure both the
+ # DiT world and the Qwen3-VL encoder TP explicitly.
+ "tensor_parallel_size": 2,
+ "text_encoder_tp_size": 2,
"vae_use_tiling": True,
+ # MiniMax H3's VAE is implemented by code shipped with the checkpoint.
+ "trust_remote_code": True,
}
omni_bl = Omni(**common_kwargs)
diff --git a/tests/diffusion/models/wan2_2/test_create_transformer_quant.py b/tests/diffusion/models/wan2_2/test_create_transformer_quant.py
index bb95363fd..b01f275eb 100644
--- a/tests/diffusion/models/wan2_2/test_create_transformer_quant.py
+++ b/tests/diffusion/models/wan2_2/test_create_transformer_quant.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Regression tests for transformer quant-config auto-detection.
The loader path at pipeline_wan2_2.py carries two quantization contracts:
@@ -27,6 +27,7 @@ from types import SimpleNamespace
import pytest
import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module
+from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import create_transformer_from_config
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -381,3 +382,213 @@ def test_pipeline_cascade_mxfp4_dualscale_each_transformer_gets_correct_ignored_
# od_config must remain unchanged — _create_transformer does not modify it.
assert od_config.quantization_config is None
+
+
+def _construct_quantized_pipeline(pipeline_kind, active, disk_quant_config, tmp_path, monkeypatch):
+ import json
+
+ import torch
+
+ import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v as i2v_module
+ import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace as vace_module
+
+ monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_npu", lambda: True)
+ for component in ("transformer", "transformer_2"):
+ directory = tmp_path / component
+ directory.mkdir()
+ config = dict(_MIN_CFG)
+ expert_quant_config = (
+ disk_quant_config[0 if component == "transformer" else 1]
+ if isinstance(disk_quant_config, tuple)
+ else disk_quant_config
+ )
+ if expert_quant_config is not None:
+ config["quantization_config"] = expert_quant_config
+ (directory / "config.json").write_text(json.dumps(config))
+ (tmp_path / "model_index.json").write_text(json.dumps({"transformer_2": ["diffusers", "WanTransformer3DModel"]}))
+
+ class FakeComponent(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.config = SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8)
+
+ class FakeTransformer(FakeComponent):
+ def __init__(self, quant_config=None, **kwargs):
+ super().__init__()
+ self.received_quant_config = quant_config
+
+ for module in (wan22_module, i2v_module):
+ monkeypatch.setattr(module, "get_local_device", lambda: torch.device("cpu"))
+ monkeypatch.setattr(module, "prefetch_subfolders", lambda *args, **kwargs: None)
+ monkeypatch.setattr(module, "from_pretrained_with_prefetch", lambda *args, **kwargs: FakeComponent())
+ monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+ monkeypatch.setattr(vace_module, "WanVACETransformer3DModel", FakeTransformer)
+ pipeline_cls = {
+ "t2v": wan22_module.Wan22Pipeline,
+ "i2v": i2v_module.Wan22I2VPipeline,
+ "vace": vace_module.Wan22VACEPipeline,
+ }[pipeline_kind]
+ monkeypatch.setattr(pipeline_cls, "setup_diffusion_pipeline_profiler", lambda *args, **kwargs: None)
+ od_config = OmniDiffusionConfig(
+ model=str(tmp_path),
+ model_class_name=pipeline_cls.__name__,
+ dtype=torch.bfloat16,
+ quantization_config=active,
+ boundary_ratio=0.875,
+ flow_shift=12.0,
+ enable_diffusion_pipeline_profiler=False,
+ )
+ od_config.enrich_config()
+ pipeline = pipeline_cls(od_config=od_config)
+ return pipeline
+
+
+@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
+@pytest.mark.parametrize("method, flag", [("mxfp4", "is_checkpoint_mxfp4_serialized")])
+@pytest.mark.parametrize("offline", [False, True])
+def test_pipeline_constructor_routes_each_expert_quantization(
+ pipeline_kind, method, flag, offline, tmp_path, monkeypatch, mocker
+):
+ """Run the actual pipeline constructors, expert factories and quant-method dispatch."""
+ from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+
+ from vllm_omni.quantization import build_quant_config
+
+ active = build_quant_config(
+ {
+ "transformer": {
+ "method": method,
+ "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"],
+ "w4a8_fallback_steps": [0, 1],
+ },
+ "transformer_2": {
+ "method": method,
+ "w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"],
+ "w4a8_fallback_steps": [],
+ },
+ }
+ )
+ disk_quant_config = None
+ if offline:
+ disk_quant_config = {
+ "quant_method": method,
+ flag: True,
+ "ignored_layers": ["blocks.12.ffn.net_2"],
+ "w4a8_fallback_layers": ["blocks.9.attn1.to_qkv"],
+ "w4a8_fallback_steps": [9],
+ }
+ pipeline = _construct_quantized_pipeline(pipeline_kind, active, disk_quant_config, tmp_path, monkeypatch)
+ for index, transformer in enumerate((pipeline.transformer, pipeline.transformer_2)):
+ received = transformer.received_quant_config
+ assert received.get_name() == method
+ assert getattr(received, flag) == offline
+ assert received.w4a8_fallback_steps == ([0, 1] if index == 0 else [])
+ selected = f"blocks.{10 + index}.attn1.to_qkv"
+ other = f"blocks.{11 - index}.attn1.to_qkv"
+ assert received.w4a8_fallback_layers == [selected]
+ assert received.get_quant_method(mocker.Mock(spec=LinearBase), selected).is_w4a8_fallback_layer
+ assert not received.get_quant_method(mocker.Mock(spec=LinearBase), other).is_w4a8_fallback_layer
+ if offline:
+ assert isinstance(
+ received.get_quant_method(mocker.Mock(spec=LinearBase), "blocks.12.ffn.net_2"), UnquantizedLinearMethod
+ )
+ assert pipeline.od_config.quantization_config is active
+
+
+@pytest.mark.parametrize("explicit_none", [False, True])
+def test_pipeline_disabled_component_requires_unquantized_checkpoint(explicit_none, monkeypatch):
+ from vllm_omni.quantization import build_quant_config
+
+ FakeTransformer, captured = _make_fake_transformer()
+ monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+ components = {"transformer": {"method": "mxfp4"}, "transformer_2": None} if explicit_none else {"vae": None}
+ active = build_quant_config(components)
+ pipeline = _FakePipeline(OmniDiffusionConfig(model="", quantization_config=active))
+ pipeline._create_transformer(_MIN_CFG, component="transformer_2")
+ assert captured[-1].get("quant_config") is None
+ with pytest.raises(ValueError, match="disabled for component 'transformer_2'"):
+ pipeline._create_transformer(
+ {**_MIN_CFG, "quantization_config": {"quant_method": "mxfp4", "is_checkpoint_mxfp4_serialized": True}},
+ component="transformer_2",
+ )
+
+
+def test_pipeline_component_default_reaches_expert_factory(monkeypatch):
+ from vllm_omni.quantization import build_quant_config
+
+ FakeTransformer, captured = _make_fake_transformer()
+ monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+ active = build_quant_config({"default": {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}})
+ pipeline = _FakePipeline(OmniDiffusionConfig(model="", quantization_config=active))
+ pipeline._create_transformer(_MIN_CFG, component="transformer_2")
+ assert captured[-1]["quant_config"] is active.default_config
+
+
+@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
+@pytest.mark.parametrize(
+ "configured_expert, with_default",
+ [
+ ("transformer", False),
+ ("transformer", True),
+ ("transformer_2", False),
+ ("transformer_2", True),
+ ],
+)
+def test_pipeline_partial_expert_config_does_not_match_sibling(
+ pipeline_kind, configured_expert, with_default, tmp_path, monkeypatch
+):
+ from vllm_omni.quantization import build_quant_config
+
+ component_spec = {
+ configured_expert: {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]},
+ }
+ if with_default:
+ component_spec["default"] = {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"]}
+ active = build_quant_config(component_spec)
+ pipeline = _construct_quantized_pipeline(pipeline_kind, active, None, tmp_path, monkeypatch)
+ sibling = "transformer_2" if configured_expert == "transformer" else "transformer"
+ selected = getattr(pipeline, configured_expert).received_quant_config
+ other = getattr(pipeline, sibling).received_quant_config
+ assert selected is active.component_configs[configured_expert]
+ assert other is active.default_config
+ if with_default:
+ assert other.w4a8_fallback_layers == ["blocks.11.attn1.to_qkv"]
+ else:
+ assert other is None
+
+
+@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
+@pytest.mark.parametrize("different_ignored_layers", [False, True])
+def test_enriched_auto_quantization_preserves_each_expert_policy(
+ pipeline_kind, different_ignored_layers, tmp_path, monkeypatch
+):
+ """Startup auto-detection must not turn the first expert into a user override."""
+ high = {
+ "quant_method": "mxfp4",
+ "is_checkpoint_mxfp4_serialized": True,
+ "ignored_layers": ["blocks.0.ffn.net_2"],
+ "w4a8_fallback_steps": [],
+ "w4a8_fallback_layers": [],
+ "mxfp4_scale_alg": 0,
+ "require_smooth_scale": False,
+ }
+ low = {
+ **high,
+ "ignored_layers": ["blocks.1.ffn.net_2"] if different_ignored_layers else high["ignored_layers"],
+ "w4a8_fallback_steps": [37],
+ "w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"],
+ "mxfp4_scale_alg": 2,
+ "require_smooth_scale": True,
+ }
+ pipeline = _construct_quantized_pipeline(pipeline_kind, None, (high, low), tmp_path, monkeypatch)
+ assert pipeline.od_config.quantization_config is not None # enrich_config really detected the first expert.
+ for transformer, expected in [(pipeline.transformer, high), (pipeline.transformer_2, low)]:
+ config = transformer.received_quant_config
+ for policy in (
+ "ignored_layers",
+ "w4a8_fallback_steps",
+ "w4a8_fallback_layers",
+ "mxfp4_scale_alg",
+ "require_smooth_scale",
+ ):
+ assert getattr(config, policy) == expected[policy], policy
diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py
index 526eeeb15..58a5e961b 100644
--- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py
+++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py
@@ -20,9 +20,11 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@pytest.fixture(autouse=True)
-def _cpu_platform(monkeypatch):
+def _cpu_pipeline_runtime(monkeypatch):
+ # These lightweight pipelines have no accelerator or distributed groups.
module = importlib.import_module("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2")
monkeypatch.setattr(module.current_omni_platform, "is_available", lambda: False)
+ monkeypatch.setattr("vllm_omni.diffusion.distributed.parallel_state._PP", SimpleNamespace(world_size=1))
class _StubTransformer(nn.Module):
@@ -477,6 +479,100 @@ def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None:
assert torch.equal(result, torch.full_like(latents, 10.0))
+@pytest.mark.parametrize("fail_first_request", [False, True])
+def test_runner_publishes_wan_step_context_across_requests(fail_first_request, monkeypatch):
+ from vllm_omni.diffusion.data import OmniDiffusionConfig
+ from vllm_omni.diffusion.forward_context import (
+ ForwardContext,
+ get_forward_context,
+ is_forward_context_available,
+ override_forward_context,
+ )
+ from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+ from vllm_omni.quantization.mxfp4_config import _is_w4a8_fallback_step
+
+ contexts: list[ForwardContext] = []
+ events = []
+
+ class InjectedPredictionError(Exception):
+ pass
+
+ class RecordingTransformer(_StubTransformer):
+ def __init__(self, expert):
+ super().__init__()
+ self.expert = expert
+
+ def forward(self, hidden_states, encoder_hidden_states, **kwargs):
+ context = get_forward_context()
+ assert context is contexts[-1]
+ request_index = len(contexts) - 1
+ step = context.denoise_step_idx
+ positive = bool(encoder_hidden_states[0, 0, 0] > 0)
+ events.append((request_index, step, self.expert, positive, _is_w4a8_fallback_step([0, 2])))
+ if fail_first_request and request_index == 0 and step == 1 and not positive:
+ raise InjectedPredictionError
+ return (torch.ones_like(hidden_states) * (1 if positive else -1),)
+
+ def encode_prompt(**kwargs):
+ context = get_forward_context()
+ assert context.denoise_step_idx is None
+ contexts.append(context)
+ embeds = torch.ones(1, 8, 8)
+ return embeds, -embeds
+
+ pipeline = _make_pipeline()
+ pipeline.transformer = RecordingTransformer("high")
+ pipeline.transformer_2 = RecordingTransformer("low")
+ pipeline.scheduler = _StubScheduler([900, 500, 100])
+ monkeypatch.setattr(pipeline, "encode_prompt", encode_prompt)
+ monkeypatch.setattr(pipeline, "scheduler_step_maybe_with_cfg", lambda pred, t, latents, cfg: latents)
+
+ # Use the real runner context manager, Wan forward/diffuse, and CFG
+ # dispatch. Only model computation, scheduler math and unrelated KV I/O
+ # are replaced; no test code publishes a denoise step.
+ runner = object.__new__(DiffusionModelRunner)
+ runner.pipeline = pipeline
+ runner.od_config = OmniDiffusionConfig(model="", dtype=torch.float32)
+ runner.vllm_config = None
+ runner.cache_backend = None
+ monkeypatch.setattr(runner, "_prepare_request_for_forward", lambda *args, **kwargs: None)
+ with override_forward_context(None):
+ for request_index in range(2):
+ request = OmniDiffusionRequest(
+ prompt="a fox walks",
+ request_id=f"request-{request_index}",
+ sampling_params=OmniDiffusionSamplingParams(
+ num_frames=1, num_inference_steps=3, guidance_scale=2.0, output_type="latent"
+ ),
+ )
+
+ def execute():
+ return runner._execute_request_list(
+ [request],
+ od_config=runner.od_config,
+ allow_single_output=True,
+ require_request_batch_support=False,
+ record_name="test_wan_step_context",
+ record_output_peak_memory=False,
+ )
+
+ if fail_first_request and request_index == 0:
+ with pytest.raises(InjectedPredictionError):
+ execute()
+ else:
+ execute()
+ assert not is_forward_context_available()
+
+ assert len(contexts) == 2 and contexts[0] is not contexts[1]
+ expected = [
+ (request, step, "high" if step == 0 else "low", positive, step in (0, 2))
+ for request in range(2)
+ for step in range(2 if fail_first_request and request == 0 else 3)
+ for positive in (True, False)
+ ]
+ assert events == expected
+
+
class _StubDMDScheduler:
def __init__(self) -> None:
self.predict_clean_calls: list[tuple[float, float, float]] = []
diff --git a/tests/diffusion/quantization/test_mxfp4_checkpoint.py b/tests/diffusion/quantization/test_mxfp4_checkpoint.py
new file mode 100644
index 000000000..e183df078
--- /dev/null
+++ b/tests/diffusion/quantization/test_mxfp4_checkpoint.py
@@ -0,0 +1,294 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CPU coverage for MXFP4 checkpoint loading and the W4A8 dtype contract."""
+
+import sys
+from types import ModuleType
+
+import pytest
+import torch
+from vllm.config.load import LoadConfig
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.forward_context import (
+ ForwardContext,
+ override_forward_context,
+ set_forward_context_denoise_step_idx,
+)
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.quantization.mxfp4_config import (
+ DiffusionMXFP4Config,
+ DiffusionMXFP4DualScaleMixedConfig,
+ NPUMxfp4DualScaleLinearMethod,
+ NPUMxfp4LinearMethod,
+ NPUMxfp4OnlineLinearMethod,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.fixture(autouse=True)
+def _patch_tp_state(monkeypatch):
+ monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0)
+ monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1)
+
+
+class _CheckpointLayer(torch.nn.Module):
+ def __init__(self, method, dtype=torch.bfloat16):
+ super().__init__()
+ method.create_weights(
+ self,
+ input_size_per_partition=512,
+ output_partition_sizes=[2],
+ input_size=512,
+ output_size=2,
+ params_dtype=dtype,
+ weight_loader=default_weight_loader,
+ )
+
+ def load_weights(self, weights):
+ params = dict(self.named_parameters())
+ loaded = set()
+ for name, value in weights:
+ if name in params:
+ params[name].weight_loader(params[name], value)
+ loaded.add(name)
+ return loaded
+
+
+def _load_checkpoint(layer, config, weights, monkeypatch):
+ od_config = OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=config)
+ loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+ loader.counter_before_loading_weights = 0.0
+ monkeypatch.setattr(loader, "get_all_weights", lambda model: iter(weights.items()))
+ loader.load_weights(layer)
+
+
+@pytest.mark.parametrize("serialized", [False, True])
+def test_single_scale_requires_checkpoint_scale_only_when_serialized(serialized, monkeypatch):
+ config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized)
+ # Exercise the shared implementation directly as well as the online MRO
+ # below: checkpoint requirements must follow the configured storage mode.
+ layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
+ weights = {"weight": torch.ones((2, 512), dtype=torch.bfloat16)}
+ if serialized:
+ with pytest.raises(ValueError, match="Required weights.*weight_scale"):
+ _load_checkpoint(layer, config, weights, monkeypatch)
+ else:
+ _load_checkpoint(layer, config, weights, monkeypatch)
+
+
+@pytest.mark.parametrize("scale_alg", [0, 2])
+def test_online_mxfp4_loads_bf16_checkpoint_without_scale(scale_alg, monkeypatch):
+ config = DiffusionMXFP4Config(mxfp4_scale_alg=scale_alg)
+ method = NPUMxfp4OnlineLinearMethod(config)
+ layer = _CheckpointLayer(method)
+ assert set(dict(layer.named_parameters())) == {"weight"}
+ assert layer.weight.is_meta
+
+ weight = torch.arange(1024, dtype=torch.bfloat16).reshape(2, 512)
+ packed_weight = torch.zeros((2, 256), dtype=torch.uint8)
+ generated_scale = torch.full((2, 16), 127, dtype=torch.uint8)
+ quantized_inputs = []
+
+ def quantize(value, **kwargs):
+ quantized_inputs.append(value.clone())
+ expected = dict(dst_type="fp4", axis=-1, block_size=32, round_mode="rint", scale_alg=scale_alg)
+ if scale_alg == 2:
+ expected["dst_type_max"] = 7.25
+ assert kwargs == expected
+ return packed_weight, generated_scale
+
+ # Keep real creation, lazy loading, processing, and strict loader checks;
+ # substitute only the device quantization operator for this CPU test.
+ npu = ModuleType("torch_npu")
+ npu.__dict__.update(float4_e2m1fn_x2="fp4", npu_dynamic_mx_quant=quantize)
+ monkeypatch.setitem(sys.modules, "torch_npu", npu)
+ _load_checkpoint(layer, config, {"weight": weight}, monkeypatch)
+
+ assert len(quantized_inputs) == 1
+ torch.testing.assert_close(quantized_inputs[0], weight, rtol=0, atol=0)
+ torch.testing.assert_close(layer.weight, packed_weight, rtol=0, atol=0)
+ torch.testing.assert_close(layer.weight_scale, generated_scale.reshape(2, 8, 2), rtol=0, atol=0)
+ assert not getattr(layer.weight_scale, "is_checkpoint_required", False)
+ method.process_weights_after_loading(layer)
+ assert len(quantized_inputs) == 1
+
+
+@pytest.mark.parametrize("require_smooth", [False, True])
+@pytest.mark.parametrize("has_smooth", [False, True])
+def test_single_scale_checkpoint_smooth_default_and_required(require_smooth, has_smooth, monkeypatch):
+ config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=require_smooth)
+ layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
+ weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
+ scale = torch.linspace(0.5, 3.0, 512)
+ if has_smooth:
+ weights["mul_scale"] = scale
+ else:
+ weights.pop("mul_scale")
+
+ if require_smooth and not has_smooth:
+ with pytest.raises(ValueError, match="mul_scale"):
+ _load_checkpoint(layer, config, weights, monkeypatch)
+ else:
+ _load_checkpoint(layer, config, weights, monkeypatch)
+ torch.testing.assert_close(layer.mul_scale, scale if has_smooth else torch.ones(512), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("bad_scale", ["shape", "name"])
+def test_declared_single_scale_smooth_rejects_malformed_checkpoint(bad_scale, monkeypatch):
+ config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True)
+ layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
+ weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
+ if bad_scale == "shape":
+ weights["mul_scale"] = torch.ones(511)
+ with pytest.raises(AssertionError):
+ _load_checkpoint(layer, config, weights, monkeypatch)
+ else:
+ weights["multiply_scale"] = weights.pop("mul_scale")
+ with pytest.raises(ValueError, match="mul_scale"):
+ _load_checkpoint(layer, config, weights, monkeypatch)
+
+
+@pytest.mark.parametrize("required", [False, True])
+@pytest.mark.parametrize("bad_value", [float("nan"), float("inf"), 0.0, -1.0])
+def test_single_scale_rejects_invalid_smooth_values(required, bad_value, monkeypatch):
+ config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=required)
+ layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
+ weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
+ weights["mul_scale"][7] = bad_value
+ with pytest.raises(ValueError, match="finite, strictly positive"):
+ _load_checkpoint(layer, config, weights, monkeypatch)
+
+
+@pytest.mark.parametrize("bad_scale", [torch.ones(512, dtype=torch.int32), torch.ones(1, 512)])
+def test_single_scale_rejects_invalid_smooth_type_or_rank(bad_scale, monkeypatch):
+ config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)
+ layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
+ weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
+ weights["mul_scale"] = bad_scale
+ with pytest.raises(ValueError, match="one-dimensional floating-point"):
+ _load_checkpoint(layer, config, weights, monkeypatch)
+
+
+@pytest.mark.parametrize("bad_value", [1e10, 1e-10])
+def test_single_scale_rejects_smooth_overflow_or_underflow_after_cast(bad_value, monkeypatch):
+ config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)
+ method = NPUMxfp4LinearMethod(config)
+ layer = _CheckpointLayer(method, torch.float16)
+ layer.mul_scale.data.fill_(bad_value)
+ # Validation must fail before any device packing or GEMM.
+ monkeypatch.setitem(sys.modules, "torch_npu", ModuleType("torch_npu"))
+ with pytest.raises(ValueError, match="finite, strictly positive"):
+ method.process_weights_after_loading(layer)
+
+
+@pytest.mark.parametrize(
+ ("mode", "missing"),
+ [
+ ("single", "weight"),
+ ("single", "weight_scale"),
+ ("dual", "weight"),
+ ("dual", "weight_scale"),
+ ("dual", "weight_dual_scale"),
+ ("dual", "mul_scale"),
+ ],
+)
+def test_offline_mxfp4_required_tensors_do_not_use_legacy_scale_tolerance(mode, missing, monkeypatch):
+ if mode == "single":
+ config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)
+ method = NPUMxfp4LinearMethod(config)
+ else:
+ config = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=True)
+ method = NPUMxfp4DualScaleLinearMethod(config)
+ layer = _CheckpointLayer(method)
+ weights = {name: torch.ones_like(param) for name, param in layer.named_parameters() if name != missing}
+ with pytest.raises(ValueError, match=missing):
+ _load_checkpoint(layer, config, weights, monkeypatch)
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_single_scale_smooth_uses_real_row_parallel_loader(rank):
+ from vllm.model_executor.layers.linear import RowParallelLinear
+
+ layer = _CheckpointLayer(NPUMxfp4LinearMethod(DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)))
+ layer.tp_rank = rank
+ checkpoint_scale = torch.arange(1024, dtype=torch.float32)
+ RowParallelLinear.weight_loader(layer, layer.mul_scale, checkpoint_scale)
+ torch.testing.assert_close(layer.mul_scale, checkpoint_scale[rank * 512 : (rank + 1) * 512], rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("fallback", [False, True])
+@pytest.mark.parametrize("with_bias", [False, True])
+@pytest.mark.parametrize("scale_alg", [0, 2])
+def test_single_scale_smooth_precedes_quantization_and_preserves_output(
+ dtype, fallback, with_bias, scale_alg, monkeypatch
+):
+ config = DiffusionMXFP4Config(
+ is_checkpoint_mxfp4_serialized=True, w4a8_fallback_steps=[0], mxfp4_scale_alg=scale_alg
+ )
+ method = NPUMxfp4LinearMethod(config)
+ layer = _CheckpointLayer(method, dtype)
+ weights = {
+ "weight": torch.linspace(-0.5, 1.5, 1024, dtype=dtype).reshape(2, 512),
+ "weight_scale": torch.full((2, 16), 127, dtype=torch.uint8),
+ "mul_scale": torch.linspace(0.5, 3.0, 512),
+ }
+ _load_checkpoint(layer, config, weights, monkeypatch)
+ x = torch.linspace(-1.0, 2.0, 3072, dtype=dtype).reshape(2, 3, 512)
+ bias = torch.tensor([0.5, -1.0], dtype=dtype) if with_bias else None
+ quantized_inputs = []
+
+ # Only the two device operators are substituted. Exercise real apply(),
+ # step selection, Smooth multiplication, bias handling, and reshape logic.
+ npu = ModuleType("torch_npu")
+
+ def quantize(value, **kwargs):
+ quantized_inputs.append(value.clone())
+ expected = dict(
+ dst_type=torch.float8_e4m3fn if fallback else "fp4",
+ axis=-1,
+ block_size=32,
+ round_mode="rint",
+ scale_alg=0 if fallback else scale_alg,
+ )
+ if not fallback and scale_alg == 2:
+ expected["dst_type_max"] = 7.25
+ assert kwargs == expected
+ return value, torch.ones((6, 16), dtype=torch.uint8)
+
+ def matmul(x_q, weight, weight_scale, *, bias, output_dtype, **kwargs):
+ if fallback:
+ assert output_dtype == torch.bfloat16
+ if bias is not None:
+ assert bias.dtype == output_dtype
+ assert bias.shape == (1, 2)
+ result = x_q.float() @ weight.float()
+ if bias is not None:
+ result += bias.float()
+ return result.to(output_dtype)
+
+ npu.__dict__.update(
+ float8_e4m3fn=torch.float8_e4m3fn,
+ float8_e8m0fnu=torch.uint8,
+ float4_e2m1fn_x2="fp4",
+ npu_dynamic_mx_quant=quantize,
+ npu_quant_matmul=matmul,
+ )
+ monkeypatch.setitem(sys.modules, "torch_npu", npu)
+ with override_forward_context(ForwardContext()):
+ set_forward_context_denoise_step_idx(0 if fallback else 1)
+ result = method.apply(layer, x, bias)
+
+ expected_input = x.reshape(-1, 512) * weights["mul_scale"].to(dtype)
+ assert len(quantized_inputs) == 1
+ torch.testing.assert_close(quantized_inputs[0], expected_input, rtol=0, atol=0)
+ expected = expected_input.float() @ weights["weight"].float().T
+ if bias is not None:
+ expected += bias.float()
+ if fallback:
+ expected = expected.bfloat16()
+ torch.testing.assert_close(result, expected.to(dtype).reshape(2, 3, 2), rtol=0, atol=0)
+ assert result.dtype == dtype
diff --git a/tests/diffusion/quantization/test_mxfp4_config.py b/tests/diffusion/quantization/test_mxfp4_config.py
index 240ee2509..c51b2013b 100644
--- a/tests/diffusion/quantization/test_mxfp4_config.py
+++ b/tests/diffusion/quantization/test_mxfp4_config.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for MXFP4 quantization configs and the MXFP4 DualScale + BF16 mixed config."""
import pytest
@@ -27,12 +27,42 @@ def test_mxfp4_config_get_name():
assert DiffusionMXFP4Config.get_name() == "mxfp4"
+@pytest.mark.parametrize("value", [True, False, None, "2", 2.0, 1, -1, 3])
+def test_mxfp4_scale_alg_rejects_invalid_values(value):
+ from vllm_omni.quantization import build_quant_config
+
+ with pytest.raises(ValueError, match="mxfp4_scale_alg"):
+ build_quant_config({"method": "mxfp4", "mxfp4_scale_alg": value})
+
+
+@pytest.mark.parametrize("runtime_alg", [0, 2])
+@pytest.mark.parametrize("serialized", [False, True])
+def test_runtime_algorithm_survives_offline_storage_rebuild(runtime_alg, serialized):
+ from vllm_omni.quantization.factory import resolve_quant_config_from_disk
+ from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+ active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized, mxfp4_scale_alg=runtime_alg)
+ resolved = resolve_quant_config_from_disk(
+ active,
+ {
+ "quant_method": "mxfp4",
+ "is_checkpoint_mxfp4_serialized": True,
+ "mxfp4_scale_alg": 2,
+ "ignored_layers": ["proj_out"],
+ },
+ )
+ assert resolved.is_checkpoint_mxfp4_serialized
+ assert resolved.mxfp4_scale_alg == runtime_alg
+ assert resolved.w4a8_fallback_layers == resolved.w4a8_fallback_steps == []
+
+
def test_mxfp4_config_from_config_defaults():
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
cfg = DiffusionMXFP4Config.from_config({})
assert cfg.is_checkpoint_mxfp4_serialized is False
assert cfg.ignored_layers == []
+ assert cfg.w4a8_fallback_layers == []
def test_mxfp4_config_from_config_serialized():
@@ -42,6 +72,98 @@ def test_mxfp4_config_from_config_serialized():
assert cfg.is_checkpoint_mxfp4_serialized is True
+def test_mxfp4_config_requires_offline_smooth():
+ from vllm_omni.quantization import build_quant_config
+ from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+ values = {"is_checkpoint_mxfp4_serialized": True, "require_smooth_scale": True}
+ assert DiffusionMXFP4Config.from_config(values).require_smooth_scale
+ assert build_quant_config({"method": "mxfp4", **values}).require_smooth_scale
+ with pytest.raises(ValueError, match="offline"):
+ build_quant_config({"method": "mxfp4", "require_smooth_scale": True})
+ with pytest.raises(ValueError, match="boolean"):
+ DiffusionMXFP4Config.from_config({**values, "require_smooth_scale": "false"})
+
+
+@pytest.mark.parametrize("serialized", [False, True])
+def test_mxfp4_disk_smooth_requirement_preserves_step_policy(serialized):
+ from vllm_omni.quantization.factory import resolve_quant_config_from_disk
+ from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+ active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized, w4a8_fallback_steps=[1])
+ resolved = resolve_quant_config_from_disk(
+ active,
+ {
+ "quant_method": "mxfp4",
+ "is_checkpoint_mxfp4_serialized": True,
+ "require_smooth_scale": True,
+ },
+ )
+ assert resolved.require_smooth_scale
+ assert resolved.is_checkpoint_mxfp4_serialized
+ assert resolved.w4a8_fallback_steps == [1]
+
+
+@pytest.mark.parametrize("disk_requirement", [None, False, True])
+def test_mxfp4_expert_rebuild_cannot_weaken_explicit_smooth_requirement(disk_requirement):
+ from vllm_omni.quantization.factory import resolve_quant_config_from_disk
+ from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+ active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True)
+ disk = {
+ "quant_method": "mxfp4",
+ "is_checkpoint_mxfp4_serialized": True,
+ "ignored_layers": ["blocks.0.attn2.to_q"],
+ }
+ if disk_requirement is not None:
+ disk["require_smooth_scale"] = disk_requirement
+ resolved = resolve_quant_config_from_disk(active, disk)
+ assert resolved.require_smooth_scale
+ assert resolved.ignored_layers == disk["ignored_layers"]
+
+
+@pytest.mark.parametrize("disk_qc", ["mxfp4", {"quant_method": "mxfp4"}])
+def test_mxfp4_method_only_disk_config_preserves_caller_layer_policy(disk_qc):
+ from vllm_omni.quantization.factory import resolve_quant_config_from_disk
+ from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+ active = DiffusionMXFP4Config(
+ ignored_layers=["blocks.0.attn2.to_q"],
+ w4a8_fallback_steps=[1],
+ w4a8_fallback_layers=["blocks.1.attn1.to_qkv"],
+ )
+
+ resolved = resolve_quant_config_from_disk(active, disk_qc)
+
+ assert resolved is active
+ assert resolved.ignored_layers == ["blocks.0.attn2.to_q"]
+ assert resolved.w4a8_fallback_steps == [1]
+ assert resolved.w4a8_fallback_layers == ["blocks.1.attn1.to_qkv"]
+
+
+def test_quantization_disk_string_rejects_mismatched_active_method():
+ from vllm_omni.quantization import build_quant_config
+ from vllm_omni.quantization.factory import resolve_quant_config_from_disk
+
+ active = build_quant_config("mxfp4")
+
+ with pytest.raises(ValueError, match="mxfp4_dualscale.*active quantization config is 'mxfp4'"):
+ resolve_quant_config_from_disk(active, "mxfp4_dualscale")
+
+
+def test_quantization_disk_string_accepts_equivalent_method_alias():
+ from vllm_omni.quantization.factory import resolve_quant_config_from_disk
+
+ class ActiveAliasConfig:
+ @staticmethod
+ def get_name() -> str:
+ return "inc"
+
+ active = ActiveAliasConfig()
+
+ assert resolve_quant_config_from_disk(active, "auto-round") is active
+
+
def test_mxfp4_config_from_config_ignored_layers():
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
@@ -57,6 +179,202 @@ def test_mxfp4_config_from_config_modules_to_not_convert_fallback():
assert cfg.ignored_layers == ["proj_out"]
+@pytest.mark.parametrize("serialized", [False, True])
+@pytest.mark.parametrize(
+ "policy",
+ [
+ {"w4a8_fallback_steps": [0]},
+ {"w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]},
+ {"w4a8_fallback_steps": [1], "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]},
+ ],
+)
+def test_dualscale_rejects_w4a8_configuration(serialized, policy):
+ from vllm_omni.quantization import build_quant_config
+ from vllm_omni.quantization.factory import resolve_quant_config_from_disk
+ from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+ config = {"quant_method": "mxfp4_dualscale", "is_checkpoint_serialized": serialized, **policy}
+ with pytest.raises(ValueError, match="does not support W4A8"):
+ DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=serialized, **policy)
+ with pytest.raises(ValueError, match="does not support W4A8"):
+ DiffusionMXFP4DualScaleMixedConfig.from_config(config)
+ with pytest.raises(ValueError, match="does not support W4A8"):
+ build_quant_config({"transformer": config})
+ with pytest.raises(ValueError, match="does not support W4A8"):
+ resolve_quant_config_from_disk(None, config)
+
+
+@pytest.mark.parametrize("method", ["mxfp4", "mxfp4_dualscale"])
+@pytest.mark.parametrize("steps", [[-1], [True], [1.5], ["2"], "0,2", (0, 2)])
+def test_mxfp4_rejects_invalid_fallback_steps(method, steps):
+ from vllm_omni.quantization import build_quant_config
+
+ with pytest.raises(ValueError, match="non-negative integer"):
+ build_quant_config({"method": method, "w4a8_fallback_steps": steps})
+
+
+def test_mxfp4_fallback_steps_follow_request_context():
+ from vllm_omni.diffusion.forward_context import (
+ ForwardContext,
+ override_forward_context,
+ set_forward_context_denoise_step_idx,
+ )
+ from vllm_omni.quantization import build_quant_config
+ from vllm_omni.quantization.mxfp4_config import _is_w4a8_fallback_step
+
+ steps = [2, 0, 2]
+ cfg = build_quant_config({"method": "mxfp4", "w4a8_fallback_steps": steps})
+ steps.append(1)
+ assert cfg.w4a8_fallback_steps == [0, 2]
+ with override_forward_context(None):
+ assert not _is_w4a8_fallback_step([])
+ with pytest.raises(RuntimeError, match="publishes its denoise step"):
+ _is_w4a8_fallback_step(cfg.w4a8_fallback_steps)
+ for _request in range(2):
+ with override_forward_context(ForwardContext()):
+ for step, expected in [(0, True), (1, False), (2, True), (3, False)]:
+ set_forward_context_denoise_step_idx(step)
+ # Both CFG passes and both experts read the same index.
+ assert [_is_w4a8_fallback_step(cfg.w4a8_fallback_steps) for _cfg_pass in range(2)] == [
+ expected,
+ expected,
+ ]
+ with pytest.raises(RuntimeError, match="publishes its denoise step"):
+ _is_w4a8_fallback_step(cfg.w4a8_fallback_steps)
+
+
+@pytest.mark.parametrize("method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized")])
+@pytest.mark.parametrize("steps", [[], [0, 2]])
+@pytest.mark.parametrize("already_serialized", [False, True])
+def test_mxfp4_disk_rebuild_preserves_runtime_fallback(method, flag, steps, already_serialized):
+ from vllm_omni.quantization import build_quant_config
+ from vllm_omni.quantization.factory import resolve_quant_config_from_disk
+
+ layers = ["blocks.10.attn1.to_qkv"] if steps else []
+ active = build_quant_config(
+ {
+ "method": method,
+ flag: already_serialized,
+ "w4a8_fallback_steps": steps,
+ "w4a8_fallback_layers": layers,
+ }
+ )
+ # Separate experts may have distinct high precision layers. Both must
+ # retain the requested step policy during online/offline reconciliation.
+ for ignored in [["blocks.0.attn1.to_q"], ["blocks.1.attn1.to_q"]]:
+ resolved = resolve_quant_config_from_disk(
+ active,
+ {
+ "quant_method": method,
+ flag: True,
+ "ignored_layers": ignored,
+ "w4a8_fallback_steps": [9],
+ "w4a8_fallback_layers": ["blocks.9.ffn.net_2"],
+ },
+ )
+ assert getattr(resolved, flag)
+ assert resolved.ignored_layers == ignored
+ assert resolved.w4a8_fallback_steps == steps
+ assert resolved.w4a8_fallback_layers == layers
+ assert active.ignored_layers == []
+
+
+@pytest.mark.parametrize(
+ "runtime_policy",
+ [{}, {"w4a8_fallback_steps": [1]}, {"w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}],
+)
+def test_mxfp4_explicit_omission_disables_saved_policies_for_both_experts(runtime_policy):
+ from vllm_omni.quantization import build_quant_config
+ from vllm_omni.quantization.factory import resolve_quant_config_from_disk
+
+ active = build_quant_config({"method": "mxfp4", **runtime_policy})
+ resolved_experts = []
+ for ignored in (["blocks.0.attn2.to_q"], ["blocks.1.ffn.net_2"]):
+ disk = {
+ "quant_method": "mxfp4",
+ "is_checkpoint_mxfp4_serialized": True,
+ "ignored_layers": ignored,
+ "w4a8_fallback_steps": [9],
+ "w4a8_fallback_layers": ["blocks.9.ffn.net_2"],
+ }
+ resolved = resolve_quant_config_from_disk(active, disk)
+ assert resolved.w4a8_fallback_steps == runtime_policy.get("w4a8_fallback_steps", [])
+ assert resolved.w4a8_fallback_layers == runtime_policy.get("w4a8_fallback_layers", [])
+ assert resolved.ignored_layers == ignored
+ detected = resolve_quant_config_from_disk(None, disk)
+ assert detected.w4a8_fallback_steps == [9]
+ assert detected.w4a8_fallback_layers == ["blocks.9.ffn.net_2"]
+ resolved_experts.append(resolved)
+ assert resolved_experts[0] is not resolved_experts[1]
+ assert active.is_checkpoint_mxfp4_serialized is False
+ assert active.ignored_layers == []
+
+
+@pytest.mark.parametrize(
+ "method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized"), ("mxfp4_dualscale", "is_checkpoint_serialized")]
+)
+def test_serialized_mxfp4_disk_omission_clears_active_ignored_layers(method, flag):
+ from vllm_omni.quantization import build_quant_config
+ from vllm_omni.quantization.factory import resolve_quant_config_from_disk
+
+ active_spec = {
+ "method": method,
+ flag: True,
+ "ignored_layers": ["blocks.0.attn2.to_q"],
+ "w4a8_fallback_steps": [1] if method == "mxfp4" else [],
+ "w4a8_fallback_layers": ["blocks.1.attn1.to_qkv"] if method == "mxfp4" else [],
+ }
+ if method == "mxfp4":
+ active_spec["require_smooth_scale"] = True
+ active = build_quant_config(active_spec)
+
+ resolved = resolve_quant_config_from_disk(active, {"quant_method": method, flag: True})
+
+ assert resolved is not active
+ assert resolved.ignored_layers == []
+ assert resolved.w4a8_fallback_steps == active.w4a8_fallback_steps
+ assert resolved.w4a8_fallback_layers == active.w4a8_fallback_layers
+ if method == "mxfp4":
+ assert resolved.require_smooth_scale
+
+
+@pytest.mark.parametrize(
+ "method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized"), ("mxfp4_dualscale", "is_checkpoint_serialized")]
+)
+@pytest.mark.parametrize("with_active", [False, True])
+def test_serialized_mxfp4_disk_uses_modules_to_not_convert(method, flag, with_active):
+ from vllm_omni.quantization import build_quant_config
+ from vllm_omni.quantization.factory import resolve_quant_config_from_disk
+
+ active = None
+ disk_steps = [9] if method == "mxfp4" else []
+ active_steps = [1] if method == "mxfp4" else []
+ expected_steps = disk_steps
+ if with_active:
+ active = build_quant_config(
+ {
+ "method": method,
+ flag: True,
+ "ignored_layers": ["blocks.0.attn2.to_q"],
+ "w4a8_fallback_steps": active_steps,
+ }
+ )
+ expected_steps = active_steps
+
+ resolved = resolve_quant_config_from_disk(
+ active,
+ {
+ "quant_method": method,
+ flag: True,
+ "modules_to_not_convert": ["blocks.2.ffn.net_2"],
+ "w4a8_fallback_steps": disk_steps,
+ },
+ )
+
+ assert resolved.ignored_layers == ["blocks.2.ffn.net_2"]
+ assert resolved.w4a8_fallback_steps == expected_steps
+
+
# ---------------------------------------------------------------------------
# build_quant_config integration
# ---------------------------------------------------------------------------
@@ -594,18 +912,31 @@ def test_rocm_online_dispatch_returns_rocm_method(_rocm_platform, monkeypatch):
assert isinstance(method, ROCmMxfp4OnlineLinearMethod)
-def test_rocm_ignored_layer_returns_unquantized(_rocm_platform, monkeypatch):
+@pytest.mark.parametrize("steps", [[], [0, 2]])
+def test_rocm_ignored_layer_returns_unquantized(_rocm_platform, monkeypatch, steps):
"""A prefix in ignored_layers must return UnquantizedLinearMethod before the gfx950 probe."""
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
_patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-")
- cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False, ignored_layers=["proj_out"])
+ cfg = DiffusionMXFP4Config(
+ is_checkpoint_mxfp4_serialized=False, ignored_layers=["proj_out"], w4a8_fallback_steps=steps
+ )
assert isinstance(cfg.get_quant_method(_fake_linear_layer(), "proj_out"), UnquantizedLinearMethod)
+def test_rocm_rejects_w4a8_step_fallback(_rocm_platform, monkeypatch):
+ from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+ _patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-")
+ cfg = DiffusionMXFP4Config(w4a8_fallback_steps=[0, 2])
+
+ with pytest.raises(NotImplementedError, match="only supported on NPU"):
+ cfg.get_quant_method(_fake_linear_layer(), "blocks.0.attn1.to_q")
+
+
def test_rocm_non_gfx950_raises(_rocm_platform, monkeypatch):
"""MXFP4 on ROCm requires gfx950; any other arch must raise."""
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
@@ -652,3 +983,221 @@ def test_rocm_create_weights_row_parallel_tp2(_rocm_platform):
assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2)
assert layer.input_size_per_partition == _TP2_K // _TP2
assert layer.output_size_per_partition == _TP2_N
+
+
+@pytest.mark.parametrize("method", ["mxfp4", "mxfp4_dualscale"])
+@pytest.mark.parametrize("layers", ["blocks.1.attn1.to_qkv", [1], [True], [""], [" a"], ["blocks.*"], ["a..b"]])
+def test_mxfp4_rejects_invalid_fallback_layers(method, layers):
+ from vllm_omni.quantization import build_quant_config
+
+ with pytest.raises(ValueError, match="exact runtime Linear paths"):
+ build_quant_config({"method": method, "w4a8_fallback_layers": layers})
+
+
+@pytest.mark.parametrize("method", ["mxfp4"])
+def test_mxfp4_layer_policy_from_checkpoint_and_component_config(method):
+ from vllm_omni.quantization import build_quant_config
+ from vllm_omni.quantization.component_config import resolve_component_quant_config
+ from vllm_omni.quantization.factory import resolve_quant_config_from_disk
+
+ names = ["blocks.10.attn1.to_qkv", "blocks.10.attn1.to_qkv"]
+ components = build_quant_config(
+ {"transformer": {"method": method, "w4a8_fallback_layers": names}, "transformer_2": {"method": method}}
+ )
+ names.append("blocks.11.attn1.to_qkv")
+ assert resolve_component_quant_config(components, "transformer").w4a8_fallback_layers == [names[0]]
+ assert resolve_component_quant_config(components, "transformer_2").w4a8_fallback_layers == []
+ detected = resolve_quant_config_from_disk(None, {"quant_method": method, "w4a8_fallback_layers": [names[0]]})
+ assert detected.w4a8_fallback_layers == [names[0]]
+ assert type(detected).from_config({"w4a8_fallback_layers": [names[0]]}).w4a8_fallback_layers == [names[0]]
+
+
+@pytest.mark.parametrize(
+ "kind, serialized_flag",
+ [("mxfp4", "is_checkpoint_mxfp4_serialized")],
+)
+@pytest.mark.parametrize("offline", [False, True])
+@pytest.mark.parametrize("scale_alg", [0, 2])
+def test_mxfp4_layer_and_step_union_routes_actual_methods(
+ kind, serialized_flag, offline, scale_alg, monkeypatch, mocker
+):
+ from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+
+ from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context
+ from vllm_omni.quantization import build_quant_config, mxfp4_config
+
+ monkeypatch.setattr(mxfp4_config.current_omni_platform, "is_npu", lambda: True)
+ selected = "blocks.10.attn1.to_qkv"
+ other = "blocks.11.attn2.to_q"
+ preserved = "blocks.12.ffn.net_2"
+ config = build_quant_config(
+ {
+ "method": kind,
+ "mxfp4_scale_alg": scale_alg,
+ serialized_flag: offline,
+ "w4a8_fallback_layers": [selected, preserved],
+ "ignored_layers": [preserved],
+ "w4a8_fallback_steps": [1],
+ }
+ )
+ # Real config dispatch creates separate methods for selected and other layers.
+ linear = mocker.Mock(spec=LinearBase)
+ assert isinstance(config.get_quant_method(linear, preserved), UnquantizedLinearMethod)
+ x = torch.ones(2, 64, dtype=torch.bfloat16)
+ layer = torch.nn.Module()
+ for name in ("weight", "weight_scale", "mul_scale"):
+ layer.register_buffer(name, torch.ones(64, dtype=x.dtype))
+ calls = []
+
+ def record_a8(*args):
+ calls.append("A8")
+ return x
+
+ def record_a4(*args):
+ calls.append("A4")
+ return x
+
+ monkeypatch.setattr(mxfp4_config, "_npu_w4a8_matmul", record_a8)
+ for name in (selected, other):
+ method = config.get_quant_method(linear, name)
+ assert method.is_w4a8_fallback_layer == (name == selected)
+ scales = (x, x)
+ monkeypatch.setattr(method, "_quantize_activation", lambda *args: scales)
+ monkeypatch.setattr(method, "_quant_matmul", record_a4)
+ for step in (0, 1, 2):
+ with override_forward_context(ForwardContext(denoise_step_idx=step)):
+ method.apply(layer, x)
+ assert calls == ["A8", "A8", "A8", "A4", "A8", "A4"]
+ # A permanently selected layer works without a request context, even if steps are also configured.
+ with override_forward_context(None):
+ config.get_quant_method(linear, selected).apply(layer, x)
+ assert calls[-1] == "A8"
+
+
+@pytest.mark.parametrize("online", [False, True])
+@pytest.mark.parametrize("selected", [False, True])
+@pytest.mark.parametrize("scale_alg", [0, 2])
+def test_single_scale_prepares_one_weight_for_both_activation_precisions(online, selected, scale_alg, monkeypatch):
+ import sys
+ from types import SimpleNamespace
+
+ from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context
+ from vllm_omni.quantization import mxfp4_config
+
+ monkeypatch.setattr(torch.Tensor, "npu", lambda self: self, raising=False)
+ prepared = []
+ packed = []
+
+ def quantize(x, **kwargs):
+ prepared.append(kwargs)
+ return x.to(torch.int8), torch.full((x.shape[0], 2), 127, dtype=torch.uint8)
+
+ def pack(x, dtype):
+ packed.append(x)
+ return x.to(dtype)
+
+ fake_npu = SimpleNamespace(
+ float4_e2m1fn_x2=torch.int8,
+ float8_e8m0fnu=torch.uint8,
+ float8_e4m3fn=torch.float8_e4m3fn,
+ npu_dtype_cast=pack,
+ npu_dynamic_mx_quant=quantize,
+ )
+ monkeypatch.setitem(sys.modules, "torch_npu", fake_npu)
+ config = mxfp4_config.DiffusionMXFP4Config(
+ is_checkpoint_mxfp4_serialized=not online,
+ w4a8_fallback_layers=["blocks.0.attn1.to_qkv"] if selected else [],
+ w4a8_fallback_steps=[1],
+ mxfp4_scale_alg=scale_alg,
+ )
+ cls = mxfp4_config.NPUMxfp4OnlineLinearMethod if online else mxfp4_config.NPUMxfp4LinearMethod
+ method = cls(config, prefix="blocks.0.attn1.to_qkv")
+ layer = torch.nn.Linear(64, 2, bias=False, dtype=torch.bfloat16)
+ layer.orig_dtype = torch.bfloat16
+ if not online:
+ layer.register_parameter("weight_scale", torch.nn.Parameter(torch.full((2, 2), 127, dtype=torch.uint8), False))
+ layer.register_parameter("mul_scale", torch.nn.Parameter(torch.linspace(0.5, 2, 64), False))
+ method.process_weights_after_loading(layer)
+ pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr())
+ method.process_weights_after_loading(layer)
+ assert len(prepared) == int(online)
+ expected = dict(dst_type=torch.int8, axis=-1, block_size=32, round_mode="rint", scale_alg=scale_alg)
+ if scale_alg == 2:
+ expected["dst_type_max"] = 7.25
+ assert prepared == ([expected] if online else [])
+ assert len(packed) == int(not online)
+ branches = []
+
+ def a4(x_q, x_scale, target, bias, dtype):
+ assert (target.weight.data_ptr(), target.weight_scale.data_ptr()) == pointers
+ branches.append("A4")
+ return torch.zeros(x_q.shape[0], 2, dtype=dtype)
+
+ def a8(x, weight, scale, bias, dtype):
+ assert (weight.data_ptr(), scale.data_ptr()) == pointers
+ branches.append("A8")
+ return torch.zeros(x.shape[0], 2, dtype=dtype)
+
+ monkeypatch.setattr(method, "_quant_matmul", a4)
+ monkeypatch.setattr(mxfp4_config, "_npu_w4a8_matmul", a8)
+ monkeypatch.setattr(method, "_quantize_activation", lambda x: (x, x))
+ monkeypatch.setattr(fake_npu, "npu_dynamic_mx_quant", lambda *a, **kw: pytest.fail("weight requantization"))
+ monkeypatch.setattr(fake_npu, "npu_dtype_cast", lambda *a, **kw: pytest.fail("weight repacking"))
+ for step in (0, 1, 2):
+ with override_forward_context(ForwardContext(denoise_step_idx=step)):
+ method.apply(layer, torch.ones(2, 64, dtype=torch.bfloat16))
+ assert branches == (["A8"] * 3 if selected else ["A4", "A8", "A4"])
+ assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == pointers
+ assert set(layer.state_dict()) == (
+ {"weight", "weight_scale"} if online else {"weight", "weight_scale", "mul_scale"}
+ )
+ assert not hasattr(layer, "w4a8_weight")
+ assert not hasattr(layer, "w4a8_weight_scale")
+
+
+@pytest.mark.parametrize("online", [False, True])
+def test_dualscale_without_fallback_keeps_original_a4_path(online, monkeypatch):
+ import sys
+ from types import SimpleNamespace
+
+ from vllm_omni.quantization import mxfp4_config
+
+ monkeypatch.setattr(torch.Tensor, "npu", lambda self: self, raising=False)
+ original_ones = torch.ones
+ monkeypatch.setattr(torch, "ones", lambda *args, **kw: original_ones(*args, **{**kw, "device": "cpu"}))
+ calls = []
+
+ def quantize_dual(x, **kw):
+ calls.append(kw.get("smooth_scale"))
+ return x.to(torch.int8), torch.ones(x.shape[0], 1), torch.ones(x.shape[0], 16, dtype=torch.uint8)
+
+ fake_npu = SimpleNamespace(
+ float4_e2m1fn_x2=torch.int8,
+ npu_dtype_cast=lambda x, dtype: x.to(dtype),
+ npu_format_cast=lambda x, *args, **kw: x,
+ npu_dynamic_dual_level_mx_quant=quantize_dual,
+ )
+ monkeypatch.setitem(sys.modules, "torch_npu", fake_npu)
+ config = mxfp4_config.DiffusionMXFP4DualScaleMixedConfig(
+ is_checkpoint_serialized=not online, w4a8_fallback_steps=[], w4a8_fallback_layers=[]
+ )
+ cls = mxfp4_config.NPUMxfp4DualScaleOnlineLinearMethod if online else mxfp4_config.NPUMxfp4DualScaleLinearMethod
+ method = cls(config)
+ layer = torch.nn.Module()
+ layer.input_size_per_partition = 512
+ for name, value in {
+ "weight": torch.ones(2, 512, dtype=torch.bfloat16),
+ "weight_scale": torch.full((2, 16), 127, dtype=torch.uint8),
+ "weight_dual_scale": torch.ones(2, 1, 1),
+ "mul_scale": torch.ones(512),
+ }.items():
+ layer.register_parameter(name, torch.nn.Parameter(value, requires_grad=False))
+ method.process_weights_after_loading(layer)
+ method.process_weights_after_loading(layer)
+ assert len(calls) == int(online)
+ assert set(layer.state_dict()) == {"weight", "weight_scale", "weight_dual_scale", "mul_scale"}
+ assert not hasattr(layer, "w4a8_weight")
+ monkeypatch.setattr(method, "_quant_matmul", lambda *args: torch.zeros(2, 2, dtype=torch.bfloat16))
+ method.apply(layer, torch.ones(2, 512, dtype=torch.bfloat16))
+ assert len(calls) == 1 + int(online)
+ torch.testing.assert_close(calls[-1], layer.mul_scale)
diff --git a/tests/diffusion/quantization/test_wan_mxfp4_loading.py b/tests/diffusion/quantization/test_wan_mxfp4_loading.py
new file mode 100644
index 000000000..fc7a43d43
--- /dev/null
+++ b/tests/diffusion/quantization/test_wan_mxfp4_loading.py
@@ -0,0 +1,237 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Exercise Wan's real name mapping and vLLM's QKV parameter loaders on CPU."""
+
+import pytest
+import torch
+from vllm.config.load import LoadConfig
+from vllm.model_executor.layers.linear import QKVParallelLinear
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
+from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, DiffusionMXFP4DualScaleMixedConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.fixture(autouse=True)
+def _patch_tp_state(monkeypatch):
+ monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0)
+ monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1)
+ monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_npu", lambda: True)
+ prefix = "vllm_omni.diffusion.models.wan2_2.wan2_2_transformer"
+ monkeypatch.setattr(f"{prefix}.get_tensor_model_parallel_rank", lambda: 0)
+ monkeypatch.setattr(f"{prefix}.get_tensor_model_parallel_world_size", lambda: 1)
+
+
+class _WanCheckpoint(torch.nn.Module):
+ load_weights = WanTransformer3DModel.load_weights
+ _log_w4a8_fallback_load_summary = WanTransformer3DModel._log_w4a8_fallback_load_summary
+
+ def __init__(self, config, *, num_blocks=1):
+ super().__init__()
+ blocks = []
+ for block_idx in range(num_blocks):
+ block = torch.nn.Module()
+ block.attn1 = torch.nn.Module()
+ block.attn1.to_qkv = QKVParallelLinear(
+ hidden_size=512,
+ head_size=32,
+ total_num_heads=4,
+ bias=False,
+ params_dtype=torch.bfloat16,
+ quant_config=config,
+ prefix=f"blocks.{block_idx}.attn1.to_qkv",
+ disable_tp=True,
+ )
+ blocks.append(block)
+ self.blocks = torch.nn.ModuleList(blocks)
+
+
+def _checkpoint(mode, fused):
+ if mode == "single":
+ config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True)
+ elif mode == "dual":
+ config = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=True)
+ else:
+ config = None
+ model = _WanCheckpoint(config)
+ weights = {}
+ expected = {}
+ for name, param in model.named_parameters():
+ suffix = name.rsplit(".", 1)[1]
+ output_dim = getattr(param, "output_dim", None)
+ if output_dim is None:
+ value = torch.full_like(param, 2.5)
+ shards = [value] * 3
+ else:
+ parts = param.chunk(3, dim=output_dim)
+ shards = [
+ torch.full_like(part, (121 if suffix == "weight_scale" else 1) + i) for i, part in enumerate(parts)
+ ]
+ value = torch.cat(shards, dim=output_dim)
+ expected[name] = value
+ if fused:
+ weights[name] = value
+ else:
+ for shard_id, shard in zip(("q", "k", "v"), shards):
+ weights[f"blocks.0.attn1.to_{shard_id}.{suffix}"] = shard
+ return model, config, weights, expected
+
+
+def _load(model, config, weights, monkeypatch):
+ loader = DiffusersPipelineLoader(
+ LoadConfig(), OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=config)
+ )
+ loader.counter_before_loading_weights = 0.0
+ monkeypatch.setattr(loader, "get_all_weights", lambda model: iter(weights.items()))
+ loader.load_weights(model)
+
+
+@pytest.mark.parametrize("mode", ["single", "dual", "bf16"])
+@pytest.mark.parametrize("fused", [False, True])
+def test_wan_qkv_loads_actual_weights_and_scales(mode, fused, monkeypatch):
+ model, config, weights, expected = _checkpoint(mode, fused)
+ _load(model, config, weights, monkeypatch)
+ for name, param in model.named_parameters():
+ torch.testing.assert_close(param, expected[name], rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("mode", ["single", "dual", "bf16"])
+def test_wan_incomplete_qkv_weight_is_not_marked_fully_loaded(mode, monkeypatch):
+ model, config, weights, _ = _checkpoint(mode, fused=False)
+ weights.pop("blocks.0.attn1.to_k.weight")
+ weights.pop("blocks.0.attn1.to_v.weight")
+ with pytest.raises(ValueError, match="to_qkv.weight"):
+ _load(model, config, weights, monkeypatch)
+
+
+@pytest.mark.parametrize("mode", ["single", "dual"])
+def test_wan_incomplete_qkv_scale_is_not_marked_fully_loaded(mode, monkeypatch):
+ model, config, weights, _ = _checkpoint(mode, fused=False)
+ weights.pop("blocks.0.attn1.to_k.weight_scale")
+ weights.pop("blocks.0.attn1.to_v.weight_scale")
+ with pytest.raises(ValueError, match="to_qkv.weight_scale"):
+ _load(model, config, weights, monkeypatch)
+
+
+@pytest.mark.parametrize("mode", ["single", "dual"])
+@pytest.mark.parametrize("reverse", [False, True])
+def test_wan_split_qkv_rejects_different_smooth_scales(mode, reverse, monkeypatch):
+ model, config, weights, _ = _checkpoint(mode, fused=False)
+ weights["blocks.0.attn1.to_k.mul_scale"] = weights["blocks.0.attn1.to_k.mul_scale"] * 2
+ if reverse:
+ weights = dict(reversed(weights.items()))
+ with pytest.raises(ValueError, match="must share the same Smooth tensor"):
+ _load(model, config, weights, monkeypatch)
+
+
+@pytest.mark.parametrize("name", ["to_q.missing", "to_qkv.missing", "to_qkv_extra.weight", "to_query.weight"])
+def test_wan_unknown_qkv_key_is_not_reported_as_loaded(name):
+ model, _, _, _ = _checkpoint("single", fused=True)
+ loaded = model.load_weights([(f"blocks.0.attn1.{name}", torch.ones(2))])
+ assert loaded == set()
+
+
+@pytest.mark.parametrize("fused", [False, True])
+def test_wan_single_scale_rejects_relabelled_dualscale_checkpoint(fused, monkeypatch):
+ model, config, weights, _ = _checkpoint("single", fused=fused)
+ projection = "to_qkv" if fused else "to_q"
+ weights[f"blocks.0.attn1.{projection}.weight_dual_scale"] = torch.ones(128, 1, 1)
+ with pytest.raises(ValueError, match="cannot load DualScale tensor"):
+ _load(model, config, weights, monkeypatch)
+
+
+def test_wan_layer_policy_validates_runtime_qkv_and_bf16_override():
+ config = DiffusionMXFP4Config(w4a8_fallback_layers=["blocks.0.attn1.to_qkv"])
+ model = _WanCheckpoint(config)
+ # Uses the same validation called at the end of WanTransformer3DModel construction.
+ WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
+ for bad in ("blocks.0.attn1.to_q", "blocks.0.attn1", "blocks.99.attn1.to_qkv", "transformer.blocks.0.attn1.to_qkv"):
+ config.w4a8_fallback_layers = [bad]
+ with pytest.raises(ValueError, match="exact runtime Linear paths"):
+ WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
+
+
+def test_wan_layer_policy_defers_only_other_pp_stage_paths():
+ from vllm.model_executor.models.utils import PPMissingLayer
+
+ config = DiffusionMXFP4Config(w4a8_fallback_layers=["blocks.0.attn1.to_qkv", "blocks.1.attn1.to_qkv"])
+ model = _WanCheckpoint(config)
+ model.blocks.append(PPMissingLayer())
+ WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
+ config.w4a8_fallback_layers = ["blocks.1"]
+ with pytest.raises(ValueError, match="exact runtime Linear paths"):
+ WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
+
+
+@pytest.mark.parametrize("is_checkpoint_serialized", [False, True], ids=["online", "offline"])
+def test_loader_logs_wan_expert_summaries_after_weight_processing(is_checkpoint_serialized, mocker):
+ selected = "blocks.0.attn1.to_qkv"
+ bf16 = "blocks.1.attn1.to_qkv"
+ transformer_config = DiffusionMXFP4Config(
+ is_checkpoint_mxfp4_serialized=is_checkpoint_serialized,
+ w4a8_fallback_layers=[selected],
+ )
+ transformer_2_config = DiffusionMXFP4Config(
+ is_checkpoint_mxfp4_serialized=is_checkpoint_serialized,
+ ignored_layers=[bf16],
+ w4a8_fallback_layers=[selected, bf16],
+ )
+ pipeline = torch.nn.Module()
+ pipeline.transformer = _WanCheckpoint(transformer_config)
+ pipeline.transformer_2 = _WanCheckpoint(transformer_2_config, num_blocks=2)
+ WanTransformer3DModel._validate_w4a8_fallback_layers(pipeline.transformer, transformer_config)
+ WanTransformer3DModel._validate_w4a8_fallback_layers(pipeline.transformer_2, transformer_2_config)
+
+ loader = DiffusersPipelineLoader(
+ LoadConfig(),
+ OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=transformer_config),
+ )
+ events = []
+ summaries = []
+
+ def process_weights(model, _target_device):
+ events.append("process")
+ for module in model.modules():
+ method = getattr(module, "quant_method", None)
+ if getattr(method, "is_w4a8_fallback_layer", False):
+ module._already_called_process_weights_after_loading = True
+
+ def capture_summary(message, *args):
+ rendered = message % args
+ summaries.append(rendered)
+ events.append(f"summary:{args[0]}")
+
+ loader._init_from_load_format = lambda *_args, **_kwargs: pipeline # type: ignore[method-assign]
+ loader.load_weights = lambda _model: events.append("load") # type: ignore[method-assign]
+ loader._process_weights_after_loading = process_weights # type: ignore[method-assign]
+ loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
+ mocker.patch(
+ "vllm_omni.diffusion.models.wan2_2.wan2_2_transformer.logger.info",
+ side_effect=capture_summary,
+ )
+
+ assert loader.load_model(load_device="cpu") is pipeline
+ assert events == ["load", "process", "summary:transformer", "summary:transformer_2"]
+ assert "component=transformer; selected=1 ['transformer.blocks.0.attn1.to_qkv']" in summaries[0]
+ assert "BF16 overrides=0 []" in summaries[0]
+ assert "weight/A8 processing state=ready (1/1 selected layers processed)" in summaries[0]
+ assert "component=transformer_2; selected=1 ['transformer_2.blocks.0.attn1.to_qkv']" in summaries[1]
+ assert "BF16 overrides=1 ['transformer_2.blocks.1.attn1.to_qkv']" in summaries[1]
+ assert "weight/A8 processing state=ready (1/1 selected layers processed)" in summaries[1]
+
+
+def test_wan_load_summary_does_not_claim_unprocessed_layer_is_ready(mocker):
+ selected = "blocks.0.attn1.to_qkv"
+ config = DiffusionMXFP4Config(w4a8_fallback_layers=[selected])
+ model = _WanCheckpoint(config)
+ WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
+ logger_info = mocker.patch("vllm_omni.diffusion.models.wan2_2.wan2_2_transformer.logger.info")
+
+ WanTransformer3DModel._log_w4a8_fallback_load_summary(model, "transformer")
+
+ message, *args = logger_info.call_args.args
+ rendered = message % tuple(args)
+ assert "weight/A8 processing state=not-ready (0/1 selected layers processed)" in rendered
diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py
index 23ff796e2..6f37f823d 100644
--- a/tests/diffusion/test_diffusion_config_propagation.py
+++ b/tests/diffusion/test_diffusion_config_propagation.py
@@ -25,6 +25,12 @@ from vllm_omni.diffusion.model_metadata import (
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+@pytest.fixture(autouse=True)
+def _local_model_paths(monkeypatch):
+ # These tests exercise config transport, not model repository resolution.
+ monkeypatch.setattr("vllm_omni.diffusion.data.get_model_path", lambda model, revision: model)
+
+
def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig:
"""Simulate the real path: create_default_diffusion → OmniDiffusionConfig.
diff --git a/tests/e2e/features/comfyui/test_comfyui_integration.py b/tests/e2e/features/comfyui/test_comfyui_integration.py
index 3c6f0492a..480381f0a 100644
--- a/tests/e2e/features/comfyui/test_comfyui_integration.py
+++ b/tests/e2e/features/comfyui/test_comfyui_integration.py
@@ -534,7 +534,7 @@ def mock_async_omni(
# 1. The API layer uses its stage_list and stage_configs attributes
# 2. Its __init__ method has slow side effects (model & config loading).
# These cases use the non-duplex mock engine without loading model configs.
- mocker.patch("vllm_omni.entrypoints.openai.api_server._is_duplex_model", return_value=False)
+ mocker.patch("vllm_omni.entrypoints.openai.api_server._should_serve_duplex", return_value=False)
mock_async_omni_cls = mocker.patch("vllm_omni.entrypoints.openai.api_server.AsyncOmni")
monkeypatch.setattr(
"vllm_omni.entrypoints.openai.serving_chat.OmniOpenAIServingChat._preprocess_chat",
diff --git a/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py b/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py
index 56a960121..4bea6563b 100644
--- a/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py
+++ b/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py
@@ -21,6 +21,7 @@ import torch.distributed as dist
from transformers import AutoTokenizer
from tests.helpers.mark import hardware_test
+from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig
from vllm_omni.diffusion.executor.uniproc_executor import UniProcDiffusionExecutor
from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
from vllm_omni.entrypoints.async_omni import AsyncOmni
@@ -136,6 +137,15 @@ async def _generate_once(
def _assert_live_uniproc_worker(engine: AsyncOmni) -> None:
"""Default single-GPU path must use the in-process executor, not mp IPC."""
+ stage_configs = engine.engine.stage_configs
+ assert len(stage_configs) == 1
+ stage_config = stage_configs[0]
+ assert isinstance(stage_config, VllmOmniDiffusionStageConfig)
+ assert stage_config.model_config.model == MODEL
+ assert stage_config.model_config.enforce_eager is True
+ assert stage_config.scheduler_config.max_num_seqs == 1
+ assert stage_config.diffusion_config.custom_pipeline_args == {"pipeline_class": CUSTOM_PIPELINE_CLASS}
+ assert stage_config.diffusion_config.worker_extension_cls == WORKER_EXTENSION_CLASS
clients = engine.engine.stage_clients
assert len(clients) == 1
client = clients[0]
diff --git a/tests/e2e/offline_inference/test_hunyuanimage3.py b/tests/e2e/offline_inference/test_hunyuanimage3.py
index 8655bd63f..21bfe9b39 100644
--- a/tests/e2e/offline_inference/test_hunyuanimage3.py
+++ b/tests/e2e/offline_inference/test_hunyuanimage3.py
@@ -1,5 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
# ruff: noqa: E501
+import copy
import importlib.util
+import os
import sys
from collections.abc import Generator
from pathlib import Path
@@ -10,20 +15,37 @@ import torch.nn.functional as F
from PIL import Image
from transformers import CLIPModel, CLIPProcessor
+from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
-from tests.helpers.stage_config import get_deploy_config_path
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
from vllm_omni import Omni
+from vllm_omni.config.omni_config import (
+ VllmOmniARStageConfig,
+ VllmOmniDiffusionStageConfig,
+)
+from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens, resolve_stop_token_ids
from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.model_extras import get_ar_input_builder, get_ar_tokenizer_validator, get_model_class_name
from vllm_omni.platforms import current_omni_platform
PROMPT = "A brown and white dog is running on the grass"
-MODEL_NAME = "tencent/HunyuanImage-3.0"
+MODEL_NAME = os.environ.get("HUNYUAN_IMAGE3_MODEL", "tencent/HunyuanImage-3.0-Instruct")
LOCAL_CLIP_PATH = "openai/clip-vit-base-patch32"
-DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image_3_moe.yaml")
+_PRODUCTION_DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image_3_moe.yaml")
+# Keep the production Mooncake connector, but use its TCP transport for this
+# single-host test. The shared server's RDMA device is not usable reliably;
+# TCP still exercises the Mooncake AR→DiT data plane without changing deploy
+# semantics in production.
+DEPLOY_CONFIG_PATH = modify_stage_config(
+ _PRODUCTION_DEPLOY_CONFIG_PATH,
+ updates={"connectors.rdma_connector.extra.protocol": "tcp"},
+)
-pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion]
+pytestmark = [
+ pytest.mark.advanced_model,
+ pytest.mark.diffusion,
+]
# System prompt type. Options: None, dynamic, en_vanilla, en_recaption, en_think_recaption, en_unified
# Below are the CLIP embedding tensors from the official HunyuanImage model (seed=1234, prompt: "A brown and white dog is running on the grass").
@@ -285,25 +307,97 @@ def omni() -> Generator[Omni, None, None]:
yield runner.omni
+@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Need at least 4 CUDA GPUs for this test.")
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+def test_structured_mixed_pipeline_config_reaches_runtime(omni: Omni) -> None:
+ """Mixed-pipeline deploy settings reach both terminal engine configs."""
+ engine = omni.engine
+
+ # Topology, connector selection, and placement remain Omni-owned.
+ resolved_stages = engine.stage_configs
+ assert len(resolved_stages) == 2
+ ar_stage, diffusion_stage = resolved_stages
+ assert isinstance(ar_stage, VllmOmniARStageConfig)
+ assert ar_stage.stage_id == 0
+ assert ar_stage.model_stage == "AR"
+ assert ar_stage.final_output_type == "text"
+ assert ar_stage.runtime_config.devices == "0,1"
+ assert ar_stage.connector_config.output_connectors == {"to_stage_1": "rdma_connector"}
+ assert isinstance(diffusion_stage, VllmOmniDiffusionStageConfig)
+ assert diffusion_stage.stage_id == 1
+ assert diffusion_stage.model_stage == "dit"
+ assert diffusion_stage.final_output_type == "image"
+ assert diffusion_stage.input_sources == [0]
+ assert diffusion_stage.runtime_config.devices == "2,3"
+ assert diffusion_stage.connector_config.input_connectors == {"from_stage_0": "rdma_connector"}
+
+ # Stage 0 is materialized as the VllmConfig actually passed to its engine.
+ ar_vllm_config, diffusion_vllm_config = engine.stage_vllm_configs
+ assert ar_vllm_config is not None
+ assert diffusion_vllm_config is None
+ assert ar_vllm_config.model_config.enforce_eager is True
+ assert ar_vllm_config.model_config.trust_remote_code is True
+ assert ar_vllm_config.parallel_config.tensor_parallel_size == 2
+ assert ar_vllm_config.scheduler_config.max_num_seqs == 1
+ assert ar_vllm_config.scheduler_config.max_num_batched_tokens == 32768
+ assert ar_vllm_config.cache_config.gpu_memory_utilization == 0.95
+ ar_omni_config = ar_vllm_config.model_config.omni_config
+ assert ar_omni_config.omni_kv_config["need_send_cache"] is True
+
+ expected_ar_sampling = {
+ "temperature": 0.0,
+ "top_p": 1,
+ "top_k": -1,
+ "max_tokens": 8192,
+ "detokenize": True,
+ "skip_special_tokens": False,
+ "include_stop_str_in_output": True,
+ }
+ ar_sampling, diffusion_sampling = engine.default_sampling_params_list
+ assert all(getattr(ar_sampling, name) == value for name, value in expected_ar_sampling.items())
+ assert diffusion_sampling.num_inference_steps == 50
+ assert diffusion_sampling.guidance_scale == 0
+
+ # Stage 1 is a diffusion engine, so its terminal owner is the launched
+ # OmniDiffusionConfig rather than a VllmConfig.
+ od_config = omni.get_diffusion_od_config()
+ assert od_config is not None
+ assert od_config.model == MODEL_NAME
+ assert od_config.enforce_eager is True
+ assert od_config.distributed_executor_backend == "mp"
+ assert od_config.parallel_config.tensor_parallel_size == 2
+ assert od_config.parallel_config.enable_expert_parallel is True
+ assert od_config.max_num_seqs == 1
+ diffusion_kv = od_config.omni_kv_config
+ assert diffusion_kv["need_recv_cache"] is True
+ assert diffusion_kv["enable_kv_async_prefetch"] is True
+ assert diffusion_kv["kv_prefetch_min_free_mem_ratio"] == 0.2
+
+
def _extract_generated_image(outputs: list[object]) -> Image.Image:
if not outputs:
raise AssertionError("No outputs were returned from Omni.generate()")
- first_output = outputs[0]
- if images := getattr(first_output, "images", None):
- return images[0]
-
- request_output = first_output
- if request_output is not None and (images := getattr(request_output, "images", None)):
- return images[0]
+ for output in outputs:
+ if images := getattr(output, "images", None):
+ return images[0]
+ multimodal_output = getattr(output, "multimodal_output", None)
+ if isinstance(multimodal_output, dict) and (images := multimodal_output.get("images")):
+ return images[0]
raise AssertionError("No generated image found in Omni output")
def extract_embedding(image: Image.Image, clip_model: CLIPModel, clip_processor: CLIPProcessor) -> torch.Tensor:
- inputs = clip_processor(images=image.convert("RGB"), return_tensors="pt")
+ # The full CLIP forward API is stable across supported Transformers
+ # versions, but it expects text inputs as well as pixels. An empty text
+ # prompt is sufficient when only the projected image embedding is needed.
+ inputs = clip_processor(text=[""], images=[image.convert("RGB")], return_tensors="pt", padding=True)
with torch.inference_mode():
- features = clip_model.get_image_features(**inputs)
+ # Use the public CLIP forward output, matching the shared CLIPScorer
+ # helper. ``get_image_features`` changed from Tensor to
+ # BaseModelOutputWithPooling in Transformers 5.x.
+ features = clip_model(**inputs).image_embeds
features = F.normalize(features, p=2, dim=-1)
return features.squeeze(0)
@@ -320,16 +414,36 @@ def compare_semantic(
def _generate_image(omni: Omni, use_system_prompt: str | None) -> Image.Image:
- generator_device = current_omni_platform.device_type or "cuda"
- sampling_params = OmniDiffusionSamplingParams(
- seed=1234,
- generator=torch.Generator(device=generator_device).manual_seed(1234),
- num_outputs_per_prompt=1,
- )
- if use_system_prompt is not None:
- sampling_params.extra_args = {"use_system_prompt": use_system_prompt}
+ from transformers import AutoTokenizer
- outputs = omni.generate({"prompt": PROMPT}, sampling_params)
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+ prompt_tokens = build_prompt_tokens(
+ PROMPT,
+ tokenizer,
+ task="t2i",
+ bot_task="think",
+ sys_type=use_system_prompt or "None",
+ )
+ stop_token_ids = resolve_stop_token_ids(task="t2i", bot_task="think", tokenizer=tokenizer, image_size="1024x1024")
+ params_list = copy.deepcopy(omni.default_sampling_params_list)
+ generator_device = current_omni_platform.device_type or "cuda"
+ for params in params_list:
+ if isinstance(params, OmniDiffusionSamplingParams):
+ params.seed = 1234
+ params.generator = torch.Generator(device=generator_device).manual_seed(1234)
+ params.num_outputs_per_prompt = 1
+ elif hasattr(params, "stop_token_ids"):
+ params.stop_token_ids = stop_token_ids
+
+ prompt = {
+ "prompt": PROMPT,
+ "prompt_token_ids": prompt_tokens.token_ids,
+ "use_system_prompt": prompt_tokens.system_prompt_type,
+ "modalities": ["image"],
+ "height": 1024,
+ "width": 1024,
+ }
+ outputs = omni.generate(prompt, params_list)
return _extract_generated_image(outputs)
@@ -342,7 +456,7 @@ def _generate_image(omni: Omni, use_system_prompt: str | None) -> Image.Image:
MIN_SEMANTIC_SIMILARITY = 0.5
-@pytest.mark.skipif(torch.accelerator.device_count() < 8, reason="Need at least 8 CUDA GPUs for this test.")
+@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Need at least 4 CUDA GPUs for this test.")
@pytest.mark.parametrize("system_prompt_name,use_system_prompt,expected_embedding", SYSTEM_PROMPT_CASES)
def test_system_prompt_scores(
omni: Omni,
@@ -378,7 +492,7 @@ def _load_text_to_image_module():
return module
-@pytest.mark.skipif(torch.accelerator.device_count() < 8, reason="Need at least 8 CUDA GPUs for this test.")
+@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Need at least 4 CUDA GPUs for this test.")
def test_shared_script_ar_path_reaches_generation(
omni: Omni,
clip_bundle: tuple[CLIPModel, CLIPProcessor],
diff --git a/tests/e2e/offline_inference/test_qwen3_omni.py b/tests/e2e/offline_inference/test_qwen3_omni.py
index 32c345194..f5937265c 100644
--- a/tests/e2e/offline_inference/test_qwen3_omni.py
+++ b/tests/e2e/offline_inference/test_qwen3_omni.py
@@ -17,6 +17,10 @@ import pytest
from tests.helpers.mark import hardware_test
from tests.helpers.media import generate_synthetic_video
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from vllm_omni.config.omni_config import (
+ VllmOmniARStageConfig,
+ VllmOmniGenerationStageConfig,
+)
from vllm_omni.platforms import current_omni_platform
models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
@@ -25,6 +29,7 @@ thinker_only_models = ["Qwen/Qwen3-Omni-30B-A3B-Captioner"]
# Single CI deploy YAML; rocm/xpu deltas are picked automatically via the
# platforms: section. Only CUDA needs an extra enforce_eager tweak.
_CI_DEPLOY = get_deploy_config_path("ci/qwen3_omni_moe.yaml")
+_PRODUCTION_DEPLOY = get_deploy_config_path("qwen3_omni_moe.yaml")
def get_cuda_graph_config():
@@ -46,6 +51,7 @@ else:
# Create parameter combinations for model and stage config
test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
+production_test_params = [(model, _PRODUCTION_DEPLOY) for model in models]
# we can use the same config for a model that only has thinker (i.e., does not
# enable audio output) because the resolver should figure out that it doesn't
# need the full pipeline based on the HF config.
@@ -59,6 +65,85 @@ def get_question(prompt_type="video"):
return prompts.get(prompt_type, prompts["video"])
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_runner", production_test_params, indirect=True)
+def test_structured_multistage_config_reaches_runtime(omni_runner, offline_client) -> None:
+ """Deploy settings reach the materialized configs used by live stages."""
+ engine = omni_runner.omni.engine
+
+ # The resolved typed stages remain the Omni-owned source for topology and
+ # placement, which are not fields of the terminal vLLM config.
+ resolved_stages = engine.stage_configs
+ assert len(resolved_stages) == 3
+ thinker, talker, code2wav = resolved_stages
+ assert isinstance(thinker, VllmOmniARStageConfig)
+ assert isinstance(talker, VllmOmniARStageConfig)
+ assert isinstance(code2wav, VllmOmniGenerationStageConfig)
+ assert [stage.stage_id for stage in resolved_stages] == [0, 1, 2]
+ assert [stage.model_stage for stage in resolved_stages] == ["thinker", "talker", "code2wav"]
+ assert code2wav.final_output_type == "audio"
+ assert thinker.runtime_config.devices == "0"
+ assert talker.runtime_config.devices == code2wav.runtime_config.devices == "1"
+ assert talker.connector_config.input_connectors == {"from_stage_0": "connector_of_shared_memory"}
+ assert code2wav.connector_config.input_connectors == {"from_stage_1": "connector_of_shared_memory"}
+
+ # Engine-owned deploy settings must survive the startup projection into
+ # the actual VllmConfig objects retained by the stage pools.
+ assert len(engine.stage_vllm_configs) == 3
+ thinker_vllm, talker_vllm, code2wav_vllm = engine.stage_vllm_configs
+ assert thinker_vllm is not None
+ assert talker_vllm is not None
+ assert code2wav_vllm is not None
+ assert thinker_vllm.scheduler_config.max_num_seqs == 64
+ assert talker_vllm.scheduler_config.max_num_seqs == 64
+ assert code2wav_vllm.scheduler_config.max_num_seqs == 64
+ assert thinker_vllm.scheduler_config.max_num_batched_tokens == 32768
+ assert talker_vllm.scheduler_config.max_num_batched_tokens == 32768
+ assert code2wav_vllm.scheduler_config.max_num_batched_tokens == 65536
+ assert thinker_vllm.cache_config.gpu_memory_utilization == 0.9
+ assert talker_vllm.cache_config.gpu_memory_utilization == 0.6
+ assert code2wav_vllm.cache_config.gpu_memory_utilization == 0.1
+ assert all(config.cache_config.enable_prefix_caching is False for config in engine.stage_vllm_configs)
+ assert all(config.model_config.trust_remote_code is True for config in engine.stage_vllm_configs)
+ assert code2wav_vllm.scheduler_config.enable_chunked_prefill is False
+ assert code2wav_vllm.scheduler_config.async_scheduling is False
+
+ # Sampling defaults are consumed from StageClient metadata, rather than
+ # from either the resolver output or VllmConfig.
+ expected_sampling = (
+ {"temperature": 0.0, "max_tokens": 2048},
+ {"temperature": 0.9, "top_k": 50, "max_tokens": 4096, "repetition_penalty": 1.05},
+ {
+ "temperature": 0.0,
+ "top_p": 1.0,
+ "top_k": -1,
+ "max_tokens": 65536,
+ "repetition_penalty": 1.1,
+ },
+ )
+ assert len(engine.default_sampling_params_list) == len(expected_sampling)
+ for runtime_params, expected in zip(engine.default_sampling_params_list, expected_sampling, strict=True):
+ # SamplingParams normalization may add backend defaults (for example
+ # ``detokenize=True``). Verify deploy-owned values without rejecting
+ # those non-lossy normalized fields.
+ assert all(getattr(runtime_params, name) == value for name, value in expected.items())
+
+ if current_omni_platform.is_cuda():
+ assert thinker_vllm.model_config.enforce_eager is False
+ assert talker_vllm.model_config.enforce_eager is False
+
+ # Exercise the materialized stages and their connector path with a real
+ # request; the assertions above are not merely resolver-shape checks.
+ offline_client.send_omni_request(
+ {
+ "prompts": "Answer with one short sentence: what is the capital of China?",
+ "modalities": ["text"],
+ }
+ )
+
+
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
diff --git a/tests/e2e/offline_inference/test_wan22_t2v.py b/tests/e2e/offline_inference/test_wan22_t2v.py
index 805fbd1fe..b3c7d6aab 100644
--- a/tests/e2e/offline_inference/test_wan22_t2v.py
+++ b/tests/e2e/offline_inference/test_wan22_t2v.py
@@ -5,6 +5,8 @@ import pytest
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OfflineOmniClient
+from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig
+from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
MODEL = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
@@ -12,6 +14,33 @@ PROMPT = "Two anthropomorphic cats in boxing gear on a spotlighted stage."
NEGATIVE_PROMPT = "low quality, blurry, watermark, text"
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards={"cuda": 1, "rocm": 1})
+@pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True)
+def test_structured_diffusion_config_reaches_runtime(omni_runner) -> None:
+ """The default WAN resolver launches the typed diffusion config itself."""
+ omni = omni_runner.omni
+ resolved_stages = omni.engine.stage_configs
+ stage_clients = omni.engine.stage_clients
+ assert len(resolved_stages) == 1
+ assert len(stage_clients) == 1
+ stage_config = resolved_stages[0]
+ od_config = stage_clients[0].od_config
+
+ assert isinstance(stage_config, VllmOmniDiffusionStageConfig)
+ assert isinstance(od_config, OmniDiffusionConfig)
+ assert stage_config.stage_id == 0
+ assert stage_config.model_stage == "dit"
+ assert stage_config.final_output_type == "video"
+
+ assert omni.engine.stage_vllm_configs == [None]
+ assert od_config.stage_id == 0
+ assert od_config.model == MODEL
+ assert od_config.model_class_name == "WanPipeline"
+ assert od_config.output_type == "pil"
+
+
@pytest.mark.advanced_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards={"cuda": 1, "rocm": 1})
diff --git a/tests/e2e/online_serving/personaplex_realtime_duplex.py b/tests/e2e/online_serving/personaplex_realtime_duplex.py
index 20da1598d..78f012548 100644
--- a/tests/e2e/online_serving/personaplex_realtime_duplex.py
+++ b/tests/e2e/online_serving/personaplex_realtime_duplex.py
@@ -10,24 +10,24 @@ import base64
import hashlib
import json
import math
+import time
import uuid
import wave
-from collections.abc import Sequence
+from collections.abc import Awaitable, Callable, Sequence
from pathlib import Path
+from typing import Protocol
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
import numpy as np
try:
import websockets
- from websockets.exceptions import ConnectionClosed
+ from websockets.exceptions import ConnectionClosed, WebSocketException
except ImportError as exc: # pragma: no cover - driver dependency.
raise SystemExit("Install websockets first: pip install websockets") from exc
-from vllm_omni.clients.duplex import (
- EventCollector,
- write_pcm16_wav,
-)
+from vllm_omni.benchmarks.audio_continuity import compute_continuity_stats
+from vllm_omni.clients.duplex import EventCollector, write_pcm16_wav
from vllm_omni.clients.duplex import (
wait_for_condition as wait_for,
)
@@ -35,6 +35,24 @@ from vllm_omni.clients.duplex import (
SAMPLE_RATE_HZ = 24_000
FRAME_SAMPLES = 1_920
FRAME_PERIOD_S = FRAME_SAMPLES / SAMPLE_RATE_HZ
+AUDIO_DELTA_EVENT_TYPES = frozenset({"response.audio.delta", "response.output_audio.delta"})
+
+
+class _AudioAcceptanceError(AssertionError):
+ """A local acceptance failure with numeric diagnostics safe to report."""
+
+ def __init__(self, check: str, message: str, stats: dict[str, object]) -> None:
+ super().__init__(message)
+ self.check = check
+ self.stats = dict(stats)
+
+
+class _ProbeTransport(Protocol):
+ async def recv(self) -> str | bytes: ...
+
+ async def send(self, message: str) -> None: ...
+
+ async def close(self) -> None: ...
class RawRealtimeProbe:
@@ -46,29 +64,33 @@ class RawRealtimeProbe:
keeps a transport-only client and reads everything from ``.events``.
"""
- def __init__(self, url: str, *, max_size: int = 64 * 1024 * 1024) -> None:
+ def __init__(self, url: str, *, max_size: int = 64 * 1024 * 1024, close_timeout_s: float = 10.0) -> None:
self.url = url
self.max_size = max_size
+ self.close_timeout_s = close_timeout_s
self.events = EventCollector()
- self._ws = None
+ self._ws: _ProbeTransport | None = None
self._reader_task: asyncio.Task[None] | None = None
async def __aenter__(self) -> RawRealtimeProbe:
- self._ws = await websockets.connect(self.url, max_size=self.max_size)
+ self._ws = await websockets.connect(self.url, max_size=self.max_size, close_timeout=self.close_timeout_s)
self._reader_task = asyncio.create_task(self._read_events())
return self
async def __aexit__(self, exc_type, exc, traceback) -> None:
- if self._ws is not None:
- await self._ws.close()
- if self._reader_task is not None:
- self._reader_task.cancel()
- try:
- await self._reader_task
- except asyncio.CancelledError:
- pass
+ try:
+ if self._ws is not None:
+ await self._ws.close()
+ finally:
+ if self._reader_task is not None:
+ self._reader_task.cancel()
+ try:
+ await self._reader_task
+ except asyncio.CancelledError:
+ pass
async def _read_events(self) -> None:
+ assert self._ws is not None
try:
while True:
raw = await self._ws.recv()
@@ -81,6 +103,7 @@ class RawRealtimeProbe:
return
async def send(self, event: dict[str, object]) -> None:
+ assert self._ws is not None
await self._ws.send(json.dumps(event))
@@ -142,6 +165,38 @@ def _events(client: RawRealtimeProbe, event_type: str) -> list[dict[str, object]
return [event for event in client.events.events if event.get("type") == event_type]
+def _audio_events(client: RawRealtimeProbe) -> list[dict[str, object]]:
+ """Return audio deltas from both current and legacy Realtime event names."""
+ return [event for event in client.events.events if event.get("type") in AUDIO_DELTA_EVENT_TYPES]
+
+
+def _validated_audio_event(event: dict[str, object]) -> bytes:
+ """Decode one wire audio delta without depending on client event aliases."""
+ encoded = event.get("delta")
+ if not isinstance(encoded, str):
+ raise ValueError("missing_audio_delta")
+ try:
+ chunk = base64.b64decode(encoded, validate=True)
+ except ValueError:
+ raise ValueError("invalid_audio_base64") from None
+ if chunk:
+ if len(chunk) % 2:
+ raise ValueError("unaligned_pcm16_packet")
+ if event.get("sample_rate_hz") != SAMPLE_RATE_HZ:
+ raise ValueError("invalid_audio_sample_rate")
+ if not isinstance(event.get("response_id"), str) or not event.get("response_id"):
+ raise ValueError("missing_audio_response_id")
+ return chunk
+
+
+def _validated_audio_chunks(client: RawRealtimeProbe) -> list[tuple[dict[str, object], bytes]]:
+ return [(event, _validated_audio_event(event)) for event in _audio_events(client)]
+
+
+def _audio_bytes(client: RawRealtimeProbe) -> bytes:
+ return b"".join(chunk for _, chunk in _validated_audio_chunks(client) if chunk)
+
+
async def _open_session(
args: argparse.Namespace,
*,
@@ -151,22 +206,7 @@ async def _open_session(
) -> tuple[RawRealtimeProbe, dict[str, object]]:
client = RawRealtimeProbe(_realtime_url(args.url, args.model, session_id))
await client.__aenter__()
- await client.send(
- {
- "type": "session.update",
- "session": {
- "session_id": session_id,
- "model": args.model,
- "modalities": ["audio", "text"],
- "input_audio_format": "pcm_f32le",
- "output_audio_format": "pcm16",
- "voice": args.voice,
- "instructions": persona,
- "turn_detection": None,
- "extra_body": {"auto_response": True},
- },
- }
- )
+ await client.send(_session_update(args, session_id=session_id, persona=persona))
event_type = "error" if expect_error else "session.created"
await wait_for(
lambda: client.events.count(event_type) > 0,
@@ -176,6 +216,23 @@ async def _open_session(
return client, _events(client, event_type)[-1]
+def _session_update(args: argparse.Namespace, *, session_id: str, persona: str) -> dict[str, object]:
+ return {
+ "type": "session.update",
+ "session": {
+ "session_id": session_id,
+ "model": args.model,
+ "modalities": ["audio", "text"],
+ "input_audio_format": "pcm_f32le",
+ "output_audio_format": "pcm16",
+ "voice": args.voice,
+ "instructions": persona,
+ "turn_detection": None,
+ "extra_body": {"auto_response": True},
+ },
+ }
+
+
async def _close_session(client: RawRealtimeProbe, *, timeout_s: float) -> None:
await client.send({"type": "session.close"})
await wait_for(
@@ -200,23 +257,24 @@ async def _stream_frames(
frame_count = min(frame_count, max_frames)
for seq in range(frame_count):
frame = pcm[seq * FRAME_SAMPLES : (seq + 1) * FRAME_SAMPLES]
- frame = np.pad(frame, (0, FRAME_SAMPLES - frame.size))
- frame = np.ascontiguousarray(frame, dtype="= 1e-4),
- }
- )
+ await client.send(_frame_event(frame, seq))
await asyncio.sleep(FRAME_PERIOD_S)
return frame_count
+def _frame_event(frame: np.ndarray, seq: int) -> dict[str, object]:
+ frame = np.ascontiguousarray(np.pad(frame, (0, FRAME_SAMPLES - frame.size)), dtype="= 1e-4),
+ }
+
+
def _audio_frame_stats(
raw: bytes,
*,
@@ -242,11 +300,7 @@ def _audio_frame_stats(
def _response_ids(client: RawRealtimeProbe) -> set[str]:
- return {
- str(event["response_id"])
- for event in _events(client, "response.output_audio.delta")
- if isinstance(event.get("response_id"), str)
- }
+ return {str(event["response_id"]) for event in _audio_events(client) if isinstance(event.get("response_id"), str)}
def _session_result(
@@ -256,40 +310,66 @@ def _session_result(
args: argparse.Namespace,
minimum_chunks: int,
) -> tuple[bytes, set[str], dict[str, object]]:
- chunks = sum(len(items) for items in client.events.response_audio.values())
- raw = client.events.audio_bytes()
- if chunks < minimum_chunks or not raw or len(raw) % 2:
- raise AssertionError(f"invalid PCM16 output: chunks={chunks}, bytes={len(raw)}")
+ if isinstance(client, RawRealtimeProbe):
+ audio_chunks = _validated_audio_chunks(client)
+ chunks = sum(bool(chunk) for _, chunk in audio_chunks)
+ raw = b"".join(chunk for _, chunk in audio_chunks if chunk)
+ else:
+ # Unit fixtures may provide an already-decoded EventCollector without
+ # retaining the base64 wire field. Real websocket probes always take
+ # the strict wire-validation path above.
+ chunks = sum(len(items) for items in client.events.response_audio.values())
+ raw = client.events.audio_bytes()
+ stats: dict[str, object] = {"audio_chunks": chunks, "audio_samples": len(raw) // 2}
+ if not raw or len(raw) % 2:
+ raise _AudioAcceptanceError("pcm16_output", f"invalid PCM16 output: chunks={chunks}, bytes={len(raw)}", stats)
pcm = np.frombuffer(raw, dtype="
async def run(args: argparse.Namespace) -> dict[str, object]:
+ if getattr(args, "sessions", None) is not None:
+ return await run_load(args)
input_path = Path(args.input_wav)
input_identity = _input_identity(
input_path,
@@ -381,14 +464,12 @@ async def run(args: argparse.Namespace) -> dict[str, object]:
)
await asyncio.gather(
wait_for(
- lambda: len(primary.events.audio_bytes()) // (2 * FRAME_SAMPLES) >= primary_frames - args.max_frame_deficit,
+ lambda: len(_audio_bytes(primary)) // (2 * FRAME_SAMPLES) >= primary_frames - args.max_frame_deficit,
timeout_s=args.timeout_s,
label="primary frame coverage",
),
wait_for(
- lambda: (
- len(secondary.events.audio_bytes()) // (2 * FRAME_SAMPLES) >= secondary_frames - args.max_frame_deficit
- ),
+ lambda: (len(_audio_bytes(secondary)) // (2 * FRAME_SAMPLES) >= secondary_frames - args.max_frame_deficit),
timeout_s=args.timeout_s,
label="secondary frame coverage",
),
@@ -426,16 +507,14 @@ async def run(args: argparse.Namespace) -> dict[str, object]:
await asyncio.gather(
wait_for(
lambda: (
- len(secondary.events.audio_bytes()) // (2 * FRAME_SAMPLES)
- >= secondary_total_frames - args.max_frame_deficit
+ len(_audio_bytes(secondary)) // (2 * FRAME_SAMPLES) >= secondary_total_frames - args.max_frame_deficit
),
timeout_s=args.timeout_s,
label="survivor frame coverage after slot recycle",
),
wait_for(
lambda: (
- len(replacement.events.audio_bytes()) // (2 * FRAME_SAMPLES)
- >= replacement_frames - args.max_frame_deficit
+ len(_audio_bytes(replacement)) // (2 * FRAME_SAMPLES) >= replacement_frames - args.max_frame_deficit
),
timeout_s=args.timeout_s,
label="replacement frame coverage",
@@ -493,6 +572,284 @@ async def run(args: argparse.Namespace) -> dict[str, object]:
return result
+def _distribution(values: list[float]) -> dict[str, int | float | None]:
+ if not values:
+ return {"count": 0, "median": None, "p99": None, "max": None, "argmax_index": None}
+ return {
+ "count": len(values),
+ "median": float(np.median(values)),
+ "p99": float(np.quantile(values, 0.99)),
+ "max": max(values),
+ "argmax_index": int(np.argmax(values)),
+ }
+
+
+def _load_metrics(client: RawRealtimeProbe, sends: list[tuple[float, float, float]]) -> dict[str, object]:
+ # Use the collector's local clock, never a timestamp supplied by the server.
+ packets: list[dict[str, object]] = []
+ response_ids: set[str] = set()
+ packet_times: list[float] = []
+ packet_bytes: list[int] = []
+ invalid_packets: list[dict[str, object]] = []
+ samples = 0
+ for index, (raw, received) in enumerate(zip(client.events.events, client.events.event_received_at_s, strict=True)):
+ if raw.get("type") not in AUDIO_DELTA_EVENT_TYPES:
+ continue
+ try:
+ chunk = _validated_audio_event(raw)
+ except ValueError as exc:
+ # Only fixed diagnostic codes from the validator are exported,
+ # never audio payloads or details from the remote endpoint.
+ invalid_packets.append({"event_index": index, "received_at_s": received, "code": str(exc)})
+ continue
+ if chunk:
+ samples += len(chunk) // 2
+ packet_times.append(received)
+ packet_bytes.append(len(chunk))
+ response_id = raw.get("response_id")
+ packets.append({"received_at_s": received, "samples": len(chunk) // 2, "response_id": response_id})
+ if isinstance(response_id, str) and response_id:
+ response_ids.add(response_id)
+ intervals = [(b - a) * 1000 for a, b in zip(packet_times, packet_times[1:])]
+ first_send = sends[0][1] if sends else None
+ audio_duration = samples / SAMPLE_RATE_HZ
+ first_audio_ms = (packet_times[0] - first_send) * 1000 if packet_times and first_send is not None else None
+ client_rtf = (packet_times[-1] - first_send) / audio_duration if audio_duration and first_send is not None else None
+ continuity = compute_continuity_stats(packet_times, packet_bytes, SAMPLE_RATE_HZ) if packet_times else None
+ lateness = _distribution([max(0.0, sent - planned) * 1000 for planned, sent, _ in sends])
+ p99_lateness = lateness["p99"]
+ return {
+ "input_frames": len(sends),
+ "output_samples": samples,
+ # Populate frame accounting from _session_result, including on failure.
+ "frame_deficit": None,
+ "response_ids": sorted(response_ids),
+ "client_first_audio_after_stream_start_ms": first_audio_ms,
+ "client_stream_rtf": client_rtf,
+ "client_audio_packet_interval_ms": _distribution(intervals),
+ "client_send_lateness_ms": lateness,
+ "client_pacing_warning": isinstance(p99_lateness, float) and p99_lateness > FRAME_PERIOD_S * 1000,
+ "audio_underrun_s": continuity.max_underrun_s if continuity else None,
+ "audio_underrun_event_count": continuity.underrun_event_count if continuity else None,
+ "audio_continuity_ok": continuity.is_continuous if continuity else None,
+ "input_timeline": [
+ {"frame_index": i, "planned_at_s": planned, "send_started_at_s": sent, "send_completed_at_s": completed}
+ for i, (planned, sent, completed) in enumerate(sends)
+ ],
+ "audio_packet_timeline": packets,
+ "invalid_audio_packets": invalid_packets,
+ }
+
+
+def _check_load_connection(client: RawRealtimeProbe) -> None:
+ if client.events.errors():
+ raise RuntimeError("server emitted a Realtime error (see the session error count)")
+ if client._reader_task is not None and client._reader_task.done():
+ client._reader_task.result() # Propagate malformed JSON/reader failures.
+ raise RuntimeError("connection closed before session completion")
+
+
+async def _wait_load_event(client: RawRealtimeProbe, event_type: str, timeout_s: float) -> None:
+ def received() -> bool:
+ if client.events.count(event_type):
+ return True
+ _check_load_connection(client)
+ return False
+
+ await wait_for(received, timeout_s=timeout_s, label=event_type)
+
+
+async def _paced_load_frames(
+ client: RawRealtimeProbe,
+ pcm: np.ndarray,
+ *,
+ epoch: float,
+ timeout_s: float,
+ sends: list[tuple[float, float, float]],
+ clock: Callable[[], float] = time.monotonic,
+ sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
+) -> None:
+ next_send = epoch
+ for seq in range(math.ceil(pcm.size / FRAME_SAMPLES)):
+ await sleep(max(0.0, next_send - clock()))
+ _check_load_connection(client)
+ event = _frame_event(pcm[seq * FRAME_SAMPLES : (seq + 1) * FRAME_SAMPLES], seq)
+ started = clock()
+ await asyncio.wait_for(client.send(event), timeout=timeout_s)
+ completed = clock()
+ sends.append((epoch + seq * FRAME_PERIOD_S, started, completed))
+ # Do not burst to catch up after a stalled send; report accumulated lag.
+ next_send = max(epoch + (seq + 1) * FRAME_PERIOD_S, completed + FRAME_PERIOD_S)
+
+
+async def _request_load_close(client: RawRealtimeProbe, timeout_s: float) -> None:
+ await client.send({"type": "session.close"})
+ await _wait_load_event(client, "session.closed", timeout_s)
+
+
+async def _run_load_session(
+ args: argparse.Namespace,
+ pcm: np.ndarray,
+ index: int,
+ ready: asyncio.Future[None],
+ start: asyncio.Future[float],
+) -> dict[str, object]:
+ session_id = f"personaplex-load-{index}-{uuid.uuid4().hex}"
+ client = RawRealtimeProbe(_realtime_url(args.url, args.model, session_id), close_timeout_s=args.cleanup_timeout_s)
+ sends: list[tuple[float, float, float]] = []
+ error: str | None = None
+ cleanup_error: str | None = None
+ metrics: dict[str, object] | None = None
+ session_stats: dict[str, object] = {}
+ acceptance_check: str | None = None
+ phase = "admission"
+ try:
+ async with client:
+ try:
+ await asyncio.wait_for(
+ client.send(_session_update(args, session_id=session_id, persona=args.persona)), args.timeout_s
+ )
+ await _wait_load_event(client, "session.created", args.timeout_s)
+ capabilities = _capabilities(_events(client, "session.created")[-1])
+ if (
+ capabilities.get("chunk_period_ms") != 80
+ or capabilities.get("supports_multi_session_same_replica") is not True
+ ):
+ raise AssertionError("server did not advertise PersonaPlex multi-session capabilities")
+ ready.set_result(None)
+ phase = "stream"
+ epoch = await asyncio.shield(start)
+ await _paced_load_frames(client, pcm, epoch=epoch, timeout_s=args.timeout_s, sends=sends)
+ phase = "drain"
+ # Fixed, reported observation window. Closing must not extend
+ # it and let a slow session pass by flushing after the cutoff.
+ await asyncio.sleep(args.drain_s)
+ _check_load_connection(client)
+ _, _, session_stats = _session_result(
+ client, input_frames=len(sends), args=args, minimum_chunks=args.minimum_audio_chunks
+ )
+ finally:
+ metrics = _load_metrics(client, sends)
+ if client.events.count("session.created") and not client.events.count("session.closed"):
+ try:
+ # Send and acknowledgement share one short cleanup budget,
+ # independent of the potentially long inference timeout.
+ await asyncio.wait_for(
+ _request_load_close(client, args.cleanup_timeout_s), args.cleanup_timeout_s
+ )
+ except (OSError, asyncio.TimeoutError, ValueError, RuntimeError, WebSocketException) as exc:
+ cleanup_error = f"cleanup: {type(exc).__name__}"
+ except (OSError, asyncio.TimeoutError, ValueError, RuntimeError, AssertionError, WebSocketException) as exc:
+ if isinstance(exc, _AudioAcceptanceError):
+ acceptance_check = exc.check
+ session_stats = exc.stats
+ error = f"{phase}: AssertionError: {exc}"
+ else:
+ # Other assertions/errors may embed handshake credentials or wire
+ # metadata. Only our local numeric acceptance diagnostics are safe.
+ error = f"{phase}: {type(exc).__name__}"
+ finally:
+ # Admission failure must not leave the other sessions at the barrier.
+ if not ready.done():
+ ready.set_result(None)
+ if metrics is None:
+ metrics = _load_metrics(client, sends)
+ if cleanup_error:
+ error = f"{error}; {cleanup_error}" if error else cleanup_error
+ if client.events.errors():
+ error = error or "server emitted a Realtime error"
+ if metrics["invalid_audio_packets"]:
+ error = error or "invalid audio packets"
+ first_audio = metrics["client_first_audio_after_stream_start_ms"]
+ if isinstance(first_audio, float) and first_audio < 0:
+ error = error or "audio arrived before the first input frame"
+ rtf = metrics["client_stream_rtf"]
+ if args.max_client_rtf is not None and isinstance(rtf, float) and rtf > args.max_client_rtf:
+ error = error or f"client_stream_rtf exceeds {args.max_client_rtf}"
+ return {
+ "index": index,
+ "session_id": session_id,
+ "ok": error is None,
+ "error": error,
+ "acceptance_check": acceptance_check,
+ "server_error_count": len(client.events.errors()),
+ "server_error_codes": [
+ str(body.get("code", "unknown")) if isinstance(body := event.get("error"), dict) else "unknown"
+ for event in client.events.errors()
+ ],
+ **metrics,
+ **session_stats,
+ }
+
+
+async def run_load(args: argparse.Namespace) -> dict[str, object]:
+ input_path = Path(args.input_wav)
+ identity = _input_identity(input_path, expected_sha256=args.expected_input_sha256)
+ pcm = _read_wav_as_float32(input_path)
+ if args.tail_s:
+ pcm = np.concatenate([pcm, np.zeros(round(args.tail_s * SAMPLE_RATE_HZ), dtype=np.float32)])
+ pcm = pcm[: args.load_frames * FRAME_SAMPLES]
+ output_dir = Path(args.output_dir)
+ output_dir.mkdir(parents=True, exist_ok=True)
+ # Refuse to silently overwrite evidence from a previous sweep.
+ result_path = output_dir / "load-result.json"
+ if result_path.exists():
+ raise FileExistsError(result_path)
+ loop = asyncio.get_running_loop()
+ start: asyncio.Future[float] = loop.create_future()
+ ready: list[asyncio.Future[None]] = [loop.create_future() for _ in range(args.sessions)]
+ tasks = [asyncio.create_task(_run_load_session(args, pcm, i, r, start)) for i, r in enumerate(ready)]
+ try:
+ await asyncio.gather(*ready)
+ start.set_result(time.monotonic())
+ rows = await asyncio.gather(*tasks)
+ finally:
+ for task in tasks:
+ if not task.done():
+ task.cancel()
+ await asyncio.gather(*tasks, return_exceptions=True)
+ seen: set[str] = set()
+ shared: set[str] = set()
+ for row in rows:
+ assert isinstance(row["response_ids"], list)
+ response_ids = set(row["response_ids"])
+ shared.update(seen & response_ids)
+ seen.update(response_ids)
+ for row in rows:
+ assert isinstance(row["response_ids"], list)
+ if shared.intersection(row["response_ids"]):
+ row.update(ok=False, error="concurrent sessions shared a response id")
+ result = {
+ "schema_version": 1,
+ "mode": "paced_multi_session",
+ "model": args.model,
+ "requested_sessions": args.sessions,
+ "passed_sessions": sum(row["ok"] is True for row in rows),
+ "ok": all(row["ok"] is True for row in rows),
+ "input": {**identity, "tail_s": args.tail_s, "load_frames": args.load_frames},
+ "frame_period_s": FRAME_PERIOD_S,
+ "drain_s": args.drain_s,
+ "cleanup_timeout_s": args.cleanup_timeout_s,
+ "client_pacing_warning_threshold_ms": FRAME_PERIOD_S * 1000,
+ "client_pacing_warning_sessions": [row["index"] for row in rows if row["client_pacing_warning"]],
+ "max_client_rtf": args.max_client_rtf,
+ "acceptance": {
+ "max_frame_deficit": args.max_frame_deficit,
+ "min_voiced_frames": args.min_voiced_frames,
+ "minimum_audio_chunks": args.minimum_audio_chunks,
+ "voiced_frame_rms_threshold": args.voiced_frame_rms_threshold,
+ },
+ "shared_response_ids": sorted(shared),
+ "sessions": rows,
+ "reported_server": {"revision": args.server_revision, "hardware": args.server_hardware},
+ "measurement_origin": "client monotonic clock; admission excluded; streaming and drain retained",
+ "interval_definition": "nonempty audio packet arrivals, not server ticks or per-frame inference latency",
+ }
+ with result_path.open("x", encoding="utf-8") as stream:
+ json.dump(result, stream, ensure_ascii=False, indent=2, allow_nan=False)
+ return result
+
+
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--url", default="ws://127.0.0.1:8099/v1/realtime?duplex=1")
@@ -510,6 +867,12 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
parser.add_argument("--tail-s", type=float, default=2.0)
parser.add_argument("--drain-s", type=float, default=2.0)
parser.add_argument("--timeout-s", type=float, default=180.0)
+ parser.add_argument(
+ "--cleanup-timeout-s",
+ type=float,
+ default=5.0,
+ help="Load mode: shared session-close/ack budget; also bounds the websocket close handshake.",
+ )
parser.add_argument(
"--max-frame-deficit",
type=int,
@@ -536,11 +899,39 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
default=0,
help="Frames sent after slot reuse; 0 replays the full input so the replacement must produce audible output.",
)
- return parser.parse_args(argv)
+ parser.add_argument(
+ "--sessions", type=int, help="Opt into paced load mode with N sessions; omit for the lifecycle test."
+ )
+ parser.add_argument(
+ "--load-frames", type=int, default=1000, help="Maximum input frames per load session (no repetition)."
+ )
+ parser.add_argument("--max-client-rtf", type=float, help="Optional client stream-window RTF acceptance ceiling.")
+ parser.add_argument("--server-revision", help="User-reported server commit/version; not verified by this client.")
+ parser.add_argument(
+ "--server-hardware", help="User-reported server hardware; not inferred from the client machine."
+ )
+ args = parser.parse_args(argv)
+ if args.sessions is not None:
+ if not 1 <= args.sessions <= 64 or not 1 <= args.load_frames <= 10000:
+ parser.error("load mode requires sessions in [1,64] and load-frames in [1,10000]")
+ for name in ("timeout_s", "cleanup_timeout_s", "drain_s", "tail_s"):
+ value = getattr(args, name)
+ if not math.isfinite(value) or value < 0 or (name in ("timeout_s", "cleanup_timeout_s") and value == 0):
+ parser.error(f"{name} must be finite and non-negative (timeout must be positive)")
+ if args.max_client_rtf is not None and (not math.isfinite(args.max_client_rtf) or args.max_client_rtf <= 0):
+ parser.error("max-client-rtf must be finite and positive")
+ if args.max_frame_deficit < 0 or args.min_voiced_frames < 1 or args.minimum_audio_chunks < 1:
+ parser.error("frame-deficit must be non-negative; voiced frames and audio chunks must be positive")
+ if not math.isfinite(args.voiced_frame_rms_threshold) or args.voiced_frame_rms_threshold < 0:
+ parser.error("voiced-frame-rms-threshold must be finite and non-negative")
+ return args
def main() -> None:
- print(json.dumps(asyncio.run(run(parse_args())), ensure_ascii=False, indent=2))
+ result = asyncio.run(run(parse_args()))
+ print(json.dumps(result, ensure_ascii=False, indent=2, allow_nan=False))
+ if result.get("mode") == "paced_multi_session" and not result["ok"]:
+ raise SystemExit(1)
if __name__ == "__main__":
diff --git a/tests/e2e/online_serving/test_personaplex_load_driver.py b/tests/e2e/online_serving/test_personaplex_load_driver.py
new file mode 100644
index 000000000..8e9a6a28d
--- /dev/null
+++ b/tests/e2e/online_serving/test_personaplex_load_driver.py
@@ -0,0 +1,568 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import asyncio
+import base64
+import json
+import wave
+from contextlib import asynccontextmanager
+from typing import TypedDict
+
+import numpy as np
+import pytest
+import websockets
+
+from tests.e2e.online_serving import personaplex_realtime_duplex as driver
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _ServerState(TypedDict):
+ index: int
+ closed: bool
+ frames: int
+ close_requested: bool
+
+
+def _args(tmp_path, *extra):
+ wav_path = tmp_path / "input.wav"
+ with wave.open(str(wav_path), "wb") as stream:
+ stream.setnchannels(1)
+ stream.setsampwidth(2)
+ stream.setframerate(driver.SAMPLE_RATE_HZ)
+ stream.writeframes(np.full(driver.FRAME_SAMPLES * 2, 2000, dtype="= 1
+ assert report["client_pacing_warning"] is False
+ assert "secret" not in json.dumps(report)
+ assert "-999" not in json.dumps(report)
+
+
+def test_empty_output_has_missing_not_zero_latency():
+ report = driver._load_metrics(driver.RawRealtimeProbe("ws://unused"), [])
+ assert report["client_stream_rtf"] is None
+ assert report["audio_continuity_ok"] is None
+ assert report["audio_underrun_s"] is None
+ intervals = report["client_audio_packet_interval_ms"]
+ assert isinstance(intervals, dict)
+ assert intervals["p99"] is None
+
+
+def test_failed_acceptance_retains_frame_counts_without_remote_data():
+ client = driver.RawRealtimeProbe("ws://unused")
+ client.events.add(_audio(), received_at_s=1.0)
+ args = driver.parse_args(
+ ["--model", "unused", "--input-wav", "unused.wav", "--max-frame-deficit", "0", "--min-voiced-frames", "1"]
+ )
+ with pytest.raises(AssertionError) as raised:
+ driver._session_result(client, input_frames=2, args=args, minimum_chunks=1)
+ assert getattr(raised.value, "check", None) == "frame_deficit"
+ stats = getattr(raised.value, "stats", {})
+ assert stats["frame_deficit"] == 1 and type(stats["frame_deficit"]) is int
+ assert stats["output_frames"] == stats["voiced_frames"] == 1
+ assert stats["frame_coverage_ratio"] == 0.5
+
+
+@pytest.mark.parametrize("event_type", sorted(driver.AUDIO_DELTA_EVENT_TYPES))
+def test_session_result_accepts_current_and_legacy_audio_event_names(tmp_path, event_type):
+ client = driver.RawRealtimeProbe("ws://unused")
+ client.events.add(_audio(samples=driver.FRAME_SAMPLES, event_type=event_type), received_at_s=1.0)
+ raw, response_ids, stats = driver._session_result(
+ client,
+ input_frames=1,
+ args=_args(tmp_path, "--sessions", "1"),
+ minimum_chunks=1,
+ )
+ assert len(raw) == driver.FRAME_SAMPLES * 2
+ assert response_ids == {"r"}
+ assert stats["frame_deficit"] == 0
+
+
+def test_legacy_audio_event_name_contributes_load_metrics():
+ client = driver.RawRealtimeProbe("ws://unused")
+ client.events.add(_audio(samples=driver.FRAME_SAMPLES * 5, event_type="response.audio.delta"), received_at_s=10.4)
+ report = driver._load_metrics(client, [(10.0, 10.0, 10.01)])
+ assert report["output_samples"] == driver.FRAME_SAMPLES * 5
+ assert report["response_ids"] == ["r"]
+ assert report["invalid_audio_packets"] == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sessions", [1, 2, 4])
+async def test_real_websocket_load_success(tmp_path, sessions):
+ args = _args(tmp_path, "--sessions", str(sessions))
+ async with _server() as (url, states):
+ args.url = url
+ result = await driver.run(args)
+ assert result["ok"] is True
+ assert result["passed_sessions"] == sessions
+ rows = result["sessions"]
+ assert isinstance(rows, list)
+ assert len(rows) == sessions
+ assert all(row["frame_deficit"] == 0 for row in rows)
+ for row in rows:
+ assert type(row["frame_deficit"]) is int
+ assert row["output_frames"] == row["voiced_frames"] == row["audio_chunks"] == 2
+ assert row["silent_frames"] == 0
+ assert row["frame_coverage_ratio"] == 1.0
+ assert row["audio_rms"] == pytest.approx(2000 / 32768)
+ assert row["acceptance_check"] is None
+ assert isinstance(row["audio_continuity_ok"], bool)
+ assert all(s["closed"] and s["close_requested"] for s in states)
+ assert "must-not-be-exported" not in json.dumps(result)
+ saved = json.loads((tmp_path / "out/load-result.json").read_text())
+ assert saved == result
+ origins = [row["input_timeline"][0]["planned_at_s"] for row in rows]
+ assert len(set(origins)) == 1
+
+
+@pytest.mark.asyncio
+async def test_real_websocket_load_accepts_legacy_audio_event_name(tmp_path):
+ args = _args(tmp_path, "--sessions", "2")
+ async with _server("legacy_audio_event") as (url, states):
+ args.url = url
+ result = await driver.run(args)
+ assert result["ok"] is True
+ assert result["passed_sessions"] == 2
+ assert result["shared_response_ids"] == []
+ assert all(s["closed"] and s["close_requested"] for s in states)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "mode",
+ [
+ "reject_second",
+ "silent",
+ "disconnect",
+ "no_close_ack",
+ "shared_response",
+ "late_on_close",
+ "error_on_close",
+ "bad_json",
+ "missing_capabilities",
+ ],
+)
+async def test_real_websocket_failures_remain_in_results(tmp_path, mode):
+ args = _args(tmp_path)
+ async with _server(mode) as (url, states):
+ args.url = url
+ result = await asyncio.wait_for(driver.run(args), timeout=4)
+ assert result["ok"] is False
+ assert "must-not-be-exported" not in json.dumps(result)
+ assert result["requested_sessions"] == 2
+ rows = result["sessions"]
+ assert isinstance(rows, list)
+ assert len(rows) == 2
+ assert all(s["closed"] for s in states)
+ if mode == "reject_second":
+ assert result["passed_sessions"] == 1
+ rejected = next(row for row in rows if not row["ok"])
+ assert rejected["input_frames"] == 0
+ assert rejected["client_stream_rtf"] is None
+ assert rejected["server_error_codes"] == ["resource_exhausted"]
+ if mode == "silent":
+ # A failed measurement must still release its admitted session.
+ assert all(s["close_requested"] for s in states)
+ if mode == "late_on_close":
+ assert all(row["output_samples"] == 0 for row in rows)
+ if mode == "shared_response":
+ assert result["passed_sessions"] == 0
+
+
+@pytest.mark.asyncio
+async def test_existing_report_is_not_overwritten(tmp_path):
+ args = _args(tmp_path)
+ path = tmp_path / "out/load-result.json"
+ path.parent.mkdir()
+ path.write_text("previous measurement")
+ with pytest.raises(FileExistsError):
+ await driver.run(args)
+ assert path.read_text() == "previous measurement"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("mode", ["success", "no_close_ack"])
+async def test_cancellation_closes_all_connections(tmp_path, mode):
+ args = _args(tmp_path, "--drain-s", "10", "--timeout-s", "30")
+ async with _server(mode) as (url, states):
+ args.url = url
+ task = asyncio.create_task(driver.run(args))
+
+ async def started():
+ while len(states) < 2 or not all(state["frames"] for state in states):
+ await asyncio.sleep(0.01)
+
+ await asyncio.wait_for(started(), timeout=3)
+ task.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await asyncio.wait_for(task, timeout=3)
+ assert all(s["closed"] for s in states)
+
+
+@pytest.mark.asyncio
+async def test_client_rtf_ceiling_marks_slow_sessions_failed(tmp_path):
+ args = _args(tmp_path, "--max-client-rtf", "0.001")
+ async with _server() as (url, _):
+ args.url = url
+ result = await driver.run(args)
+ assert result["passed_sessions"] == 0
+ rows = result["sessions"]
+ assert isinstance(rows, list)
+ assert all("client_stream_rtf exceeds" in row["error"] for row in rows)
+
+
+@pytest.mark.asyncio
+async def test_stalled_send_does_not_cause_catchup_burst():
+ class Clock:
+ now = 10.0
+
+ def monotonic(self):
+ return self.now
+
+ async def sleep(self, delay):
+ self.now += delay
+
+ class Client(driver.RawRealtimeProbe):
+ async def send(self, event):
+ if event["audio_end_ms"] == 80:
+ clock.now += 0.25
+
+ clock = Clock()
+ sends: list[tuple[float, float, float]] = []
+ await driver._paced_load_frames(
+ Client("ws://unused"),
+ np.zeros(driver.FRAME_SAMPLES * 3, dtype="= sends[0][2] + driver.FRAME_PERIOD_S
+ assert sends[1][1] - sends[1][0] == pytest.approx(0.25)
+ assert sends[2][1] >= sends[1][2] + driver.FRAME_PERIOD_S
+ metrics = driver._load_metrics(Client("ws://unused"), sends)
+ assert metrics["client_pacing_warning"] is True
+ assert asyncio.sleep != clock.sleep
+
+
+@pytest.mark.asyncio
+async def test_failed_frame_acceptance_preserves_numeric_diagnostics(tmp_path):
+ args = _args(tmp_path, "--sessions", "1")
+ async with _server("short_audio") as (url, states):
+ args.url = url
+ result = await driver.run(args)
+ rows = result["sessions"]
+ assert isinstance(rows, list)
+ row = rows[0]
+ assert row["ok"] is False
+ assert row["acceptance_check"] == "frame_deficit"
+ assert row["frame_deficit"] == 1 and type(row["frame_deficit"]) is int
+ assert row["output_frames"] == row["voiced_frames"] == row["audio_chunks"] == 1
+ assert row["frame_coverage_ratio"] == 0.5
+ assert row["audio_rms"] == pytest.approx(2000 / 32768)
+ assert "input=2" in row["error"] and "maximum_deficit=0" in row["error"]
+ assert all(state["closed"] and state["close_requested"] for state in states)
+ assert json.loads((tmp_path / "out/load-result.json").read_text()) == result
+
+
+@pytest.mark.asyncio
+async def test_acceptance_failure_does_not_export_remote_metadata(tmp_path):
+ args = _args(tmp_path, "--sessions", "1")
+ async with _server("private_metadata") as (url, _):
+ args.url = url
+ result = await driver.run(args)
+ rows = result["sessions"]
+ assert isinstance(rows, list)
+ row = rows[0]
+ assert row["ok"] is False
+ assert row["acceptance_check"] == "scheduler_data_plane"
+ assert "remote-detail" not in json.dumps(result)
+ assert "must-not-be-exported" not in json.dumps(result)
+
+
+@pytest.mark.parametrize("ok", [True, False])
+def test_load_main_exit_status(tmp_path, monkeypatch, ok):
+ args = _args(tmp_path)
+
+ async def completed(_args):
+ return {"mode": "paced_multi_session", "ok": ok}
+
+ monkeypatch.setattr(driver, "parse_args", lambda: args)
+ monkeypatch.setattr(driver, "run", completed)
+ if ok:
+ driver.main()
+ else:
+ with pytest.raises(SystemExit) as exc:
+ driver.main()
+ assert exc.value.code == 1
+
+
+@pytest.mark.asyncio
+async def test_handshake_rejection_is_retained_without_server_details(tmp_path, monkeypatch):
+ from websockets.exceptions import InvalidHandshake
+
+ async def rejected(_self):
+ raise InvalidHandshake("credentials-in-remote-handshake-detail")
+
+ monkeypatch.setattr(driver.RawRealtimeProbe, "__aenter__", rejected)
+ result = await asyncio.wait_for(driver.run(_args(tmp_path)), timeout=3)
+ assert result["ok"] is False
+ rows = result["sessions"]
+ assert isinstance(rows, list) and len(rows) == 2
+ assert all(row["error"] == "admission: InvalidHandshake" for row in rows)
+ assert "credentials-in-remote-handshake-detail" not in json.dumps(result)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sessions", [1, 2])
+@pytest.mark.parametrize(
+ "mode",
+ [
+ "missing_packet_rate",
+ "string_packet_rate",
+ "invalid_base64",
+ "orphan_audio",
+ "odd_packet",
+ ],
+)
+async def test_load_rejects_invalid_individual_audio_packets(tmp_path, sessions, mode):
+ # The concatenated PCM can still meet coverage/voicing thresholds, so the
+ # validator must also check each packet before accepting the whole run.
+ args = _args(tmp_path, "--sessions", str(sessions))
+ async with _server(mode) as (url, states):
+ args.url = url
+ result = await asyncio.wait_for(driver.run(args), timeout=4)
+ assert result["ok"] is False
+ assert result["passed_sessions"] == 0
+ rows = result["sessions"]
+ assert isinstance(rows, list) and len(rows) == sessions
+ assert all(row["error"] is not None for row in rows)
+ expected_code = {
+ "missing_packet_rate": "invalid_audio_sample_rate",
+ "string_packet_rate": "invalid_audio_sample_rate",
+ "invalid_base64": "invalid_audio_base64",
+ "orphan_audio": "missing_audio_response_id",
+ "odd_packet": "unaligned_pcm16_packet",
+ }[mode]
+ assert all({packet["code"] for packet in row["invalid_audio_packets"]} == {expected_code} for row in rows)
+ assert all(
+ sum(packet["samples"] for packet in row["audio_packet_timeline"]) == row["output_samples"] for row in rows
+ )
+ assert all(state["closed"] and state["close_requested"] for state in states)
+ assert json.loads((tmp_path / "out/load-result.json").read_text()) == result
+
+
+def test_load_allows_empty_and_multi_frame_audio_packets():
+ client = driver.RawRealtimeProbe("ws://unused")
+ client.events.add(_audio(samples=0), received_at_s=1.0)
+ client.events.add(_audio(samples=driver.FRAME_SAMPLES * 5), received_at_s=1.4)
+ metrics = driver._load_metrics(client, [(1.0, 1.0, 1.001)])
+ assert metrics["invalid_audio_packets"] == []
+ assert metrics["output_samples"] == driver.FRAME_SAMPLES * 5
+ packets = metrics["audio_packet_timeline"]
+ assert isinstance(packets, list) and len(packets) == 1
+
+
+def test_load_packet_diagnostics_do_not_copy_remote_payload():
+ client = driver.RawRealtimeProbe("ws://unused")
+ event = _audio()
+ event["delta"] = "private-remote-payload!"
+ client.events.add(event, received_at_s=1.4)
+ metrics = driver._load_metrics(client, [])
+ assert metrics["output_samples"] == 0
+ assert metrics["invalid_audio_packets"] == [
+ {
+ "event_index": 0,
+ "received_at_s": 1.4,
+ "code": "invalid_audio_base64",
+ }
+ ]
+ assert "private-remote-payload" not in json.dumps(metrics)
diff --git a/tests/engine/test_async_omni_engine_abort_ack.py b/tests/engine/test_async_omni_engine_abort_ack.py
index 43a29bf3a..366e640f6 100644
--- a/tests/engine/test_async_omni_engine_abort_ack.py
+++ b/tests/engine/test_async_omni_engine_abort_ack.py
@@ -66,9 +66,10 @@ def test_abort_async_waits_for_ack(mocker: MockerFixture):
assert msg.rpc_id == "abort-rpc-1"
assert not task.done()
- rpc_q.put(AbortResultMessage(rpc_id="abort-rpc-1", success=True))
+ partial_output = SimpleNamespace(request_id="req-a", finished=True)
+ rpc_q.put(AbortResultMessage(rpc_id="abort-rpc-1", success=True, abort_outputs=[partial_output]))
result = await task
- assert result == []
+ assert result == [partial_output]
try:
asyncio.run(_run())
diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py
index 15d31cbbc..75301ee0a 100644
--- a/tests/engine/test_async_omni_engine_stage_init.py
+++ b/tests/engine/test_async_omni_engine_stage_init.py
@@ -8,25 +8,66 @@ import json
import os
import time
import types
+from dataclasses import dataclass, field
import pytest
from omegaconf import OmegaConf
+from vllm.v1.engine.utils import EngineZmqAddresses
+from vllm_omni.config.omni_config import OmniStageRuntimeConfig
from vllm_omni.diffusion.data import AttentionConfig
from vllm_omni.engine import omni_engine_base as async_omni_engine_module
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+from vllm_omni.engine.stage_engine_startup import StageReplicaResources
from vllm_omni.engine.stage_init_utils import (
LogicalStageInitPlan,
ReplicaInitPlan,
build_stage0_input_processor,
compute_replica_layout,
split_devices_for_replicas,
+ stage_runtime_env,
)
from vllm_omni.engine.stage_runtime import StageRuntime
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+@dataclass
+class _FakeParallelConfig:
+ enable_fault_tolerance: bool = False
+ enable_elastic_ep: bool = False
+ data_parallel_size: int = 1
+ use_ray: bool = False
+
+
+@dataclass
+class _FakeVllmConfig:
+ parallel_config: _FakeParallelConfig = field(default_factory=_FakeParallelConfig)
+
+
+@dataclass
+class _FakeRuntimeConfig:
+ devices: str
+
+
+@dataclass
+class _FakeStageConfig:
+ stage_id: int
+ stage_type: str
+ engine_args: dict[str, object]
+ runtime: _FakeRuntimeConfig
+
+
+def test_stage_runtime_env_accepts_typed_runtime_config(monkeypatch):
+ env_key = "VLLM_OMNI_TEST_TYPED_STAGE_ENV"
+ monkeypatch.delenv(env_key, raising=False)
+
+ with stage_runtime_env(0, OmniStageRuntimeConfig(env={env_key: "typed-value"})):
+ assert os.environ[env_key] == "typed-value"
+
+ assert env_key not in os.environ
+
+
def test_orchestrator_startup_timeout_warns_how_to_raise_limits(monkeypatch):
engine = object.__new__(AsyncOmniEngine)
engine.orchestrator_thread = types.SimpleNamespace(is_alive=lambda: True)
@@ -188,7 +229,7 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred():
)
-def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeypatch):
+def test_async_omni_engine_initialize_stages_passes_log_stats_and_client_config_to_runtime(monkeypatch):
import vllm_omni.engine.omni_engine_base as engine_mod
engine = object.__new__(AsyncOmniEngine)
@@ -206,6 +247,7 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp
engine._omni_lb_policy = "random"
engine.request_queue = types.SimpleNamespace()
engine._log_stats = True
+ engine._client_config = engine_mod.OmniClientConfig(client_count=2, client_index=1, stage_addresses={})
engine._parallel_stage_init = False
captured: dict[str, object] = {}
@@ -221,6 +263,51 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp
assert captured["stage_init_timeout"] == 7
assert captured["log_stats"] is True
+ assert captured["client_config"] is engine._client_config
+
+
+def test_async_omni_engine_initialize_stages_retains_stage0_prompt_transform(monkeypatch):
+ import vllm_omni.engine.omni_engine_base as engine_mod
+
+ engine = object.__new__(AsyncOmniEngine)
+ engine.stage_configs = [types.SimpleNamespace()]
+ engine.model = "dummy-model"
+ engine.config_path = "dummy-config"
+ engine.single_stage_mode = False
+ engine.async_chunk = False
+ engine.tokenizer = None
+ engine._single_stage_id_filter = None
+ engine._omni_master_address = None
+ engine._omni_master_port = None
+ engine._omni_dp_size_local = 1
+ engine._omni_heartbeat_timeout = 30.0
+ engine._omni_lb_policy = "random"
+ engine.request_queue = types.SimpleNamespace()
+ engine._log_stats = False
+ engine._parallel_stage_init = False
+
+ prompt_transform = object()
+ client = types.SimpleNamespace(
+ prompt_transform_func=prompt_transform,
+ prompt_expand_func=None,
+ default_sampling_params=types.SimpleNamespace(),
+ final_output=True,
+ final_output_type="text",
+ stage_type="llm",
+ model_stage="text_encoder",
+ is_comprehension=True,
+ )
+ pool = types.SimpleNamespace(
+ stage_client=client,
+ stage_vllm_config=None,
+ output_processor=None,
+ )
+ runtime = types.SimpleNamespace(stage_pools=[pool], initialize=lambda: None)
+ monkeypatch.setattr(engine_mod, "create_stage_runtime", lambda **_kwargs: runtime)
+
+ engine._initialize_stages(stage_init_timeout=7)
+
+ assert engine.prompt_transform_func is prompt_transform
def test_compute_replica_layout_splits_diffusion_devices_by_world_size():
@@ -777,6 +864,289 @@ def test_stage_runtime_passes_log_stats_to_llm_replica_launch(monkeypatch):
assert captured["client_log_stats"] is True
+def test_stage_runtime_attaches_external_llm_client_without_launching_engine(monkeypatch):
+ import vllm_omni.engine.stage_runtime as runtime_mod
+
+ client_addresses = {
+ "input_address": "ipc://stage0-input-1",
+ "output_address": "ipc://stage0-output-1",
+ }
+ client_config = {
+ "client_count": 2,
+ "client_index": 1,
+ "stage_addresses": {0: {0: client_addresses}},
+ }
+ runtime = StageRuntime(
+ stage_configs=[],
+ model="dummy-model",
+ config_path="dummy-config",
+ stage_init_timeout=1,
+ async_chunk=False,
+ client_config=client_config,
+ )
+ parallel_config = _FakeParallelConfig()
+ vllm_config = _FakeVllmConfig(parallel_config=parallel_config)
+ plan = _make_llm_plan(0, stage_id=0, vllm_config=vllm_config).replicas[0]
+ captured: dict[str, object] = {}
+ stage_client = object()
+
+ def _capture_client(**kwargs):
+ captured.update(kwargs)
+ return stage_client
+
+ monkeypatch.setattr(runtime_mod.StageEngineCoreClientBase, "make_async_mp_client", _capture_client)
+ monkeypatch.setattr(
+ runtime_mod,
+ "launch_stage_replica",
+ lambda **_kwargs: pytest.fail("external client must not launch a stage engine"),
+ )
+
+ assert runtime._initialize_local_llm_replica(plan, stage_init_timeout=1) is stage_client
+ assert captured["client_addresses"] == client_addresses
+ assert captured["client_count"] == 2
+ assert captured["client_index"] == 1
+ assert not hasattr(parallel_config, "_api_process_count")
+ assert not hasattr(parallel_config, "_api_process_rank")
+
+
+@pytest.mark.parametrize("stage_ids", [(0,), (0, 1)], ids=["single-stage", "multi-stage"])
+@pytest.mark.parametrize("client_count", [1, 2])
+def test_stage_runtime_launches_shared_engines_with_per_client_addresses(monkeypatch, stage_ids, client_count):
+ import vllm_omni.engine.stage_runtime as runtime_mod
+
+ runtime = _make_stage_runtime()
+ stage_plans = []
+ for stage_id in stage_ids:
+ parallel_config = _FakeParallelConfig()
+ plan = _make_llm_plan(
+ stage_id,
+ stage_id=stage_id,
+ vllm_config=_FakeVllmConfig(parallel_config=parallel_config),
+ )
+ plan.replicas[0].engine_args_dict = {}
+ stage_plans.append(plan)
+ stage_plans[0].replicas[0].metadata.runtime_cfg = {"env": {"VLLM_OMNI_TEST_STAGE_RUNTIME_ENV": "enabled"}}
+
+ monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans)
+ monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda stage_id, _cfg: str(stage_id))
+ monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args: [])
+
+ events: list[tuple[str, int]] = []
+ captured_launch_kwargs: list[dict[str, object]] = []
+ captured_launch_env: list[str | None] = []
+
+ @contextlib.contextmanager
+ def _fake_launch_stage_replica(**kwargs):
+ stage_id = kwargs["stage_id"]
+ events.append(("enter", stage_id))
+ captured_launch_kwargs.append(kwargs)
+ captured_launch_env.append(os.environ.get("VLLM_OMNI_TEST_STAGE_RUNTIME_ENV"))
+ yield StageReplicaResources(
+ addresses=EngineZmqAddresses(
+ inputs=[f"ipc://stage{stage_id}-input-{idx}" for idx in range(client_count)],
+ outputs=[f"ipc://stage{stage_id}-output-{idx}" for idx in range(client_count)],
+ ),
+ )
+ events.append(("exit", stage_id))
+
+ monkeypatch.setattr(runtime_mod, "launch_stage_replica", _fake_launch_stage_replica)
+
+ with runtime.launch_stage_engines(client_count) as launch:
+ expected = (
+ [("enter", stage_id) for stage_id in stage_ids]
+ if client_count > 1
+ else [(event, stage_id) for stage_id in stage_ids for event in ("enter", "exit")]
+ )
+ assert events == expected
+ for client_index, config in enumerate(launch.client_configs):
+ assert config["client_count"] == client_count
+ assert config["client_index"] == client_index
+ for stage_id in stage_ids:
+ assert config["stage_addresses"][stage_id][0] == {
+ "input_address": f"ipc://stage{stage_id}-input-{client_index}",
+ "output_address": f"ipc://stage{stage_id}-output-{client_index}",
+ }
+
+ assert events == (
+ [(event, stage_id) for event in ("enter", "exit") for stage_id in stage_ids] if client_count > 1 else expected
+ )
+ assert not hasattr(stage_plans[0].replicas[0].stage_vllm_config.parallel_config, "_api_process_count")
+ assert not hasattr(stage_plans[0].replicas[0].stage_vllm_config.parallel_config, "_api_process_rank")
+ assert all(
+ kwargs["watched_frontend_processes"] is (launch.watched_frontend_processes if client_count > 1 else None)
+ for kwargs in captured_launch_kwargs
+ )
+ assert captured_launch_env == ["enabled" if stage_id == 0 else None for stage_id in stage_ids]
+ assert os.environ.get("VLLM_OMNI_TEST_STAGE_RUNTIME_ENV") is None
+
+
+@pytest.mark.parametrize("client_count", [1, 2])
+def test_stage_runtime_overlapping_devices_acquire_real_locks_once(monkeypatch, tmp_path, client_count):
+ import fcntl
+
+ import vllm_omni.engine.stage_init_utils as init_utils
+ import vllm_omni.engine.stage_runtime as runtime_mod
+
+ runtime = _make_stage_runtime()
+ plans = [_make_llm_plan(stage_id, stage_id=stage_id, vllm_config=_FakeVllmConfig()) for stage_id in (0, 1)]
+ for stage_id, plan in enumerate(plans):
+ plan.replicas[0].engine_args_dict = {"tensor_parallel_size": stage_id + 1}
+ monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: plans)
+ monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda sid, _cfg: "0" if sid == 0 else "0,1")
+ monkeypatch.setattr(
+ init_utils, "device_init_lock_path", lambda device_id: str(tmp_path / f"device-{device_id}.lock")
+ )
+ monkeypatch.setattr(
+ init_utils.time, "sleep", lambda _: pytest.fail("Overlapping stages must not wait on their own lock")
+ )
+
+ @contextlib.contextmanager
+ def launch(**kwargs):
+ # Real flock on a separate descriptor proves the runtime holds each lock.
+ for device_id in range(kwargs["stage_id"] + 1):
+ with open(tmp_path / f"device-{device_id}.lock", "a") as probe:
+ with pytest.raises(BlockingIOError):
+ fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB)
+ yield StageReplicaResources(
+ addresses=EngineZmqAddresses(
+ inputs=[f"ipc://input-{i}" for i in range(client_count)],
+ outputs=[f"ipc://output-{i}" for i in range(client_count)],
+ )
+ )
+
+ monkeypatch.setattr(runtime_mod, "launch_stage_replica", launch)
+ with runtime.launch_stage_engines(client_count):
+ pass
+ for device_id in (0, 1):
+ with open(tmp_path / f"device-{device_id}.lock", "a") as probe:
+ fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB)
+ fcntl.flock(probe, fcntl.LOCK_UN)
+
+
+@pytest.mark.parametrize("diffusion_stage_id", [0, 1], ids=["diffusion-only", "enginecore-to-diffusion"])
+def test_stage_runtime_multi_api_rejects_diffusion_before_launch(monkeypatch, diffusion_stage_id):
+ import vllm_omni.engine.stage_runtime as runtime_mod
+
+ runtime = _make_stage_runtime()
+ stage_plans = [
+ _make_llm_plan(stage_id, stage_id=stage_id, vllm_config=_FakeVllmConfig())
+ for stage_id in range(diffusion_stage_id)
+ ]
+ stage_plans.append(_make_diffusion_plan(diffusion_stage_id, stage_id=diffusion_stage_id))
+ monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans)
+ monkeypatch.setattr(
+ runtime_mod,
+ "acquire_device_locks",
+ lambda *_args: pytest.fail("unsupported topology must fail before acquiring devices"),
+ )
+ monkeypatch.setattr(
+ runtime_mod,
+ "launch_stage_replica",
+ lambda **_kwargs: pytest.fail("unsupported topology must not partially launch EngineCore stages"),
+ )
+
+ with pytest.raises(ValueError, match="diffusion stage\\(s\\) are not supported"):
+ with runtime.launch_stage_engines(2):
+ pytest.fail("unsupported topology must not yield a launch handle")
+
+
+def test_stage_runtime_multi_api_maps_stage_devices_from_launcher_visibility(monkeypatch):
+ import vllm_omni.engine.stage_runtime as runtime_mod
+
+ runtime = _make_stage_runtime()
+ stage_plan = _make_llm_plan(
+ 0,
+ stage_id=0,
+ vllm_config=_FakeVllmConfig(),
+ )
+ replica = stage_plan.replicas[0]
+ replica.engine_args_dict = {}
+ replica.metadata.runtime_cfg = {"devices": "0"}
+
+ device_env = runtime_mod.current_omni_platform.device_control_env_var
+ monkeypatch.setenv(device_env, "5")
+ monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [stage_plan])
+ monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args: [])
+ captured_devices: list[str | None] = []
+
+ @contextlib.contextmanager
+ def _fake_launch_stage_replica(**kwargs):
+ captured_devices.append(kwargs["stage_visible_devices"])
+ yield StageReplicaResources(
+ addresses=EngineZmqAddresses(
+ inputs=["ipc://input-0", "ipc://input-1"],
+ outputs=["ipc://output-0", "ipc://output-1"],
+ ),
+ )
+
+ monkeypatch.setattr(runtime_mod, "launch_stage_replica", _fake_launch_stage_replica)
+
+ with runtime.launch_stage_engines(2):
+ pass
+
+ assert captured_devices == ["5"]
+ assert os.environ[device_env] == "5"
+
+
+def test_stage_runtime_multi_api_rejects_elastic_ep(monkeypatch):
+ runtime = _make_stage_runtime()
+ stage_plan = _make_llm_plan(
+ 0,
+ stage_id=0,
+ vllm_config=_FakeVllmConfig(parallel_config=_FakeParallelConfig(enable_elastic_ep=True)),
+ )
+ stage_plan.replicas[0].engine_args_dict = {}
+ monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [stage_plan])
+
+ with pytest.raises(ValueError, match="enable-elastic-ep"):
+ with runtime.launch_stage_engines(2):
+ pass
+
+
+def test_stage_runtime_multi_api_failure_shuts_down_before_exceptional_context_exit(monkeypatch):
+ import vllm_omni.engine.stage_runtime as runtime_mod
+
+ runtime = _make_stage_runtime()
+ stage_plan = _make_llm_plan(0, stage_id=0, vllm_config=_FakeVllmConfig())
+ stage_plan.replicas[0].engine_args_dict = {}
+ monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [stage_plan])
+ monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda *_args: None)
+ monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args: [])
+
+ events: list[str] = []
+
+ class _Manager:
+ def shutdown(self) -> None:
+ events.append("shutdown")
+
+ @contextlib.contextmanager
+ def _fake_launch_stage_replica(**_kwargs):
+ try:
+ yield StageReplicaResources(
+ manager=_Manager(),
+ addresses=EngineZmqAddresses(
+ inputs=["tcp://127.0.0.1:0", "tcp://127.0.0.1:1"],
+ outputs=["tcp://127.0.0.1:2", "tcp://127.0.0.1:3"],
+ ),
+ )
+ except RuntimeError as exc:
+ events.append(f"exceptional-exit:{exc}")
+ raise
+ else:
+ events.append("normal-exit")
+
+ monkeypatch.setattr(runtime_mod, "launch_stage_replica", _fake_launch_stage_replica)
+
+ with pytest.raises(RuntimeError, match="deferred TCP addresses"):
+ with runtime.launch_stage_engines(2):
+ pass
+
+ assert events == [
+ "shutdown",
+ "exceptional-exit:Stage 0 returned deferred TCP addresses; multi-API launch requires fixed ports or IPC addresses",
+ ]
+
+
def test_stage_runtime_passes_log_stats_to_output_processor(monkeypatch):
import vllm_omni.engine.stage_runtime as runtime_mod
@@ -1023,7 +1393,7 @@ def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stag
num_replicas=1,
launch_mode="local",
stage_cfg=types.SimpleNamespace(engine_args={}, runtime=types.SimpleNamespace(devices="0")),
- metadata=types.SimpleNamespace(stage_id=0, runtime_cfg={"devices": "0"}),
+ metadata=types.SimpleNamespace(stage_id=0, stage_type="llm", runtime_cfg={"devices": "0"}),
stage_connector_spec={},
omni_kv_connector=(None, None, None),
stage_vllm_config=fake_vllm_config,
@@ -1715,6 +2085,51 @@ def test_inject_kv_stage_info_infers_receiver_tp_topology():
assert stage1.engine_args["omni_kv_config"]["rank_mapping"] == {"from_tp": 4, "to_tp": 2}
+def test_inject_kv_stage_info_updates_typed_connector_config():
+ from vllm_omni.config.omni_config import (
+ OmniStageConnectorConfig,
+ OmniStageDiffusionParallelConfig,
+ VllmOmniDiffusionStageConfig,
+ )
+ from vllm_omni.config.stage_config import StageExecutionType, StagePipelineConfig
+ from vllm_omni.engine.stage_init_utils import inject_kv_stage_info
+ from vllm_omni.entrypoints.utils import inject_omni_kv_config
+
+ stage0 = VllmOmniDiffusionStageConfig(
+ stage_pipeline_config=StagePipelineConfig(
+ stage_id=0,
+ model_stage="diffusion",
+ execution_type=StageExecutionType.DIFFUSION,
+ ),
+ connector_config=OmniStageConnectorConfig(
+ omni_kv_config={
+ "need_send_cache": True,
+ "omni_from_stage": "0",
+ "omni_to_stage": "1",
+ }
+ ),
+ parallel_config=OmniStageDiffusionParallelConfig(tensor_parallel_size=4),
+ )
+ stage1 = VllmOmniDiffusionStageConfig(
+ stage_pipeline_config=StagePipelineConfig(
+ stage_id=1,
+ model_stage="diffusion",
+ execution_type=StageExecutionType.DIFFUSION,
+ input_sources=(0,),
+ ),
+ connector_config=OmniStageConnectorConfig(omni_kv_config={"need_recv_cache": True}),
+ parallel_config=OmniStageDiffusionParallelConfig(tensor_parallel_size=2),
+ )
+
+ inject_omni_kv_config(stage0, {"kv_connector": "P2pNcclConnector"}, "0", "1")
+ inject_kv_stage_info(stage0, 0, [stage0, stage1])
+
+ assert stage0.connector_config.omni_kv_config["stage_id"] == 0
+ assert stage0.connector_config.omni_kv_config["connector_config"] == {"kv_connector": "P2pNcclConnector"}
+ assert stage0.connector_config.omni_kv_config["engine_input_source"] == []
+ assert stage0.connector_config.omni_kv_config["rank_mapping"] == {"from_tp": 4, "to_tp": 2}
+
+
def test_extract_legacy_stage_metadata_rocm_does_not_inject_diffusion_attention(monkeypatch):
"""ROCm default attention logic only applies to LLM stages, not diffusion."""
from vllm_omni.engine.stage_init_utils import extract_legacy_stage_metadata
@@ -1832,6 +2247,102 @@ def test_port_from_zmq_address_parsing():
assert _port_from_zmq_address("tcp://host:not-a-port") is None
+@pytest.mark.parametrize("parallel", [False, True])
+@pytest.mark.parametrize("failure", ["ready", "attach", None])
+def test_single_api_common_launch_ownership_and_rollback(monkeypatch, parallel, failure):
+ import vllm_omni.engine.stage_runtime as runtime_mod
+
+ runtime = _make_stage_runtime()
+ runtime._parallel_stage_init = parallel
+ plan = _make_llm_plan(0, stage_id=0, vllm_config=_FakeVllmConfig()).replicas[0]
+ plan.engine_args_dict = {}
+ events = []
+
+ class Manager:
+ def shutdown(self):
+ events.append("shutdown")
+
+ manager = Manager()
+ addresses = EngineZmqAddresses(inputs=["ipc://input"], outputs=["ipc://output"])
+ monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda *_: "0")
+
+ def acquire(*_):
+ events.append("lock")
+ return [42]
+
+ monkeypatch.setattr(runtime_mod, "acquire_device_locks", acquire)
+ monkeypatch.setattr(runtime_mod, "release_device_locks", lambda _: events.append("unlock"))
+
+ @contextlib.contextmanager
+ def launch(**kwargs):
+ assert kwargs["num_api_servers"] == 1
+ assert kwargs["omni_parallel_stage_init"] is parallel
+ events.append("spawn")
+ yield StageReplicaResources(manager=manager, addresses=addresses)
+ assert runtime._replica_launch_lock.locked() is False
+ events.append("ready")
+ if failure == "ready":
+ raise RuntimeError("ready")
+
+ def attach(**kwargs):
+ assert "ready" in events
+ assert kwargs["engine_manager"] is manager
+ events.append("attach")
+ if failure == "attach":
+ raise RuntimeError("attach")
+ return manager
+
+ monkeypatch.setattr(runtime_mod, "launch_stage_replica", launch)
+ monkeypatch.setattr(runtime_mod.StageEngineCoreClientBase, "make_async_mp_client", attach)
+ if failure:
+ with pytest.raises(RuntimeError, match=failure):
+ runtime._initialize_local_llm_replica(plan, 1)
+ assert events.count("shutdown") == 1
+ else:
+ assert runtime._initialize_local_llm_replica(plan, 1) is manager
+ assert "shutdown" not in events
+ assert events.count("lock") == events.count("unlock") == (0 if parallel else 1)
+
+
+@pytest.mark.parametrize("client_count", [1, 2])
+def test_common_launch_parallel_admission_before_spawn(monkeypatch, client_count):
+ import vllm_omni.engine.stage_runtime as runtime_mod
+
+ runtime = _make_stage_runtime()
+ runtime._parallel_stage_init = True
+ plan = _make_llm_plan(0, stage_id=0, vllm_config=_FakeVllmConfig())
+ plan.replicas[0].engine_args_dict = {}
+ events = []
+ monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [plan])
+ monkeypatch.setattr(runtime, "_reject_unguardable_executors", lambda _: events.append("guard"))
+ monkeypatch.setattr(runtime, "_run_stage_admission", lambda _: events.append("admit"))
+ monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda *_: "0")
+
+ def forbidden_lock(*_):
+ pytest.fail("parallel initialization must use child phase locks")
+
+ monkeypatch.setattr(runtime_mod, "acquire_device_locks", forbidden_lock)
+
+ @contextlib.contextmanager
+ def launch(**kwargs):
+ assert events == ["guard", "admit"]
+ assert kwargs["omni_parallel_stage_init"] is True
+ events.append("spawn")
+ yield StageReplicaResources(
+ addresses=EngineZmqAddresses(
+ inputs=[f"ipc://in{i}" for i in range(client_count)],
+ outputs=[f"ipc://out{i}" for i in range(client_count)],
+ )
+ )
+ assert not runtime._replica_launch_lock.locked()
+ events.append("ready")
+
+ monkeypatch.setattr(runtime_mod, "launch_stage_replica", launch)
+ with runtime.launch_stage_engines(client_count):
+ pass
+ assert events == ["guard", "admit", "spawn", "ready"]
+
+
@pytest.mark.parametrize(
"stage_modes,expected", [([False, True, True], True), ([False, False], False), ([True, False], True)]
)
@@ -1852,3 +2363,80 @@ def test_engine_async_chunk_includes_downstream_stages(monkeypatch, stage_modes,
with pytest.raises(ConfigResolvedError):
engine.__init__("test-model")
assert engine.async_chunk is expected
+
+
+@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"])
+def test_dist_stage_runtime_applies_local_dp_to_stage_config(typed):
+ from vllm_omni.config.omni_config import VllmOmniARStageConfig
+ from vllm_omni.config.stage_config import StagePipelineConfig
+ from vllm_omni.engine.stage_runtime import DistStageRuntime
+
+ if typed:
+ stage = VllmOmniARStageConfig(stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="ar"))
+ runtime_cfg = stage.runtime_config
+ else:
+ runtime_cfg = types.SimpleNamespace(num_replicas=1)
+ stage = types.SimpleNamespace(stage_id=0, runtime=runtime_cfg)
+ runtime = DistStageRuntime(
+ stage_configs=[stage],
+ model="dummy-model",
+ config_path="dummy-config",
+ stage_init_timeout=1,
+ async_chunk=False,
+ single_stage_id_filter=0,
+ omni_master_address="127.0.0.1",
+ omni_master_port=12345,
+ omni_dp_size_local=2,
+ )
+
+ runtime._validate_single_stage_mode_replica_constraints()
+
+ assert runtime_cfg.num_replicas == 2
+
+
+@pytest.mark.parametrize("num_replicas", [1, 2, 3])
+def test_typed_diffusion_replicas_share_one_config_between_planning_and_launch(mocker, num_replicas):
+ from vllm_omni.config.config_factory import StageConfigFactory
+ from vllm_omni.engine import stage_runtime as runtime_module
+
+ stage = StageConfigFactory.create_typed_default_diffusion(
+ "generic-diffusion", {"model_class_name": "QwenImagePipeline"}
+ ).stage_configs[0]
+ devices = ",".join(str(i) for i in range(num_replicas))
+ stage.runtime_config.devices = devices
+ stage.runtime_config.num_replicas = num_replicas
+ runtime = StageRuntime(
+ stage_configs=[stage],
+ model="dummy-model",
+ config_path="dummy-config",
+ stage_init_timeout=1,
+ async_chunk=False,
+ )
+ mocker.patch.object(runtime_module, "get_stage_connector_spec", return_value={})
+ mocker.patch.object(runtime_module, "resolve_omni_kv_config_for_stage", return_value=(None, None, None))
+ mocker.patch.object(runtime, "_stage_device_scope", side_effect=lambda *_: contextlib.nullcontext())
+ client = mocker.Mock()
+ launch = mocker.patch.object(runtime_module, "launch_diffusion_stage_replica", return_value=(client, None))
+
+ # Replanning must not inherit the first replica's narrowed device slice.
+ for _ in range(2):
+ counts, device_map = compute_replica_layout([stage])
+ plans = runtime._build_logical_stage_init_plans(
+ omni_transfer_config=None,
+ replicas_per_stage=counts,
+ replica_devices_map=device_map,
+ )
+ replicas = plans[0].replicas
+ assert stage.runtime_config.devices == devices
+ assert len({id(plan.stage_cfg) for plan in replicas}) == num_replicas
+ for i, plan in enumerate(replicas):
+ assert plan.metadata.runtime_cfg is plan.stage_cfg.runtime_config
+ assert plan.stage_cfg.runtime_config.devices == str(i)
+ if num_replicas > 1:
+ assert plan.stage_cfg is not stage
+ assert runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=1) is client
+ assert launch.call_args.kwargs["stage_config"] is plan.stage_cfg
+ assert launch.call_args.kwargs["metadata"] is plan.metadata
+ assert launch.call_args.kwargs["use_inline"] is (num_replicas == 1)
+
+ assert launch.call_count == 2 * num_replicas
diff --git a/tests/engine/test_parallel_stage_init.py b/tests/engine/test_parallel_stage_init.py
index d88d49825..9f89aabca 100644
--- a/tests/engine/test_parallel_stage_init.py
+++ b/tests/engine/test_parallel_stage_init.py
@@ -306,6 +306,23 @@ def test_check_admission_admits_diffusion_with_explicit_util():
assert ledgers[0].kv_budget_bytes == int(cap * 0.3)
+def test_check_admission_reads_typed_diffusion_cache_config():
+ cap = 80 * 1024**3
+ diff = _local_diffusion_replica(0, 0, "0", util=None)
+ diff.stage_cfg.cache_config = types.SimpleNamespace(gpu_memory_utilization=0.35)
+ plan = LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=[diff])
+
+ ledgers = check_admission(
+ [plan],
+ resolve_physical_devices=lambda _r: [0],
+ device_total_memory=lambda _d: cap,
+ external_reserve_bytes=0,
+ safety_margin_bytes=0,
+ )
+
+ assert ledgers[0].kv_budget_bytes == int(cap * 0.35)
+
+
def test_check_admission_unresolved_local_raises():
"""A local replica invisible to the ledger must fail admission, not skip it:
it would otherwise initialize concurrently with unbounded demand."""
diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py
index 41d63af9e..9ce707dbc 100644
--- a/tests/engine/test_single_stage_mode.py
+++ b/tests/engine/test_single_stage_mode.py
@@ -16,7 +16,8 @@ from pytest_mock import MockerFixture
from vllm.v1.engine.utils import EngineZmqAddresses
from vllm_omni.config.config_factory import StageConfigFactory
-from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
+from vllm_omni.config.omni_config import VllmOmniARStageConfig, VllmOmniDiffusionStageConfig
+from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig, StageExecutionType, StagePipelineConfig
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase
from vllm_omni.engine.stage_engine_startup import (
@@ -57,6 +58,10 @@ def _make_llm_plan(
vllm_config: Any | None = None,
) -> LogicalStageInitPlan:
stage_cfg = _make_stage_cfg(stage_id)
+ if launch_mode == "remote":
+ stage_cfg = VllmOmniARStageConfig(
+ stage_pipeline_config=StagePipelineConfig(stage_id=stage_id, model_stage="thinker"),
+ )
metadata = SimpleNamespace(
stage_id=stage_id,
stage_type="llm",
@@ -98,6 +103,16 @@ def _make_diffusion_plan(
launch_mode: str,
) -> LogicalStageInitPlan:
stage_cfg = _make_stage_cfg(stage_id, stage_type="diffusion")
+ if launch_mode == "remote":
+ stage_cfg = VllmOmniDiffusionStageConfig(
+ stage_pipeline_config=StagePipelineConfig(
+ stage_id=stage_id,
+ model_stage="diffusion",
+ execution_type=StageExecutionType.DIFFUSION,
+ final_output=True,
+ final_output_type="image",
+ ),
+ )
metadata = SimpleNamespace(
stage_id=stage_id,
stage_type="diffusion",
@@ -960,6 +975,7 @@ class TestSingleStageReplicaInitialization:
assert mock_connect.call_args.kwargs["stage_id"] == 1
assert mock_connect.call_args.kwargs["replica_id"] == 0
assert client_kwargs["log_stats"] is True
+ assert client_kwargs["metadata"].model_stage == "thinker"
assert events == ["enter", "exit", "attach"]
def test_initialize_llm_replica_remote_missing_registered_stage_config_raises(self, mocker: MockerFixture):
@@ -1125,11 +1141,9 @@ class TestSingleStageReplicaInitialization:
return _ctx()
- remote_metadata = _make_diffusion_plan(1, stage_id=1, launch_mode="remote").replicas[0].metadata
plan = _make_diffusion_plan(1, stage_id=1, launch_mode="remote").replicas[0]
sentinel_client = SimpleNamespace()
- mocker.patch.object(runtime_mod, "extract_legacy_stage_metadata", return_value=remote_metadata)
mock_connect = mocker.patch.object(runtime_mod, "connect_remote_diffusion_proc", side_effect=_fake_connect)
mock_from_addresses = mocker.patch(
"vllm_omni.diffusion.stage_diffusion_client.StageDiffusionClient.from_addresses",
@@ -1146,6 +1160,7 @@ class TestSingleStageReplicaInitialization:
replica_id=0,
)
mock_from_addresses.assert_called_once()
+ assert mock_from_addresses.call_args.args[0].final_output_type == "image"
def test_initialize_local_diffusion_replica_registers_with_master(self, mocker: MockerFixture):
import vllm_omni.engine.stage_runtime as runtime_mod
diff --git a/tests/engine/test_stage_device_layout.py b/tests/engine/test_stage_device_layout.py
index 0473dc341..6dd30373d 100644
--- a/tests/engine/test_stage_device_layout.py
+++ b/tests/engine/test_stage_device_layout.py
@@ -18,6 +18,7 @@ import types
from unittest import mock
import pytest
+from pytest_mock import MockerFixture
from vllm_omni.engine import stage_init_utils
from vllm_omni.engine.stage_init_utils import (
@@ -156,6 +157,43 @@ def test_replica_guard_and_splitter_share_local_world_size(devices, expected):
assert replica_devices_map == {0: expected}
+@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"])
+@pytest.mark.parametrize("local_dp", [None, 0, 1, 2])
+@pytest.mark.parametrize("pool_devices", [False, True], ids=["template", "pool"])
+def test_replica_layout_preserves_local_dp_semantics(typed, local_dp, pool_devices):
+ from vllm_omni.config.omni_config import VllmOmniConfig
+ from vllm_omni.config.stage_config import PipelineConfig, StagePipelineConfig
+ from vllm_omni.engine.stage_engine_startup import get_headless_replica_devices
+
+ engine_args = {
+ "tensor_parallel_size": 2,
+ "data_parallel_size": 4,
+ "data_parallel_size_local": local_dp,
+ "pipeline_parallel_size": 2,
+ }
+ width = 2 * max(1, local_dp if local_dp is not None else 4) * 2
+ devices = ",".join(str(i) for i in range(width * (2 if pool_devices else 1)))
+ if typed:
+ pipeline = PipelineConfig(
+ model_type="test",
+ model_arch="TestModel",
+ stages=(StagePipelineConfig(stage_id=0, model_stage="ar", final_output=True),),
+ )
+ stage = VllmOmniConfig.from_pipeline_config(
+ pipeline,
+ cli_overrides={
+ f"stage_0_{key}": value for key, value in {**engine_args, "devices": devices, "num_replicas": 2}.items()
+ },
+ ).stage_configs[0]
+ else:
+ stage = _stage(0, devices=devices, num_replicas=2, engine_args=engine_args)
+ expected = [",".join(str(i) for i in range(r * width, (r + 1) * width)) for r in range(2)]
+
+ assert get_stage_devices_per_replica(stage) == width
+ assert compute_replica_layout([stage]) == ([2], {0: expected})
+ assert get_headless_replica_devices(stage, stage_id=0, omni_dp_size_local=2) == expected
+
+
def test_non_tp_layout_error_uses_generic_guidance():
"""A PP mismatch must not be reported as a top-level TP broadcast."""
stage = _stage(0, devices="0")
@@ -203,7 +241,7 @@ def test_build_vllm_config_fails_before_engine_config_on_mismatch():
get_class.assert_not_called()
-def test_build_vllm_config_proceeds_on_consistent_layout():
+def test_build_vllm_config_proceeds_on_consistent_layout(mocker: MockerFixture) -> None:
"""The guard must not false-positive: a consistent single-GPU stage flows past
it and reaches ``create_engine_config`` / ``Executor.get_class`` as usual."""
stage = _stage(1, devices="1")
@@ -212,18 +250,30 @@ def test_build_vllm_config_proceeds_on_consistent_layout():
model_config=types.SimpleNamespace(hf_config=types.SimpleNamespace()),
)
sentinel_executor = object()
- with (
- mock.patch.object(
- stage_init_utils.OmniEngineArgs, "create_engine_config", return_value=fake_config
- ) as create_engine_config,
- mock.patch.object(stage_init_utils.Executor, "get_class", return_value=sentinel_executor),
- mock.patch.object(stage_init_utils.OmniINCConfig, "maybe_upgrade", side_effect=lambda quant: quant),
- ):
- vllm_config, executor_class = build_vllm_config(
- stage,
- model="dummy-model",
- engine_args_dict={"tensor_parallel_size": 1},
- )
+ captured_api_process_config: dict[str, int] = {}
+
+ def capture_create_engine_config(engine_args, *args, **kwargs):
+ captured_api_process_config["count"] = engine_args._api_process_count
+ captured_api_process_config["rank"] = engine_args._api_process_rank
+ return fake_config
+
+ create_engine_config = mocker.patch.object(
+ stage_init_utils.OmniEngineArgs,
+ "create_engine_config",
+ autospec=True,
+ side_effect=capture_create_engine_config,
+ )
+ mocker.patch.object(stage_init_utils.Executor, "get_class", return_value=sentinel_executor)
+ mocker.patch.object(stage_init_utils.OmniINCConfig, "maybe_upgrade", side_effect=lambda quant: quant)
+
+ vllm_config, executor_class = build_vllm_config(
+ stage,
+ model="dummy-model",
+ engine_args_dict={"tensor_parallel_size": 1},
+ api_process_count=3,
+ api_process_rank=2,
+ )
create_engine_config.assert_called_once()
assert vllm_config is fake_config
assert executor_class is sentinel_executor
+ assert captured_api_process_config == {"count": 3, "rank": 2}
diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py
index 7ac5b6b69..6650c8e45 100644
--- a/tests/engine/test_stage_engine_args.py
+++ b/tests/engine/test_stage_engine_args.py
@@ -95,6 +95,7 @@ _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset(
"num_weight_load_threads",
"omni_kv_config",
"parallel_config",
+ "pooling_output_decoder",
"silence_ban_frames",
"subtalker_sampling_params",
"task_type",
@@ -366,6 +367,21 @@ def test_typed_llm_projection_does_not_emit_inherited_upstream_defaults():
assert inherited_defaults.isdisjoint(engine_args)
+def test_typed_llm_projection_omits_diffusion_only_and_process_only_defaults():
+ stage_config = VllmOmniARStageConfig(
+ stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"),
+ )
+
+ engine_args = stage_init_utils._project_omni_stage_engine_args(stage_config)
+
+ assert {
+ "enable_multithread_weight_load",
+ "num_weight_load_threads",
+ "disable_autocast",
+ "log_level",
+ }.isdisjoint(engine_args)
+
+
def test_mammoth_fp8_kv_deploy_projects_only_ar_stage(monkeypatch):
monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda _engine_args: None)
diff --git a/tests/entrypoints/openai_api/test_api_server_guards.py b/tests/entrypoints/openai_api/test_api_server_guards.py
index 37e2dc430..ebc08662f 100644
--- a/tests/entrypoints/openai_api/test_api_server_guards.py
+++ b/tests/entrypoints/openai_api/test_api_server_guards.py
@@ -42,7 +42,10 @@ from __future__ import annotations
import asyncio
import json
+import signal
+from argparse import Namespace
from contextlib import asynccontextmanager
+from dataclasses import dataclass
from types import SimpleNamespace
import pytest
@@ -59,11 +62,78 @@ from vllm_omni.entrypoints.serve.utils import errors as serve_errors
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+@dataclass
+class _FakeParallelConfig:
+ _api_process_rank: int
+
+
+@dataclass
+class _FakeVllmConfig:
+ parallel_config: _FakeParallelConfig
+
+
+def test_omni_api_worker_sets_parent_death_signal_before_validation(mocker):
+ args = _minimal_args(_omni_stage_client_configs=[{"stage_addresses": {}}])
+ death_signal = mocker.patch.object(api_server, "set_death_signal")
+
+ with pytest.raises(RuntimeError, match="client configuration"):
+ api_server.run_omni_api_server_worker_proc(
+ "127.0.0.1:8000",
+ mocker.Mock(),
+ args,
+ {"client_count": 2, "client_index": -1},
+ )
+
+ death_signal.assert_called_once_with(signal.SIGTERM)
+
+
+@pytest.mark.asyncio
+async def test_omni_run_server_worker_forwards_client_config(monkeypatch):
+ captured: dict[str, object] = {}
+
+ @asynccontextmanager
+ async def fake_build_async_omni(*_args, **kwargs):
+ captured.update(kwargs)
+ yield _FakeEngineClient(vllm_config=_FakeVllmConfig(_FakeParallelConfig(_api_process_rank=1)))
+
+ monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni)
+ monkeypatch.setattr(api_server, "build_openai_app", lambda *_args: FastAPI())
+ monkeypatch.setattr(api_server, "omni_init_app_state", lambda *_args: asyncio.sleep(0))
+ monkeypatch.setattr(api_server, "shutdown_unsupported_routes", lambda *_args: None)
+
+ async def fake_storage_start() -> None:
+ return None
+
+ monkeypatch.setattr(api_server.STORAGE_MANAGER, "start", fake_storage_start)
+
+ async def fake_serve_http(app, **_kwargs):
+ assert app.state.api_server_count == 2
+ task = asyncio.create_task(asyncio.sleep(0))
+ await asyncio.sleep(0)
+ return task
+
+ monkeypatch.setattr(api_server, "serve_http", fake_serve_http)
+
+ args = _minimal_args(model="dummy")
+ config = {"client_count": 2, "client_index": 1, "stage_addresses": {}}
+
+ await api_server.omni_run_server_worker("127.0.0.1:8000", _FakeSocket(), args, config)
+ assert captured["client_config"] is config
+
+
# FastAPI/Starlette auto-adds HEAD on every GET route. Including HEAD here
# would invent ("HEAD", path) entries for every GET and fail the current census.
# TODO: extend this set and the expected census if @router.head / @router.options appear.
_HTTP_METHODS = {"GET", "POST", "DELETE", "PUT", "PATCH"}
+
+def test_profiler_endpoints_detect_typed_stage_config():
+ stage = SimpleNamespace(profiler_config=SimpleNamespace(profiler="torch"))
+
+ assert api_server._should_enable_profiler_endpoints([stage])
+
+
# Exact Omni router census on main. Update intentionally when routes change.
_EXPECTED_ROUTER_ROUTES = {
("POST", "/v1/chat/completions"),
@@ -619,6 +689,48 @@ async def test_realtime_route_defaults_to_configured_duplex_handler(
assert calls == [expected_handler]
+@pytest.mark.parametrize("path", ["/v1/duplex", "/v1/realtime", "/v1/realtime?duplex=1"])
+def test_multi_api_duplex_reconnect_rejected_on_each_worker(path: str) -> None:
+ class SessionHandler:
+ async def handle_realtime_session(self, _websocket) -> None:
+ pytest.fail("Multi-API requests must not access process-local sessions")
+
+ class PendingWarmup:
+ def is_set(self) -> bool:
+ raise AssertionError("Multi-API rejection must not wait for duplex warmup")
+
+ # Separate apps represent workers with independent frontend registries.
+ # Opening on A then reconnecting to B must reject at the same boundary.
+ for worker_index in (0, 1):
+ app = FastAPI()
+ app.state.api_server_count = 2
+ app.state.api_server_index = worker_index
+ app.state.openai_serving_duplex = SessionHandler()
+ app.state.duplex_warmup_done = PendingWarmup()
+ app.include_router(api_server.router)
+ with TestClient(app) as client, client.websocket_connect(path) as websocket:
+ assert websocket.receive_json()["code"] == "multi_api_duplex_unsupported"
+ with pytest.raises(WebSocketDisconnect) as exc:
+ websocket.receive_text()
+ assert exc.value.code == 1008
+
+
+@pytest.mark.parametrize("path", ["/v1/duplex", "/v1/realtime?duplex=1"])
+def test_single_api_duplex_session_still_reaches_handler(path: str) -> None:
+ class SessionHandler:
+ async def handle_realtime_session(self, websocket) -> None:
+ await websocket.accept()
+ await websocket.send_json({"type": "session.created"})
+ await websocket.close()
+
+ app = FastAPI()
+ app.state.api_server_count = 1
+ app.state.openai_serving_duplex = SessionHandler()
+ app.include_router(api_server.router)
+ with TestClient(app) as client, client.websocket_connect(path) as websocket:
+ assert websocket.receive_json() == {"type": "session.created"}
+
+
def test_health_without_engine_returns_stable_unhealthy_response() -> None:
"""Lock ``/health`` with no engine initialized.
@@ -708,6 +820,38 @@ def test_speech_without_handler_preserves_not_found_http_error() -> None:
assert exc_info.value.detail == "The model does not support Speech API"
+@pytest.mark.asyncio
+async def test_multi_api_rejects_runtime_voice_upload() -> None:
+ app = FastAPI()
+ app.state.api_server_count = 2
+ raw_request = _request_for(app, method="POST", path="/v1/audio/voices")
+
+ with pytest.raises(HTTPException) as exc_info:
+ await api_server.upload_voice(
+ raw_request,
+ audio_sample=None,
+ speaker_embedding=None,
+ consent="consent-id",
+ name="probe",
+ )
+
+ assert exc_info.value.status_code == 409
+ assert "process-local frontend state" in exc_info.value.detail
+
+
+@pytest.mark.asyncio
+async def test_multi_api_rejects_sleep_route() -> None:
+ app = FastAPI()
+ app.state.api_server_count = 2
+ raw_request = _request_for(app, method="POST", path="/v1/omni/sleep")
+
+ with pytest.raises(HTTPException) as exc_info:
+ await api_server.omni_sleep(api_server.OmniSleepRequest(stage_ids=[0]), raw_request)
+
+ assert exc_info.value.status_code == 409
+ assert "process-local frontend state" in exc_info.value.detail
+
+
def test_engine_error_json_response_includes_request_and_stage_fields(monkeypatch) -> None:
"""Lock shared ``EngineGenerateError`` JSON fields used by Omni handlers.
@@ -930,3 +1074,57 @@ async def test_multistage_app_state_key_snapshot(monkeypatch) -> None:
must_be_wired=_MULTISTAGE_MUST_BE_WIRED,
must_be_none=_MULTISTAGE_MUST_BE_NONE,
)
+
+
+@pytest.mark.parametrize("count", [None, 0, -1, "2", True])
+def test_process_local_guard_rejects_unknown_topology(count):
+ app = FastAPI()
+ if count is not None:
+ app.state.api_server_count = count
+ # CLI args cannot substitute for missing/invalid actual runtime state.
+ app.state.args = Namespace(api_server_count=1)
+ with pytest.raises(HTTPException) as exc:
+ api_server._reject_process_local_state_with_multiple_api_workers(_request_for(app), "Sleep")
+ assert exc.value.status_code == 503
+
+
+@pytest.mark.parametrize(
+ "requested,config,expected",
+ [(None, None, 1), (1, None, 1), (2, {"client_count": 2}, 2), (None, {"client_count": 2}, 2)],
+)
+def test_api_topology_resolved_from_launch(requested, config, expected):
+ assert api_server._resolve_api_server_count(Namespace(api_server_count=requested), config) == expected
+
+
+@pytest.mark.parametrize("requested,config", [(2, None), (2, {}), (2, {"client_count": 1}), (2, {"client_count": "2"})])
+def test_api_topology_rejects_inconsistent_launch(requested, config):
+ with pytest.raises(ValueError):
+ api_server._resolve_api_server_count(Namespace(api_server_count=requested), config)
+
+
+@pytest.mark.parametrize("count,status", [(2, 409), (None, 503)])
+@pytest.mark.parametrize(
+ "method,path",
+ [
+ ("POST", "/v1/videos"),
+ ("GET", "/v1/videos"),
+ ("GET", "/v1/videos/job"),
+ ("GET", "/v1/videos/job/content"),
+ ("DELETE", "/v1/videos/job"),
+ ],
+)
+def test_video_lifecycle_guard_precedes_processing(monkeypatch, count, status, method, path):
+ from fastapi.testclient import TestClient
+
+ app = FastAPI()
+ if count is not None:
+ app.state.api_server_count = count
+ app.include_router(api_server.router)
+
+ async def forbidden_form():
+ pytest.fail("video form processing must not run before topology rejection")
+
+ app.dependency_overrides[api_server._parse_video_form] = forbidden_form
+ with TestClient(app) as client:
+ response = client.request(method, path)
+ assert response.status_code == status
diff --git a/tests/entrypoints/openai_api/test_audex_serving_guards.py b/tests/entrypoints/openai_api/test_audex_serving_guards.py
index 87524825e..141dd8cac 100644
--- a/tests/entrypoints/openai_api/test_audex_serving_guards.py
+++ b/tests/entrypoints/openai_api/test_audex_serving_guards.py
@@ -15,6 +15,7 @@ from types import SimpleNamespace
import pytest
import torch
+from vllm_omni.config.stage_config import StagePipelineConfig
from vllm_omni.entrypoints.openai.serving_speech import (
_AUDEX_NO_AUDIO_GUARD_MODEL_TYPES,
OmniOpenAIServingSpeech,
@@ -251,9 +252,28 @@ class TestAudexStageDetection:
stage = serving._find_tts_stage()
assert stage is not None and stage.engine_args.model_stage == "audex_thinker"
+ def test_typed_pipeline_detected_without_legacy_engine_args(self):
+ serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech)
+ stages = [
+ SimpleNamespace(
+ stage_pipeline_config=StagePipelineConfig(stage_id=index, model_stage=model_stage),
+ model_config=SimpleNamespace(model_arch=None),
+ worker_type="generation",
+ )
+ for index, model_stage in enumerate(("audex_thinker", "audex_code2wav"))
+ ]
+ serving.engine_client = SimpleNamespace(stage_configs=stages)
+
+ stage = serving._find_tts_stage()
+ serving._tts_stage = stage
+
+ assert stage is stages[0]
+ assert serving._detect_tts_model_type() == "audex"
+
def test_guided_injection_sets_guided_temperature():
"""Guided requests decode at the measured CFG temperature (0.05)."""
+
serving = TestAudexCfgInjection()._serving_with_fake_tokenizer()
cond = TestAudexCfgInjection()._cond_prompt()
params = SimpleNamespace(extra_args={"cfg_scale": 1.5}, temperature=0.1)
@@ -265,6 +285,31 @@ def test_guided_injection_sets_guided_temperature():
assert out.temperature == 0.1
+def test_typed_full_pipeline_uses_full_checkpoint_tokenizer(mocker):
+ from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext
+
+ serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech)
+ serving._tts_stage = SimpleNamespace(
+ stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="audex_omni"),
+ model_config=SimpleNamespace(model_arch=None),
+ worker_type="ar",
+ )
+ serving.engine_client = SimpleNamespace(
+ model_config=SimpleNamespace(model="/models/audex"),
+ )
+ ensure_snapshot = mocker.patch(
+ "vllm_omni.model_executor.models.audex.checkpoint.ensure_audex_snapshot",
+ return_value="/snapshots/audex",
+ )
+ from_pretrained = mocker.patch("transformers.AutoTokenizer.from_pretrained", return_value=object())
+
+ adapter = AudexAdapter(SpeechServingContext(server=serving, engine_client=serving.engine_client))
+ adapter._get_tokenizer()
+
+ ensure_snapshot.assert_called_once_with("/models/audex", profile="full")
+ from_pretrained.assert_called_once_with("/snapshots/audex/checkpoint_folder_full")
+
+
# ---------------------------------------------------------------- TTA serving surface
from vllm_omni.entrypoints.openai.tts_adapters import resolve_adapter # noqa: E402
diff --git a/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py b/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py
index 615248fcf..ba5925d8f 100644
--- a/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py
+++ b/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py
@@ -20,6 +20,7 @@ def _make_app(engine_client):
from vllm_omni.entrypoints.openai.api_server import router
app = FastAPI()
+ app.state.api_server_count = 1
app.include_router(router)
app.state.engine_client = engine_client
app.state.sleeping_stages = set()
@@ -253,6 +254,7 @@ def pure_diffusion_app(pure_diffusion_engine, mocker):
)
app = FastAPI()
+ app.state.api_server_count = 1
app.include_router(router)
mocker.patch(
diff --git a/tests/entrypoints/openai_api/test_openpi_serving.py b/tests/entrypoints/openai_api/test_openpi_serving.py
index 46ff0f822..d274eadda 100644
--- a/tests/entrypoints/openai_api/test_openpi_serving.py
+++ b/tests/entrypoints/openai_api/test_openpi_serving.py
@@ -138,6 +138,36 @@ def test_policy_server_config_reads_stage_config_model_config():
assert serving.policy_server_config.to_dict() == policy_config
+@pytest.mark.parametrize("inline", [False, True], ids=["out_of_process", "inline"])
+def test_policy_server_config_reads_typed_runtime(inline):
+ from tests.helpers.stage_config import get_deploy_config_path
+ from vllm_omni.config.omni_config import VllmOmniConfig
+ from vllm_omni.diffusion.models.pi0_pipeline_config import PI0_PIPELINE
+ from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+ typed_stage = VllmOmniConfig.from_pipeline_config(
+ PI0_PIPELINE, deploy_config_path=get_deploy_config_path("pi0.yaml")
+ ).stage_configs[0]
+ policy_config = typed_stage.diffusion_config.model_config["policy_server_config"]
+ stage_client = SimpleNamespace(stage_type="diffusion")
+ if inline:
+ # Worker-enriched inline metadata remains authoritative over the
+ # startup stage config; only the process client needs the fallback.
+ policy_config = {**policy_config, "worker_metadata": "loaded"}
+ stage_client.od_config = SimpleNamespace(model_config={"policy_server_config": policy_config})
+ engine_client = AsyncOmni.__new__(AsyncOmni)
+ engine_client.engine = SimpleNamespace(
+ stage_clients=[stage_client],
+ stage_configs=[typed_stage],
+ stage_vllm_configs=[None],
+ get_diffusion_od_config=lambda: SimpleNamespace(model_class_name="Pi0Pipeline"),
+ )
+
+ serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine_client)
+
+ assert serving.policy_server_config.to_dict() == policy_config
+
+
def test_policy_server_config_reads_omegaconf_stage_config():
engine_client = SimpleNamespace(
get_diffusion_od_config=lambda: None,
diff --git a/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py b/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py
index 59d070a39..417f457d3 100644
--- a/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py
+++ b/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py
@@ -27,7 +27,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@dataclass
class _EngineClientStub:
- stage_configs: list[dict[str, object]]
+ stage_configs: list[object]
@dataclass
@@ -149,6 +149,37 @@ async def test_chat_reference_audio_reaches_code2wav_transport(
assert code2wav_payload.meta.ref_audio_sr == 16000
+@pytest.mark.asyncio
+async def test_typed_stage_enables_reference_audio(
+ serving_chat: OmniOpenAIServingChat,
+ mocker: MockerFixture,
+) -> None:
+ reference_audio_url = "data:audio/wav;base64,AAAA"
+ reference_waveform = np.array([0.25, -0.5], dtype=np.float32)
+ serving_chat.engine_client = _EngineClientStub(
+ [
+ SimpleNamespace(
+ model_config=SimpleNamespace(model_arch="MiniCPMO45OmniForConditionalGeneration"),
+ model_stage="llm",
+ )
+ ]
+ )
+ serving_chat.model_config = _ModelConfigStub()
+ engine_prompt: TokensPrompt = {"type": "tokens", "prompt_token_ids": [1]}
+ media_connector = mocker.patch(
+ "vllm_omni.entrypoints.openai.serving_chat.MediaConnector",
+ ).return_value
+ media_connector.fetch_audio_async = mocker.AsyncMock(return_value=(reference_waveform, 16000))
+
+ await serving_chat._attach_minicpmo45_reference_audio(
+ engine_prompt,
+ _chat_request(reference_audio_url),
+ )
+
+ media_connector.fetch_audio_async.assert_awaited_once_with(reference_audio_url)
+ assert MINICPMO45_REFERENCE_AUDIO_KEY in cast(dict[str, object], engine_prompt)
+
+
@pytest.mark.parametrize("request_shape", ["extra_body", "model_extra"])
@pytest.mark.asyncio
async def test_reference_audio_accepts_request_extra_mappings(
diff --git a/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_stream_output_kinds.py b/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_stream_output_kinds.py
index 1a3a4f22e..5ce129d15 100644
--- a/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_stream_output_kinds.py
+++ b/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_stream_output_kinds.py
@@ -1,4 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Regression tests for MiniCPM-o 4.5 audio-streaming output-kind coercion.
Chat ``stream=true`` coerces every AR stage to DELTA. For MiniCPM-o 4.5 that
@@ -33,6 +34,13 @@ def _stage(model_arch: str, model_stage: str) -> SimpleNamespace:
return SimpleNamespace(engine_args=SimpleNamespace(model_arch=model_arch, model_stage=model_stage))
+def _typed_stage(model_arch: str, model_stage: str) -> SimpleNamespace:
+ return SimpleNamespace(
+ model_config=SimpleNamespace(model_arch=model_arch),
+ model_stage=model_stage,
+ )
+
+
def _minicpmo45_stages():
return [_stage(MINICPMO45_ARCH, "llm"), _stage(MINICPMO45_ARCH, "tts")]
@@ -50,6 +58,20 @@ def test_thinker_forced_final_only_talker_stays_delta(serving_chat):
assert out[1].output_kind is RequestOutputKind.DELTA
+def test_typed_stages_force_thinker_final_only(serving_chat):
+ serving_chat.engine_client = SimpleNamespace(
+ stage_configs=[
+ _typed_stage(MINICPMO45_ARCH, "llm"),
+ _typed_stage(MINICPMO45_ARCH, "tts"),
+ ]
+ )
+
+ out = serving_chat._fix_minicpmo45_audio_stream_output_kinds(_delta_params(), ["text", "audio"])
+
+ assert out[0].output_kind is RequestOutputKind.FINAL_ONLY
+ assert out[1].output_kind is RequestOutputKind.DELTA
+
+
@pytest.mark.parametrize("modalities", [None, [], ["text"]])
def test_non_audio_request_is_untouched(serving_chat, modalities):
serving_chat.engine_client = SimpleNamespace(stage_configs=_minicpmo45_stages())
diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py
index f8dfaafb6..7cdeebc39 100644
--- a/tests/entrypoints/openai_api/test_serving_speech.py
+++ b/tests/entrypoints/openai_api/test_serving_speech.py
@@ -27,6 +27,7 @@ from pydantic import ValidationError
from pytest_mock import MockerFixture
from vllm.entrypoints.serve.engine.protocol import ErrorInfo, ErrorResponse
+from vllm_omni.config.stage_config import StagePipelineConfig
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler
from vllm_omni.diffusion.sched.step_scheduler import StepScheduler
@@ -1039,6 +1040,38 @@ class TestTTSMethods:
asyncio.run(server._prepare_speech_generation(request))
server.shutdown()
+ def test_prepare_speech_rejects_typed_non_tts_omni_model(self, mocker: MockerFixture):
+ """Typed Qwen3-Omni stages must hit the same talker guard as legacy stages."""
+ mock_engine_client = mocker.MagicMock()
+ mock_engine_client.errored = False
+ mock_engine_client.tts_max_instructions_length = None
+ mock_engine_client.stage_configs = [
+ SimpleNamespace(
+ stage_pipeline_config=StagePipelineConfig(stage_id=index, model_stage=model_stage),
+ model_config=SimpleNamespace(model_arch=None),
+ worker_type=worker_type,
+ )
+ for index, (model_stage, worker_type) in enumerate(
+ (
+ ("thinker", "ar"),
+ ("talker", "ar"),
+ ("code2wav", "generation"),
+ )
+ )
+ ]
+
+ mock_models = mocker.MagicMock()
+ mock_models.is_base_model.return_value = True
+ server = OmniOpenAIServingSpeech(
+ engine_client=mock_engine_client,
+ models=mock_models,
+ request_logger=mocker.MagicMock(),
+ )
+ request = OpenAICreateSpeechRequest(input="Hello world")
+ with pytest.raises(ValueError, match="only supported for dedicated TTS models"):
+ asyncio.run(server._prepare_speech_generation(request))
+ server.shutdown()
+
def test_estimate_prompt_len_fallback(self, speech_server):
"""Test prompt length estimation falls back to 2048 when model is unavailable."""
tts_params = {"text": ["Hello"], "task_type": ["CustomVoice"]}
@@ -3793,6 +3826,7 @@ def test_api_server_create_speech_wraps_error_response_status(mocker: MockerFixt
def _make_api_server_request(handler, *, method: str = "POST", path: str = "/v1/audio/voices") -> Request:
app = FastAPI()
app.state.openai_serving_speech = handler
+ app.state.api_server_count = 1
scope = {
"type": "http",
"app": app,
diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py
index e0a165a55..500feab35 100644
--- a/tests/entrypoints/openai_api/test_tts_adapter.py
+++ b/tests/entrypoints/openai_api/test_tts_adapter.py
@@ -22,6 +22,7 @@ from vllm_omni.entrypoints.openai.tts_adapters import (
detect_tts_model_type,
resolve_adapter,
)
+from vllm_omni.entrypoints.openai.tts_adapters.base import resolve_stage_model_path
from vllm_omni.entrypoints.openai.tts_adapters.covo_audio import CovoAudioAdapter
from vllm_omni.entrypoints.openai.tts_adapters.higgs_audio_v2 import HiggsAudioV2Adapter
from vllm_omni.entrypoints.openai.tts_adapters.indextts2 import (
@@ -99,6 +100,31 @@ def test_resolve_unknown_returns_none():
assert resolve_adapter(None) is None
+def test_stage_model_path_prefers_typed_override():
+ engine_client = SimpleNamespace(
+ stage_configs=[
+ SimpleNamespace(engine_args=SimpleNamespace(model="legacy-stage-model")),
+ SimpleNamespace(
+ model_config=SimpleNamespace(model="typed-stage-model"),
+ ),
+ ],
+ model="served-model",
+ )
+
+ assert resolve_stage_model_path(engine_client) == "typed-stage-model"
+
+
+def test_stage_model_path_falls_back_to_legacy_then_served_model():
+ legacy_client = SimpleNamespace(
+ stage_configs=[SimpleNamespace(engine_args=SimpleNamespace(model="legacy-stage-model"))],
+ model="served-model",
+ )
+ served_client = SimpleNamespace(stage_configs=[SimpleNamespace()], model="served-model")
+
+ assert resolve_stage_model_path(legacy_client) == "legacy-stage-model"
+ assert resolve_stage_model_path(served_client) == "served-model"
+
+
def test_voxcpm2_resolves():
"""VoxCPM2 (the served ``latent_generator`` model) resolves cleanly.
diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py
index 2de5c4056..715c7ca04 100644
--- a/tests/entrypoints/openai_api/test_video_server.py
+++ b/tests/entrypoints/openai_api/test_video_server.py
@@ -463,6 +463,7 @@ async def test_server_worker_keeps_engine_alive_until_http_shutdown(monkeypatch)
@pytest.fixture
def test_client():
app = FastAPI()
+ app.state.api_server_count = 1
app.include_router(router)
app.state.openai_serving_video = OmniOpenAIServingVideo.for_diffusion(
diffusion_engine=FakeAsyncOmni(),
@@ -1134,6 +1135,32 @@ def test_mixed_reference_capability_uses_model_metadata_when_config_defaults_fal
assert handler.supports_mixed_reference_inputs
+def test_typed_stage_drives_video_capability_checks():
+ from vllm_omni.config.config_factory import StageConfigFactory
+
+ minimax_stage = StageConfigFactory.create_typed_default_diffusion(
+ "minimax-h3",
+ {"model_class_name": "MiniMaxH3Pipeline"},
+ ).stage_configs[0]
+ minimax_handler = OmniOpenAIServingVideo.for_diffusion(
+ SimpleNamespace(od_config=SimpleNamespace(model_class_name=None)),
+ model_name="minimax-h3",
+ stage_configs=[minimax_stage],
+ )
+ assert minimax_handler.supports_mixed_reference_inputs
+
+ cosmos_stage = StageConfigFactory.create_typed_default_diffusion(
+ "cosmos3",
+ {"model_class_name": "Cosmos3OmniDiffusersPipeline"},
+ ).stage_configs[0]
+ cosmos_handler = OmniOpenAIServingVideo.for_diffusion(
+ SimpleNamespace(od_config=SimpleNamespace(model_class_name=None)),
+ model_name="cosmos3",
+ stage_configs=[cosmos_stage],
+ )
+ assert cosmos_handler.supported_control_upload_types == frozenset({"edge", "blur", "depth", "seg", "wsm"})
+
+
@pytest.mark.parametrize("model_class_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"])
def test_control_upload_capability_is_declared_only_by_cosmos3(test_client, model_class_name):
handler = test_client.app.state.openai_serving_video
@@ -1191,6 +1218,35 @@ def test_cosmos3_reference_video_limit_uses_v2v_condition_frames():
assert spec.keep == "first"
+@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"])
+@pytest.mark.parametrize(
+ ("num_frames", "extra_params", "expected"),
+ [
+ (189, {"condition_frame_indexes_vision": [0, 2], "condition_video_keep": "last"}, (9, "last")),
+ (189, {"condition_frame_indexes_vision": [0, 2]}, (9, "first")),
+ (5, {"condition_frame_indexes_vision": [0, 20]}, (5, "first")),
+ (None, {"action_mode": "inverse_dynamics", "action_chunk_size": 16}, (17, "first")),
+ ],
+)
+def test_cosmos3_reference_video_decode_policy_with_runtime_configs(typed, num_frames, extra_params, expected):
+ from vllm_omni.config.config_factory import StageConfigFactory
+
+ if typed:
+ stages = list(
+ StageConfigFactory.create_typed_default_diffusion(
+ "cosmos3",
+ {"model_class_name": "Cosmos3OmniDiffusersPipeline"},
+ ).stage_configs
+ )
+ else:
+ stages = _cosmos3_stage_configs()
+ request = VideoGenerationRequest(prompt="Continue this motion.", num_frames=num_frames, extra_params=extra_params)
+
+ spec = _reference_video_decode_spec(request, stages)
+
+ assert (spec.max_frames, spec.keep) == expected
+
+
def test_cosmos3_reference_video_limit_preserves_action_frames():
request = VideoGenerationRequest(
prompt="Predict the action.",
@@ -1749,6 +1805,7 @@ def test_action_extraction_accepts_multimodal_actions_payload():
def test_missing_handler_returns_503():
app = FastAPI()
+ app.state.api_server_count = 1
app.include_router(router)
app.state.openai_serving_video = None
client = TestClient(app)
@@ -2735,6 +2792,7 @@ def test_cosmos3_control_upload_rejects_invalid_size(control_bytes, message, tes
def test_sync_missing_handler_returns_503():
app = FastAPI()
+ app.state.api_server_count = 1
app.include_router(router)
app.state.openai_serving_video = None
client = TestClient(app)
diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py
index 3b378930a..9b99af712 100644
--- a/tests/entrypoints/test_async_omni_diffusion_config.py
+++ b/tests/entrypoints/test_async_omni_diffusion_config.py
@@ -2,13 +2,15 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
+from typing import Any
import pytest
import torch
from pydantic import ValidationError
from vllm_omni.config.config_factory import StageConfigFactory
-from vllm_omni.config.resolver import OmniConfigResolution
+from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig
+from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config
from vllm_omni.diffusion.data import AttentionConfig, OmniDiffusionConfig
from vllm_omni.engine import stage_init_utils
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
@@ -117,7 +119,7 @@ def test_stage_override_preserves_model_extras_for_default_diffusion_stage(mocke
trust_remote_code=False,
)
- assert stage_configs[0]["engine_args"]["extras"]["ltx2_use_conv_vae"] is True
+ assert stage_configs[0].diffusion_config.extras["ltx2_use_conv_vae"] is True
def test_default_stage_rejects_unknown_nested_parallel_config_key():
@@ -239,6 +241,33 @@ def test_default_stage_config_includes_default_sampling_params():
}
+@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"])
+@pytest.mark.parametrize("sampling_defaults", [{"0": {"guidance_scale": 7.5}}, '{"0":{"guidance_scale":7.5}}'])
+def test_generic_diffusion_sampling_defaults_remain_overridable(typed, sampling_defaults):
+ from vllm_omni.config.yaml_util import create_config
+ from vllm_omni.entrypoints.omni_base import OmniBase
+ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+ kwargs = {"model_class_name": "QwenImagePipeline", "default_sampling_params": sampling_defaults}
+ if typed:
+ stage = StageConfigFactory.create_typed_default_diffusion("generic-diffusion", kwargs).stage_configs[0]
+ metadata = stage_init_utils.extract_stage_metadata_from_omni_stage_config(stage)
+ else:
+ stage = create_config(StageConfigFactory.create_default_diffusion(kwargs))[0]
+ metadata = stage_init_utils.extract_legacy_stage_metadata(stage)
+ base = OmniBase.__new__(OmniBase)
+ base.engine = SimpleNamespace(num_stages=1, stage_configs=[stage])
+ base.default_sampling_params_list = [metadata.default_sampling_params]
+ base.sampling_constraints_list = base._get_sampling_constraints_list([stage])
+
+ assert base.resolve_sampling_params_list(None)[0].guidance_scale == 7.5
+ requested = OmniDiffusionSamplingParams(guidance_scale=2.0)
+ assert base.resolve_sampling_params_list(requested)[0].guidance_scale == 2.0
+ assert requested.guidance_scale == 2.0
+ assert base.default_sampling_params_list[0].guidance_scale == 7.5
+ assert base.sampling_constraints_list == [{}]
+
+
def test_default_stage_config_includes_diffusion_attention_backend():
"""Ensure diffusion attention shorthand lands in engine_args.diffusion_attention_config."""
stage_cfg = StageConfigFactory.create_default_diffusion(
@@ -726,3 +755,106 @@ def test_default_stage_config_includes_quantization_config():
stage_cfg = StageConfigFactory.create_default_diffusion({"quantization_config": quantization_config})[0]
assert stage_cfg["engine_args"]["quantization_config"] == quantization_config
+
+
+@pytest.mark.parametrize("model_class_name", ["HeliosPipeline", "HunyuanVideo15Pipeline"])
+def test_generic_diffusion_uses_canonical_video_output_type(model_class_name):
+ config = StageConfigFactory.create_typed_default_diffusion(
+ "generic-video",
+ {"model_class_name": model_class_name},
+ )
+
+ assert config.stage_configs[0].final_output_type == "video"
+
+
+def test_generic_diffusion_resolves_structured_stage_without_legacy_conversion(mocker):
+ """Generic diffusion reaches runtime as the structured stage itself."""
+ mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None)
+ mocker.patch(
+ "vllm_omni.config.resolver._resolve_generic_diffusion_model_class",
+ return_value=(True, "FakeDiffusionPipeline"),
+ )
+
+ resolved = resolve_omni_config(
+ "generic-diffusion",
+ trust_remote_code=False,
+ deploy_config_path=None,
+ cli_overrides={
+ "num_gpus": 4,
+ "tensor_parallel_size": 2,
+ "default_sampling_params": '{"0": {"guidance_scale": 7.5}}',
+ },
+ stage_overrides=None,
+ strategy_config_path=None,
+ )
+
+ assert resolved.pipeline_config is not None
+ assert resolved.pipeline_config.model_type == "generic_diffusion"
+ assert len(resolved.stage_configs) == 1
+ stage = resolved.stage_configs[0]
+ assert isinstance(stage, VllmOmniDiffusionStageConfig)
+ assert stage.model_config.model == "generic-diffusion"
+ assert stage.model_config.default_sampling_params == {"guidance_scale": 7.5}
+ assert stage.parallel_config.tensor_parallel_size == 2
+ assert stage.parallel_config.data_parallel_size == 2
+ assert stage.parallel_config.world_size == 4
+ assert stage.runtime_config.devices == "0,1,2,3"
+
+
+def test_generic_diffusion_structured_stage_reaches_standard_startup(mocker):
+ """Standard runtime resolves and starts the typed stage through the real launcher."""
+ from vllm_omni.engine import stage_engine_startup as startup_module
+ from vllm_omni.engine import stage_runtime as runtime_module
+ from vllm_omni.engine.stage_runtime import StageRuntime
+
+ mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None)
+ mocker.patch(
+ "vllm_omni.config.resolver._resolve_generic_diffusion_model_class",
+ return_value=(True, "FakeDiffusionPipeline"),
+ )
+ resolved = resolve_omni_config(
+ "generic-diffusion",
+ trust_remote_code=False,
+ deploy_config_path=None,
+ cli_overrides={"num_gpus": 1},
+ stage_overrides=None,
+ strategy_config_path=None,
+ )
+ stage = resolved.stage_configs[0]
+ launched: dict[str, Any] = {}
+ client = SimpleNamespace(input_address=None, shutdown=mocker.Mock())
+
+ mocker.patch.object(runtime_module, "prepare_engine_environment")
+ mocker.patch.object(runtime_module, "load_omni_transfer_config_for_model", return_value=None)
+ mocker.patch.object(runtime_module, "get_stage_connector_spec", return_value={})
+ mocker.patch.object(runtime_module, "resolve_omni_kv_config_for_stage", return_value=(None, None, None))
+
+ def _initialize_typed_stage(stage_id, model, stage_config, metadata, **kwargs):
+ launched.update(
+ stage_id=stage_id,
+ model=model,
+ stage_config=stage_config,
+ metadata=metadata,
+ **kwargs,
+ )
+ return client
+
+ mocker.patch.object(startup_module, "initialize_diffusion_stage", side_effect=_initialize_typed_stage)
+
+ runtime = StageRuntime(
+ stage_configs=list(resolved.stage_configs),
+ model="generic-diffusion",
+ config_path="",
+ stage_init_timeout=10,
+ async_chunk=False,
+ )
+ runtime.initialize()
+
+ assert launched["stage_config"] is stage
+ assert isinstance(launched["stage_config"], VllmOmniDiffusionStageConfig)
+ assert launched["metadata"].stage_type == "diffusion"
+ assert launched["metadata"].model_stage == "diffusion"
+ assert launched["use_inline"] is True
+ assert launched["model"] == "generic-diffusion"
+ assert launched["stage_id"] == 0
+ assert runtime.stage_pools[0].clients == [client]
diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py
index ea57d07ab..969bbaf91 100644
--- a/tests/entrypoints/test_omni_entrypoints.py
+++ b/tests/entrypoints/test_omni_entrypoints.py
@@ -17,6 +17,8 @@ from vllm.entrypoints.openai.models.serving import OpenAIServingModels
from vllm.sampling_params import RequestOutputKind, SamplingParams
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
+from vllm_omni.config.omni_config import VllmOmniConfig
+from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
from vllm_omni.config.stage_config import StageConfig
from vllm_omni.engine.async_omni_engine import StageRuntimeInfo
from vllm_omni.engine.messages import ErrorMessage, OutputMessage
@@ -202,6 +204,49 @@ def _make_base():
return obj
+@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"])
+def test_resolve_sampling_params_list_preserves_stage_constraints(typed):
+ from vllm_omni.config.omni_config import VllmOmniARStageConfig
+ from vllm_omni.config.stage_config import StagePipelineConfig
+
+ base = _make_base()
+ base.engine.num_stages = 1
+ base.default_sampling_params_list = [SamplingParams(max_tokens=1000, detokenize=False, stop_token_ids=[42])]
+ if typed:
+ stage = VllmOmniARStageConfig(
+ stage_pipeline_config=StagePipelineConfig(
+ stage_id=0,
+ model_stage="dummy-model",
+ sampling_constraints={"detokenize": False, "stop_token_ids": [42]},
+ )
+ )
+ else:
+ stage = StageConfig(
+ stage_id=0,
+ model_stage="dummy-model",
+ sampling_constraints={"detokenize": False, "stop_token_ids": [42]},
+ ).to_omegaconf()
+ base.engine.stage_configs = [stage]
+ base.sampling_constraints_list = base._get_sampling_constraints_list(base.engine.stage_configs)
+ assert base.sampling_constraints_list == [{"detokenize": False, "stop_token_ids": [42]}]
+
+ resolved_defaults = base.resolve_sampling_params_list(None)
+ assert resolved_defaults[0].stop_token_ids == [42]
+
+ caller_params = SamplingParams(seed=1234, max_tokens=7, detokenize=True, stop_token_ids=[7])
+ resolved = base.resolve_sampling_params_list(caller_params)
+
+ assert resolved[0] is not caller_params
+ assert resolved[0].seed == 1234
+ assert resolved[0].max_tokens == 7
+ assert resolved[0].detokenize is False
+ assert resolved[0].stop_token_ids == [7, 42]
+ assert 7 in resolved[0]._all_stop_token_ids
+ assert 42 in resolved[0]._all_stop_token_ids
+ assert caller_params.detokenize is True
+ assert caller_params.stop_token_ids == [7]
+
+
def test_resolve_sampling_params_list_merges_required_stop_tokens():
base = _make_base()
base.engine.num_stages = 1
@@ -219,22 +264,13 @@ def test_resolve_sampling_params_list_merges_required_stop_tokens():
base.sampling_constraints_list = base._get_sampling_constraints_list(base.engine.stage_configs)
assert base.sampling_constraints_list == [{"detokenize": False, "stop_token_ids": required_stop_ids}]
- resolved_defaults = base.resolve_sampling_params_list(None)
- assert resolved_defaults[0].stop_token_ids == [99, 151704, 151645]
+def test_sampling_constraints_are_forwarded_by_typed_stage_configs():
+ config = VllmOmniConfig.from_pipeline_config(OMNI_PIPELINES["qwen3_tts"])
- caller_params = SamplingParams(seed=1234, max_tokens=7, detokenize=True, stop_token_ids=[100, 151704])
+ constraints = OmniBase._get_sampling_constraints_list(config.stage_configs)
- resolved = base.resolve_sampling_params_list(caller_params)
-
- assert resolved[0] is not caller_params
- assert resolved[0].seed == 1234
- assert resolved[0].max_tokens == 7
- assert resolved[0].detokenize is False
- assert resolved[0].stop_token_ids == [100, 151704, 151645]
- assert {100, 151704, 151645}.issubset(resolved[0]._all_stop_token_ids)
- assert caller_params.detokenize is True
- assert caller_params.stop_token_ids == [100, 151704]
+ assert constraints == [dict(stage.stage_pipeline_config.sampling_constraints) for stage in config.stage_configs]
@pytest.mark.parametrize("use_defaults", [False, True])
diff --git a/tests/entrypoints/test_resolve_dreamzero_config.py b/tests/entrypoints/test_resolve_dreamzero_config.py
index 450e1be39..3cfa0f0e5 100644
--- a/tests/entrypoints/test_resolve_dreamzero_config.py
+++ b/tests/entrypoints/test_resolve_dreamzero_config.py
@@ -32,15 +32,18 @@ def test_dreamzero_resolves_through_registry_with_model_defaults(monkeypatch):
stage_overrides=None,
strategy_config_path=None,
)
- engine_args = resolved.stage_configs[0].engine_args
+ stage = resolved.stage_configs[0]
assert resolved.config_path is not None
assert resolved.config_path.endswith("vllm_omni/deploy/dreamzero.yaml")
- assert engine_args.model_class_name == "DreamZeroPipeline"
- assert engine_args.model_config.policy_server_config.action_space == "joint_position"
+ assert stage.diffusion_config.model_class_name == "DreamZeroPipeline"
+ assert stage.diffusion_config.model_config["policy_server_config"]["action_space"] == "joint_position"
def test_dreamzero_enrich_config_preserves_explicit_model_class_name(monkeypatch):
+ def _get_hf_config(path, _model):
+ return None if path == "model_index.json" else {"model_type": "vla", "architectures": ["VLA"]}
+
monkeypatch.setattr(
"vllm.transformers_utils.config.get_hf_file_to_dict",
lambda path, _model, **_kwargs: (
diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py
index 1f3b58170..7f8b6a4a9 100644
--- a/tests/entrypoints/test_serve.py
+++ b/tests/entrypoints/test_serve.py
@@ -6,12 +6,19 @@
from __future__ import annotations
import argparse
+from dataclasses import dataclass
+from multiprocessing.reduction import ForkingPickler
from types import SimpleNamespace
+from typing import Any
import pytest
from pytest_mock import MockerFixture
+from vllm.v1.engine.utils import EngineZmqAddresses
+from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig
from vllm_omni.config.resolver import OmniConfigResolution
+from vllm_omni.engine.stage_engine_startup import StageReplicaResources
+from vllm_omni.engine.stage_runtime import StageEngineLaunch
from vllm_omni.entrypoints.cli.serve import (
OmniServeCommand,
_parse_stage_overrides,
@@ -23,6 +30,31 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingName
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+@dataclass(frozen=True)
+class _FakeProcess:
+ sentinel: str
+ exitcode: int | None
+ name: str
+ pid: int
+
+
+@dataclass
+class _FakeProcessOwner:
+ processes: list[_FakeProcess]
+
+
+@dataclass
+class _FakeStageEngineArgs:
+ async_chunk: bool
+ enable_sleep_mode: bool
+
+
+@dataclass
+class _FakeStageConfig:
+ stage_id: int
+ engine_args: _FakeStageEngineArgs
+
+
def _resolved(*stages: SimpleNamespace) -> OmniConfigResolution:
return OmniConfigResolution(config_path="/fake/stages.yaml", stage_configs=tuple(stages))
@@ -42,6 +74,8 @@ def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None:
explicit = args.get_explicit_kwargs_dict()
assert args.get_explicit_kwargs_dict()
assert not explicit["async_chunk"]
+ assert (args.api_server_count or 1) == 1
+ assert "api_server_count" not in explicit
def _parse_serve_args(argv: list[str]) -> TrackingNamespace:
@@ -167,6 +201,244 @@ def test_serve_parser_rejects_invalid_video_output_transport(value: str) -> None
parser.parse_args(["serve", "fake-model", "--omni", "--video-output-transport", value])
+def test_tracking_namespace_is_picklable_for_spawned_api_workers() -> None:
+ parser = TrackingArgumentParser()
+ subparsers = parser.add_subparsers(dest="subcommand")
+ OmniServeCommand().subparser_init(subparsers)
+ args = parser.parse_args(["serve", "fake-model", "--omni", "--api-server-count", "2"])
+ args._omni_stage_client_configs = [{"stage_addresses": {0: {0: {"input_address": "ipc://input"}}}}]
+
+ restored = ForkingPickler.loads(ForkingPickler.dumps(args))
+
+ assert restored.api_server_count == 2
+ assert restored.get_explicit_kwargs_dict()["api_server_count"] == 2
+ assert restored._omni_stage_client_configs == args._omni_stage_client_configs
+
+
+def test_serve_validate_rejects_multiple_api_servers_for_diffusion(mocker: MockerFixture) -> None:
+ parser = TrackingArgumentParser()
+ subparsers = parser.add_subparsers(dest="subcommand")
+ cmd = OmniServeCommand()
+ cmd.subparser_init(subparsers)
+ args = parser.parse_args(["serve", "fake-diffusion-model", "--omni", "--api-server-count", "2"])
+
+ mocker.patch("vllm_omni.diffusion.utils.hf_utils.is_diffusion_model", return_value=True)
+ validate = mocker.patch("vllm_omni.entrypoints.cli.serve.validate_parsed_serve_args")
+
+ with pytest.raises(ValueError, match="not supported for diffusion"):
+ cmd.validate(args)
+
+ validate.assert_not_called()
+
+
+def test_serve_validate_rejects_sleep_mode_with_multiple_api_servers(mocker: MockerFixture) -> None:
+ parser = TrackingArgumentParser()
+ subparsers = parser.add_subparsers(dest="subcommand")
+ cmd = OmniServeCommand()
+ cmd.subparser_init(subparsers)
+ args = parser.parse_args(["serve", "fake-model", "--omni", "--api-server-count", "2", "--enable-sleep-mode"])
+
+ mocker.patch("vllm_omni.diffusion.utils.hf_utils.is_diffusion_model", return_value=False)
+ mocker.patch("vllm_omni.entrypoints.cli.serve.validate_parsed_serve_args")
+
+ with pytest.raises(ValueError, match="enable-sleep-mode"):
+ cmd.validate(args)
+
+
+@pytest.mark.parametrize("typed", [False, True])
+@pytest.mark.parametrize("downstream_async_chunk", [False, True])
+def test_build_multi_api_stage_runtime_matches_current_constructor(
+ mocker: MockerFixture, typed: bool, downstream_async_chunk: bool
+) -> None:
+ from vllm_omni.entrypoints.cli import serve as serve_module
+
+ args = TrackingNamespace(
+ argparse.Namespace(
+ model="dummy-model",
+ model_tag=None,
+ stage_init_timeout=1,
+ tokenizer=None,
+ disable_log_stats=False,
+ ),
+ frozenset({"model"}),
+ )
+ if typed:
+ from vllm_omni.config.omni_config import VllmOmniARStageConfig
+ from vllm_omni.config.stage_config import StagePipelineConfig
+
+ stages = [
+ VllmOmniARStageConfig(stage_pipeline_config=StagePipelineConfig(stage_id=i, model_stage="ar"))
+ for i in range(2)
+ ]
+ stages[1].connector_config.async_chunk = downstream_async_chunk
+ else:
+ stages = [
+ _FakeStageConfig(
+ stage_id=i,
+ engine_args=_FakeStageEngineArgs(
+ async_chunk=bool(i and downstream_async_chunk), enable_sleep_mode=False
+ ),
+ )
+ for i in range(2)
+ ]
+ mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", return_value="dummy-model")
+ mocker.patch(
+ "vllm_omni.config.config_factory.with_trust_remote_code_override",
+ side_effect=lambda kwargs, _trust: kwargs,
+ )
+ mocker.patch("vllm_omni.entrypoints.utils.parse_stage_overrides", return_value={})
+ mocker.patch(
+ "vllm_omni.config.resolver.resolve_omni_config",
+ return_value=OmniConfigResolution(config_path="dummy-config", stage_configs=tuple(stages)),
+ )
+
+ runtime = serve_module._build_multi_api_stage_runtime(args, 2)
+
+ assert runtime._client_count == 1
+ assert runtime._client_index == 0
+ assert runtime._async_chunk is downstream_async_chunk
+
+
+@pytest.mark.parametrize("typed", [False, True])
+def test_build_multi_api_stage_runtime_rejects_sleep_enabled_in_stage_config(
+ mocker: MockerFixture, typed: bool
+) -> None:
+ from vllm_omni.entrypoints.cli import serve as serve_module
+
+ args = TrackingNamespace(
+ argparse.Namespace(
+ model="dummy-model",
+ model_tag=None,
+ stage_init_timeout=1,
+ tokenizer=None,
+ disable_log_stats=False,
+ ),
+ frozenset({"model"}),
+ )
+ stage_config = _FakeStageConfig(
+ stage_id=3,
+ engine_args=_FakeStageEngineArgs(async_chunk=False, enable_sleep_mode=True),
+ )
+ if typed:
+ from vllm_omni.config.omni_config import VllmOmniARStageConfig
+ from vllm_omni.config.stage_config import StagePipelineConfig
+
+ typed_stage = VllmOmniARStageConfig(stage_pipeline_config=StagePipelineConfig(stage_id=3, model_stage="ar"))
+ typed_stage.model_config.enable_sleep_mode = True
+ stage_config = typed_stage
+ mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", return_value="dummy-model")
+ mocker.patch(
+ "vllm_omni.config.config_factory.with_trust_remote_code_override",
+ side_effect=lambda kwargs, _trust: kwargs,
+ )
+ mocker.patch("vllm_omni.entrypoints.utils.parse_stage_overrides", return_value={})
+ mocker.patch(
+ "vllm_omni.config.resolver.resolve_omni_config",
+ return_value=OmniConfigResolution(config_path="dummy-config", stage_configs=(stage_config,)),
+ )
+
+ with pytest.raises(ValueError, match=r"sleep mode.*stage\(s\) \[3\]"):
+ serve_module._build_multi_api_stage_runtime(args, 2)
+
+
+def test_run_multi_api_server_omni_starts_workers_after_shared_engine_launch(mocker: MockerFixture) -> None:
+ from vllm_omni.entrypoints.cli import serve as serve_module
+
+ args = TrackingNamespace(
+ argparse.Namespace(api_server_count=2, shutdown_timeout=1),
+ frozenset({"api_server_count", "shutdown_timeout"}),
+ )
+ socket = mocker.Mock()
+ primary_addresses = EngineZmqAddresses(
+ inputs=["ipc://input-0", "ipc://input-1"],
+ outputs=["ipc://output-0", "ipc://output-1"],
+ )
+ engine_launch = StageEngineLaunch(
+ client_configs=[{"client_count": 2, "client_index": index, "stage_addresses": {}} for index in range(2)],
+ resources=[StageReplicaResources(addresses=primary_addresses)],
+ )
+ runtime = mocker.MagicMock()
+ runtime.launch_stage_engines.return_value.__enter__.return_value = engine_launch
+ runtime.launch_stage_engines.return_value.__exit__.return_value = False
+
+ manager = mocker.Mock()
+ manager.processes = [mocker.Mock()]
+ manager.gather_actual_addresses.return_value = (primary_addresses.inputs, primary_addresses.outputs)
+
+ mocker.patch.object(serve_module, "_build_multi_api_stage_runtime", return_value=runtime)
+ mocker.patch.object(serve_module, "_wait_for_multi_api_server_completion")
+ mocker.patch("signal.signal")
+ mocker.patch("vllm.entrypoints.openai.api_server.setup_server", return_value=("127.0.0.1:8000", socket))
+ mocker.patch("vllm.v1.metrics.prometheus.setup_multiprocess_prometheus")
+ start_manager = mocker.patch.object(serve_module, "_start_api_server_process_manager", return_value=manager)
+
+ serve_module.run_multi_api_server_omni(args)
+
+ assert args._omni_stage_client_configs is engine_launch.client_configs
+ start_manager.assert_called_once()
+ manager_kwargs = start_manager.call_args.kwargs
+ assert manager_kwargs["cleanup_timeout"] == 1
+ assert manager_kwargs["num_servers"] == 2
+ assert manager_kwargs["input_addresses"] == ["ipc://input-0", "ipc://input-1"]
+ assert manager_kwargs["target_server_fn"] is serve_module.run_omni_api_server_worker_proc
+ assert engine_launch.watched_frontend_processes == manager.processes
+ runtime.shutdown.assert_called_once_with()
+ socket.close.assert_called_once_with()
+
+
+def test_start_api_server_process_manager_cleans_up_partial_start(mocker: MockerFixture) -> None:
+ from vllm_omni.entrypoints.cli import serve as serve_module
+
+ first_process = mocker.Mock(name="first_process")
+ second_process = mocker.Mock(name="second_process")
+ first_process.pid = 1234
+ second_process.pid = None
+ second_process.start.side_effect = RuntimeError("spawn failed")
+ parent_pipes = [mocker.Mock(name="parent_pipe_0"), mocker.Mock(name="parent_pipe_1")]
+ child_pipes = [mocker.Mock(name="child_pipe_0"), mocker.Mock(name="child_pipe_1")]
+ spawn_context = mocker.Mock()
+ spawn_context.Process.side_effect = [first_process, second_process]
+ spawn_context.Pipe.side_effect = list(zip(parent_pipes, child_pipes))
+ mocker.patch("vllm.v1.utils.multiprocessing.get_context", return_value=spawn_context)
+ shutdown = mocker.patch("vllm.v1.utils.shutdown")
+
+ with pytest.raises(RuntimeError, match="spawn failed"):
+ serve_module._start_api_server_process_manager(
+ cleanup_timeout=1.5,
+ listen_address="127.0.0.1:8000",
+ sock=mocker.Mock(),
+ args=argparse.Namespace(),
+ num_servers=2,
+ input_addresses=["ipc://input-0", "ipc://input-1"],
+ output_addresses=["ipc://output-0", "ipc://output-1"],
+ target_server_fn=mocker.Mock(),
+ )
+
+ first_process.start.assert_called_once_with()
+ second_process.start.assert_called_once_with()
+ for pipe in parent_pipes:
+ pipe.close.assert_called_once_with()
+ shutdown.assert_called_once_with([first_process], timeout=1.5)
+
+
+def test_wait_for_multi_api_server_completion_rejects_engine_failure(
+ mocker: MockerFixture,
+) -> None:
+ from vllm_omni.entrypoints.cli import serve as serve_module
+
+ api_process = _FakeProcess(sentinel="api-0", exitcode=None, name="api-0", pid=10)
+ engine_process = _FakeProcess(sentinel="engine-0", exitcode=4, name="engine-0", pid=20)
+ manager = _FakeProcessOwner(processes=[api_process])
+ engine_launch = StageEngineLaunch(
+ client_configs=[],
+ resources=[StageReplicaResources(manager=_FakeProcessOwner(processes=[engine_process]))],
+ )
+ mocker.patch("multiprocessing.connection.wait", return_value=["engine-0"])
+
+ with pytest.raises(RuntimeError, match=r"Shared stage engine process engine-0 .* exited with code 4"):
+ serve_module._wait_for_multi_api_server_completion(manager, engine_launch)
+
+
def test_serve_parser_rejects_stage_configs_path() -> None:
parser = TrackingArgumentParser()
subparsers = parser.add_subparsers(dest="subcommand")
@@ -334,6 +606,26 @@ def test_parse_stage_overrides_invalid_json_raises() -> None:
assert f"Got: {bad!r}" in message
+@pytest.mark.parametrize(
+ "payload",
+ ['{"abc": {}}', '{"-1": {}}', '{"1.5": {}}', '{"\uff10": {}}'],
+)
+def test_parse_stage_overrides_rejects_invalid_stage_ids(payload: str) -> None:
+ with pytest.raises(argparse.ArgumentTypeError, match="non-negative integer stage ids"):
+ _parse_stage_overrides(payload)
+
+
+def test_parse_stage_overrides_preserves_arbitrary_override_fields() -> None:
+ parsed = _parse_stage_overrides(
+ '{"0": {"extras": {"ltx2_use_conv_vae": true}, "kv_cache_dtype": "fp8", "typo_field_xyz": 1}}'
+ )
+ assert parsed["0"] == {
+ "extras": {"ltx2_use_conv_vae": True},
+ "kv_cache_dtype": "fp8",
+ "typo_field_xyz": 1,
+ }
+
+
def test_parse_stage_overrides_rejects_non_dict_top_level() -> None:
"""Top-level must be a JSON object (dict). A list, scalar, or non-dict
mapping is rejected with a ValueError naming ``--stage-overrides`` and
@@ -706,6 +998,77 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture)
assert manager_kwargs["omni_replica_id"] == 0
+def test_run_headless_generic_diffusion_launches_structured_stage(mocker: MockerFixture) -> None:
+ """Headless resolution starts the typed stage through the real group launcher."""
+ from vllm_omni.engine import stage_engine_startup as startup_module
+
+ mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None)
+ mocker.patch(
+ "vllm_omni.config.resolver._resolve_generic_diffusion_model_class",
+ return_value=(True, "FakeDiffusionPipeline"),
+ )
+ mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
+ mocker.patch.object(startup_module.stage_init_utils, "load_omni_transfer_config_for_model", return_value=None)
+ mocker.patch.object(
+ startup_module.initialization,
+ "resolve_omni_kv_config_for_stage",
+ return_value=(None, None, None),
+ )
+ captured: dict[str, Any] = {}
+ od_config = SimpleNamespace()
+
+ def _build_diffusion_config(model, stage_config, metadata):
+ captured.update(model=model, stage_config=stage_config, metadata=metadata)
+ return od_config
+
+ def _launch_replica_group(**kwargs):
+ captured.update(group_kwargs=kwargs)
+ captured["manager"] = kwargs["launch_one"](0)
+
+ def _launch_replica(**kwargs):
+ captured.update(replica_kwargs=kwargs)
+ return SimpleNamespace(exitcode=None)
+
+ mocker.patch.object(startup_module.stage_init_utils, "build_diffusion_config", side_effect=_build_diffusion_config)
+ mocker.patch.object(startup_module, "launch_headless_replica_group", side_effect=_launch_replica_group)
+ mocker.patch.object(startup_module, "launch_headless_diffusion_replica", side_effect=_launch_replica)
+
+ explicit_keys = frozenset(
+ {
+ "model",
+ "stage_id",
+ "omni_master_address",
+ "omni_master_port",
+ "worker_backend",
+ "model_class_name",
+ "num_gpus",
+ }
+ )
+ args = _make_headless_args(
+ explicit_keys=explicit_keys,
+ model="generic-diffusion",
+ model_class_name="FakeDiffusionPipeline",
+ num_gpus=1,
+ )
+
+ run_headless(args)
+
+ stage = captured["stage_config"]
+ assert isinstance(stage, VllmOmniDiffusionStageConfig)
+ assert captured["metadata"].stage_type == "diffusion"
+ assert captured["metadata"].model_stage == "diffusion"
+ assert captured["model"] == "generic-diffusion"
+ assert captured["group_kwargs"]["stage_id"] == 0
+ assert captured["group_kwargs"]["omni_dp_size_local"] == 1
+ assert captured["group_kwargs"]["per_replica_devices"] == ["0"]
+ assert captured["replica_kwargs"]["stage_config"] is stage
+ assert captured["replica_kwargs"]["stage_id"] == 0
+ assert captured["replica_kwargs"]["omni_master_address"] == "127.0.0.1"
+ assert captured["replica_kwargs"]["omni_master_port"] == 26000
+ assert captured["replica_kwargs"]["od_config"] is od_config
+ assert stage.runtime_config.devices == "0"
+
+
def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixture) -> None:
"""A diffusion replica that exits with a non-zero code must surface as a
RuntimeError from ``run_headless`` (the head needs the signal to roll
@@ -754,3 +1117,22 @@ def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixtur
with pytest.raises(RuntimeError, match=r"exited with code 137"):
run_headless(_make_headless_args(stage_id=1))
+
+
+@pytest.mark.parametrize("name", ["unfiltered_ns", "explicit_keys", "missing", "__dict__"])
+def test_tracking_namespace_uninitialized_access_raises_attribute_error(name):
+ namespace = object.__new__(TrackingNamespace)
+ with pytest.raises(AttributeError):
+ getattr(namespace, name)
+
+
+@pytest.mark.parametrize("deep", [False, True])
+def test_tracking_namespace_copy_preserves_tracking(deep):
+ import copy
+
+ namespace = TrackingNamespace(
+ argparse.Namespace(model="example", api_server_count=2), frozenset({"api_server_count"})
+ )
+ restored = copy.deepcopy(namespace) if deep else copy.copy(namespace)
+ assert restored.get_explicit_kwargs_dict() == {"api_server_count": 2}
+ assert restored.model == "example"
diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py
index 680bb110c..65395265b 100644
--- a/tests/entrypoints/test_utils.py
+++ b/tests/entrypoints/test_utils.py
@@ -13,6 +13,14 @@ import torch
from pytest_mock import MockerFixture
from vllm.sampling_params import RequestOutputKind, SamplingParams
+from vllm_omni.config.composable_parallel import (
+ Broadcast,
+ FanInByStage,
+ MeshAxisSpec,
+ RouteByStage,
+ StrategySpec,
+ TakeRank,
+)
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
from vllm_omni.config.resolver import (
OmniConfigResolution,
@@ -21,7 +29,6 @@ from vllm_omni.config.resolver import (
resolve_omni_config,
)
from vllm_omni.config.stage_config import PipelineConfig
-from vllm_omni.config.yaml_util import create_config
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.engine.arg_utils import OmniEngineArgs
from vllm_omni.entrypoints.utils import (
@@ -356,16 +363,12 @@ class TestResolveOmniConfig:
)
assert resolved.config_path is None
assert len(resolved.stage_configs) == 1
- engine_args = resolved.stage_configs[0]["engine_args"]
- assert "dtype" in engine_args
- assert engine_args["engine_backend"] == engine_backend
- assert engine_args["revision"] == "pinned-revision"
+ assert resolved.stage_configs[0].diffusion_config.dtype is torch.float32
+ assert resolved.stage_configs[0].diffusion_config.engine_backend == engine_backend
+ assert resolved.stage_configs[0].diffusion_config.revision == "pinned-revision"
def test_generic_diffusion_uses_registered_model_metadata(self, mocker: MockerFixture):
- mocker.patch(
- "vllm_omni.config.resolver.StageConfigFactory.create_from_model",
- return_value=None,
- )
+ mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None)
resolve_model_class = mocker.patch(
"vllm_omni.config.resolver.resolve_model_class_name",
return_value="WanImageToVideoPipeline",
@@ -374,59 +377,98 @@ class TestResolveOmniConfig:
"vllm_omni.config.resolver.DiffusionModelRegistry.get_supported_archs",
return_value={"WanImageToVideoPipeline"},
)
+ load_model_class = mocker.patch("vllm_omni.config.resolver.DiffusionModelRegistry._try_load_model_cls")
resolved = resolve_omni_config(
"/models/Wan2.2-I2V",
trust_remote_code=False,
deploy_config_path=None,
- cli_overrides={
- "diffusion_load_format": "diffusers",
- "revision": "pinned-revision",
- },
+ cli_overrides={"diffusion_load_format": "diffusers", "revision": "pinned-revision"},
stage_overrides=None,
strategy_config_path=None,
)
+ resolve_model_class.assert_called_once_with("/models/Wan2.2-I2V", "diffusers", "pinned-revision")
+ load_model_class.assert_not_called()
stage = resolved.stage_configs[0]
- resolve_model_class.assert_called_once_with(
- "/models/Wan2.2-I2V",
- "diffusers",
- "pinned-revision",
- )
- assert stage.engine_args.model_class_name == "WanImageToVideoPipeline"
+ assert stage.diffusion_config.model_class_name == "WanImageToVideoPipeline"
+ assert stage.diffusion_config.revision == "pinned-revision"
assert stage.final_output_type == "video"
+ def test_generic_diffusion_stage_overrides_reach_typed_backend(self, mocker: MockerFixture):
+ from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
+
+ mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None)
+ mocker.patch(
+ "vllm_omni.config.resolver._resolve_generic_diffusion_model_class",
+ return_value=(True, "LTX2Pipeline"),
+ )
+ resolved = resolve_omni_config(
+ "/models/LTX-2.5-Diffusers",
+ trust_remote_code=False,
+ deploy_config_path=None,
+ cli_overrides={
+ "extras": {"ltx2_use_conv_vae": False, "keep": "global"},
+ "tensor_parallel_size": 1,
+ "diffusion_streaming_output": True,
+ },
+ stage_overrides={
+ "0": {
+ "extras": {"ltx2_use_conv_vae": True},
+ "tensor_parallel_size": 2,
+ "devices": "2,3",
+ "engine_backend": "custom.backend",
+ },
+ },
+ strategy_config_path=None,
+ )
+ stage = resolved.stage_by_id(0)
+ engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="/models/LTX-2.5-Diffusers")
+
+ assert stage.runtime_config.devices == "2,3"
+ assert stage.parallel_config.world_size == 2
+ assert engine_args["engine_backend"] == "custom.backend"
+ assert engine_args["extras"]["ltx2_use_conv_vae"] is True
+ assert engine_args["extras"]["keep"] == "global"
+ assert engine_args["streaming_output"] is True
+
def test_registered_pipeline_uses_structured_metadata_and_preserves_override_trust(self, mocker: MockerFixture):
endpoint_restriction = SimpleNamespace(name="chat")
+ typed_stage = SimpleNamespace(stage_id=1)
structured_config = SimpleNamespace(
- orchestrator_config=SimpleNamespace(deploy_config_path="/resolved/deploy.yaml"),
+ orchestrator_config=SimpleNamespace(
+ deploy_config_path="/resolved/deploy.yaml",
+ omni_lb_policy="round_robin",
+ ),
pipeline_config=SimpleNamespace(endpoint_restrictions=(endpoint_restriction,)),
+ stage_configs=(typed_stage,),
+ strategy_omni_lb_policy="round-robin",
)
- runtime_stage = create_config(
- {
- "stage_id": 1,
- "runtime": {"devices": "1,2,3", "num_replicas": 3},
- "engine_args": {"model": "dummy-model"},
- }
- )
- legacy_stage = SimpleNamespace(to_omegaconf=lambda: runtime_stage)
create_structured = mocker.patch(
"vllm_omni.config.resolver.StageConfigFactory.create_from_model",
return_value=structured_config,
)
create_legacy = mocker.patch(
- "vllm_omni.config.resolver.StageConfigFactory._resolve_legacy_from_registry",
- return_value=SimpleNamespace(
- stage_configs=[legacy_stage],
- pipeline_config=structured_config.pipeline_config,
- omni_lb_policy="round_robin",
- ),
+ "vllm_omni.config.resolver.StageConfigFactory.create_legacy_stage_configs_from_model",
)
- strategy_specs = {"stage_1": {"dp": 3}}
+ strategy_specs = {
+ "stage_1": [
+ StrategySpec(
+ "stage_replica",
+ MeshAxisSpec("stage_replica", 2),
+ RouteByStage("round_robin"),
+ FanInByStage(),
+ )
+ ]
+ }
load_strategy = mocker.patch(
"vllm_omni.config.resolver._load_strategy_specs",
return_value=strategy_specs,
)
+ apply_strategy = mocker.patch(
+ "vllm_omni.config.composable_parallel.apply_strategy_specs",
+ return_value=SimpleNamespace(omni_lb_policy="round_robin"),
+ )
resolved = resolve_omni_config(
"dummy-model",
@@ -447,31 +489,63 @@ class TestResolveOmniConfig:
trust_remote_code=None,
cli_overrides=expected_overrides,
deploy_config_path="deploy.yaml",
- )
- create_legacy.assert_called_once_with(
- structured_config.pipeline_config,
- expected_overrides,
- "/resolved/deploy.yaml",
strategy_specs=strategy_specs,
)
+ create_legacy.assert_not_called()
load_strategy.assert_called_once_with("strategy.yaml")
+ apply_strategy.assert_not_called()
assert resolved.config_path == "/resolved/deploy.yaml"
assert resolved.pipeline_config is structured_config.pipeline_config
- assert resolved.omni_lb_policy == "round_robin"
+ assert resolved.omni_lb_policy == "round-robin"
assert resolved.endpoint_restrictions == (endpoint_restriction,)
- assert resolved.stage_configs == (runtime_stage,)
+ assert resolved.stage_configs == (typed_stage,)
- def test_registered_resolution_exposes_forced_aligner_topology(self, mocker: MockerFixture):
- pipeline = OMNI_PIPELINES["qwen3_tts"]
- assert isinstance(pipeline, PipelineConfig)
+ def test_tp_only_strategy_does_not_report_default_lb_policy_as_derived(self, mocker: MockerFixture):
structured_config = SimpleNamespace(
- orchestrator_config=SimpleNamespace(deploy_config_path=None),
- pipeline_config=pipeline,
+ orchestrator_config=SimpleNamespace(
+ deploy_config_path="/resolved/deploy.yaml",
+ omni_lb_policy="round-robin",
+ ),
+ pipeline_config=SimpleNamespace(endpoint_restrictions=()),
+ stage_configs=(SimpleNamespace(stage_id=0),),
+ strategy_omni_lb_policy=None,
)
mocker.patch(
"vllm_omni.config.resolver.StageConfigFactory.create_from_model",
return_value=structured_config,
)
+ mocker.patch(
+ "vllm_omni.config.resolver._load_strategy_specs",
+ return_value={
+ "thinker": [
+ StrategySpec(
+ "tp",
+ MeshAxisSpec("tp", 2),
+ Broadcast(),
+ TakeRank(),
+ )
+ ]
+ },
+ )
+
+ resolved = resolve_omni_config(
+ "dummy-model",
+ trust_remote_code=True,
+ deploy_config_path="deploy.yaml",
+ cli_overrides={"omni_lb_policy": "round-robin"},
+ stage_overrides=None,
+ strategy_config_path="strategy.yaml",
+ )
+
+ assert resolved.omni_lb_policy is None
+
+ def test_registered_resolution_exposes_forced_aligner_topology(self, mocker: MockerFixture):
+ pipeline = OMNI_PIPELINES["qwen3_tts"]
+ assert isinstance(pipeline, PipelineConfig)
+ mocker.patch(
+ "vllm_omni.config.config_factory.StageConfigFactory.get_pipeline_config",
+ return_value=pipeline,
+ )
resolved = resolve_omni_config(
"Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
@@ -483,9 +557,9 @@ class TestResolveOmniConfig:
)
assert resolved.pipeline_config is not None
- pipeline_stage_ids = [stage.stage_id for stage in resolved.pipeline_config.stages]
- runtime_stage_ids = [stage.stage_id for stage in resolved.stage_configs]
- assert pipeline_stage_ids == runtime_stage_ids
+ assert [stage.stage_id for stage in resolved.pipeline_config.stages] == [
+ stage.stage_id for stage in resolved.stage_configs
+ ]
assert resolved.pipeline_config.stages[-1].model_stage == "forced_aligner"
assert len(resolved.pipeline_config.stages) == len(pipeline.stages) + 1
@@ -529,3 +603,21 @@ class TestCumulativeStreamingCoercion:
assert all([isinstance(r, SamplingParams) for r in result])
assert result[0].output_kind == RequestOutputKind.DELTA
assert result[1].output_kind == RequestOutputKind.DELTA
+
+
+@pytest.mark.parametrize("api_server_count", [1, 2])
+def test_prepare_stage_config_inputs_filters_frontend_options(api_server_count):
+ from vllm_omni.entrypoints.utils import prepare_stage_config_inputs
+
+ kwargs = {
+ "api_server_count": api_server_count,
+ "disable_log_stats": True,
+ "model_tag": "example",
+ "max_num_seqs": 8,
+ "stage_overrides": '{"0": {"max_num_seqs": 4}}',
+ }
+ inputs = prepare_stage_config_inputs("example", kwargs, trust_remote_code=None)
+
+ assert inputs.kwargs == {"max_num_seqs": 8}
+ assert inputs.stage_overrides == {"0": {"max_num_seqs": 4}}
+ assert kwargs["api_server_count"] == api_server_count
diff --git a/tests/examples/offline_inference/test_text_to_video_prompts.py b/tests/examples/offline_inference/test_text_to_video_prompts.py
index c588b7a38..e0c381bf3 100644
--- a/tests/examples/offline_inference/test_text_to_video_prompts.py
+++ b/tests/examples/offline_inference/test_text_to_video_prompts.py
@@ -1,10 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import functools
import importlib.util
+import sys
+from types import SimpleNamespace
import pytest
import torch
@@ -24,6 +26,51 @@ def _load_text_to_video():
return module
+@pytest.mark.parametrize(
+ ("arguments", "expected"),
+ [
+ (["--quantization", "mxfp4"], {"quantization": "mxfp4"}),
+ (
+ ["--quantization-config", '{"method":"mxfp4","w4a8_fallback_steps":[0,2]}'],
+ {"quantization_config": {"method": "mxfp4", "w4a8_fallback_steps": [0, 2]}},
+ ),
+ ],
+)
+def test_quantization_cli_reaches_omni(arguments, expected, monkeypatch):
+ mod = _load_text_to_video()
+ captured = {}
+
+ class StopBeforeModelLoadError(Exception):
+ pass
+
+ def capture_omni(**kwargs):
+ captured.update(kwargs)
+ raise StopBeforeModelLoadError
+
+ monkeypatch.setattr(sys, "argv", ["text_to_video.py", "--model", "Wan-AI/Wan2.2-T2V-A14B-Diffusers", *arguments])
+ monkeypatch.setattr(mod, "current_omni_platform", SimpleNamespace(device_type="cpu"))
+ monkeypatch.setattr(mod, "Omni", capture_omni)
+
+ with pytest.raises(StopBeforeModelLoadError):
+ mod.main()
+
+ assert {key: captured[key] for key in ("quantization", "quantization_config") if key in captured} == expected
+
+
+def test_quantization_cli_options_are_mutually_exclusive(monkeypatch):
+ mod = _load_text_to_video()
+ monkeypatch.setattr(
+ sys,
+ "argv",
+ ["text_to_video.py", "--quantization", "mxfp4", "--quantization-config", '{"method":"mxfp4"}'],
+ )
+
+ with pytest.raises(SystemExit) as exc:
+ mod.parse_args()
+
+ assert exc.value.code == 2
+
+
@pytest.mark.parametrize(
("negative_prompt", "expected_negative_prompt"),
[(None, None), ("", ""), ("blurry", "blurry")],
diff --git a/tests/platforms/npu/quant/test_mxfp4_step_fallback.py b/tests/platforms/npu/quant/test_mxfp4_step_fallback.py
new file mode 100644
index 000000000..8431ebf17
--- /dev/null
+++ b/tests/platforms/npu/quant/test_mxfp4_step_fallback.py
@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Real A5 kernels: step selection, bias, shape, and reuse of prepared W4 weights."""
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.npu]
+
+torch_npu = pytest.importorskip("torch_npu")
+
+
+@torch.inference_mode()
+def _exercise_mxfp4_fallback(with_bias, dtype, input_shape, mode, monkeypatch, *, policy, scale_alg=2):
+ from vllm_omni.diffusion.forward_context import (
+ ForwardContext,
+ override_forward_context,
+ set_forward_context_denoise_step_idx,
+ )
+ from vllm_omni.quantization.mxfp4_config import (
+ DiffusionMXFP4Config,
+ NPUMxfp4LinearMethod,
+ NPUMxfp4OnlineLinearMethod,
+ )
+
+ prefix = "blocks.10.attn1.to_qkv"
+ steps = [1] if policy in ("step", "layer_and_step") else []
+ layers = [prefix] if policy in ("layer", "layer_and_step") else []
+ fallback = {"w4a8_fallback_steps": steps, "w4a8_fallback_layers": layers, "mxfp4_scale_alg": scale_alg}
+ torch.manual_seed(42)
+ layer = torch.nn.Linear(1024, 128, bias=with_bias, device="npu", dtype=dtype)
+ if mode == "single":
+ method = NPUMxfp4OnlineLinearMethod(DiffusionMXFP4Config(**fallback), prefix=prefix)
+ elif mode == "single_offline":
+ method = NPUMxfp4LinearMethod(
+ DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, **fallback), prefix=prefix
+ )
+ layer.orig_dtype = dtype
+ values = torch.tensor([-4.0, -2.0, -1.0, -0.5, 0.0, 0.5, 1.0, 2.0])
+ layer.weight.copy_(values[torch.randint(8, (128, 1024))].to(layer.weight))
+ layer.register_parameter(
+ "weight_scale", torch.nn.Parameter(torch.full((128, 32), 124, dtype=torch.uint8, device="npu"), False)
+ )
+ layer.register_parameter("mul_scale", torch.nn.Parameter((torch.rand(1024) * 1.5 + 0.5).npu(), False))
+ else:
+ raise ValueError(f"Unsupported single-scale test mode: {mode}")
+ method.process_weights_after_loading(layer)
+ x = torch.randn(input_shape, device="npu", dtype=dtype)
+ pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr())
+ method.process_weights_after_loading(layer)
+ quantize = torch_npu.npu_dynamic_mx_quant
+ activation_types = []
+
+ def record_quantize(x, **kwargs):
+ dst = kwargs["dst_type"]
+ expected_alg = 0 if dst == torch_npu.float8_e4m3fn else scale_alg
+ assert kwargs["scale_alg"] == expected_alg
+ assert kwargs["axis"] == -1 and kwargs["block_size"] == 32 and kwargs["round_mode"] == "rint"
+ if expected_alg == 2:
+ assert kwargs["dst_type_max"] == 7.25
+ else:
+ assert "dst_type_max" not in kwargs
+ activation_types.append(dst)
+ return quantize(x, **kwargs)
+
+ monkeypatch.setattr(torch_npu, "npu_dynamic_mx_quant", record_quantize)
+ results = []
+ for _request in range(2):
+ with override_forward_context(ForwardContext()):
+ for step in range(3):
+ set_forward_context_denoise_step_idx(step)
+ output = method.apply(layer, x, layer.bias)
+ torch.npu.synchronize()
+ assert output.shape == (*input_shape[:-1], 128)
+ assert output.dtype == dtype
+ assert torch.isfinite(output).all()
+ results.append(output.cpu())
+ expected_types = [
+ torch_npu.float8_e4m3fn if layers or step in steps else torch_npu.float4_e2m1fn_x2 for step in range(3)
+ ]
+ assert activation_types == expected_types * 2
+ assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == pointers
+ assert not hasattr(layer, "w4a8_weight")
+ assert not hasattr(layer, "w4a8_weight_scale")
+ for first, second in zip(results[:3], results[3:]):
+ torch.testing.assert_close(first, second, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("with_bias", [False, True])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("input_shape", [(16, 1024), (2, 8, 1024)])
+@pytest.mark.parametrize("mode", ["single", "single_offline"])
+def test_mxfp4_step_fallback(with_bias, dtype, input_shape, mode, monkeypatch):
+ _exercise_mxfp4_fallback(with_bias, dtype, input_shape, mode, monkeypatch, policy="step")
+
+
+@pytest.mark.cards_1
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@torch.inference_mode()
+def test_single_scale_step_fallback_matches_independent_dense_reference(dtype):
+ if not torch.npu.is_available() or "Ascend950" not in torch.npu.get_device_name():
+ pytest.skip("MXFP4 kernels require an Ascend950 NPU")
+
+ from vllm_omni.diffusion.forward_context import (
+ ForwardContext,
+ override_forward_context,
+ set_forward_context_denoise_step_idx,
+ )
+ from vllm_omni.quantization.mxfp4_config import (
+ DiffusionMXFP4Config,
+ NPUMxfp4LinearMethod,
+ )
+
+ # Synthetic serialized single-scale values and E8M0 exponent bytes.
+ # Powers of two make both activation precisions exact: this independent
+ # CPU dense oracle needs no production packing, dequantization or apply.
+ # Real UOS-quantized checkpoint validation remains a separate checkpoint/E2E gate.
+ values = torch.tensor([1.0, -1.0, 2.0, -2.0]).repeat(32).unsqueeze(1).expand(128, 512)
+ weight = torch.cat((values, values), dim=1)
+ fine = torch.cat((torch.full((128, 16), 127), torch.full((128, 16), 128)), dim=1).to(torch.uint8)
+ smooth = torch.cat((torch.full((512,), 0.5), torch.full((512,), 2.0)))
+ bias = torch.tensor([0.25, -0.25]).repeat(64)
+ x = torch.tensor([1.0, -1.0, 0.5, -0.5]).repeat(4).unsqueeze(1).expand(16, 1024) / 64
+ expected = torch.nn.functional.linear(x * smooth, torch.cat((values, 2 * values), dim=1), bias)
+
+ layer = torch.nn.Linear(1024, 128, bias=True, device="npu", dtype=dtype)
+ layer.weight.copy_(weight)
+ layer.bias.copy_(bias)
+ layer.orig_dtype = dtype
+ for name, value in {"weight_scale": fine, "mul_scale": smooth}.items():
+ layer.register_parameter(name, torch.nn.Parameter(value.npu(), requires_grad=False))
+ method = NPUMxfp4LinearMethod(
+ DiffusionMXFP4Config(
+ is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True, w4a8_fallback_steps=[1], mxfp4_scale_alg=2
+ )
+ )
+ method.process_weights_after_loading(layer)
+ prepared_pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr())
+ assert set(layer.state_dict()) == {"weight", "bias", "weight_scale", "mul_scale"}
+ assert not hasattr(layer, "w4a8_weight")
+ assert not hasattr(layer, "w4a8_weight_scale")
+ x = x.to(device="npu", dtype=dtype)
+
+ with override_forward_context(ForwardContext()):
+ for step in range(3):
+ set_forward_context_denoise_step_idx(step)
+ actual = method.apply(layer, x, layer.bias).cpu()
+ # A8's documented GEMM output is BF16 before restoring the caller
+ # dtype; the reference does not reuse quantization or unpacking.
+ reference = expected.bfloat16().to(dtype) if step == 1 else expected.to(dtype)
+ torch.testing.assert_close(actual, reference, rtol=0, atol=0)
+ assert actual.dtype == dtype
+ assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == prepared_pointers
+
+
+@pytest.mark.cards_1
+@pytest.mark.parametrize("mode", ["single", "single_offline"])
+@pytest.mark.parametrize("policy", ["layer", "layer_and_step"])
+def test_mxfp4_layer_fallback(mode, policy, monkeypatch):
+ """Reuse the kernel/shape/cache assertions without multiplying the full dtype/bias matrix."""
+ _exercise_mxfp4_fallback(False, torch.bfloat16, (16, 1024), mode, monkeypatch, policy=policy)
+
+
+@pytest.mark.cards_1
+@pytest.mark.parametrize("mode", ["single", "single_offline"])
+def test_mxfp4_ocp_compatibility(mode, monkeypatch):
+ """OCP compatibility: A4 -> A8 -> A4 with the same W4 across requests."""
+ _exercise_mxfp4_fallback(False, torch.bfloat16, (16, 1024), mode, monkeypatch, policy="step", scale_alg=0)
diff --git a/vllm_omni/config/composable_parallel/apply.py b/vllm_omni/config/composable_parallel/apply.py
index 879c7aec3..acfb9e229 100644
--- a/vllm_omni/config/composable_parallel/apply.py
+++ b/vllm_omni/config/composable_parallel/apply.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Apply declarative parallel strategies onto merged stage configs.
This is the *override-after-merge* seam. ``merge_pipeline_deploy`` first fuses
@@ -198,10 +198,41 @@ def _check_devices(runtime: dict[str, Any], cfg: OmniParallelConfig, *, role: Ro
def _apply_to_stage(stage: Any, cfg: OmniParallelConfig, *, role: RoleKey) -> None:
+ declared = set(cfg.l1_owners.keys())
+ if hasattr(stage, "parallel_config") and hasattr(stage, "runtime_config"):
+ parallel = stage.parallel_config
+ for kind, field_name in _ENGINE_FIELD_BY_KIND.items():
+ if kind in declared:
+ existing = getattr(parallel, field_name, None)
+ value = getattr(cfg, field_name)
+ if existing not in (None, 1, value):
+ _fail(
+ f"role {role!r}: strategy derives {field_name}={value} "
+ f"but structured config already set {existing}"
+ )
+ setattr(parallel, field_name, value)
+ if "ep" in declared and cfg.enable_expert_parallel:
+ setattr(parallel, "enable_expert_parallel", True)
+ if "stage_replica" in declared:
+ existing = stage.runtime_config.num_replicas
+ if existing not in (None, 1, cfg.stage_replica_size):
+ _fail(
+ f"role {role!r}: strategy derives "
+ f"num_replicas={cfg.stage_replica_size} but structured config already set {existing}"
+ )
+ stage.runtime_config.num_replicas = cfg.stage_replica_size
+ check_device_layout(
+ stage.runtime_config.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(stage.runtime_config.num_replicas or 1),
+ role=role,
+ )
+ return
+
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)
@@ -209,7 +240,6 @@ def _apply_to_stage(stage: Any, cfg: OmniParallelConfig, *, role: RoleKey) -> No
_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)
diff --git a/vllm_omni/config/config_factory.py b/vllm_omni/config/config_factory.py
index 90108e7df..94decce9a 100644
--- a/vllm_omni/config/config_factory.py
+++ b/vllm_omni/config/config_factory.py
@@ -25,6 +25,9 @@ from vllm_omni.config.stage_config import (
DeployConfig,
PipelineConfig,
StageConfig,
+ StageDeployConfig,
+ StageExecutionType,
+ StagePipelineConfig,
StageType,
build_stage_runtime_overrides,
load_deploy_config,
@@ -373,6 +376,7 @@ class StageConfigFactory:
trust_remote_code: bool | None,
cli_overrides: dict[str, Any],
deploy_config_path: str | None,
+ strategy_specs: Mapping[Any, Any] | None = None,
) -> VllmOmniConfig | None:
"""Build the structured Omni config for a model/deploy pair."""
user_deploy_config = cls._load_user_deploy_config(deploy_config_path)
@@ -387,6 +391,20 @@ class StageConfigFactory:
if pipeline_cfg is None:
return None
+ # Keep the typed representation's topology in lockstep with the
+ # compatibility launch representation. The aligner is a real extra
+ # stage, not an engine-only setting, so it must be injected before
+ # either representation is built.
+ if user_deploy_config is None:
+ if pipeline_cfg.default_deploy_config_name is not None:
+ default_deploy_path = _DEPLOY_DIR / pipeline_cfg.default_deploy_config_name
+ user_deploy_config = load_deploy_config(default_deploy_path)
+ deploy_config_path = str(default_deploy_path)
+ else:
+ user_deploy_config = DeployConfig()
+ from vllm_omni.utils.forced_aligner import inject_forced_aligner_stage
+
+ pipeline_cfg, user_deploy_config = inject_forced_aligner_stage(pipeline_cfg, user_deploy_config, cli_overrides)
registry_cli_overrides = with_trust_remote_code_override(
{**cli_overrides, "model": model},
trust_remote_code,
@@ -396,6 +414,7 @@ class StageConfigFactory:
user_deploy_config=user_deploy_config,
deploy_config_path=deploy_config_path,
cli_overrides=registry_cli_overrides,
+ strategy_specs=strategy_specs,
)
@classmethod
@@ -621,13 +640,11 @@ class StageConfigFactory:
)
@classmethod
- def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]:
- """Build the temporary runtime ABI for a generic diffusion stage.
-
- The terminal diffusion config owns engine defaults and normalization;
- this compatibility builder only adds Omni stage topology, request
- defaults, and device placement.
- """
+ def _normalize_default_diffusion(
+ cls,
+ kwargs: dict[str, Any],
+ ) -> tuple[dict[str, Any], DiffusionParallelConfig, dict[str, Any], str]:
+ """Normalize inputs shared by typed and compatibility diffusion builders."""
raw_sampling_params = kwargs.get("default_sampling_params")
if isinstance(raw_sampling_params, str):
try:
@@ -643,8 +660,6 @@ class StageConfigFactory:
if kwargs.get("num_gpus") is not None:
parallel_config.resolve_data_parallel_size(int(kwargs["num_gpus"]))
engine_args = OmniDiffusionConfig.normalize_init_kwargs(kwargs)
- engine_args["parallel_config"] = asdict(parallel_config)
- engine_args["model_stage"] = "diffusion"
extras = dict(engine_args.get("extras") or {})
for key, default in (
@@ -665,6 +680,21 @@ class StageConfigFactory:
final_output_type,
model_class_name,
)
+ return engine_args, parallel_config, default_sampling_params, final_output_type
+
+ @classmethod
+ def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]:
+ """Build the temporary runtime ABI for a generic diffusion stage.
+
+ The terminal diffusion config owns engine defaults and normalization;
+ this compatibility builder only adds Omni stage topology, request
+ defaults, and device placement.
+ """
+ engine_args, parallel_config, default_sampling_params, final_output_type = cls._normalize_default_diffusion(
+ kwargs
+ )
+ engine_args["parallel_config"] = asdict(parallel_config)
+ engine_args["model_stage"] = "diffusion"
return [
{
@@ -681,6 +711,50 @@ class StageConfigFactory:
}
]
+ @classmethod
+ def create_typed_default_diffusion(
+ cls,
+ model: str,
+ kwargs: dict[str, Any],
+ ) -> VllmOmniConfig:
+ """Build generic diffusion directly into the structured runtime config."""
+ engine_overrides, parallel_config, default_sampling_params, final_output_type = (
+ cls._normalize_default_diffusion(kwargs)
+ )
+ engine_overrides.update(asdict(parallel_config))
+ model_class_name = engine_overrides.get("model_class_name")
+ pipeline = PipelineConfig(
+ model_type="generic_diffusion",
+ model_arch=str(model_class_name or ""),
+ stages=(
+ StagePipelineConfig(
+ stage_id=0,
+ model_stage="diffusion",
+ execution_type=StageExecutionType.DIFFUSION,
+ final_output=True,
+ final_output_type=final_output_type,
+ ),
+ ),
+ )
+
+ # These values have already been normalized for this diffusion stage.
+ # Scope them explicitly so diffusion-only fields such as engine_backend
+ # and extras are not filtered by the global LLM CLI argument surface.
+ stage_overrides = {
+ f"stage_0_{key}": value for key, value in engine_overrides.items() if key not in {"model", "stage_id"}
+ }
+ stage_overrides["stage_0_devices"] = (
+ kwargs.get("stage_0_devices")
+ or kwargs.get("devices")
+ or ",".join(str(i) for i in range(parallel_config.world_size))
+ )
+ # Caller defaults belong to deployment, not immutable pipeline constraints:
+ # explicit request sampling parameters must still be able to override them.
+ deploy = DeployConfig(stages=[StageDeployConfig(stage_id=0, default_sampling_params=default_sampling_params)])
+ return VllmOmniConfig.from_pipeline_config(
+ pipeline, user_deploy_config=deploy, cli_overrides={"model": model, **stage_overrides}
+ )
+
@classmethod
def _merge_cli_overrides(
cls,
diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py
index 56b40b836..1407615ab 100644
--- a/vllm_omni/config/omni_config.py
+++ b/vllm_omni/config/omni_config.py
@@ -29,7 +29,9 @@ from vllm.config import ProfilerConfig as VllmProfilerConfig
from vllm.config import SchedulerConfig as VllmSchedulerConfig
from vllm.config.utils import config
from vllm.engine.arg_utils import EngineArgs as VllmEngineArgs
+from vllm.logger import init_logger
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.pooling_params import PoolingParams
from vllm_omni.config.stage_config import (
_DEPLOY_DIR,
@@ -42,11 +44,13 @@ from vllm_omni.config.stage_config import (
StagePipelineConfig,
StageType,
_apply_platform_overrides,
+ _get_recursively_merged_dict,
_resolve_scheduler,
_scheduler_path,
_select_processor_funcs,
build_stage_runtime_overrides,
load_deploy_config,
+ merge_pipeline_deploy,
merge_sampling_constraints,
normalize_pipeline_cli_overrides,
reconcile_diffusion_attention_overrides,
@@ -55,6 +59,8 @@ from vllm_omni.config.stage_config import (
)
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+logger = init_logger(__name__)
+
_EXECUTION_TYPE_TO_STAGE_WORKER: dict[StageExecutionType, tuple[StageType, str | None]] = {
StageExecutionType.LLM_AR: (StageType.LLM, "ar"),
StageExecutionType.LLM_GENERATION: (StageType.LLM, "generation"),
@@ -186,6 +192,11 @@ class _ModelEngineOverrides(TypedDict, total=False):
enable_prompt_embeds: bool
+class _PoolingEngineOverrides(TypedDict, total=False):
+ runner: str
+ pooling_output_decoder: str
+
+
class _LoadEngineOverrides(TypedDict, total=False):
tokenizer: str
download_dir: str
@@ -261,6 +272,7 @@ class _StageEngineValues:
quantization: _QuantizationEngineOverrides
model: _ModelEngineOverrides
load: _LoadEngineOverrides
+ pooling: _PoolingEngineOverrides
cache: _CacheEngineOverrides
scheduler: _SchedulerEngineOverrides
connector: _ConnectorEngineOverrides
@@ -476,6 +488,15 @@ class OmniStageModelConfig(_TrackExplicitConfigFields):
enable_prompt_embeds: bool | None = None
+@config(config=ConfigDict(arbitrary_types_allowed=True))
+class OmniStagePoolingConfig:
+ """Typed inputs owned by vLLM pooling stages."""
+
+ runner: str | None = None
+ pooling_output_decoder: str | None = None
+ default_pooling_params: PoolingParams | None = None
+
+
@_enforce_keyword_only_init
@config(kw_only=True)
class OmniStageLoadConfig(_TrackExplicitConfigFields, VllmLoadConfig):
@@ -817,6 +838,8 @@ class _DiffusionConfigProjection:
enable_stage_verification: bool = True
prompt_file_path: str | None = None
quantization_config: _QuantizationConfigType = None
+ # Internal provenance, retained across config projection and worker transport.
+ quantization_config_is_auto_detected: bool = False
extras: dict[str, Any] = field(default_factory=dict)
@field_validator("kv_transfer_config", mode="before")
@@ -831,12 +854,12 @@ class _DiffusionConfigProjection:
from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs
from vllm_omni.diffusion.offloader.config import parse_diffusion_offload_config
- normalized_kwargs = normalize_omni_diffusion_kwargs(kwargs)
+ config_kwargs = normalize_omni_diffusion_kwargs(kwargs)
# Validate before stage construction while retaining the raw mapping
# needed by dataclass/config serialization across process boundaries.
- parse_diffusion_offload_config(normalized_kwargs.get("diffusion_offload_config"))
+ parse_diffusion_offload_config(config_kwargs.get("diffusion_offload_config"))
valid_fields = {f.name for f in fields(cast(Any, cls))}
- return cls(**{k: v for k, v in normalized_kwargs.items() if k in valid_fields})
+ return cls(**{k: v for k, v in config_kwargs.items() if k in valid_fields})
def __post_init__(self) -> None:
# Keep diffusion imports lazy so importing vllm_omni.config does not
@@ -964,6 +987,8 @@ class _DiffusionConfigProjection:
or (is_checkpoint_nvfp4 and self._is_generic_nvfp4_quant_config(self.quantization_config))
)
if should_use_checkpoint_config:
+ if self.quantization_config is None:
+ self.quantization_config_is_auto_detected = True
self.quantization_config = quant_config
if quant_method is not None:
self.additional_config.setdefault("auto_detected_quant_method", quant_method)
@@ -1141,6 +1166,7 @@ _MODEL_ENGINE_FIELDS = frozenset(_ModelEngineOverrides.__annotations__)
_LOAD_ENGINE_FIELDS = frozenset(_LoadEngineOverrides.__annotations__)
_CACHE_ENGINE_FIELDS = frozenset(_CacheEngineOverrides.__annotations__)
_SCHEDULER_ENGINE_FIELDS = frozenset(_SchedulerEngineOverrides.__annotations__)
+_POOLING_ENGINE_FIELDS = frozenset(_PoolingEngineOverrides.__annotations__)
_CONNECTOR_ENGINE_FIELDS = frozenset(_ConnectorEngineOverrides.__annotations__)
_RUNTIME_ENGINE_FIELDS = frozenset(_RuntimeEngineOverrides.__annotations__)
_DIRECT_VLLM_CONFIG_ENGINE_FIELDS = frozenset({"compilation_config", "profiler_config"})
@@ -1171,6 +1197,7 @@ _LLM_STAGE_ENGINE_FIELDS = (
| _LLM_CACHE_ENGINE_FIELDS
| _LLM_SCHEDULER_ENGINE_FIELDS
| _LLM_PARALLEL_CONFIG_ENGINE_FIELDS
+ | _POOLING_ENGINE_FIELDS
| {"parallel_config"}
)
_DIFFUSION_OWNED_STAGE_ENGINE_FIELDS = (
@@ -1304,7 +1331,12 @@ def _stage_engine_values(
if topology.execution_type == StageExecutionType.DIFFUSION:
# Mirror StageConfig.to_omegaconf so both projections resolve alike.
reconcile_diffusion_attention_overrides(engine, stage_cli_overrides)
- engine.update(_copy_value(stage_cli_overrides))
+ for key, value in stage_cli_overrides.items():
+ existing = engine.get(key)
+ if key != "omni_kv_config" and isinstance(existing, dict) and isinstance(value, Mapping):
+ engine[key] = _get_recursively_merged_dict(existing, dict(value))
+ else:
+ engine[key] = _copy_value(value)
_validate_stage_engine_override_ownership(
topology.stage_id,
topology.execution_type,
@@ -1330,6 +1362,10 @@ def _stage_engine_values(
),
model=cast(_ModelEngineOverrides, _select_engine_overrides(engine, _MODEL_ENGINE_FIELDS)),
load=cast(_LoadEngineOverrides, _select_engine_overrides(engine, load_engine_fields)),
+ pooling=cast(
+ _PoolingEngineOverrides,
+ _select_engine_overrides(engine, _POOLING_ENGINE_FIELDS),
+ ),
cache=cast(_CacheEngineOverrides, _select_engine_overrides(engine, cache_engine_fields)),
scheduler=cast(
_SchedulerEngineOverrides,
@@ -1402,6 +1438,7 @@ class BaseVllmOmniStageConfig:
cache_config: OmniStageCacheConfig = field(default_factory=OmniStageCacheConfig)
scheduler_config: OmniStageSchedulerConfig = field(default_factory=OmniStageSchedulerConfig)
connector_config: OmniStageConnectorConfig = field(default_factory=OmniStageConnectorConfig)
+ pooling_config: OmniStagePoolingConfig = field(default_factory=OmniStagePoolingConfig)
runtime_config: OmniStageRuntimeConfig = field(default_factory=OmniStageRuntimeConfig)
parallel_config: OmniStageParallelConfig = field(default_factory=OmniStageParallelConfig)
compilation_config: VllmCompilationConfig | None = None
@@ -1486,6 +1523,10 @@ class BaseVllmOmniStageConfig:
def prompt_transform_func(self) -> str | None:
return self.stage_pipeline_config.prompt_transform_func
+ @property
+ def sampling_constraints(self) -> dict[str, Any]:
+ return dict(self.stage_pipeline_config.sampling_constraints)
+
@property
def cfg_kv_collect_func(self) -> str | None:
return self.stage_pipeline_config.cfg_kv_collect_func
@@ -1539,6 +1580,7 @@ def _build_common_stage_config_kwargs(
model=model,
),
"load_config": _build_load_config(topology, engine.load),
+ "pooling_config": _build_pooling_config(stage_deploy, engine.pooling),
"cache_config": _build_cache_config(
deploy,
engine.cache,
@@ -1752,6 +1794,20 @@ def _build_model_config(
)
+def _build_pooling_config(
+ stage_deploy: StageDeployConfig | None,
+ engine: _PoolingEngineOverrides,
+) -> OmniStagePoolingConfig:
+ default_pooling_params = None
+ if stage_deploy is not None and stage_deploy.default_pooling_params:
+ default_pooling_params = PoolingParams(**dict(stage_deploy.default_pooling_params))
+ return cast(Any, OmniStagePoolingConfig)(
+ runner=_copy_value(engine.get("runner")),
+ pooling_output_decoder=_copy_value(engine.get("pooling_output_decoder")),
+ default_pooling_params=default_pooling_params,
+ )
+
+
def _build_load_config(
topology: StagePipelineConfig,
engine: _LoadEngineOverrides,
@@ -1930,6 +1986,10 @@ class VllmOmniConfig:
pipeline_config: PipelineConfig
stage_configs: tuple[StageConfigType, ...]
orchestrator_config: VllmOmniOrchestratorConfig = field(default_factory=VllmOmniOrchestratorConfig)
+ # Keep strategy provenance separate from the effective orchestrator value:
+ # an explicit CLI policy may be present even when the strategy declares no
+ # stage-replica axis.
+ strategy_omni_lb_policy: str | None = None
def stage_by_id(self, stage_id: int) -> StageConfigType:
for stage in self.stage_configs:
@@ -1945,6 +2005,7 @@ class VllmOmniConfig:
user_deploy_config: DeployConfig | None = None,
deploy_config_path: str | None = None,
cli_overrides: dict[str, Any] | None = None,
+ strategy_specs: Mapping[Any, Any] | None = None,
) -> VllmOmniConfig:
"""Create a structured config from a resolved pipeline and deploy YAML."""
if cli_overrides is None:
@@ -1969,6 +2030,62 @@ class VllmOmniConfig:
deploy.async_chunk = False
_validate_async_chunk_support(pipeline_cfg, deploy)
validate_stage_async_chunk_edges(pipeline_cfg, deploy)
+
+ strategy_result = None
+ if strategy_specs:
+ from vllm_omni.config.composable_parallel import apply_strategy_specs
+
+ strategy_stages = merge_pipeline_deploy(pipeline_cfg, copy.deepcopy(deploy), {})
+ strategy_result = apply_strategy_specs(strategy_stages, strategy_specs)
+ strategy_overrides: dict[str, Any] = {}
+ axis_fields = {
+ "tp": "tensor_parallel_size",
+ "dp": "data_parallel_size",
+ "pp": "pipeline_parallel_size",
+ }
+ for stage_id, parallel in strategy_result.per_stage_config.items():
+ declared = set(parallel.l1_owners)
+ explicit = build_stage_runtime_overrides(stage_id, cli_overrides)
+ for axis, field_name in axis_fields.items():
+ if axis not in declared:
+ continue
+ derived = getattr(parallel, field_name)
+ value = explicit.get(field_name)
+ if value is None:
+ value = derived
+ elif value != 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_id,
+ field_name,
+ value,
+ field_name,
+ derived,
+ )
+ strategy_overrides[f"stage_{stage_id}_{field_name}"] = value
+ if "ep" in declared and parallel.enable_expert_parallel:
+ strategy_overrides[f"stage_{stage_id}_enable_expert_parallel"] = bool(
+ explicit.get("enable_expert_parallel", True)
+ )
+ if "stage_replica" in declared:
+ derived = parallel.stage_replica_size
+ value = explicit.get("num_replicas")
+ if value is None:
+ value = derived
+ elif value != derived:
+ 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_id,
+ value,
+ derived,
+ )
+ strategy_overrides[f"stage_{stage_id}_num_replicas"] = value
+ cli_overrides = {**strategy_overrides, **cli_overrides}
+ if strategy_result.omni_lb_policy is not None and cli_overrides.get("omni_lb_policy") is None:
+ cli_overrides["omni_lb_policy"] = strategy_result.omni_lb_policy
+
deploy_by_id = {stage.stage_id: stage for stage in deploy.stages}
model = cli_overrides.get("model")
@@ -1991,6 +2108,20 @@ class VllmOmniConfig:
)
for topology in pipeline_cfg.stages
)
+ if strategy_result is not None:
+ from vllm_omni.config.composable_parallel import check_device_layout
+
+ by_id = {stage.stage_id: stage for stage in stage_configs}
+ for stage_id in strategy_result.per_stage_config:
+ stage = by_id[stage_id]
+ check_device_layout(
+ stage.runtime_config.devices,
+ tensor_parallel_size=stage.parallel_config.tensor_parallel_size,
+ data_parallel_size=stage.parallel_config.data_parallel_size,
+ pipeline_parallel_size=stage.parallel_config.pipeline_parallel_size,
+ num_replicas=stage.runtime_config.num_replicas,
+ role=stage.model_stage,
+ )
orchestrator_config = cast(Any, VllmOmniOrchestratorConfig)(
deploy_config_path=loaded_deploy_config_path,
@@ -2000,6 +2131,7 @@ class VllmOmniConfig:
pipeline_config=pipeline_cfg,
stage_configs=stage_configs,
orchestrator_config=orchestrator_config,
+ strategy_omni_lb_policy=(strategy_result.omni_lb_policy if strategy_result is not None else None),
)
diff --git a/vllm_omni/config/resolver.py b/vllm_omni/config/resolver.py
index 4f5ee8a6c..0db436768 100644
--- a/vllm_omni/config/resolver.py
+++ b/vllm_omni/config/resolver.py
@@ -14,7 +14,6 @@ from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remot
from vllm_omni.config.endpoint_policy import EndpointRestriction
from vllm_omni.config.omni_config import VllmOmniConfig
from vllm_omni.config.stage_config import PipelineConfig
-from vllm_omni.config.yaml_util import create_config
from vllm_omni.diffusion.data import resolve_model_class_name
from vllm_omni.diffusion.registry import DiffusionModelRegistry
from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
@@ -24,19 +23,10 @@ logger = init_logger(__name__)
@dataclass(frozen=True)
class OmniConfigResolution:
- """Migration envelope returned by the production config resolver.
-
- ``stage_configs`` intentionally carries the current OmegaConf-compatible
- runtime ABI only until stage startup consumes ``VllmOmniConfig`` directly.
- ``pipeline_config`` is the effective topology after runtime stage injection,
- so both fields always describe the same set of stages.
- It is not a stable authoring or extension API; new production callers
- should resolve through :func:`resolve_omni_config` and must not construct or
- merge this compatibility shape themselves.
- """
+ """Canonical typed result returned by the production config resolver."""
config_path: str | None
- stage_configs: tuple[Any, ...] # Temporary StageConfig/OmegaConf bridge.
+ stage_configs: tuple[Any, ...]
pipeline_config: PipelineConfig | None = None
omni_lb_policy: str | None = None
@@ -176,27 +166,13 @@ def _load_strategy_specs(strategy_config_path: str | None) -> Mapping[Any, Any]
def _build_registered_resolution(
structured_config: VllmOmniConfig,
- *,
- cli_overrides: dict[str, Any],
- strategy_config_path: str | None,
) -> OmniConfigResolution:
- """Build the temporary runtime view for an already-resolved pipeline."""
- # Runtime consumers have not yet moved to typed per-stage configs. Use the
- # factory-owned compatibility bridge instead of reimplementing legacy YAML
- # discovery and merging in this resolver.
- effective_deploy_path = structured_config.orchestrator_config.deploy_config_path
- legacy_resolution = StageConfigFactory._resolve_legacy_from_registry(
- structured_config.pipeline_config,
- cli_overrides,
- effective_deploy_path,
- strategy_specs=_load_strategy_specs(strategy_config_path),
- )
-
+ """Return the canonical structured runtime view for a resolved pipeline."""
return OmniConfigResolution(
- config_path=effective_deploy_path,
- stage_configs=tuple(stage.to_omegaconf() for stage in legacy_resolution.stage_configs),
- pipeline_config=legacy_resolution.pipeline_config,
- omni_lb_policy=legacy_resolution.omni_lb_policy,
+ config_path=structured_config.orchestrator_config.deploy_config_path,
+ stage_configs=tuple(structured_config.stage_configs),
+ pipeline_config=structured_config.pipeline_config,
+ omni_lb_policy=getattr(structured_config, "strategy_omni_lb_policy", None),
)
@@ -231,18 +207,16 @@ def resolve_omni_config(
registry_overrides = dict(normalized_overrides)
_flatten_stage_overrides(registry_overrides, stage_overrides)
+ strategy_specs = _load_strategy_specs(strategy_config_path)
structured_config = StageConfigFactory.create_from_model(
model,
trust_remote_code=trust_remote_code,
cli_overrides=registry_overrides,
deploy_config_path=deploy_config_path,
+ strategy_specs=strategy_specs,
)
if structured_config is not None:
- return _build_registered_resolution(
- structured_config,
- cli_overrides=registry_overrides,
- strategy_config_path=strategy_config_path,
- )
+ return _build_registered_resolution(structured_config)
_apply_generic_stage_overrides(normalized_overrides, stage_overrides)
supported, model_class_name = _resolve_generic_diffusion_model_class(model, normalized_overrides)
@@ -252,11 +226,12 @@ def resolve_omni_config(
)
if model_class_name is not None:
normalized_overrides.setdefault("model_class_name", model_class_name)
- default_stages = StageConfigFactory.create_default_diffusion(normalized_overrides)
+ default_config = StageConfigFactory.create_typed_default_diffusion(model, normalized_overrides)
return OmniConfigResolution(
config_path=deploy_config_path,
- stage_configs=tuple(create_config(_convert_dataclasses_to_dict(default_stages))),
+ stage_configs=tuple(default_config.stage_configs),
+ pipeline_config=default_config.pipeline_config,
)
diff --git a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml
index 7914ecd52..36c85e41d 100644
--- a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml
+++ b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml
@@ -9,6 +9,15 @@
# Profile validated for the c64 / PROMPTS=512 performance experiments:
# Stage 0 talker on GPU 0 with S0=64, Stage 1 Code2Wav on GPU 1 with S1=10.
#
+# For CPU-bound frontend or orchestration workloads, this profile can be used
+# with ``--api-server-count 2`` or higher to run multiple API frontends over
+# one shared set of local stage engines. This can improve request parsing,
+# multimodal/audio preprocessing, and response handling under concurrency;
+# it does not duplicate GPU engines or make a single GPU-bound request faster.
+# Each frontend keeps process-local state, so runtime voice mutations, sleep /
+# wakeup, async video jobs, and resumable duplex sessions are not supported
+# with more than one API worker. See docs/cli/serve.md for the full boundary.
+#
async_chunk: true
connectors:
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index e1ce550d4..0ffd0c54e 100644
--- a/vllm_omni/diffusion/data.py
+++ b/vllm_omni/diffusion/data.py
@@ -1021,6 +1021,8 @@ class OmniDiffusionConfig:
# str is resolved to {"method": } internally.
# Per-component: {"transformer": {"method": "fp8"}, "vae": None}
quantization_config: str | QuantizationConfig | dict[str, Any] | None = None
+ # Internal provenance, retained across config projection and worker transport.
+ quantization_config_is_auto_detected: bool = False
# Explicit runtime override for ModelOpt FP8 diffusion checkpoints. This
# does not enable FP8 by itself; it only selects CUTLASS once the checkpoint
# has already resolved to vLLM's ModelOpt FP8 linear method.
@@ -1358,6 +1360,8 @@ class OmniDiffusionConfig:
or (is_checkpoint_nvfp4 and self._is_generic_nvfp4_quant_config(self.quantization_config))
)
if should_use_checkpoint_config:
+ if self.quantization_config is None:
+ self.quantization_config_is_auto_detected = True
self.quantization_config = tf_config.quant_config
logger.info(
"Auto-detected quantization '%s' from model config",
diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py
index 78247e7b9..834887f41 100644
--- a/vllm_omni/diffusion/model_loader/diffusers_loader.py
+++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py
@@ -7,7 +7,7 @@ import json
import os
import re
import time
-from collections.abc import Generator, Iterable, Sequence
+from collections.abc import Callable, Generator, Iterable, Sequence
from pathlib import Path
from typing import cast
@@ -481,8 +481,14 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
)
def _get_expected_parameter_names(self, model: nn.Module) -> set[str]:
- """Return parameter names that should be covered by strict load checks."""
- all_parameter_names = {name for name, _ in model.named_parameters()}
+ """Return parameter names that should be covered by strict load checks.
+
+ A parameter with an initialized checkpoint default can explicitly set
+ ``is_checkpoint_optional``. It is still loaded when present on disk.
+ """
+ all_parameter_names = {
+ name for name, param in model.named_parameters() if not getattr(param, "is_checkpoint_optional", False)
+ }
sources = self._get_weight_sources(model)
# Keep strict behavior if no source metadata exists.
@@ -853,9 +859,24 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
del model
return self.load_fresh_canonical_model()
raise
+ self._log_w4a8_fallback_load_summaries(model)
self._attach_offload_startup_state(model)
return model
+ @staticmethod
+ def _log_w4a8_fallback_load_summaries(model: nn.Module) -> None:
+ """Ask discovered DiTs to report W4A8 state at the common load exit.
+
+ Each DiT derives readiness from its processed layers, so a deferred
+ weight plan cannot be mistaken for the ordinary or HSDP load path.
+ """
+ components = ModuleDiscovery.discover(model)
+ for component_name, dit in zip(components.dit_names, components.dits):
+ candidate = getattr(dit, "_log_w4a8_fallback_load_summary", None)
+ if callable(candidate):
+ reporter = cast(Callable[[str], None], candidate)
+ reporter(component_name)
+
@staticmethod
def _request_offload_after_quant(model: nn.Module) -> int:
"""Ask online-quant layers to return to host memory once quantized.
@@ -1065,6 +1086,16 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
# that have loaded weights tracking currently.
if loaded_weights is not None:
weights_not_loaded = weights_to_load - loaded_weights
+ # Offline formats can require scales that older online loaders
+ # were allowed to synthesize. Do not apply that legacy tolerance
+ # to an explicitly required checkpoint tensor.
+ required_missing = {
+ name
+ for name, param in model.named_parameters()
+ if name in weights_not_loaded and getattr(param, "is_checkpoint_required", False)
+ }
+ if required_missing:
+ raise ValueError(f"Required weights were not initialized from checkpoint: {required_missing}")
# NOTE: if the model is quantized, ignore not_loaded check for scale
# weights. ModelOpt FP8 carries a per-tensor `weight_scale` and a
# static activation `input_scale`, which the quant method may
diff --git a/vllm_omni/diffusion/models/minimax_h3/fasth3.py b/vllm_omni/diffusion/models/minimax_h3/fasth3.py
index c4860c57e..dcc513028 100644
--- a/vllm_omni/diffusion/models/minimax_h3/fasth3.py
+++ b/vllm_omni/diffusion/models/minimax_h3/fasth3.py
@@ -102,6 +102,7 @@ _MODEL_LEVEL_TARGETS = {
"time_embedder.linear_2": "time_embedder.proj_out",
"norm_out.linear": "final_layer.adaln_proj.linear",
"norm_out.norm": "final_layer.norm",
+ "token_refiner.final_norm": "token_refiner.final_norm",
}
# Per-block adapter suffix -> (native suffix, how a delta enters the native
@@ -115,6 +116,8 @@ _BLOCK_TARGETS = {
"attn.to_v": ("attn.qkv_proj", _QKV_V),
"attn.to_out.0": ("attn.out_proj", _PLAIN),
"attn.to_gate_compress": ("attn.to_gate_compress", _PLAIN),
+ "attn.norm_q": ("attn.q_norm", _PLAIN),
+ "attn.norm_k": ("attn.k_norm", _PLAIN),
"ff.net.0.proj": ("mlp.fc1", _SWAP_HALVES),
"ff.net.2": ("mlp.fc2", _PLAIN),
"adaln_proj.linear": ("adaln_proj.linear", _PLAIN),
diff --git a/vllm_omni/diffusion/models/minimax_h3/fasth3_checkpoint.py b/vllm_omni/diffusion/models/minimax_h3/fasth3_checkpoint.py
new file mode 100644
index 000000000..0382e9220
--- /dev/null
+++ b/vllm_omni/diffusion/models/minimax_h3/fasth3_checkpoint.py
@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Sampling contract for FastVideo FastH3 full checkpoints."""
+
+from __future__ import annotations
+
+import json
+import math
+from collections.abc import Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule
+from vllm_omni.errors import OmniClientError
+from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+
+from .fasth3 import FASTH3_BASE_MODEL, _resolve_dit_attention_backend
+
+if TYPE_CHECKING:
+ from vllm_omni.diffusion.data import OmniDiffusionConfig
+ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+FASTH3_V2_MODEL_ID = "FastVideo/FastVideo-FastH3-8-Step-V2"
+FASTH3_V2_BASE_SCHEDULE = DMD2SigmaSchedule.from_positions((0.999, 0.874, 0.749, 0.624, 0.5, 0.375, 0.25, 0.125, 0.0))
+
+
+@dataclass(frozen=True)
+class FastH3CheckpointSpec:
+ """The full V2 release has its own schedule and trained attention policy.
+
+ This is independent of ``FastH3WeightFusion``: the released weights already
+ contain the entire student, including its learned compression gates.
+ """
+
+ vsa_sparsity: float = 0.8
+
+ @classmethod
+ def from_metadata(cls, metadata: Mapping[str, object]) -> FastH3CheckpointSpec:
+ """Validate the release's own fastvideo_inference.json."""
+ required = {
+ "model_id": FASTH3_V2_MODEL_ID,
+ "schema_version": "fasth3-inference-contract-v1",
+ "dmd_denoising_steps": [999, 874, 749, 624, 500, 375, 250, 125],
+ "transformer_forwards": 8,
+ "video_scheduler_shift": 10.0,
+ "audio_scheduler_shift": 3.0,
+ "guidance_scale": 1.0,
+ "attention_backend": "VIDEO_SPARSE_ATTN_H3",
+ "vsa_sparsity": 0.8,
+ "vsa_tile_size": 64,
+ "task": "t2av",
+ }
+ for key, expected in required.items():
+ if metadata.get(key) != expected:
+ raise ValueError(f"unsupported FastH3 contract: {key}={metadata.get(key)!r}, expected {expected!r}")
+ return cls()
+
+ def release_metadata(self) -> dict[str, Any]:
+ """Express the sampling policy in the existing H3 pipeline schema."""
+ return {
+ "partition": "fl2va",
+ "tasks": ["t2va"],
+ "sigma_shift_scales": {"video": 10.0, "audio": 3.0},
+ "base_schedule": list(FASTH3_V2_BASE_SCHEDULE.base_schedule),
+ }
+
+ def resolve_native_vaes(self, model_root: Path) -> Path:
+ """Reuse the frozen base VAEs with Omni's native tiled/parallel runtime.
+
+ Only VAE components are fetched; the student and text encoder come
+ directly from the FastVideo release. The release pins the base revision.
+ """
+ provenance = json.loads((model_root / "provenance.json").read_text(encoding="utf-8"))
+ base = provenance.get("base_model", "")
+ prefix = f"hf://{FASTH3_BASE_MODEL}@"
+ if not isinstance(base, str) or not base.startswith(prefix) or not base[len(prefix) :]:
+ raise ValueError("FastH3 V2 provenance must pin a MiniMaxAI/MiniMax-H3 base revision")
+ return (
+ Path(
+ download_weights_from_hf_specific(
+ model_name_or_path=FASTH3_BASE_MODEL,
+ cache_dir=None,
+ allow_patterns=["FL2VA/video_vae/**", "FL2VA/audio_vae/**"],
+ revision=base[len(prefix) :],
+ require_all=True,
+ )
+ )
+ / "FL2VA"
+ )
+
+ def check_serving_contract(self, *, partition: str, od_config: OmniDiffusionConfig) -> None:
+ if partition != "fl2va":
+ raise ValueError("FastH3 V2 requires --task-type fl2va (T2VA requests only)")
+ if getattr(od_config, "lora_path", None):
+ raise ValueError("FastH3 V2 is a full checkpoint; additional LoRA adapters are unsupported")
+ if _resolve_dit_attention_backend(od_config) != "FASTVIDEO_VSA":
+ raise ValueError("FastH3 V2 requires --diffusion-attention-backend FASTVIDEO_VSA")
+ attention_config = getattr(od_config, "diffusion_attention_config", None)
+ per_role = getattr(attention_config, "per_role", None) or {}
+ spec = per_role.get("self") or getattr(attention_config, "default", None)
+ if getattr(spec, "fastvideo_vsa_topk", None) is not None:
+ raise ValueError("FastH3 V2 pins VSA sparsity=0.8; remove the fixed fastvideo_vsa_topk override")
+ parallel = getattr(od_config, "parallel_config", None)
+ if any(int(getattr(parallel, key, 1) or 1) != 1 for key in ("ring_degree", "allgather_degree")):
+ raise ValueError("FastH3 V2 supports local attention or pure Ulysses sequence parallelism")
+
+ def check_request(self, sampling: OmniDiffusionSamplingParams, *, step_execution: bool = False) -> None:
+ if step_execution:
+ # StepScheduler admits requests before the pipeline hook runs and
+ # derives its lifetime from these fields. The pinned schedule must
+ # therefore be explicit and cannot be overridden by custom arrays.
+ if sampling.num_inference_steps != FASTH3_V2_BASE_SCHEDULE.num_inference_steps:
+ raise OmniClientError("FastH3 V2 step execution requires num_inference_steps=8")
+ if sampling.timesteps is not None or sampling.sigmas is not None:
+ raise OmniClientError("FastH3 V2 step execution does not support custom timesteps or sigmas")
+ if sampling.lora_request is not None:
+ raise OmniClientError("FastH3 V2 does not support per-request LoRA adapters")
+ steps = sampling.num_inference_steps
+ if steps is not None and steps != FASTH3_V2_BASE_SCHEDULE.num_inference_steps:
+ raise OmniClientError("FastH3 V2 requires num_inference_steps=8 (nine sigma points), or omitted")
+ extra = sampling.extra_args or {}
+ for key, expected in (("flow_shift", 10.0), ("audio_flow_shift", 3.0)):
+ try:
+ value = float(extra.get(key, expected))
+ except (TypeError, ValueError) as exc:
+ raise OmniClientError(f"FastH3 V2 requires {key}={expected:g}") from exc
+ if not math.isclose(value, expected):
+ raise OmniClientError(f"FastH3 V2 requires {key}={expected:g}, got {value:g}")
+ if sampling.guidance_scale is not None and sampling.guidance_scale != 1.0:
+ raise OmniClientError("FastH3 V2 requires guidance_scale=1")
diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py
index e26493931..5662a9641 100644
--- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py
+++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py
@@ -55,6 +55,8 @@ from vllm_omni.diffusion.layers.rope import RotaryEmbedding
from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
from vllm_omni.platforms import current_omni_platform
+from .fasth3 import _resolve_native_target
+
if TYPE_CHECKING:
from vllm.model_executor.layers.quantization.base_config import (
QuantizationConfig,
@@ -118,6 +120,18 @@ class MiniMaxH3DiTArchConfig:
@classmethod
def from_mapping(cls, config: Mapping[str, Any]) -> MiniMaxH3DiTArchConfig:
+ # The modular Diffusers checkpoint uses these spellings for the same
+ # architecture. Normalize before constructing the existing native DiT.
+ aliases = {
+ "num_refiner_layers": "token_refiner_num_layers",
+ "ffn_dim": "ffn_hidden_size",
+ "in_channels": "latents_dim",
+ "audio_in_channels": "audio_latents_dim",
+ "freq_dim": "timestep_input_dim",
+ "time_embed_hidden_dim": "time_embed_hidden_size",
+ "rope_freq_dim": "rope_inv_freq_len",
+ }
+ config = {aliases.get(name, name): value for name, value in config.items()}
fields = cls.__dataclass_fields__
values = {name: config[name] for name in fields if name in config}
if "patch_size" in values:
@@ -394,6 +408,9 @@ def _sdpa_varlen_attention(
class MiniMaxH3Attention(nn.Module):
+ # Full sparse checkpoints pin a ratio; legacy adapters use backend top-k.
+ vsa_sparsity: float | None = None
+
def __init__(
self,
arch: MiniMaxH3DiTArchConfig,
@@ -588,6 +605,7 @@ class MiniMaxH3Attention(nn.Module):
# (see MINIMAX_H3_LASER_INPUT_SCALE). Ignored by every other
# backend/path.
"laser_input_scale": MINIMAX_H3_LASER_INPUT_SCALE,
+ **({"vsa_h3_sparsity": self.vsa_sparsity} if self.vsa_sparsity is not None else {}),
# Present only for a VSA artifact; the VSA backend reads it as
# the learned compression gate and every other backend ignores it.
**({"gate_compress": gate_compress.unsqueeze(0)} if gate_compress is not None else {}),
@@ -1127,10 +1145,23 @@ class MiniMaxH3DiTModel(nn.Module):
self,
od_config: OmniDiffusionConfig,
quant_config: QuantizationConfig | None = None,
+ *,
+ diffusers_weights: bool | None = None,
) -> None:
super().__init__()
tf_config = od_config.tf_model_config
config_mapping = tf_config.to_dict() if hasattr(tf_config, "to_dict") else dict(tf_config)
+ # The native MiniMax-H3 Hub snapshot advertises the Diffusers
+ # transformer class in its root config, while its FL2VA/Ref2VA
+ # components still contain native weights. The pipeline has already
+ # resolved the actual source format, so let it override the
+ # class-name heuristic. Keep the heuristic for standalone callers.
+ self._diffusers_weights = (
+ config_mapping.get("_class_name") == "MiniMaxH3Transformer3DModel"
+ if diffusers_weights is None
+ else diffusers_weights
+ )
+ self._rope_theta = float(config_mapping.get("rope_theta", 10000.0))
arch = MiniMaxH3DiTArchConfig.from_mapping(config_mapping)
self.arch = arch
self.od_config = od_config
@@ -1211,7 +1242,7 @@ class MiniMaxH3DiTModel(nn.Module):
)
self._mark_missing_params_required()
- def enable_vsa_gates(self) -> None:
+ def enable_vsa_gates(self, *, sparsity: float | None = None) -> None:
"""Give every DiT block's attention a VSA compression gate.
A FastH3 VSA artifact assigns these projections rather than adding to
@@ -1223,6 +1254,9 @@ class MiniMaxH3DiTModel(nn.Module):
return
for block in self.blocks:
block.attn.enable_vsa_gate()
+ block.attn.vsa_sparsity = sparsity
+ if sparsity is not None:
+ block.attn.to_gate_compress.weight.missing_param_init = "error"
self.vsa_gates_enabled = True
def _mark_missing_params_required(self) -> None:
@@ -1293,17 +1327,37 @@ class MiniMaxH3DiTModel(nn.Module):
self,
weights: Iterable[tuple[str, torch.Tensor]],
) -> set[str]:
- """Load exact H3 checkpoint names with logical TP-aware loaders."""
+ """Load native or Diffusers H3 weights with the existing TP-aware loaders."""
params = dict(self.named_parameters())
params.update(dict(self.named_buffers()))
loaded: set[str] = set()
+ diffusers_weights = getattr(self, "_diffusers_weights", False)
+ qkv_parts: dict[str, set[str]] = {}
+ source_names: set[str] = set()
for name, loaded_weight in weights:
+ layout = "plain"
+ if diffusers_weights:
+ if name in source_names:
+ raise ValueError(f"duplicate Diffusers H3 weight: {name}")
+ source_names.add(name)
+ module, _, kind = name.rpartition(".")
+ target = _resolve_native_target(module)
+ if target is None or kind not in {"weight", "bias"}:
+ raise ValueError(f"unsupported Diffusers H3 weight: {name}")
+ name, layout = f"{target[0]}.{kind}", target[1]
param = params.get(name)
if param is None:
+ if diffusers_weights:
+ raise ValueError(f"Diffusers H3 weight has no model parameter: {name}")
logger.warning("Skipping MiniMax H3 weight not present in model: %s", name)
continue
weight_loader = getattr(param, "weight_loader", default_weight_loader)
- if name.endswith(".attn.qkv_proj.weight"):
+ if layout in {"q", "k", "v"}:
+ # vLLM can load each projection directly into its packed QKV
+ # parameter, including TP slicing and online quantization.
+ weight_loader(param, loaded_weight, layout)
+ qkv_parts.setdefault(name, set()).add(layout)
+ elif name.endswith(".attn.qkv_proj.weight"):
# Transform checkpoint layout before entering vLLM's loader so
# online FP8 can keep ``online_process_loader`` outermost.
loaded_weight = _reorder_grouped_qkv_to_qkv(
@@ -1319,12 +1373,26 @@ class MiniMaxH3DiTModel(nn.Module):
"MiniMax H3 fc1 checkpoint rows must split evenly into "
f"gate/up matrices, got {tuple(loaded_weight.shape)}"
)
- gate, up = loaded_weight.chunk(2, dim=0)
+ first, second = loaded_weight.chunk(2, dim=0)
+ gate, up = (second, first) if layout == "swap_halves" else (first, second)
weight_loader(param, gate, 0)
weight_loader(param, up, 1)
else:
weight_loader(param, loaded_weight)
loaded.add(name)
+ if diffusers_weights:
+ for name, parts in qkv_parts.items():
+ if parts != {"q", "k", "v"}:
+ raise ValueError(f"incomplete Diffusers H3 QKV group {name}: {sorted(parts)}")
+ # Diffusers reconstructs RoPE from config instead of storing this
+ # native checkpoint buffer. Compute on CPU for identical values.
+ freq_dim = self.arch.rope_inv_freq_len
+ rope = 1.0 / (
+ self._rope_theta
+ ** (torch.arange(0, 2 * freq_dim, 2, dtype=torch.float32, device="cpu") / (2 * freq_dim))
+ )
+ default_weight_loader(params["rope.inv_freq"], rope)
+ loaded.add("rope.inv_freq")
return loaded
@staticmethod
diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py
index 1b8bd7fea..60d6f4125 100644
--- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py
+++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py
@@ -67,6 +67,7 @@ from vllm_omni.model_executor.model_loader.weight_utils import (
download_weights_from_hf_specific,
)
from vllm_omni.model_executor.models.minimax_h3.checkpoint import (
+ is_minimax_h3_modular,
resolve_minimax_h3_partition,
)
from vllm_omni.model_executor.models.minimax_h3.conditioning import (
@@ -106,6 +107,7 @@ from .denoise_loop import (
)
from .encoder import MiniMaxH3Qwen3VLEncoder
from .fasth3 import FastH3WeightFusion, resolve_fasth3_fusion
+from .fasth3_checkpoint import FastH3CheckpointSpec
from .lora import TurboSpec, load_minimax_h3_turbo_lora
from .minimax_h3_transformer import (
MiniMaxH3Attention,
@@ -213,7 +215,11 @@ def _resolve_minimax_h3_model_root(
if path.name in {"FL2VA", "Ref2VA"}:
return path.parent
return path
- if load_text_encoder:
+ if is_minimax_h3_modular(model, revision):
+ allow_patterns = ["modular_model_index.json", "fastvideo_inference.json", "provenance.json", "transformer/**"]
+ if load_text_encoder:
+ allow_patterns += ["text_encoder/**", "tokenizer/**", "processor/**"]
+ elif load_text_encoder:
allow_patterns = (
MINIMAX_H3_DOWNLOAD_PATTERNS if partition == "combined" else MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[partition]
)
@@ -314,6 +320,8 @@ def resolve_minimax_h3_diffusion_model_path(
partition,
load_text_encoder=False,
)
+ if is_minimax_h3_modular(str(model_root), revision):
+ return str(model_root)
if partition == "combined":
return str(model_root)
subdir = "Ref2VA" if partition == "ref2va" else "FL2VA"
@@ -575,6 +583,7 @@ class MiniMaxH3Pipeline(
_base_schedule_by_partition: ClassVar[Mapping[str, DMD2SigmaSchedule | None]] = {}
# Set from --lora-path during construction; absent means no FastH3 adapter.
_fasth3: FastH3WeightFusion | None = None
+ _fasth3_checkpoint: FastH3CheckpointSpec | None = None
def _load_diffusion_lora_adapter(
self,
@@ -800,10 +809,13 @@ class MiniMaxH3Pipeline(
self._PROFILER_TARGETS.remove("encode_prompt")
if not self.load_vae_encoder:
self._PROFILER_TARGETS.remove("_encode_local_media")
+ modular = is_minimax_h3_modular(str(od_config.model), od_config.revision)
self.partition = _minimax_h3_partition_for_task(
getattr(od_config, "task_type", None),
str(od_config.model),
)
+ if modular and str(od_config.task_type or "auto").lower() == "auto":
+ self.partition = "fl2va"
self._turbo_lora_specs: dict[int, TurboSpec] = {}
self._native_lora_adapter_ids: set[int] = set()
self._lora_sigma_schedules: dict[int, DMD2SigmaSchedule] = {}
@@ -813,9 +825,19 @@ class MiniMaxH3Pipeline(
self.partition,
load_text_encoder=self.load_text_encoder,
)
- model_path = model_root / ("Ref2VA" if self.partition == "ref2va" else "FL2VA")
- model_index = json.loads((model_path / "model_index.json").read_text(encoding="utf-8"))
- release = model_index.get("_minimax_h3") or {}
+ if modular:
+ model_path = model_root
+ self._fasth3_checkpoint = FastH3CheckpointSpec.from_metadata(
+ json.loads((model_root / "fastvideo_inference.json").read_text(encoding="utf-8"))
+ )
+ self._fasth3_checkpoint.check_serving_contract(partition=self.partition, od_config=od_config)
+ release = self._fasth3_checkpoint.release_metadata()
+ vae_model_path = self._fasth3_checkpoint.resolve_native_vaes(model_root)
+ else:
+ model_path = model_root / ("Ref2VA" if self.partition == "ref2va" else "FL2VA")
+ model_index = json.loads((model_path / "model_index.json").read_text(encoding="utf-8"))
+ release = model_index.get("_minimax_h3") or {}
+ vae_model_path = model_path
partition = str(release.get("partition", "")).lower()
expected_partition = "ref2va" if self.partition == "ref2va" else "fl2va"
if partition != expected_partition:
@@ -876,11 +898,18 @@ class MiniMaxH3Pipeline(
self.transformer = MiniMaxH3DiTModel(
od_config,
quant_config=transformer_quant_config,
+ diffusers_weights=modular,
)
+ if self._fasth3_checkpoint is not None:
+ self.transformer.enable_vsa_gates(sparsity=self._fasth3_checkpoint.vsa_sparsity)
+ logger.info(
+ "FastH3 V2 full checkpoint: 8 transformer forwards, video/audio shifts 10/3, VSA sparsity=0.8 tile=64"
+ )
if ref2va_model_path is not None:
self.transformers_ref = MiniMaxH3DiTModel(
od_config,
quant_config=transformer_quant_config,
+ diffusers_weights=modular,
)
self._fasth3 = resolve_fasth3_fusion(od_config, self.transformer)
@@ -963,14 +992,14 @@ class MiniMaxH3Pipeline(
# so VAEs stay resident for new configurations.
component_load_device = torch.device("cpu") if legacy_manual_components else self.device
self.video_vae = MiniMaxH3VideoVAE(
- os.path.join(model_path, "video_vae"),
+ os.path.join(vae_model_path, "video_vae"),
device=self.device,
load_device=component_load_device,
decode_only=not self.load_vae_encoder,
trust_remote_code=od_config.trust_remote_code,
)
self.audio_vae = MiniMaxH3AudioVAE(
- os.path.join(model_path, "audio_vae"),
+ os.path.join(vae_model_path, "audio_vae"),
device=self.device,
load_device=component_load_device,
decode_only=not self.load_vae_encoder,
@@ -1046,6 +1075,12 @@ class MiniMaxH3Pipeline(
# load_weights only warns on a parameter the model does not have, so
# close the adapter against what the DiT actually consumed.
self._fasth3.validate_fully_applied(transformer_loaded)
+ if self._fasth3_checkpoint is not None:
+ required_gates = {
+ f"blocks.{i}.attn.to_gate_compress.weight" for i in range(self.transformer.arch.num_layers)
+ }
+ if missing_gates := required_gates - transformer_loaded:
+ raise ValueError(f"FastH3 V2 checkpoint is missing compression gates: {sorted(missing_gates)}")
return loaded_with_prefix
@property
@@ -2056,6 +2091,11 @@ class MiniMaxH3Pipeline(
self._validate_turbo_sampling(sampling, turbo_spec)
if has_native_lora:
self._validate_native_sampling(sampling, task=task)
+ if self._fasth3_checkpoint is not None:
+ self._fasth3_checkpoint.check_request(
+ sampling,
+ step_execution=bool(getattr(self.od_config, "step_execution", False)),
+ )
if self._fasth3 is not None:
self._fasth3.check_request(
sampling,
diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
index 15b036864..915fb45f7 100644
--- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
+++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
@@ -179,8 +179,36 @@ def load_transformer_config(model_path: str, subfolder: str = "transformer", loc
return {}
+def resolve_wan_transformer_quant_config(
+ config: dict, quant_config: QuantizationConfig | None, component: str
+) -> QuantizationConfig | None:
+ """Resolve the expert before applying its checkpoint's storage contract."""
+ from vllm_omni.quantization.component_config import ComponentQuantizationConfig
+ from vllm_omni.quantization.factory import resolve_quant_config_from_disk
+
+ # Wan experts are siblings: "transformer_2" must not inherit "transformer"
+ # through the generic layer-prefix resolver. Select the whole expert name.
+ component_quant_config = (
+ quant_config.component_configs.get(component, quant_config.default_config)
+ if isinstance(quant_config, ComponentQuantizationConfig)
+ else quant_config
+ )
+ quantization_disabled = isinstance(quant_config, ComponentQuantizationConfig) and component_quant_config is None
+ resolved_quant_config = resolve_quant_config_from_disk(component_quant_config, config.get("quantization_config"))
+ if quantization_disabled and resolved_quant_config is not None:
+ raise ValueError(
+ f"Quantization is disabled for component {component!r}, but its checkpoint declares quantization. "
+ "Use a BF16 checkpoint for this component or enable its matching quantization method."
+ )
+ return resolved_quant_config
+
+
def create_transformer_from_config(
- config: dict, quant_config: QuantizationConfig | None = None, prefix: str = ""
+ config: dict,
+ quant_config: QuantizationConfig | None = None,
+ prefix: str = "",
+ *,
+ component: str = "transformer",
) -> WanTransformer3DModel:
"""Create WanTransformer3DModel from config dict."""
kwargs: dict = {}
@@ -216,11 +244,7 @@ def create_transformer_from_config(
if "pos_embed_seq_len" in config:
kwargs["pos_embed_seq_len"] = config["pos_embed_seq_len"]
- if "quantization_config" in config:
- from vllm_omni.quantization.factory import resolve_quant_config_from_disk
-
- quant_config = resolve_quant_config_from_disk(quant_config, config["quantization_config"])
-
+ quant_config = resolve_wan_transformer_quant_config(config, quant_config, component)
if quant_config is not None:
kwargs["quant_config"] = quant_config
if prefix:
@@ -460,13 +484,13 @@ class Wan22Pipeline(
# Initialize transformers with correct config (weights loaded via load_weights)
if load_transformer:
transformer_config = load_transformer_config(model, "transformer", local_files_only)
- self.transformer = self._create_transformer(transformer_config)
+ self.transformer = self._create_transformer(transformer_config, component="transformer")
else:
self.transformer = None
if load_transformer_2:
transformer_2_config = load_transformer_config(model, "transformer_2", local_files_only)
- self.transformer_2 = self._create_transformer(transformer_2_config)
+ self.transformer_2 = self._create_transformer(transformer_2_config, component="transformer_2")
else:
self.transformer_2 = None
@@ -504,10 +528,17 @@ class Wan22Pipeline(
enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
)
- def _create_transformer(self, config: dict) -> WanTransformer3DModel:
+ def _create_transformer(self, config: dict, component: str = "transformer") -> WanTransformer3DModel:
"""Create a transformer from a config dict. Respects od_config.quantization_config."""
quant_config = getattr(self.od_config, "quantization_config", None)
- return create_transformer_from_config(config, quant_config=quant_config)
+ # Startup metadata describes the first expert, not a user policy for both.
+ if getattr(self.od_config, "quantization_config_is_auto_detected", False):
+ quant_config = None
+ return create_transformer_from_config(
+ config,
+ quant_config=quant_config,
+ component=component,
+ )
@property
def guidance_scale(self):
diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py
index 01bd53bb8..1a75c40da 100644
--- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py
+++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py
@@ -304,25 +304,10 @@ class Wan22I2VPipeline(
# Transformers (weights loaded via load_weights)
# Load config from model directory or HF Hub to get correct in_channels for I2V models
transformer_config = load_transformer_config(model, "transformer", local_files_only)
- self.transformer = create_transformer_from_config(
- transformer_config,
- quant_config=od_config.quantization_config,
- )
+ self.transformer = self._create_transformer(transformer_config, component="transformer")
if self.has_transformer_2:
transformer_2_config = load_transformer_config(model, "transformer_2", local_files_only)
- t2_quant = transformer_2_config.get("quantization_config")
- if isinstance(t2_quant, dict) and "quant_method" in t2_quant:
- from vllm_omni.quantization.factory import build_quant_config
-
- method = t2_quant["quant_method"]
- kwargs = {k: v for k, v in t2_quant.items() if k != "quant_method"}
- t2_quant = build_quant_config(method, **kwargs)
- else:
- t2_quant = None
- self.transformer_2 = create_transformer_from_config(
- transformer_2_config,
- quant_config=t2_quant,
- )
+ self.transformer_2 = self._create_transformer(transformer_2_config, component="transformer_2")
else:
self.transformer_2 = None
@@ -462,10 +447,13 @@ class Wan22I2VPipeline(
image_embeds = self.image_encoder(pixel_values, output_hidden_states=True)
return image_embeds.hidden_states[-2]
- def _create_transformer(self, config: dict) -> WanTransformer3DModel:
+ def _create_transformer(self, config: dict, component: str = "transformer") -> WanTransformer3DModel:
"""Create a transformer from a config dict. Respects od_config.quantization_config."""
quant_config = getattr(self.od_config, "quantization_config", None)
- return create_transformer_from_config(config, quant_config=quant_config)
+ # Startup metadata describes the first expert, not a user policy for both.
+ if getattr(self.od_config, "quantization_config_is_auto_detected", False):
+ quant_config = None
+ return create_transformer_from_config(config, quant_config=quant_config, component=component)
def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
sampling_params_list = req.sampling_params_list
diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py
index aeecff608..0d5f3cd67 100644
--- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py
+++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py
@@ -28,6 +28,7 @@ from vllm_omni.diffusion.models.interface import SupportImageInput
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
Wan22Pipeline,
load_wan_weights_with_optional_gate,
+ resolve_wan_transformer_quant_config,
retrieve_latents,
)
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
@@ -60,6 +61,8 @@ def create_vace_transformer_from_config(
config: dict,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
+ *,
+ component: str = "transformer",
) -> WanVACETransformer3DModel:
"""Create WanVACETransformer3DModel from config dict."""
kwargs = {}
@@ -97,6 +100,7 @@ def create_vace_transformer_from_config(
kwargs["vace_layers"] = config["vace_layers"]
if "vace_in_channels" in config:
kwargs["vace_in_channels"] = config["vace_in_channels"]
+ quant_config = resolve_wan_transformer_quant_config(config, quant_config, component)
if quant_config is not None:
kwargs["quant_config"] = quant_config
if prefix:
@@ -205,10 +209,13 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
super().__init__(od_config=od_config, prefix=prefix)
- def _create_transformer(self, config: dict) -> WanVACETransformer3DModel:
+ def _create_transformer(self, config: dict, component: str = "transformer") -> WanVACETransformer3DModel:
"""Build VACE transformer. Respects od_config.quantization_config."""
quant_config = getattr(self.od_config, "quantization_config", None)
- return create_vace_transformer_from_config(config, quant_config=quant_config)
+ # Startup metadata describes the first expert, not a user policy for both.
+ if getattr(self.od_config, "quantization_config_is_auto_detected", False):
+ quant_config = None
+ return create_vace_transformer_from_config(config, quant_config=quant_config, component=component)
def diffuse(
self,
diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py
index dbc555604..2808dcee7 100644
--- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py
+++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py
@@ -18,7 +18,7 @@ from vllm.distributed import (
)
from vllm.logger import init_logger
from vllm.model_executor.layers.conv import Conv3dLayer
-from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear
+from vllm.model_executor.layers.linear import ColumnParallelLinear, LinearBase, QKVParallelLinear, RowParallelLinear
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.models.utils import (
@@ -45,6 +45,7 @@ from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
from vllm_omni.diffusion.layers.norm import LayerNorm, RMSNorm
from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan
from vllm_omni.platforms import current_omni_platform
+from vllm_omni.quantization.mxfp4_config import NPUMxfp4LinearMethod
logger = init_logger(__name__)
@@ -985,6 +986,68 @@ class WanTransformer3DModel(nn.Module):
# ROPE helper
self._cached_rope_emb = None
self._cached_rope_resolution = None
+ self._validate_w4a8_fallback_layers(quant_config)
+
+ def _validate_w4a8_fallback_layers(self, quant_config: QuantizationConfig | None) -> None:
+ """Reject checkpoint aliases/typos before loading weights or running inference."""
+ requested = tuple(getattr(quant_config, "w4a8_fallback_layers", []))
+ self._w4a8_fallback_layers = requested
+ self._local_w4a8_fallback_layers: tuple[str, ...] = ()
+ if not requested:
+ return
+ modules = dict(self.named_modules())
+ missing_pp_prefixes = [name for name, module in modules.items() if isinstance(module, PPMissingLayer)]
+ local_layers = [
+ name for name in requested if not any(name.startswith(prefix + ".") for prefix in missing_pp_prefixes)
+ ]
+ self._local_w4a8_fallback_layers = tuple(local_layers)
+ invalid = [name for name in local_layers if not isinstance(modules.get(name), LinearBase)]
+ if invalid:
+ raise ValueError(
+ "w4a8_fallback_layers requires exact runtime Linear paths relative to each Wan transformer; "
+ f"unknown/non-Linear paths: {invalid}. Use attn1.to_qkv for fused self-attention, "
+ "not checkpoint attn1.to_q/to_k/to_v names."
+ )
+
+ def _log_w4a8_fallback_load_summary(self, component_name: str) -> None:
+ """Log post-load transforms/cache readiness, not numerical acceptance."""
+ requested = self._w4a8_fallback_layers
+ if not requested:
+ return
+ modules = dict(self.named_modules())
+ local_layers = self._local_w4a8_fallback_layers
+ selected_layers = [
+ name for name in local_layers if getattr(modules[name].quant_method, "is_w4a8_fallback_layer", False)
+ ]
+ bf16_overrides = [name for name in local_layers if name not in selected_layers]
+ processed = [
+ name
+ for name in selected_layers
+ if getattr(modules[name], "_already_called_process_weights_after_loading", False)
+ ]
+ if not selected_layers:
+ processing_state = "not-required"
+ elif len(processed) == len(selected_layers):
+ processing_state = "ready"
+ else:
+ processing_state = "not-ready"
+
+ def qualified(names: Iterable[str]) -> list[str]:
+ return [f"{component_name}.{name}" for name in names]
+
+ logger.info(
+ "Wan W4A8 post-load summary: component=%s; selected=%d %s; BF16 overrides=%d %s; "
+ "weight/A8 processing state=%s (%d/%d selected layers processed); other PP ranks=%d",
+ component_name,
+ len(selected_layers),
+ qualified(selected_layers),
+ len(bf16_overrides),
+ qualified(bf16_overrides),
+ processing_state,
+ len(processed),
+ len(selected_layers),
+ len(requested) - len(local_layers),
+ )
@property
def dtype(self) -> torch.dtype:
@@ -1155,7 +1218,22 @@ class WanTransformer3DModel(nn.Module):
}
params_dict = dict(self.named_parameters())
+ single_scale_prefixes = {
+ name
+ for name, module in self.named_modules()
+ if isinstance(getattr(module, "quant_method", None), NPUMxfp4LinearMethod)
+ }
+
+ def reject_dualscale_tensor(name: str) -> None:
+ if name.endswith(".weight_dual_scale") and name.rsplit(".", 1)[0] in single_scale_prefixes:
+ raise ValueError(
+ f"Single-scale mxfp4 cannot load DualScale tensor {name}; "
+ "changing quant_method does not convert a checkpoint."
+ )
+
loaded_params: set[str] = set()
+ loaded_qkv_shards: dict[str, set[str]] = {}
+ qkv_smooth_scales: dict[str, torch.Tensor] = {}
for name, loaded_weight in weights:
name = weight_name_remapping.get(name, name)
@@ -1166,15 +1244,32 @@ class WanTransformer3DModel(nn.Module):
# Pre-fused to_qkv tensors (from offline MXFP8 merged checkpoint) fall
# through to the else branch and are loaded directly.
for param_name, weight_name, shard_id in stacked_params_mapping:
- if weight_name not in original_name:
+ if f"{weight_name}." not in original_name:
continue
- lookup_name = original_name.replace(weight_name, param_name)
+ lookup_name = original_name.replace(f"{weight_name}.", f"{param_name}.", 1)
+ reject_dualscale_tensor(lookup_name)
# Skip weights that belong to PP stages other than this one
if is_pp_missing_parameter(lookup_name, self) or lookup_name not in params_dict:
break
param = params_dict[lookup_name]
+ if lookup_name.endswith(".mul_scale"):
+ previous_scale = qkv_smooth_scales.get(lookup_name)
+ if previous_scale is not None and not torch.equal(previous_scale, loaded_weight):
+ raise ValueError(f"Fused Q/K/V must share the same Smooth tensor: {lookup_name}")
+ if previous_scale is None:
+ qkv_smooth_scales[lookup_name] = loaded_weight.clone()
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
+ loaded_params.add(original_name)
+ if getattr(param, "output_dim", None) is None and not getattr(param, "needs_scalar_to_array", False):
+ # Shared input-channel Smooth is a complete tensor even
+ # when supplied under a single Q/K/V source name.
+ loaded_params.add(lookup_name)
+ else:
+ shards = loaded_qkv_shards.setdefault(lookup_name, set())
+ shards.add(shard_id)
+ if shards == {"q", "k", "v"}:
+ loaded_params.add(lookup_name)
break
else:
# diffusers: ffn.net.0.proj.weight -> our: ffn.net_0.proj.weight
@@ -1198,6 +1293,7 @@ class WanTransformer3DModel(nn.Module):
if is_pp_missing_parameter(lookup_name, self):
continue
+ reject_dualscale_tensor(lookup_name)
if lookup_name not in params_dict:
logger.warning(f"Skipping weight {original_name} -> {lookup_name}")
continue
@@ -1222,8 +1318,6 @@ class WanTransformer3DModel(nn.Module):
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
-
- loaded_params.add(original_name)
- loaded_params.add(lookup_name)
+ loaded_params.update((original_name, lookup_name))
return loaded_params
diff --git a/vllm_omni/engine/omni_engine_base.py b/vllm_omni/engine/omni_engine_base.py
index 85310e9dd..e22bfb241 100644
--- a/vllm_omni/engine/omni_engine_base.py
+++ b/vllm_omni/engine/omni_engine_base.py
@@ -61,11 +61,12 @@ from vllm_omni.engine.stage_client import StageClient
from vllm_omni.engine.stage_init_utils import build_stage0_input_processor
from vllm_omni.engine.stage_pool import StagePool
from vllm_omni.engine.stage_runtime import (
+ OmniClientConfig,
StageRuntimeInfo,
create_stage_runtime,
)
from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
-from vllm_omni.entrypoints.utils import parse_stage_overrides
+from vllm_omni.entrypoints.utils import prepare_stage_config_inputs
from vllm_omni.inputs.data import OmniSamplingParams
from vllm_omni.metrics.prometheus import OmniRequestCounter
@@ -122,6 +123,7 @@ class OmniEngineBase:
_transfer_emitter: Any = None
_prom_metrics: Any = None
_enable_orch_monitor: bool = False
+ _client_config: OmniClientConfig | None = None
# Lazily created by get_output_blocking_async().
_output_drain_executor: concurrent.futures.ThreadPoolExecutor | None = None
@@ -136,6 +138,7 @@ class OmniEngineBase:
log_stats: bool = False,
tokenizer: str | None = None,
trust_remote_code: bool | None = None,
+ client_config: OmniClientConfig | None = None,
**kwargs: Any,
) -> None:
self.model = model
@@ -155,6 +158,7 @@ class OmniEngineBase:
# --log-stats CLI flag set by the user via OmniBase.
self._log_stats = log_stats
self._enable_orch_monitor = bool(kwargs.pop("enable_orch_monitor", False))
+ self._client_config = client_config
logger.info(f"[OmniEngine] Initializing with model {model}")
@@ -166,6 +170,8 @@ class OmniEngineBase:
_stage_id_kwarg = kwargs.get("stage_id")
if isinstance(_stage_id_kwarg, int) and not single_stage_mode:
single_stage_mode = True
+ if client_config is not None and int(client_config.get("client_count", 1)) > 1 and single_stage_mode:
+ raise ValueError("Multiple API servers cannot be combined with single-stage distributed mode")
self.single_stage_mode: bool = single_stage_mode
self._single_stage_id_filter: int | None = (
@@ -272,7 +278,14 @@ class OmniEngineBase:
self.num_stages = len(self.stage_configs)
self.async_chunk = any(
- bool(getattr(getattr(stage, "engine_args", None), "async_chunk", False)) for stage in self.stage_configs
+ bool(
+ getattr(
+ getattr(stage, "connector_config", None),
+ "async_chunk",
+ getattr(getattr(stage, "engine_args", None), "async_chunk", False),
+ )
+ )
+ for stage in self.stage_configs
)
self.stage_pools: list[StagePool] = []
self.stage_clients: list[StageClient] = [] # logical-stage view for external readers
@@ -371,6 +384,7 @@ class OmniEngineBase:
omni_lb_policy=self._omni_lb_policy,
request_queue=self.request_queue,
log_stats=self._log_stats,
+ client_config=self._client_config,
)
self._runtime.initialize()
@@ -885,24 +899,15 @@ class OmniEngineBase:
) -> tuple[str, list[Any]]:
"""Resolve stage configs and inject defaults shared by orchestrator/headless."""
- for legacy_arg in ("stage_configs_path", "stage_configs"):
- if legacy_arg in kwargs:
- raise ValueError(f"`{legacy_arg}` is no longer supported; use `deploy_config` instead.")
-
- # log_stats is captured by __init__; its CLI-only negative alias must
- # not cross into per-stage structured config ownership validation.
- kwargs.pop("disable_log_stats", None)
- deploy_config_path = kwargs.pop("deploy_config", None)
- strategy_config_path = kwargs.pop("strategy_config", None)
- # CLI callers arrive pre-parsed; offline Python callers may use the
- # JSON-string form documented in recipes.
- stage_overrides = parse_stage_overrides(kwargs.pop("stage_overrides", None))
-
- # ``diffusion_streaming_output`` is the public AsyncOmni/serve kwarg;
- # stage configs know the field as ``streaming_output``. Mirror only a
- # truthy value so the CLI's False default does not override deploy YAML.
- if kwargs.get("diffusion_streaming_output") and kwargs.get("streaming_output") is None:
- kwargs["streaming_output"] = True
+ config_inputs = prepare_stage_config_inputs(
+ model,
+ kwargs,
+ trust_remote_code=trust_remote_code,
+ )
+ kwargs = config_inputs.kwargs
+ deploy_config_path = config_inputs.deploy_config_path
+ strategy_config_path = config_inputs.strategy_config_path
+ stage_overrides = config_inputs.stage_overrides
resolution = cast(
_ConfigResolutionResult,
diff --git a/vllm_omni/engine/stage_admission.py b/vllm_omni/engine/stage_admission.py
index 469787bca..bd57f805d 100644
--- a/vllm_omni/engine/stage_admission.py
+++ b/vllm_omni/engine/stage_admission.py
@@ -287,8 +287,13 @@ def check_admission(
def _diffusion_utilization(replica: Any) -> float | None:
- """Best-effort gpu_memory_utilization for a diffusion replica from raw args."""
+ """Resolve a diffusion replica's memory budget from typed or legacy config."""
stage_cfg = getattr(replica, "stage_cfg", None)
+ cache_config = getattr(stage_cfg, "cache_config", None)
+ util = getattr(cache_config, "gpu_memory_utilization", None)
+ if util is not None:
+ return float(util)
+
engine_args = getattr(stage_cfg, "engine_args", None)
if isinstance(engine_args, dict):
util = engine_args.get("gpu_memory_utilization")
diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py
index 71ba17c7c..1aaebc82b 100644
--- a/vllm_omni/engine/stage_engine_startup.py
+++ b/vllm_omni/engine/stage_engine_startup.py
@@ -13,10 +13,12 @@ import threading
from collections.abc import Callable, Iterator
from dataclasses import dataclass, field
from multiprocessing import connection
+from multiprocessing.process import BaseProcess
from types import SimpleNamespace
from typing import Any
import msgspec
+import torch
import zmq
from omegaconf import OmegaConf
from vllm.config import VllmConfig
@@ -33,6 +35,7 @@ from vllm.v1.engine.utils import (
)
from vllm.v1.executor import Executor
+from vllm_omni.config.omni_config import BaseVllmOmniStageConfig
from vllm_omni.distributed.omni_connectors.utils import initialization
from vllm_omni.engine import stage_init_utils
from vllm_omni.engine.stage_init_utils import (
@@ -67,11 +70,13 @@ def _serialize_stage_config(stage_config: Any) -> Any:
"""Convert a stage config to msgpack-friendly builtins."""
if stage_config is None or isinstance(stage_config, (str, bytes, int, float, bool)):
return stage_config
+ if isinstance(stage_config, torch.dtype):
+ return str(stage_config).removeprefix("torch.")
if OmegaConf.is_config(stage_config):
return _serialize_stage_config(OmegaConf.to_container(stage_config, resolve=True))
- if dataclasses.is_dataclass(stage_config):
+ if dataclasses.is_dataclass(stage_config) and not isinstance(stage_config, type):
return _serialize_stage_config(dataclasses.asdict(stage_config))
if isinstance(stage_config, dict):
@@ -1134,6 +1139,8 @@ def launch_stage_replica(
stage_visible_devices: str | None = None,
spawn_device_lock: threading.Lock | None = None,
omni_parallel_stage_init: bool = False,
+ num_api_servers: int = 1,
+ watched_frontend_processes: list[BaseProcess] | None = None,
) -> Iterator[StageReplicaResources]:
"""Launch a local LLM stage replica.
@@ -1144,6 +1151,8 @@ def launch_stage_replica(
keep the same returned resource bundle.
"""
if omni_master_server is not None:
+ if num_api_servers != 1:
+ raise ValueError("Multiple API servers are not supported with distributed stage launch")
with _launch_omni_core_engines(
vllm_config=vllm_config,
executor_class=executor_class,
@@ -1169,7 +1178,16 @@ def launch_stage_replica(
from vllm_omni.engine.stage_engine_core_proc_manager import StageEngineCoreProcManager
- addresses = get_engine_zmq_addresses(vllm_config)
+ if num_api_servers > 1:
+ # Multi-API Omni is intentionally local-only. IPC paths avoid the
+ # bind-after-allocation race of tcp://host:0 and are stable for every
+ # stage, including stages that are not used by APIServerProcessManager.
+ addresses = EngineZmqAddresses(
+ inputs=[get_open_zmq_ipc_path() for _ in range(num_api_servers)],
+ outputs=[get_open_zmq_ipc_path() for _ in range(num_api_servers)],
+ )
+ else:
+ addresses = get_engine_zmq_addresses(vllm_config)
handshake_address = get_open_zmq_ipc_path()
engines_to_handshake = [CoreEngine(index=0, local=True)]
with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock, stage_id=stage_id, replica_id=replica_id):
@@ -1193,18 +1211,21 @@ def launch_stage_replica(
manager=engine_manager,
addresses=addresses,
)
+ engine_launch = CoreEngineLaunch(
+ engine_manager=engine_manager,
+ coordinator=None,
+ addresses=addresses,
+ tensor_queue=None,
+ )
+ if watched_frontend_processes is not None:
+ engine_launch.watched_frontend_processes = watched_frontend_processes
wait_for_engine_startup(
handshake_socket,
engines_to_handshake,
vllm_config.parallel_config,
False, # coordinated_dp
vllm_config.cache_config,
- CoreEngineLaunch(
- engine_manager=engine_manager,
- coordinator=None,
- addresses=addresses,
- tensor_queue=None,
- ),
+ engine_launch,
)
@@ -1403,7 +1424,7 @@ def get_headless_replica_devices(
omni_dp_size_local: int,
) -> list[str | None]:
"""Return per-replica device slices for a headless stage."""
- runtime_cfg = getattr(stage_cfg, "runtime", None)
+ runtime_cfg = getattr(stage_cfg, "runtime_config", getattr(stage_cfg, "runtime", None))
devices_str: str | None = None
if runtime_cfg is not None:
devices_str = (
@@ -1509,10 +1530,11 @@ def launch_headless_diffusion_replicas(
stage_id,
)
- # Headless diffusion startup and its downstream helpers still consume the
- # legacy StageConfig shape; switch this with the coordinated RFC #4021
- # stage-init cutover.
- metadata = stage_init_utils.extract_legacy_stage_metadata(stage_cfg)
+ metadata = (
+ stage_init_utils.extract_stage_metadata_from_omni_stage_config(stage_cfg)
+ if isinstance(stage_cfg, BaseVllmOmniStageConfig)
+ else stage_init_utils.extract_legacy_stage_metadata(stage_cfg)
+ )
if omni_conn_cfg:
inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to)
# Headless single-stage launch must still infer cross-stage TP topology
diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py
index 9668dfc99..a10cf5861 100644
--- a/vllm_omni/engine/stage_init_utils.py
+++ b/vllm_omni/engine/stage_init_utils.py
@@ -402,6 +402,11 @@ def _tp_size_for_stage(stage_configs: Sequence[Any], stage_id: Any) -> int | Non
for stage_cfg in stage_configs:
if str(getattr(stage_cfg, "stage_id", None)) not in id_strs:
continue
+ if isinstance(stage_cfg, BaseVllmOmniStageConfig):
+ try:
+ return max(1, int(stage_cfg.parallel_config.tensor_parallel_size))
+ except (TypeError, ValueError):
+ return 1
engine_args = getattr(stage_cfg, "engine_args", None)
if engine_args is None:
return 1
@@ -489,11 +494,15 @@ def inject_kv_stage_info(stage_cfg: Any, stage_id: int, stage_configs: Sequence[
rank mappings automatically.
"""
try:
- engine_args = stage_cfg.engine_args
- if hasattr(engine_args, "get"):
- omni_kv = engine_args.get("omni_kv_config", None)
+ connector_config = getattr(stage_cfg, "connector_config", None)
+ if connector_config is not None:
+ omni_kv = connector_config.omni_kv_config
else:
- omni_kv = getattr(engine_args, "omni_kv_config", None)
+ engine_args = stage_cfg.engine_args
+ if hasattr(engine_args, "get"):
+ omni_kv = engine_args.get("omni_kv_config", None)
+ else:
+ omni_kv = getattr(engine_args, "omni_kv_config", None)
if omni_kv is None:
return
@@ -504,7 +513,11 @@ def inject_kv_stage_info(stage_cfg: Any, stage_id: int, stage_configs: Sequence[
if "stage_id" not in omni_kv:
omni_kv["stage_id"] = stage_id
- engine_input_source = getattr(stage_cfg, "engine_input_source", None)
+ engine_input_source = (
+ getattr(stage_cfg, "input_sources", None)
+ if connector_config is not None
+ else getattr(stage_cfg, "engine_input_source", None)
+ )
if engine_input_source is not None:
if hasattr(omni_kv, "setdefault"):
omni_kv.setdefault("engine_input_source", list(engine_input_source))
@@ -703,17 +716,16 @@ def _resolve_omni_metadata_hook(path: str | None) -> Callable | None:
def extract_stage_metadata_from_omni_stage_config(
stage_config: BaseVllmOmniStageConfig,
) -> StageMetadata:
- """Project one typed stage config into metadata for a future cutover.
-
- This projection is not used by production startup yet. Current replica
- layout, engine-argument, remote-diffusion, and platform setup paths still
- require the legacy StageConfig/OmegaConf shape.
- """
+ """Project one typed stage config into production runtime metadata."""
stage_type: Literal["llm", "diffusion"] = "diffusion" if stage_config.stage_type == StageType.DIFFUSION else "llm"
- sampling_params_cls = SamplingParams if stage_type == "llm" else OmniDiffusionSamplingParams
- sampling_params: OmniSamplingParams = sampling_params_cls(
- **(stage_config.model_config.default_sampling_params or {})
- )
+ pooling_config = stage_config.pooling_config
+ if stage_type == "llm" and (pooling_config.runner or "").lower() == "pooling":
+ sampling_params: OmniSamplingParams | PoolingParams = (
+ copy.deepcopy(pooling_config.default_pooling_params) or PoolingParams()
+ )
+ else:
+ sampling_params_cls = SamplingParams if stage_type == "llm" else OmniDiffusionSamplingParams
+ sampling_params = sampling_params_cls(**(stage_config.model_config.default_sampling_params or {}))
custom_process_input_func = _resolve_omni_metadata_hook(stage_config.custom_process_input_func)
if stage_type == "diffusion":
@@ -878,7 +890,9 @@ def _get_local_llm_parallel_sizes(
unset.
"""
if engine_args is None:
- engine_args = getattr(stage_cfg, "engine_args", {})
+ engine_args = getattr(stage_cfg, "parallel_config", None)
+ if engine_args is None:
+ engine_args = getattr(stage_cfg, "engine_args", {})
tp_size = int(_get_attr_or_item(engine_args, "tensor_parallel_size", 1) or 1)
pp_size = int(_get_attr_or_item(engine_args, "pipeline_parallel_size", 1) or 1)
local_dp_size = _get_attr_or_item(engine_args, "data_parallel_size_local", None)
@@ -889,9 +903,12 @@ def _get_local_llm_parallel_sizes(
def get_stage_devices_per_replica(stage_cfg: Any, engine_args: Any | None = None) -> int:
"""Return the number of devices consumed by one replica of *stage_cfg*."""
- if engine_args is None:
- engine_args = getattr(stage_cfg, "engine_args", {})
if getattr(stage_cfg, "stage_type", "llm") == "diffusion":
+ typed_parallel_config = getattr(stage_cfg, "parallel_config", None)
+ if typed_parallel_config is not None:
+ return max(1, int(typed_parallel_config.world_size))
+ if engine_args is None:
+ engine_args = getattr(stage_cfg, "engine_args", {})
parallel_config = _get_attr_or_item(engine_args, "parallel_config")
if parallel_config is None:
return 1
@@ -933,7 +950,7 @@ def compute_replica_layout(
"""
replicas_per_stage: list[int] = []
for stage_cfg in stage_configs:
- runtime_cfg = getattr(stage_cfg, "runtime", {})
+ runtime_cfg = getattr(stage_cfg, "runtime_config", getattr(stage_cfg, "runtime", {}))
num_replicas = int(
runtime_cfg.get("num_replicas", 1)
if hasattr(runtime_cfg, "get")
@@ -948,7 +965,7 @@ def compute_replica_layout(
num_replicas = replicas_per_stage[stage_id]
if num_replicas <= 1:
continue
- runtime_cfg = getattr(stage_cfg, "runtime", {})
+ runtime_cfg = getattr(stage_cfg, "runtime_config", getattr(stage_cfg, "runtime", {}))
devices_str = (
runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None)
)
@@ -1086,15 +1103,31 @@ def _project_omni_stage_engine_args(
diffusion_stage = cast(VllmOmniDiffusionStageConfig, stage_config)
engine_args.update(_project_omni_config_fields(diffusion_stage.diffusion_config))
+ model_excluded_fields = {
+ "default_sampling_params",
+ "has_sampling_extra_args",
+ }
+ runtime_excluded_fields = {"devices", "num_replicas", "env", "num_gpus"}
+ if not is_diffusion:
+ # These values configure OmniDiffusionConfig or its worker process;
+ # OmniEngineArgs has no matching fields for LLM stages.
+ model_excluded_fields.update(
+ {
+ "disable_autocast",
+ "enable_multithread_weight_load",
+ "num_weight_load_threads",
+ }
+ )
+ runtime_excluded_fields.add("log_level")
+
for config, excluded_fields in (
(
stage_config.model_config,
- frozenset({"default_sampling_params", "has_sampling_extra_args"}),
+ frozenset(model_excluded_fields),
),
(
stage_config.runtime_config,
- frozenset({"devices", "num_replicas", "env", "num_gpus"})
- | (frozenset({"additional_config"}) if is_diffusion else frozenset()),
+ frozenset(runtime_excluded_fields) | (frozenset({"additional_config"}) if is_diffusion else frozenset()),
),
):
engine_args.update(
@@ -1103,6 +1136,11 @@ def _project_omni_stage_engine_args(
exclude=excluded_fields,
)
)
+ pooling_config = stage_config.pooling_config
+ if pooling_config.runner is not None:
+ engine_args["runner"] = pooling_config.runner
+ if pooling_config.pooling_output_decoder is not None:
+ engine_args["pooling_output_decoder"] = pooling_config.pooling_output_decoder
if is_diffusion:
for config, field_map in (
@@ -1330,11 +1368,9 @@ def build_engine_args_dict(
stage_connector_spec: dict[str, Any] | None = None,
cli_tokenizer: str | None = None,
) -> dict[str, Any]:
- """Build engine arguments through the stable production entry point.
+ """Preserve the legacy engine-argument API for compatibility callers.
- Production inputs still use the legacy stage representation. Keep that
- compatibility choice behind this function so callers do not bind directly
- to a representation-specific implementation.
+ Typed runtime stages use ``build_engine_args_dict_from_omni_stage_config``.
"""
return build_legacy_engine_args_dict(
stage_config,
@@ -1350,13 +1386,7 @@ def build_engine_args_dict_from_omni_stage_config(
stage_connector_spec: dict[str, Any] | None = None,
cli_tokenizer: str | None = None,
) -> dict[str, Any]:
- """Project one typed stage config into backend engine arguments.
-
- This projection is prepared for the RFC #4021 stage-init cutover. Current
- production startup reaches the legacy implementation through
- ``build_engine_args_dict`` while strategy and startup-plan inputs still
- use the legacy representation.
- """
+ """Project one typed production stage config into backend engine arguments."""
engine_args_dict = _project_omni_stage_engine_args(stage_config)
_apply_rocm_attention_backend(engine_args_dict, stage_config.stage_type)
return _finalize_engine_args_dict(
@@ -1391,7 +1421,7 @@ def _check_stage_device_layout(stage_config: Any, engine_args_dict: dict[str, An
"""
from vllm_omni.config.composable_parallel import StrategyApplyError, check_device_layout
- runtime = getattr(stage_config, "runtime", None)
+ runtime = getattr(stage_config, "runtime_config", getattr(stage_config, "runtime", None))
devices = _get_attr_or_item(runtime, "devices", None) if runtime is not None else None
if devices is None:
# No explicit placement -> vLLM assigns devices itself; nothing to check.
@@ -1445,6 +1475,8 @@ def build_vllm_config(
stage_connector_spec: dict[str, Any] | None = None,
engine_args_dict: dict[str, Any] | None = None,
headless: bool = False,
+ api_process_count: int = 1,
+ api_process_rank: int = 0,
) -> tuple[Any, type]:
"""Build engine args, then create VllmConfig and executor_class.
@@ -1452,13 +1484,18 @@ def build_vllm_config(
(vllm_config, executor_class)
"""
if engine_args_dict is None:
- engine_args_dict = build_engine_args_dict(
- stage_config,
- model,
- stage_connector_spec=stage_connector_spec,
+ engine_args_dict = (
+ build_engine_args_dict_from_omni_stage_config(
+ stage_config, model, stage_connector_spec=stage_connector_spec
+ )
+ if isinstance(stage_config, BaseVllmOmniStageConfig)
+ else build_engine_args_dict(stage_config, model, stage_connector_spec=stage_connector_spec)
)
filtered_engine_args_dict = filter_dataclass_kwargs(OmniEngineArgs, engine_args_dict)
+ if api_process_count != 1 or api_process_rank != 0:
+ filtered_engine_args_dict["_api_process_count"] = api_process_count
+ filtered_engine_args_dict["_api_process_rank"] = api_process_rank
# _to_dict serializes dataclass fields (e.g. StructuredOutputsConfig) into
# plain dicts. When OmniEngineArgs is instantiated with the dict, these
@@ -1715,6 +1752,7 @@ def acquire_device_locks(
stage_id: int,
engine_args_dict: dict[str, Any],
stage_init_timeout: int,
+ locked_devices: set[int] | None = None,
) -> list[int]:
"""Acquire exclusive file locks on devices needed by this stage.
@@ -1769,8 +1807,7 @@ def acquire_device_locks(
f"but only {len(physical_devices)} device(s) are available: {physical_devices}"
)
- num_devices_to_lock = num_devices_per_stage
- devices_to_lock = sorted(physical_devices[:num_devices_to_lock])
+ devices_to_lock = sorted(set(physical_devices[:num_devices_per_stage]) - (locked_devices or set()))
logger.debug(
"Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d, CFG=%d; will lock %d devices: %s",
@@ -1780,7 +1817,7 @@ def acquire_device_locks(
prefill_context_parallel_size,
sequence_parallel_size,
cfg_parallel_size,
- num_devices_to_lock,
+ len(devices_to_lock),
devices_to_lock,
)
@@ -1798,6 +1835,8 @@ def acquire_device_locks(
record_lock_holder_pid(lock_fd, lock_writable)
lock_acquired = True
lock_fds.append(lock_fd)
+ if locked_devices is not None:
+ locked_devices.add(device_id)
logger.debug("Acquired exclusive lock for device %s", device_id)
except BlockingIOError:
os.close(lock_fd)
@@ -1927,7 +1966,11 @@ def build_diffusion_config(
) -> Any:
"""Build diffusion config for a stage."""
- engine_args_dict = build_engine_args_dict(stage_cfg, model)
+ engine_args_dict = (
+ build_engine_args_dict_from_omni_stage_config(stage_cfg, model)
+ if isinstance(stage_cfg, BaseVllmOmniStageConfig)
+ else build_engine_args_dict(stage_cfg, model)
+ )
od_config = OmniDiffusionConfig.from_kwargs(**engine_args_dict)
num_devices_per_stage = od_config.parallel_config.world_size
diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py
index 53400df5e..ce958ba6d 100644
--- a/vllm_omni/engine/stage_runtime.py
+++ b/vllm_omni/engine/stage_runtime.py
@@ -6,19 +6,20 @@
from __future__ import annotations
import concurrent.futures
+import contextlib
import copy
import os
import threading
-import time
from collections.abc import Callable, Iterator, Mapping, Sequence
-from contextlib import ExitStack, contextmanager
-from dataclasses import dataclass
-from typing import Any, cast
+from contextlib import AbstractContextManager, contextmanager
+from dataclasses import dataclass, field
+from multiprocessing.process import BaseProcess
+from typing import Any, TypedDict, cast
import janus
-from omegaconf import OmegaConf
from vllm.logger import init_logger
+from vllm_omni.config.omni_config import BaseVllmOmniStageConfig
from vllm_omni.distributed.omni_connectors.utils.initialization import (
resolve_omni_kv_config_for_stage,
)
@@ -49,11 +50,13 @@ from vllm_omni.engine.stage_init_utils import (
_inject_inferred_kv_tp_topology,
acquire_device_locks,
build_engine_args_dict,
+ build_engine_args_dict_from_omni_stage_config,
build_llm_stage_output_processor,
build_vllm_config,
compute_replica_layout,
device_overlap_group_keys,
extract_legacy_stage_metadata,
+ extract_stage_metadata_from_omni_stage_config,
get_stage_connector_spec,
inject_kv_stage_info,
inject_omni_kv_connector_config,
@@ -92,6 +95,44 @@ class StageRemoteFactoryContext:
executor_class: type | None = None
+class OmniClientConfig(TypedDict, total=False):
+ """Per-frontend transport/configuration passed across process boundaries."""
+
+ client_count: int
+ client_index: int
+ stage_addresses: dict[int | str, dict[int | str, dict[str, Any]]]
+ input_address: str
+ output_address: str
+ stats_update_address: str
+ actual_address_pipe: Any
+ tensor_queue: Any
+
+
+@dataclass
+class StageEngineLaunch:
+ """Engine resources handed to a local client or the serving parent.
+
+ Successful launch transfers resource ownership to the caller. Startup or
+ attachment failures are rolled back by ``launch_stage_engines``.
+ """
+
+ client_configs: list[OmniClientConfig]
+ resources: list[StageReplicaResources]
+ watched_frontend_processes: list[BaseProcess] = field(default_factory=list)
+
+ def shutdown(self) -> None:
+ seen: set[int] = set()
+ for stage_resources in reversed(self.resources):
+ for resource in (stage_resources.manager, stage_resources.coordinator):
+ if resource is None or id(resource) in seen:
+ continue
+ seen.add(id(resource))
+ try:
+ resource.shutdown()
+ except Exception:
+ logger.warning("[StageEngineLaunch] resource shutdown failed", exc_info=True)
+
+
def _build_load_balancer_factory(policy: str) -> Callable[[], LoadBalancer]:
try:
normalized = LoadBalancingPolicy(policy)
@@ -123,13 +164,14 @@ class StageRuntime:
self,
stage_configs: list[Any],
model: str,
- config_path: str,
+ config_path: str | None,
*,
stage_init_timeout: int,
async_chunk: bool,
tokenizer: str | None = None,
parallel_stage_init: bool = False,
log_stats: bool = False,
+ client_config: OmniClientConfig | None = None,
) -> None:
self._stage_configs = stage_configs
self._model = model
@@ -144,6 +186,14 @@ class StageRuntime:
self._parallel_stage_init = parallel_stage_init
self._log_stats = log_stats
self._num_stages = len(stage_configs)
+ self._client_count = int((client_config or {}).get("client_count", 1))
+ self._client_index = int((client_config or {}).get("client_index", 0))
+ if self._client_count < 1:
+ raise ValueError(f"client_count must be >= 1, got {self._client_count}")
+ if self._client_index != -1 and not 0 <= self._client_index < self._client_count:
+ raise ValueError(f"client_index must be in [0, {self._client_count}), got {self._client_index}")
+ self._external_stage_addresses = (client_config or {}).get("stage_addresses")
+ self._api_process_rank = self._client_index
# Populated by initialize()
self.stage_pools: list[StagePool] = []
@@ -262,6 +312,205 @@ class StageRuntime:
self._cleanup_after_initialize_failure()
raise exc
+ @contextmanager
+ def launch_stage_engines(
+ self,
+ num_api_servers: int = 1,
+ *,
+ stage_plans: Sequence[LogicalStageInitPlan] | None = None,
+ stage_init_timeout: int | None = None,
+ ) -> Iterator[StageEngineLaunch]:
+ """Launch EngineCore stages for local or subprocess client attachment.
+
+ The regular runtime supplies a replica plan from its device-group
+ scheduler; the serving parent resolves the whole pipeline here. Both
+ use the same device locks, spawn, readiness and rollback machinery.
+ One client completes readiness before attachment, preserving the
+ existing in-process frontend. Multiple clients attach in the caller's
+ ``with`` body before readiness completes on context exit.
+ """
+ if self._external_stage_addresses is not None:
+ raise ValueError("A runtime attached to external stage engines cannot launch engines")
+ if num_api_servers < 1:
+ raise ValueError(f"num_api_servers must be >= 1, got {num_api_servers}")
+ timeout = self._stage_init_timeout if stage_init_timeout is None else stage_init_timeout
+ resolve_plans = stage_plans is None
+ if stage_plans is None:
+ self._client_count = num_api_servers
+ self._api_process_rank = -1 if num_api_servers > 1 else self._client_index
+ self._init_visible_devices_baseline = os.environ.get(current_omni_platform.device_control_env_var)
+ stage_plans = self._prepare_stage_plans()
+
+ unsupported = [
+ plan.stage_id
+ for plan in stage_plans
+ if any(replica.metadata.stage_type == "diffusion" for replica in plan.replicas)
+ ]
+ if unsupported:
+ raise ValueError(
+ "--api-server-count currently supports EngineCore stages only; "
+ f"diffusion stage(s) are not supported: {unsupported}"
+ )
+ if any(replica.launch_mode != "local" for plan in stage_plans for replica in plan.replicas):
+ raise ValueError("--api-server-count is not supported with remote or headless stages")
+
+ if num_api_servers > 1:
+ parallel_configs: list[Any] = []
+ for plan in stage_plans:
+ for replica in plan.replicas:
+ if replica.stage_vllm_config is None or replica.executor_class is None:
+ raise RuntimeError(f"LLM stage {plan.stage_id} is missing its engine configuration")
+ parallel_configs.append(replica.stage_vllm_config.parallel_config)
+ if any(getattr(parallel_config, "enable_fault_tolerance", False) for parallel_config in parallel_configs):
+ raise ValueError("--api-server-count > 1 cannot be combined with --enable-fault-tolerance")
+ if any(getattr(parallel_config, "enable_elastic_ep", False) for parallel_config in parallel_configs):
+ raise ValueError("--api-server-count > 1 cannot be combined with --enable-elastic-ep")
+ if any(
+ getattr(parallel_config, "data_parallel_size", 1) != 1 or getattr(parallel_config, "use_ray", False)
+ for parallel_config in parallel_configs
+ ):
+ raise ValueError(
+ "--api-server-count > 1 currently supports one local process group per replica; "
+ "intra-replica data parallelism and Ray backends are not supported"
+ )
+
+ if resolve_plans and self._parallel_stage_init:
+ self._reject_unguardable_executors(stage_plans)
+ self._run_stage_admission(stage_plans)
+
+ client_configs: list[OmniClientConfig] = [
+ OmniClientConfig(
+ client_count=num_api_servers,
+ client_index=client_index,
+ stage_addresses={},
+ )
+ for client_index in range(num_api_servers)
+ ]
+ entered_contexts: list[AbstractContextManager[StageReplicaResources]] = []
+ resources: list[StageReplicaResources] = []
+ watched_frontend_processes: list[BaseProcess] = []
+ # launch_stage_replica performs its readiness handshake in __exit__,
+ # after API workers are started, so initialization locks must remain
+ # held until all launch contexts have exited.
+ lock_fds: list[int] = []
+ locked_devices: set[int] = set()
+ launch: StageEngineLaunch | None = None
+ exited_contexts = 0
+
+ try:
+ for plan in stage_plans:
+ for replica in plan.replicas:
+ if replica.stage_vllm_config is None or replica.executor_class is None:
+ raise RuntimeError(f"LLM stage {plan.stage_id} is missing its engine configuration")
+ if replica.engine_args_dict is None:
+ raise RuntimeError(f"LLM stage {plan.stage_id} is missing engine args")
+
+ physical_devices = self._resolve_replica_physical_devices(
+ replica.metadata.stage_id,
+ replica.metadata.runtime_cfg,
+ )
+ if not self._parallel_stage_init:
+ with self._scoped_spawn_device_env(physical_devices):
+ lock_fds.extend(
+ acquire_device_locks(
+ replica.metadata.stage_id,
+ replica.engine_args_dict,
+ timeout,
+ locked_devices,
+ )
+ )
+
+ launch_context = launch_stage_replica(
+ vllm_config=replica.stage_vllm_config,
+ executor_class=replica.executor_class,
+ log_stats=self._log_stats,
+ stage_id=replica.metadata.stage_id,
+ replica_id=replica.replica_id,
+ stage_config=replica.stage_cfg,
+ omni_master_server=self._get_omni_master_server(),
+ omni_coordinator_address=self._get_coordinator_address(),
+ stage_visible_devices=physical_devices,
+ spawn_device_lock=self._spawn_device_lock,
+ omni_parallel_stage_init=self._parallel_stage_init,
+ num_api_servers=num_api_servers,
+ watched_frontend_processes=watched_frontend_processes if num_api_servers > 1 else None,
+ )
+ # Environment overlays are process-global. Serialize spawn;
+ # parallel initialization waits for READY outside this lock.
+ with self._replica_launch_lock:
+ with stage_runtime_env(replica.metadata.stage_id, replica.metadata.runtime_cfg):
+ stage_resources = launch_context.__enter__()
+ entered_contexts.append(launch_context)
+ if stage_resources is None:
+ raise RuntimeError(f"LLM stage {plan.stage_id} launcher returned no resources")
+ resources.append(stage_resources)
+ if num_api_servers == 1:
+ exited_contexts += 1
+ launch_context.__exit__(None, None, None)
+ addresses = stage_resources.addresses
+ # The generic resource bundle allows no addresses (e.g.
+ # diffusion), but EngineCore attachment requires them.
+ if addresses is None:
+ raise RuntimeError(f"LLM stage {plan.stage_id} launcher returned no addresses")
+ # Unlike the single-client handshake, only stage 0 is
+ # handed to APIServerProcessManager for address updates.
+ # Other stages must already expose connectable endpoints.
+ if num_api_servers > 1 and any(
+ address.startswith("tcp://") and address.rsplit(":", 1)[-1] == "0"
+ for address in (*addresses.inputs, *addresses.outputs)
+ ):
+ raise RuntimeError(
+ f"Stage {plan.stage_id} returned deferred TCP addresses; "
+ "multi-API launch requires fixed ports or IPC addresses"
+ )
+ if len(addresses.inputs) != num_api_servers or len(addresses.outputs) != num_api_servers:
+ raise RuntimeError(
+ f"Stage {plan.stage_id} returned {len(addresses.inputs)} input and "
+ f"{len(addresses.outputs)} output addresses for {num_api_servers} API servers"
+ )
+
+ if num_api_servers == 1 and lock_fds:
+ release_device_locks(lock_fds)
+ lock_fds.clear()
+ locked_devices.clear()
+
+ for client_index, client_config in enumerate(client_configs):
+ stage_addresses = client_config["stage_addresses"]
+ replica_addresses = stage_addresses.setdefault(plan.stage_id, {})
+ replica_addresses[replica.replica_id] = {
+ "input_address": addresses.inputs[client_index],
+ "output_address": addresses.outputs[client_index],
+ }
+ if addresses.frontend_stats_publish_address is not None:
+ replica_addresses[replica.replica_id]["stats_update_address"] = (
+ addresses.frontend_stats_publish_address
+ )
+
+ launch = StageEngineLaunch(
+ client_configs=client_configs,
+ resources=resources,
+ watched_frontend_processes=watched_frontend_processes,
+ )
+ yield launch
+
+ while exited_contexts < len(entered_contexts):
+ entered_context = entered_contexts[exited_contexts]
+ exited_contexts += 1
+ entered_context.__exit__(None, None, None)
+ except BaseException as exc:
+ if launch is not None:
+ launch.shutdown()
+ else:
+ StageEngineLaunch(client_configs=client_configs, resources=resources).shutdown()
+ exc_info = (type(exc), exc, exc.__traceback__)
+ for pending_context in reversed(entered_contexts[exited_contexts:]):
+ with contextlib.suppress(BaseException):
+ pending_context.__exit__(*exc_info)
+ raise
+ finally:
+ if lock_fds:
+ release_device_locks(lock_fds)
+
def shutdown(self) -> None:
for pool in self.stage_pools:
for client in pool.clients:
@@ -345,12 +594,12 @@ class StageRuntime:
"""Build startup plans for every logical stage and replica."""
stage_plans: list[LogicalStageInitPlan] = []
- # RFC #4021 transition boundary: stage planning still relies on legacy
- # StageConfig.runtime and StageConfig.engine_args for replica and engine
- # setup. Keep metadata extraction on the legacy path until the
- # coordinated stage-init cutover.
for stage_idx, stage_cfg in enumerate(self._stage_configs):
- base_metadata = extract_legacy_stage_metadata(stage_cfg)
+ base_metadata = (
+ extract_stage_metadata_from_omni_stage_config(stage_cfg)
+ if isinstance(stage_cfg, BaseVllmOmniStageConfig)
+ else extract_legacy_stage_metadata(stage_cfg)
+ )
stage_id = int(base_metadata.stage_id)
if stage_id != stage_idx:
raise ValueError(
@@ -374,14 +623,20 @@ class StageRuntime:
executor_class = None
engine_args_dict = None
if base_metadata.stage_type != "diffusion":
- # The stable adapter entry point still receives the same
- # legacy stage object as replica planning. Its implementation
- # switches only at the coordinated RFC #4021 cutover.
- engine_args_dict = build_engine_args_dict(
- stage_cfg,
- self._model,
- stage_connector_spec=stage_connector_spec,
- cli_tokenizer=self._tokenizer,
+ engine_args_dict = (
+ build_engine_args_dict_from_omni_stage_config(
+ stage_cfg,
+ self._model,
+ stage_connector_spec=stage_connector_spec,
+ cli_tokenizer=self._tokenizer,
+ )
+ if isinstance(stage_cfg, BaseVllmOmniStageConfig)
+ else build_engine_args_dict(
+ stage_cfg,
+ self._model,
+ stage_connector_spec=stage_connector_spec,
+ cli_tokenizer=self._tokenizer,
+ )
)
inject_omni_kv_connector_config(
engine_args_dict,
@@ -398,14 +653,25 @@ class StageRuntime:
self._model,
stage_connector_spec=stage_connector_spec,
engine_args_dict=engine_args_dict,
+ api_process_count=self._client_count,
+ api_process_rank=self._api_process_rank,
)
for replica_id in range(num_replicas):
- replica_cfg = copy.deepcopy(stage_cfg) if replica_id > 0 else stage_cfg
+ # Keep the logical stage's device pool intact; each replica owns
+ # the same config throughout planning, metadata and launch.
+ replica_cfg = copy.deepcopy(stage_cfg) if num_replicas > 1 else stage_cfg
if stage_idx in replica_devices_map:
- replica_cfg.runtime.devices = replica_devices_map[stage_idx][replica_id]
+ devices = replica_devices_map[stage_idx][replica_id]
+ runtime_cfg = getattr(replica_cfg, "runtime_config", getattr(replica_cfg, "runtime", None))
+ if runtime_cfg is not None:
+ runtime_cfg.devices = devices
- replica_metadata = extract_legacy_stage_metadata(replica_cfg)
+ replica_metadata = (
+ extract_stage_metadata_from_omni_stage_config(replica_cfg)
+ if isinstance(replica_cfg, BaseVllmOmniStageConfig)
+ else extract_legacy_stage_metadata(replica_cfg)
+ )
replica_metadata.replica_id = replica_id
if launch_mode == "remote" and replica_metadata.stage_type != "diffusion":
replica_metadata.runtime_cfg = None
@@ -720,122 +986,70 @@ class StageRuntime:
stage_init_timeout: int,
) -> StagePoolClient:
"""Initialize one local LLM replica using vLLM's launch/attach pattern."""
- resources: StageReplicaResources | None = None
- stage_client = None
- lock_fds: list[int] = []
- try:
- physical_devices = self._resolve_replica_physical_devices(
- plan.metadata.stage_id,
- plan.metadata.runtime_cfg,
- )
- if physical_devices:
- logger.info(
- "[stage_init] Stage-%s set runtime devices: %s",
- plan.metadata.stage_id,
- physical_devices,
- )
- vllm_config = plan.stage_vllm_config
- executor_class = plan.executor_class
- if vllm_config is None:
- raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing vllm_config")
- if executor_class is None:
- raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing executor_class")
- if plan.engine_args_dict is None:
- raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing engine args")
- # G3 device locks: in the default (serial) path the parent holds a
- # per-device LOCK_EX across the whole child init. When parallel stage
- # init is enabled the engine-core child takes phased SH/EX locks
- # itself (see stage_phase_lock), so the parent must NOT also hold the
- # lock here — it would self-deadlock while waiting for the child's
- # READY handshake.
- if not self._parallel_stage_init:
- g3_start = time.perf_counter()
- with self._scoped_spawn_device_env(physical_devices):
- lock_fds = acquire_device_locks(
- plan.metadata.stage_id,
- plan.engine_args_dict,
- stage_init_timeout,
- )
- logger.debug(
- "[stage_init] Stage-%s G3 device-lock acquire took %.3fs",
- plan.metadata.stage_id,
- time.perf_counter() - g3_start,
- )
-
- launch_cm = launch_stage_replica(
- vllm_config=vllm_config,
- executor_class=executor_class,
- log_stats=self._log_stats,
- stage_id=plan.metadata.stage_id,
- replica_id=plan.replica_id,
- stage_config=plan.stage_cfg,
- omni_master_server=self._get_omni_master_server(),
- omni_coordinator_address=self._get_coordinator_address(),
- stage_visible_devices=physical_devices,
- spawn_device_lock=self._spawn_device_lock,
- omni_parallel_stage_init=self._parallel_stage_init,
- )
- # Only process spawning and the runtime.env overlay need the global
- # launch lock. The launch context's exit waits for READY: holding the
- # lock there serializes model loading and compilation even across
- # different GPUs. Default initialization keeps its per-device EX
- # locks until READY; parallel_stage_init uses child phase locks.
- g2_start = time.perf_counter()
- with ExitStack() as launch_stack:
- with self._replica_launch_lock:
- g2_locked = time.perf_counter()
- with stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg):
- resources = launch_stack.enter_context(launch_cm)
- g2_spawned = time.perf_counter()
- logger.debug(
- "[stage_init] Stage-%s G2 launch-lock wait=%.3fs, spawn=%.3fs, READY=%.3fs",
- plan.metadata.stage_id,
- g2_locked - g2_start,
- g2_spawned - g2_locked,
- time.perf_counter() - g2_spawned,
+ external_addresses = self._get_external_client_addresses(
+ plan.metadata.stage_id,
+ plan.replica_id,
+ )
+ if external_addresses is not None:
+ if plan.stage_vllm_config is None or plan.executor_class is None:
+ raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing its engine configuration")
+ return cast(
+ StagePoolClient,
+ StageEngineCoreClientBase.make_async_mp_client(
+ vllm_config=plan.stage_vllm_config,
+ executor_class=plan.executor_class,
+ log_stats=self._log_stats,
+ metadata=plan.metadata,
+ client_addresses=external_addresses,
+ client_count=self._client_count,
+ client_index=self._client_index,
+ ),
)
- logger.info("[StageRuntime] Stage %s engine startup completed", plan.metadata.stage_id)
- if resources is None:
- raise RuntimeError(f"LLM stage {plan.metadata.stage_id} launcher returned no resources")
- if resources.addresses is None:
- raise RuntimeError(f"LLM stage {plan.metadata.stage_id} launcher returned no addresses")
- stage_client = StageEngineCoreClientBase.make_async_mp_client(
- vllm_config=vllm_config,
- executor_class=executor_class,
- log_stats=self._log_stats,
- metadata=plan.metadata,
- client_addresses=self._client_addresses_from_zmq(resources.addresses),
- engine_manager=resources.manager,
- coordinator=resources.coordinator,
+ if plan.stage_vllm_config is None or plan.executor_class is None:
+ raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing its engine configuration")
+ stage_plan = LogicalStageInitPlan(
+ stage_idx=plan.metadata.stage_id,
+ stage_id=plan.metadata.stage_id,
+ replicas=[plan],
+ )
+ with self.launch_stage_engines(
+ stage_plans=[stage_plan],
+ stage_init_timeout=stage_init_timeout,
+ ) as launch:
+ resources = launch.resources[0]
+ return cast(
+ StagePoolClient,
+ StageEngineCoreClientBase.make_async_mp_client(
+ vllm_config=plan.stage_vllm_config,
+ executor_class=plan.executor_class,
+ log_stats=self._log_stats,
+ metadata=plan.metadata,
+ client_addresses=self._client_addresses_from_zmq(resources.addresses),
+ engine_manager=resources.manager,
+ coordinator=resources.coordinator,
+ ),
)
- logger.info("[StageRuntime] Stage %s initialized", plan.metadata.stage_id)
- return cast(StagePoolClient, stage_client)
- except Exception:
- if stage_client is not None:
- try:
- stage_client.shutdown()
- except Exception as cleanup_error:
- logger.warning(
- "[StageRuntime] Failed to cleanup stage %s after init failure: %s",
- plan.metadata.stage_id,
- cleanup_error,
- )
- else:
- self._cleanup_launched_resources(
- stage_id=plan.metadata.stage_id,
- resources=resources,
- )
- raise
- finally:
- if lock_fds:
- release_device_locks(lock_fds)
-
def _get_coordinator_address(self) -> str | None:
"""Return coordinator router address. Overridden by DistStageRuntime."""
return None
+ def _get_external_client_addresses(self, stage_id: int, replica_id: int) -> dict[str, Any] | None:
+ if self._external_stage_addresses is None:
+ return None
+ stage_addresses = self._external_stage_addresses.get(stage_id)
+ if stage_addresses is None:
+ stage_addresses = self._external_stage_addresses.get(str(stage_id))
+ if stage_addresses is None:
+ raise ValueError(f"Missing external addresses for stage {stage_id}")
+ replica_addresses = stage_addresses.get(replica_id)
+ if replica_addresses is None:
+ replica_addresses = stage_addresses.get(str(replica_id))
+ if replica_addresses is None:
+ raise ValueError(f"Missing external addresses for stage {stage_id} replica {replica_id}")
+ return dict(replica_addresses)
+
def _get_omni_master_server(self) -> OmniMasterServer | None:
"""Return the master server for local distributed launches, if any."""
return None
@@ -859,17 +1073,25 @@ class StageRuntime:
raise RuntimeError("Omni KV connector requires source and destination stages")
inject_omni_kv_config(plan.stage_cfg, omni_conn_cfg, omni_from, omni_to)
inject_kv_stage_info(plan.stage_cfg, plan.metadata.stage_id, self._stage_configs)
- engine_args = getattr(plan.stage_cfg, "engine_args", {})
- inline_diffusion = (
- engine_args.get("inline_diffusion", False)
- if hasattr(engine_args, "get")
- else getattr(engine_args, "inline_diffusion", False)
- )
- custom_pipeline_args = (
- engine_args.get("custom_pipeline_args")
- if hasattr(engine_args, "get")
- else getattr(engine_args, "custom_pipeline_args", None)
- )
+ if isinstance(plan.stage_cfg, BaseVllmOmniStageConfig):
+ inline_diffusion = plan.stage_cfg.stage_pipeline_config.inline_diffusion
+ custom_pipeline_args = getattr(
+ getattr(plan.stage_cfg, "diffusion_config", None),
+ "custom_pipeline_args",
+ None,
+ )
+ else:
+ engine_args = getattr(plan.stage_cfg, "engine_args", {})
+ inline_diffusion = (
+ engine_args.get("inline_diffusion", False)
+ if hasattr(engine_args, "get")
+ else getattr(engine_args, "inline_diffusion", False)
+ )
+ custom_pipeline_args = (
+ engine_args.get("custom_pipeline_args")
+ if hasattr(engine_args, "get")
+ else getattr(engine_args, "custom_pipeline_args", None)
+ )
client, resources = launch_diffusion_stage_replica(
model=self._model,
stage_config=plan.stage_cfg,
@@ -966,7 +1188,7 @@ class DistStageRuntime(StageRuntime):
self,
stage_configs: list[Any],
model: str,
- config_path: str,
+ config_path: str | None,
*,
stage_init_timeout: int,
async_chunk: bool,
@@ -1028,7 +1250,7 @@ class DistStageRuntime(StageRuntime):
for idx, stage_cfg in enumerate(self._stage_configs):
stage_id = int(getattr(stage_cfg, "stage_id", idx))
- runtime_cfg = getattr(stage_cfg, "runtime", None)
+ runtime_cfg = getattr(stage_cfg, "runtime_config", getattr(stage_cfg, "runtime", None))
if runtime_cfg is None:
continue
if stage_id == target_stage_id:
@@ -1100,13 +1322,9 @@ class DistStageRuntime(StageRuntime):
if registered_stage_cfg is None:
raise ValueError(f"Remote stage {plan.metadata.stage_id} registered without stage config")
- # Remote diffusion registration still transports the legacy mapping
- # shape. Reconstruct and project that shape until its RFC #4021 cutover.
- metadata = (
- extract_legacy_stage_metadata(OmegaConf.create(registered_stage_cfg))
- if plan.metadata.stage_type == "diffusion"
- else copy.deepcopy(plan.metadata)
- )
+ # Registration is transport-only: the head-side typed plan remains
+ # authoritative for metadata and topology.
+ metadata = extract_stage_metadata_from_omni_stage_config(plan.stage_cfg)
metadata.replica_id = plan.replica_id
ctx = StageRemoteFactoryContext(
stage_id=plan.metadata.stage_id,
@@ -1303,7 +1521,7 @@ class DistStageRuntime(StageRuntime):
def create_stage_runtime(
stage_configs: list[Any],
model: str,
- config_path: str,
+ config_path: str | None,
*,
single_stage_mode: bool,
stage_init_timeout: int,
@@ -1319,6 +1537,7 @@ def create_stage_runtime(
omni_lb_policy: str = "random",
request_queue: janus.Queue[EngineQueueMessage] | None = None,
log_stats: bool = False,
+ client_config: OmniClientConfig | None = None,
) -> StageRuntime:
"""Factory: select StageRuntime or DistStageRuntime."""
if single_stage_mode:
@@ -1350,4 +1569,5 @@ def create_stage_runtime(
tokenizer=tokenizer,
parallel_stage_init=parallel_stage_init,
log_stats=log_stats,
+ client_config=client_config,
)
diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py
index bc7f81132..493f2533e 100644
--- a/vllm_omni/entrypoints/cli/serve.py
+++ b/vllm_omni/entrypoints/cli/serve.py
@@ -8,13 +8,16 @@ Supports both multi-stage LLM models (e.g., Qwen2.5-Omni) and
diffusion models (e.g., Qwen-Image) through the same CLI interface.
"""
+from __future__ import annotations
+
import argparse
+import contextlib
import json
import math
import os
import signal
from types import FrameType
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast
import uvloop
from vllm.entrypoints.cli.types import CLISubcommand
@@ -23,10 +26,18 @@ from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG
from vllm.logger import init_logger
from vllm_omni.entrypoints.cli.logo import log_logo
-from vllm_omni.entrypoints.openai.api_server import omni_run_server
-from vllm_omni.entrypoints.utils import parse_stage_overrides
+from vllm_omni.entrypoints.openai.api_server import (
+ omni_run_server,
+ run_omni_api_server_worker_proc,
+)
+from vllm_omni.entrypoints.utils import parse_stage_overrides, prepare_stage_config_inputs
from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace
+if TYPE_CHECKING:
+ from vllm.v1.utils import APIServerProcessManager
+
+ from vllm_omni.engine.stage_runtime import StageEngineLaunch, StageRuntime
+
logger = init_logger(__name__)
DESCRIPTION = """Launch a local OpenAI-compatible API server to serve Omni models
@@ -141,6 +152,8 @@ class OmniServeCommand(CLISubcommand):
if args.headless:
run_headless(args)
+ elif (getattr(args, "api_server_count", None) or 1) > 1:
+ run_multi_api_server_omni(args)
else:
uvloop.run(omni_run_server(args))
@@ -198,6 +211,39 @@ class OmniServeCommand(CLISubcommand):
if getattr(args, "omni_replica_address", None) is not None and not args.headless:
raise ValueError("--omni-replica-address requires --headless to be set")
+ api_server_count = getattr(args, "api_server_count", None)
+ if api_server_count is not None and api_server_count < 1 and not args.headless:
+ raise ValueError("--api-server-count must be >= 1 unless --headless is set")
+ if args.headless and api_server_count is not None and api_server_count > 0:
+ raise ValueError("--api-server-count cannot be used with --headless")
+ if api_server_count is not None and api_server_count > 1:
+ distributed_args = (
+ "stage_id",
+ "omni_master_address",
+ "omni_master_port",
+ "omni_replica_address",
+ )
+ if any(getattr(args, name, None) is not None for name in distributed_args):
+ raise ValueError(
+ "--api-server-count > 1 cannot be combined with stage-based "
+ "or distributed stage arguments (--stage-id/--omni-master-* / "
+ "--omni-replica-address)"
+ )
+ if getattr(args, "omni_dp_size_local", 1) != 1:
+ raise ValueError("--api-server-count > 1 requires --omni-dp-size-local=1")
+ if getattr(args, "worker_backend", "multi_process") != "multi_process":
+ raise ValueError("--api-server-count > 1 requires --worker-backend=multi_process")
+ if getattr(args, "enable_fault_tolerance", False):
+ raise ValueError("--api-server-count > 1 cannot be combined with --enable-fault-tolerance")
+ if getattr(args, "enable_elastic_ep", False):
+ raise ValueError("--api-server-count > 1 cannot be combined with --enable-elastic-ep")
+ if getattr(args, "enable_sleep_mode", False):
+ raise ValueError("--api-server-count > 1 cannot be combined with --enable-sleep-mode")
+ from vllm import envs as vllm_envs
+
+ if getattr(vllm_envs, "VLLM_ALLOW_RUNTIME_LORA_UPDATING", False):
+ raise ValueError("--api-server-count > 1 cannot be combined with VLLM_ALLOW_RUNTIME_LORA_UPDATING")
+
# --omni-dp-size-local is process-local. A value other than 1 only
# makes sense when this process owns a stage (head or headless).
omni_dp_size_local = getattr(args, "omni_dp_size_local", None)
@@ -221,7 +267,6 @@ class OmniServeCommand(CLISubcommand):
"data_parallel_rpc_port": "--data-parallel-rpc-port",
"data_parallel_start_rank": "--data-parallel-start-rank",
"data_parallel_backend": "--data-parallel-backend",
- "api_server_count": "--api-server-count",
"enable_expert_parallel": "--enable-expert-parallel",
}
offenders = sorted(flag for dest, flag in prohibited_with_omni.items() if dest in explicit_cli_keys)
@@ -261,6 +306,8 @@ class OmniServeCommand(CLISubcommand):
model = getattr(args, "model_tag", None) or getattr(args, "model", None)
if model and is_diffusion_model(model):
+ if api_server_count is not None and api_server_count > 1:
+ raise ValueError("--api-server-count > 1 is not supported for diffusion models")
logger.info("Detected diffusion model: %s", model)
return
validate_parsed_serve_args(args)
@@ -1027,6 +1074,207 @@ class OmniServeCommand(CLISubcommand):
return serve_parser
+def _build_multi_api_stage_runtime(args: TrackingNamespace, num_api_servers: int) -> StageRuntime:
+ """Resolve the local EngineCore stages that the parent process owns."""
+ from vllm_omni.config.resolver import resolve_omni_config
+ from vllm_omni.engine.stage_runtime import StageRuntime
+
+ kwargs = args.get_explicit_kwargs_dict()
+ model = kwargs.pop("model", None) or args.model
+
+ trust_remote_code = kwargs.get("trust_remote_code")
+ if trust_remote_code is False:
+ trust_remote_code = None
+ config_inputs = prepare_stage_config_inputs(
+ model,
+ kwargs,
+ trust_remote_code=trust_remote_code,
+ snapshot_model=True,
+ )
+ model = config_inputs.model
+ kwargs = config_inputs.kwargs
+ resolved = resolve_omni_config(
+ model,
+ cli_overrides=kwargs,
+ trust_remote_code=config_inputs.trust_remote_code,
+ deploy_config_path=config_inputs.deploy_config_path,
+ stage_overrides=config_inputs.stage_overrides,
+ strategy_config_path=config_inputs.strategy_config_path,
+ )
+
+ config_path = resolved.config_path
+ stage_configs = list(resolved.stage_configs)
+
+ sleep_stages = [
+ int(getattr(stage_config, "stage_id", stage_index))
+ for stage_index, stage_config in enumerate(stage_configs)
+ if bool(
+ getattr(
+ getattr(stage_config, "model_config", getattr(stage_config, "engine_args", None)),
+ "enable_sleep_mode",
+ False,
+ )
+ )
+ ]
+ if sleep_stages:
+ raise ValueError(
+ "--api-server-count > 1 cannot be combined with sleep mode; "
+ f"disable enable_sleep_mode for stage(s) {sleep_stages}"
+ )
+
+ async_chunk = any(
+ bool(
+ getattr(
+ getattr(stage, "connector_config", None),
+ "async_chunk",
+ getattr(getattr(stage, "engine_args", None), "async_chunk", False),
+ )
+ )
+ for stage in stage_configs
+ )
+ return StageRuntime(
+ stage_configs=stage_configs,
+ model=model,
+ config_path=config_path,
+ stage_init_timeout=int(getattr(args, "stage_init_timeout", 300)),
+ parallel_stage_init=bool(getattr(args, "parallel_stage_init", False)),
+ async_chunk=async_chunk,
+ tokenizer=getattr(args, "tokenizer", None),
+ log_stats=not bool(getattr(args, "disable_log_stats", False)),
+ )
+
+
+def _wait_for_multi_api_server_completion(
+ api_server_manager: APIServerProcessManager,
+ engine_launch: StageEngineLaunch,
+) -> None:
+ """Wait until API workers complete or any shared stage engine fails."""
+ from multiprocessing import connection
+
+ api_processes = list(api_server_manager.processes)
+ engine_processes = [
+ process
+ for resources in engine_launch.resources
+ if resources.manager is not None
+ for process in resources.manager.processes
+ ]
+ api_by_sentinel = {process.sentinel: process for process in api_processes}
+ engine_by_sentinel = {process.sentinel: process for process in engine_processes}
+
+ while api_by_sentinel:
+ ready = connection.wait([*api_by_sentinel, *engine_by_sentinel])
+ for sentinel in ready:
+ if sentinel in engine_by_sentinel:
+ process = engine_by_sentinel[sentinel]
+ raise RuntimeError(
+ f"Shared stage engine process {process.name} (PID: {process.pid}) "
+ f"exited with code {process.exitcode}"
+ )
+ process = api_by_sentinel.pop(sentinel)
+ if process.exitcode != 0:
+ raise RuntimeError(
+ f"API server process {process.name} (PID: {process.pid}) exited with code {process.exitcode}"
+ )
+
+
+def _start_api_server_process_manager(
+ *,
+ cleanup_timeout: float,
+ **manager_kwargs: object,
+) -> APIServerProcessManager:
+ """Construct vLLM's manager with rollback for a partial ``__init__``.
+
+ ``APIServerProcessManager`` starts workers one by one in ``__init__`` and
+ installs its finalizer only after every ``Process.start`` succeeds. Keep a
+ reference to the partially initialized object so workers started before a
+ later failure are still terminated.
+ """
+ from vllm.v1.utils import APIServerProcessManager, shutdown
+
+ manager = APIServerProcessManager.__new__(APIServerProcessManager)
+ try:
+ APIServerProcessManager.__init__(manager, **manager_kwargs)
+ except BaseException:
+ try:
+ for pipe in getattr(manager, "_address_pipes", ()):
+ with contextlib.suppress(Exception):
+ pipe.close()
+ # A Process whose start() raised is still present in the upstream
+ # manager's list, but calling is_alive() on it raises because it
+ # has no Popen object. Only hand successfully started children to
+ # vLLM's shutdown helper.
+ started_processes = [process for process in getattr(manager, "processes", ()) if process.pid is not None]
+ if started_processes:
+ shutdown(started_processes, timeout=cleanup_timeout)
+ except Exception:
+ logger.exception("Failed to clean up partially started API server processes")
+ raise
+ return manager
+
+
+def run_multi_api_server_omni(args: TrackingNamespace) -> None:
+ """Launch API subprocesses that share one set of local stage engines."""
+ from vllm.entrypoints.openai.api_server import setup_server
+ from vllm.v1.metrics.prometheus import setup_multiprocess_prometheus
+
+ num_api_servers = int(args.api_server_count)
+ if num_api_servers < 2:
+ raise ValueError(f"api_server_count must be >= 2, got {num_api_servers}")
+
+ setup_multiprocess_prometheus()
+ shutdown_requested = False
+
+ def signal_handler(signum: int, frame: FrameType | None) -> None:
+ nonlocal shutdown_requested
+ logger.debug("Received %d signal.", signum)
+ if not shutdown_requested:
+ shutdown_requested = True
+ raise SystemExit
+
+ signal.signal(signal.SIGTERM, signal_handler)
+ signal.signal(signal.SIGINT, signal_handler)
+
+ listen_address, sock = setup_server(args, reuse_port=True)
+ stage_runtime = None
+ api_server_manager = None
+ engine_launch = None
+ try:
+ stage_runtime = _build_multi_api_stage_runtime(args, num_api_servers)
+ with stage_runtime.launch_stage_engines(num_api_servers) as engine_launch:
+ args._omni_stage_client_configs = engine_launch.client_configs
+ primary_addresses = engine_launch.resources[0].addresses
+ if primary_addresses is None:
+ raise RuntimeError("Primary stage engine returned no addresses")
+ timeout = float(getattr(args, "shutdown_timeout", 5) or 5)
+ api_server_manager = _start_api_server_process_manager(
+ cleanup_timeout=timeout,
+ listen_address=listen_address,
+ sock=sock,
+ args=args,
+ num_servers=num_api_servers,
+ input_addresses=primary_addresses.inputs,
+ output_addresses=primary_addresses.outputs,
+ target_server_fn=run_omni_api_server_worker_proc,
+ )
+ engine_launch.watched_frontend_processes.extend(api_server_manager.processes)
+ actual_inputs, actual_outputs = api_server_manager.gather_actual_addresses()
+ primary_addresses.inputs = actual_inputs
+ primary_addresses.outputs = actual_outputs
+
+ _wait_for_multi_api_server_completion(api_server_manager, engine_launch)
+ finally:
+ timeout = float(getattr(args, "shutdown_timeout", 5) or 5)
+ if api_server_manager is not None:
+ api_server_manager.shutdown(timeout=timeout)
+ if engine_launch is not None:
+ engine_launch.shutdown()
+ if stage_runtime is not None:
+ stage_runtime.shutdown()
+ sock.close()
+ if shutdown_requested:
+ logger.info("Shared API server shutdown completed")
+
+
def run_headless(args: TrackingNamespace) -> None:
"""Run a single stage in headless mode.
@@ -1048,6 +1296,7 @@ def run_headless(args: TrackingNamespace) -> None:
)
from vllm_omni.engine.stage_init_utils import (
build_engine_args_dict,
+ build_engine_args_dict_from_omni_stage_config,
build_vllm_config,
get_stage_connector_spec,
inject_omni_kv_connector_config,
@@ -1074,15 +1323,6 @@ def run_headless(args: TrackingNamespace) -> None:
# Filter down to a dict of things explicitly requested by the user
args_dict = args.get_explicit_kwargs_dict()
- # This controls only the API-process WebSocket and is not a stage-engine option.
- args_dict.pop("robot_openpi_idle_timeout", None)
- # This CLI-only negative alias is consumed below when selecting the
- # launcher log_stats value; it is not a per-stage config override.
- args_dict.pop("disable_log_stats", None)
-
- deploy_config_path = args_dict.pop("deploy_config", None)
- strategy_config_path = args_dict.pop("strategy_config", None)
- stage_overrides = args_dict.pop("stage_overrides", None)
# ``--replica-id`` is deprecated and ignored — replica ids are
# auto-assigned by ``OmniMasterServer`` so headless processes carry
@@ -1098,15 +1338,21 @@ def run_headless(args: TrackingNamespace) -> None:
)
args_dict.pop("replica_id")
- resolved = resolve_omni_config(
+ config_inputs = prepare_stage_config_inputs(
model,
+ args_dict,
# store_true cannot express an explicit False: absent maps to None
# ("not specified") so the deploy yaml's per-stage value applies.
trust_remote_code=getattr(args, "trust_remote_code", None) or None,
+ )
+ args_dict = config_inputs.kwargs
+ resolved = resolve_omni_config(
+ model,
+ trust_remote_code=config_inputs.trust_remote_code,
cli_overrides=args_dict,
- deploy_config_path=deploy_config_path,
- stage_overrides=stage_overrides,
- strategy_config_path=strategy_config_path,
+ deploy_config_path=config_inputs.deploy_config_path,
+ stage_overrides=config_inputs.stage_overrides,
+ strategy_config_path=config_inputs.strategy_config_path,
)
config_path = resolved.config_path
stage_configs = list(resolved.stage_configs)
@@ -1141,21 +1387,21 @@ def run_headless(args: TrackingNamespace) -> None:
stage_connector_spec = get_stage_connector_spec(
omni_transfer_config=omni_transfer_config,
stage_id=stage_id,
- async_chunk=bool(stage_cfg.engine_args.get("async_chunk", False)),
+ async_chunk=bool(
+ getattr(getattr(stage_cfg, "connector_config", None), "async_chunk", None)
+ if hasattr(stage_cfg, "connector_config")
+ else stage_cfg.engine_args.get("async_chunk", False)
+ ),
)
- # ``runtime_cfg`` is mostly inherited from the parent's
- # CUDA_VISIBLE_DEVICES; when ``--omni-dp-size-local > 1`` we additionally
- # bracket each replica's spawn below with setup_stage_devices so they
- # don't all stack on cuda:0 (see ``per_replica_devices`` above).
- # Headless startup still supplies the legacy OmegaConf stage shape through
- # the stable adapter entry point. The implementation switches only when
- # RFC #4021 threads structured stage configs through the launch plan.
- engine_args_dict = build_engine_args_dict(
- stage_cfg,
- model,
- stage_connector_spec=stage_connector_spec,
- cli_tokenizer=getattr(args, "tokenizer", None),
+ engine_args_dict = (
+ build_engine_args_dict_from_omni_stage_config(
+ stage_cfg, model, stage_connector_spec=stage_connector_spec, cli_tokenizer=getattr(args, "tokenizer", None)
+ )
+ if hasattr(stage_cfg, "connector_config")
+ else build_engine_args_dict(
+ stage_cfg, model, stage_connector_spec=stage_connector_spec, cli_tokenizer=getattr(args, "tokenizer", None)
+ )
)
inject_omni_kv_connector_config(engine_args_dict, omni_kv_connector, stage_id)
diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py
index 20706396c..94f1ab513 100644
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/vllm_omni/entrypoints/openai/api_server.py
@@ -6,6 +6,7 @@ This module owns app construction, server startup, app-state initialization,
and route bodies that have not yet moved to endpoint-owned modules."""
import asyncio
+import copy
import dataclasses
import json
import multiprocessing
@@ -14,6 +15,8 @@ import os
# Image generation API imports
import random
+import signal
+import socket
import time
from argparse import Namespace
from collections.abc import AsyncIterator
@@ -22,6 +25,7 @@ from http import HTTPStatus
from pathlib import Path
from typing import Annotated, Any, Literal
+import uvloop
import vllm.envs as envs
from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, Request, UploadFile, WebSocket
from fastapi.responses import FileResponse, JSONResponse, Response, StreamingResponse
@@ -81,12 +85,14 @@ from vllm.renderers.online_renderer import OnlineRenderer
from vllm.tasks import POOLING_TASKS
from vllm.tool_parsers import ToolParserManager
from vllm.utils import random_uuid
-from vllm.utils.system_utils import decorate_logs
+from vllm.utils.system_utils import decorate_logs, set_process_title
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
from vllm_omni.config.endpoint_policy import (
shutdown_unsupported_routes,
)
+from vllm_omni.engine.stage_init_utils import set_death_signal
+from vllm_omni.engine.stage_runtime import OmniClientConfig
from vllm_omni.entrypoints.async_omni import ABORT_TIMEOUT_S, AsyncOmni
from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler
from vllm_omni.entrypoints.duplex.warmup import _warmup_duplex_realtime
@@ -232,8 +238,55 @@ async def omni_run_server(args, **uvicorn_kwargs) -> None:
await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs)
-async def omni_run_server_worker(listen_address, sock, args, client_config=None, **uvicorn_kwargs) -> None:
+def run_omni_api_server_worker_proc(
+ listen_address: str,
+ sock: socket.socket,
+ args: TrackingNamespace,
+ client_config: dict[str, Any] | None = None,
+ **uvicorn_kwargs: object,
+) -> None:
+ """Entrypoint used by vLLM's API server process manager."""
+ set_death_signal(signal.SIGTERM)
+ manager_config = client_config or {}
+ client_index = int(manager_config.get("client_index", 0))
+ all_client_configs = getattr(args, "_omni_stage_client_configs", None)
+ if not all_client_configs or not 0 <= client_index < len(all_client_configs):
+ raise RuntimeError(f"Missing Omni stage client configuration for API server {client_index}")
+
+ omni_client_config: OmniClientConfig = copy.deepcopy(all_client_configs[client_index])
+ omni_client_config["client_count"] = int(manager_config.get("client_count", 1))
+ omni_client_config["client_index"] = client_index
+
+ stage_addresses = omni_client_config["stage_addresses"]
+ first_stage_id = min(stage_addresses)
+ first_replica_id = min(stage_addresses[first_stage_id])
+ first_addresses = stage_addresses[first_stage_id][first_replica_id]
+ for key in ("input_address", "output_address", "actual_address_pipe", "tensor_queue"):
+ if key in manager_config:
+ first_addresses[key] = manager_config[key]
+
+ set_process_title("APIServer", str(client_index))
+ decorate_logs("APIServer", skip_if_decorated=True)
+ uvloop.run(
+ omni_run_server_worker(
+ listen_address,
+ sock,
+ args,
+ client_config=omni_client_config,
+ **uvicorn_kwargs,
+ )
+ )
+
+
+async def omni_run_server_worker(
+ listen_address: str,
+ sock: socket.socket,
+ args: TrackingNamespace,
+ client_config: OmniClientConfig | None = None,
+ **uvicorn_kwargs: object,
+) -> None:
"""Run a single API server worker."""
+ api_server_count = _resolve_api_server_count(args, client_config)
if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3:
ToolParserManager.import_tool_parser(args.tool_parser_plugin)
@@ -273,6 +326,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
# OMNI: Pass supported_tasks to build_app (required by upstream vLLM)
app = build_openai_app(args, supported_tasks)
+ app.state.api_server_count = api_server_count
# OMNI: Remove upstream routes that we override with omni-specific handlers
remove_route_from_app(app, "/v1/chat/completions", {"POST"})
@@ -348,6 +402,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
if scope["type"] == "http":
scope.setdefault("state", {})
scope["state"]["request_timestamp"] = time.time()
+
await self._inner(scope, receive, send)
# Startup duplex warmup (duplex_session.warmup_frames in the deploy
@@ -406,7 +461,7 @@ async def build_async_omni(
args: TrackingNamespace,
*,
disable_frontend_multiprocessing: bool | None = None,
- client_config: dict[str, Any] | None = None,
+ client_config: OmniClientConfig | None = None,
) -> AsyncIterator[EngineClient]:
"""Build an AsyncOmni instance from command-line arguments.
@@ -437,6 +492,7 @@ async def build_async_omni(
async with build_async_omni_from_stage_config(
args,
disable_frontend_multiprocessing=disable_frontend_multiprocessing,
+ client_config=client_config,
) as async_omni:
yield async_omni
@@ -476,6 +532,7 @@ async def build_async_omni_from_stage_config(
args: TrackingNamespace,
*,
disable_frontend_multiprocessing: bool = False,
+ client_config: OmniClientConfig | None = None,
) -> AsyncIterator[EngineClient]:
"""Create AsyncOmni from stage configuration.
@@ -532,6 +589,8 @@ async def build_async_omni_from_stage_config(
kwargs.pop("robot_openpi_idle_timeout", None)
model = kwargs.pop("model", None) or args.model
kwargs.setdefault("log_stats", not args.disable_log_stats)
+ if client_config is not None:
+ kwargs["client_config"] = client_config
if _should_serve_duplex(model, kwargs):
async_omni = DuplexOmni(model=model, **kwargs)
else:
@@ -1430,6 +1489,42 @@ async def list_voices(raw_request: Request):
return JSONResponse(content={"voices": speakers, "uploaded_voices": uploaded_speakers})
+def _resolve_api_server_count(args: Namespace, client_config: OmniClientConfig | None) -> int:
+ """Resolve frontend topology once, from the launcher's client configuration."""
+ requested = getattr(args, "api_server_count", None)
+ actual = 1 if client_config is None else client_config.get("client_count")
+ if type(actual) is not int or actual < 1:
+ raise ValueError("API client configuration requires a positive integer client_count")
+ if requested is not None and (type(requested) is not int or requested < 1 or requested != actual):
+ raise ValueError("API server count does not match the frontend launch configuration")
+ return actual
+
+
+def _reject_process_local_state_with_multiple_api_workers(raw_request: Request, operation: str) -> None:
+ """Reject access to frontend state that is not shared across API workers."""
+ api_server_count = getattr(raw_request.app.state, "api_server_count", None)
+ if type(api_server_count) is not int or api_server_count < 1:
+ raise HTTPException(
+ status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
+ detail="API worker topology is not initialized.",
+ )
+ if api_server_count > 1:
+ raise HTTPException(
+ status_code=HTTPStatus.CONFLICT.value,
+ detail=(
+ f"{operation} is not supported with --api-server-count > 1 because "
+ "the operation uses process-local frontend state"
+ ),
+ )
+
+
+async def _require_single_api_video_store(raw_request: Request) -> None:
+ # VIDEO_STORE and VIDEO_TASKS are process-local, including reads and
+ # cancellation. Keep this independent of the temporary diffusion launch
+ # restriction: enabling multi-API diffusion alone cannot make jobs shared.
+ _reject_process_local_state_with_multiple_api_workers(raw_request, "Asynchronous video jobs")
+
+
@router.post(
"/v1/audio/voices",
responses={
@@ -1475,6 +1570,7 @@ async def upload_voice(
Returns:
JSON response with voice information
"""
+ _reject_process_local_state_with_multiple_api_workers(raw_request, "Runtime voice upload")
handler = Omnispeech(raw_request)
if handler is None:
return _create_speech_error_json_response(
@@ -1540,6 +1636,7 @@ async def delete_voice(name: str, raw_request: Request):
Returns:
JSON response indicating success or failure
"""
+ _reject_process_local_state_with_multiple_api_workers(raw_request, "Runtime voice deletion")
handler = Omnispeech(raw_request)
if handler is None:
return _create_speech_error_json_response(
@@ -1633,13 +1730,15 @@ async def streaming_video_output(websocket: WebSocket):
@router.websocket("/v1/realtime")
async def realtime_websocket(websocket: WebSocket):
"""WebSocket endpoint for OpenAI-style realtime interactions."""
- await _wait_for_duplex_warmup(websocket)
duplex_handler = getattr(websocket.app.state, "openai_serving_duplex", None)
duplex_query = websocket.query_params.get("duplex")
use_duplex_realtime = duplex_handler is not None and (
duplex_query is None or (isinstance(duplex_query, str) and duplex_query.lower() in {"1", "true", "on"})
)
if use_duplex_realtime and duplex_handler is not None:
+ if await _reject_multi_api_duplex(websocket):
+ return
+ await _wait_for_duplex_warmup(websocket)
await duplex_handler.handle_realtime_session(websocket)
return
@@ -1669,6 +1768,8 @@ async def _wait_for_duplex_warmup(websocket: WebSocket) -> None:
@router.websocket("/v1/duplex")
async def duplex_websocket(websocket: WebSocket):
"""Alias of ``/v1/realtime?duplex=1``: the same Realtime duplex session protocol."""
+ if await _reject_multi_api_duplex(websocket):
+ return
await _wait_for_duplex_warmup(websocket)
handler = getattr(websocket.app.state, "openai_serving_duplex", None)
if handler is None:
@@ -1709,6 +1810,24 @@ async def realtime_robot_openpi(websocket: WebSocket):
await connection.handle_connection()
+async def _reject_multi_api_duplex(websocket: WebSocket) -> bool:
+ # Session credentials and replay state are local to each frontend; a shared
+ # listening socket cannot route reconnects back to the session owner.
+ count = getattr(websocket.app.state, "api_server_count", None)
+ if count is None or (type(count) is int and count == 1):
+ return False
+ await websocket.accept()
+ await websocket.send_json(
+ {
+ "type": "error",
+ "error": "Duplex sessions require a single API worker with initialized topology.",
+ "code": "multi_api_duplex_unsupported",
+ }
+ )
+ await websocket.close(code=1008)
+ return True
+
+
# RL Rollout serving (RFC #3747, P0)
@@ -2476,6 +2595,7 @@ async def edit_images(
@router.post(
"/v1/videos",
+ dependencies=[Depends(_require_single_api_video_store)],
responses={
HTTPStatus.OK.value: {"model": VideoResponse},
HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
@@ -2617,7 +2737,7 @@ async def create_video_sync(
)
-@router.get("/v1/videos", response_model=VideoListResponse)
+@router.get("/v1/videos", response_model=VideoListResponse, dependencies=[Depends(_require_single_api_video_store)])
async def list_videos(
after: str | None = None,
limit: int | None = Query(None, ge=0, le=100),
@@ -2654,7 +2774,7 @@ async def list_videos(
return VideoListResponse(data=jobs, has_more=has_more, first_id=first_id, last_id=last_id)
-@router.get("/v1/videos/{video_id}", response_model=None)
+@router.get("/v1/videos/{video_id}", response_model=None, dependencies=[Depends(_require_single_api_video_store)])
async def retrieve_video(video_id: str) -> VideoResponse | JSONResponse:
"""Retrieve metadata for a previously created video job.
@@ -2682,7 +2802,7 @@ async def retrieve_video(video_id: str) -> VideoResponse | JSONResponse:
return job
-@router.delete("/v1/videos/{video_id}")
+@router.delete("/v1/videos/{video_id}", dependencies=[Depends(_require_single_api_video_store)])
async def delete_video(video_id: str, raw_request: Request) -> VideoDeleteResponse:
"""Delete a stored video job and any generated output.
@@ -2752,7 +2872,7 @@ async def delete_video(video_id: str, raw_request: Request) -> VideoDeleteRespon
return VideoDeleteResponse(id=job.id, deleted=True)
-@router.get("/v1/videos/{video_id}/content")
+@router.get("/v1/videos/{video_id}/content", dependencies=[Depends(_require_single_api_video_store)])
async def download_video(video_id: str) -> Response:
"""Download the generated file for a completed video job.
@@ -2845,6 +2965,7 @@ async def stop_profile(raw_request: Request, request: ProfileRequest | None = No
@router.post("/v1/omni/sleep")
async def omni_sleep(request: OmniSleepRequest, raw_request: Request):
+ _reject_process_local_state_with_multiple_api_workers(raw_request, "Sleep")
engine_client = raw_request.app.state.engine_client
sleeping_set = raw_request.app.state.sleeping_stages
if not hasattr(engine_client, "sleep"):
@@ -2857,6 +2978,7 @@ async def omni_sleep(request: OmniSleepRequest, raw_request: Request):
@router.post("/v1/omni/wakeup")
async def omni_wakeup(request: OmniWakeupRequest, raw_request: Request):
+ _reject_process_local_state_with_multiple_api_workers(raw_request, "Wakeup")
engine_client = raw_request.app.state.engine_client
sleeping_set = raw_request.app.state.sleeping_stages
if not any(sid in sleeping_set for sid in request.stage_ids):
diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py
index 484d5d170..5b3f2d54b 100644
--- a/vllm_omni/entrypoints/openai/serving_chat.py
+++ b/vllm_omni/entrypoints/openai/serving_chat.py
@@ -248,10 +248,27 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
return obj.get(key, default)
return getattr(obj, key, default)
+ @classmethod
+ def _stage_model_metadata(cls, stage: Any) -> tuple[str | None, str | None]:
+ """Return ``(model_arch, model_stage)`` for typed or legacy stages."""
+ model_config = cls._stage_get(stage, "model_config")
+ model_arch = cls._stage_get(model_config, "model_arch")
+ model_stage = cls._stage_get(stage, "model_stage")
+
+ topology = cls._stage_get(stage, "stage_pipeline_config")
+ if model_stage is None:
+ model_stage = cls._stage_get(topology, "model_stage")
+
+ engine_args = cls._stage_get(stage, "engine_args")
+ if model_arch is None:
+ model_arch = cls._stage_get(engine_args, "model_arch")
+ if model_stage is None:
+ model_stage = cls._stage_get(engine_args, "model_stage")
+ return model_arch, model_stage
+
def _has_minicpmo45_stage(self) -> bool:
for stage in getattr(self.engine_client, "stage_configs", []) or []:
- engine_args = self._stage_get(stage, "engine_args")
- model_arch = self._stage_get(engine_args, "model_arch")
+ model_arch, _ = self._stage_model_metadata(stage)
if model_arch == "MiniCPMO45OmniForConditionalGeneration":
return True
return False
@@ -324,11 +341,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
if not hasattr(sp, "output_kind"):
continue
- engine_args = self._stage_get(stage, "engine_args")
- if self._stage_get(engine_args, "model_arch") != "MiniCPMO45OmniForConditionalGeneration":
+ model_arch, model_stage = self._stage_model_metadata(stage)
+ if model_arch != "MiniCPMO45OmniForConditionalGeneration":
continue
- model_stage = self._stage_get(engine_args, "model_stage")
if model_stage == "llm":
sp.output_kind = RequestOutputKind.FINAL_ONLY
elif model_stage == "tts":
diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py
index 80f357651..8aa32899a 100644
--- a/vllm_omni/entrypoints/openai/serving_speech.py
+++ b/vllm_omni/entrypoints/openai/serving_speech.py
@@ -36,6 +36,7 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams
from vllm.utils import random_uuid
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
+from vllm_omni.config.stage_config import StagePipelineConfig
from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin, StreamingAudioResampler
from vllm_omni.entrypoints.openai.protocol.audio import (
AudioResponse,
@@ -76,6 +77,24 @@ _SPEECH_USAGE_TOTAL_TOKENS_HEADER = "X-VLLM-OMNI-TOTAL-TOKENS"
_SPEECH_USAGE_INPUT_TEXT_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-TEXT-TOKENS"
_SPEECH_USAGE_INPUT_AUDIO_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-AUDIO-TOKENS"
+
+def _stage_speech_metadata(stage: Any) -> tuple[str | None, str | None, str | None]:
+ """Return speech routing metadata from typed or legacy stage configs."""
+ topology = getattr(stage, "stage_pipeline_config", None)
+ if isinstance(topology, StagePipelineConfig):
+ return (
+ getattr(topology, "model_stage", None),
+ getattr(getattr(stage, "model_config", None), "model_arch", None),
+ getattr(stage, "worker_type", None),
+ )
+ engine_args = getattr(stage, "engine_args", None)
+ return (
+ getattr(engine_args, "model_stage", None),
+ getattr(engine_args, "model_arch", None),
+ getattr(engine_args, "worker_type", None),
+ )
+
+
# TTS Configuration
#
# The stage-key -> model-type mapping is NOT declared here: it is derived from
@@ -472,12 +491,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
"""Find and return the TTS stage config, or None if not found."""
tts_stage_keys = all_tts_stage_keys()
entry_stage_archs = tts_entry_stage_archs()
- all_stages = frozenset(stage.engine_args.model_stage for stage in self.engine_client.stage_configs)
+ all_stages = frozenset(
+ model_stage
+ for stage in self.engine_client.stage_configs
+ if (model_stage := _stage_speech_metadata(stage)[0]) is not None
+ )
for stage in self.engine_client.stage_configs:
- engine_args = stage.engine_args
- model_stage = engine_args.model_stage
- model_arch = getattr(engine_args, "model_arch", None)
- worker_type = getattr(engine_args, "worker_type", None)
+ model_stage, model_arch, worker_type = _stage_speech_metadata(stage)
if model_stage in tts_stage_keys:
# Owning the stage key is not always enough: a model may be
# speech-capable only in some deployment topologies. Ask the
@@ -501,10 +521,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
"""
if self._tts_stage is None:
return None
- return detect_tts_model_type(
- getattr(self._tts_stage.engine_args, "model_stage", None),
- getattr(self._tts_stage.engine_args, "model_arch", None),
- )
+ return detect_tts_model_type(*_stage_speech_metadata(self._tts_stage)[:2])
def _compute_max_instructions_length(self) -> int:
"""Compute max instructions length with precedence: CLI > stage config > default.
@@ -518,7 +535,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
# 2. Try to get from TTS stage config
if self._tts_stage is not None:
- tts_args = getattr(self._tts_stage, "tts_args", {})
+ topology = getattr(self._tts_stage, "stage_pipeline_config", None)
+ tts_args = (
+ getattr(topology, "extras", {}).get("tts_args", {})
+ if isinstance(topology, StagePipelineConfig)
+ else getattr(self._tts_stage, "tts_args", {})
+ )
if "max_instructions_length" in tts_args:
return tts_args["max_instructions_length"]
@@ -1896,9 +1918,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
# chat-template markers (im_start_token_id 151644). A raw-text
# prompt produces a 1-token placeholder that crashes the talker's
# prefill/decode handoff. Reject early with an actionable message.
- stage_names = {
- getattr(getattr(s, "engine_args", None), "model_stage", None) for s in self.engine_client.stage_configs
- }
+ stage_names = {_stage_speech_metadata(s)[0] for s in self.engine_client.stage_configs}
if "talker" in stage_names:
raise ValueError(
"The /v1/audio/speech endpoint is only supported for "
diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py
index 9a24a1ef9..717496362 100644
--- a/vllm_omni/entrypoints/openai/serving_video.py
+++ b/vllm_omni/entrypoints/openai/serving_video.py
@@ -53,6 +53,30 @@ logger = init_logger(__name__)
_VIDEO_RESPONSE_FRAME_CONVERSION_WORKERS = 8
+
+def _config_value(config: Any, key: str, default: Any = None) -> Any:
+ if isinstance(config, Mapping):
+ return config.get(key, default)
+ return getattr(config, key, default)
+
+
+def _stage_diffusion_model_class_name(stage_config: Any) -> str | None:
+ """Resolve a stage's diffusion class across typed and legacy shapes."""
+ diffusion_config = _config_value(stage_config, "diffusion_config")
+ model_class_name = _config_value(diffusion_config, "model_class_name")
+ if model_class_name:
+ return str(model_class_name)
+
+ model_config = _config_value(stage_config, "model_config")
+ model_arch = _config_value(model_config, "model_arch") or _config_value(stage_config, "model_arch")
+ if model_arch:
+ return str(model_arch)
+
+ engine_args = _config_value(stage_config, "engine_args", {})
+ model_class_name = _config_value(engine_args, "model_class_name")
+ return str(model_class_name) if model_class_name else None
+
+
if TYPE_CHECKING:
from vllm_omni.diffusion.data import OmniDiffusionConfig
@@ -195,18 +219,7 @@ class OmniOpenAIServingVideo:
model_class_name = getattr(od_config, "model_class_name", None)
model_archs = [model_class_name]
for stage_config in self.stage_configs or ():
- stage_get = (
- stage_config.get if isinstance(stage_config, Mapping) else lambda key: getattr(stage_config, key, None)
- )
- engine_args = stage_get("engine_args") or {}
- model_archs.extend(
- (
- stage_get("model_arch"),
- engine_args.get("model_class_name")
- if isinstance(engine_args, Mapping)
- else getattr(engine_args, "model_class_name", None),
- )
- )
+ model_archs.append(_stage_diffusion_model_class_name(stage_config))
metadata_capability = any(
get_diffusion_model_metadata(model_arch).supports_mixed_reference_inputs for model_arch in model_archs
)
@@ -223,18 +236,7 @@ class OmniOpenAIServingVideo:
od_config = self._resolve_diffusion_od_config()
model_archs = [None if od_config is None else getattr(od_config, "model_class_name", None)]
for stage_config in self.stage_configs or ():
- stage_get = (
- stage_config.get if isinstance(stage_config, Mapping) else lambda key: getattr(stage_config, key, None)
- )
- engine_args = stage_get("engine_args") or {}
- model_archs.extend(
- (
- stage_get("model_arch"),
- engine_args.get("model_class_name")
- if isinstance(engine_args, Mapping)
- else getattr(engine_args, "model_class_name", None),
- )
- )
+ model_archs.append(_stage_diffusion_model_class_name(stage_config))
supported: set[str] = set()
for model_arch in model_archs:
diff --git a/vllm_omni/entrypoints/openai/tts_adapters/audex.py b/vllm_omni/entrypoints/openai/tts_adapters/audex.py
index 0e31b633f..3b9e4a9a5 100644
--- a/vllm_omni/entrypoints/openai/tts_adapters/audex.py
+++ b/vllm_omni/entrypoints/openai/tts_adapters/audex.py
@@ -42,7 +42,11 @@ class AudexAdapter(ARTTSAdapter):
from vllm_omni.model_executor.models.audex.checkpoint import ensure_audex_snapshot
- stage = getattr(getattr(self.ctx.server._tts_stage, "engine_args", None), "model_stage", None)
+ stage_config = self.ctx.server._tts_stage
+ topology = getattr(stage_config, "stage_pipeline_config", None)
+ stage = getattr(topology, "model_stage", None)
+ if stage is None:
+ stage = getattr(getattr(stage_config, "engine_args", None), "model_stage", None)
profile, folder = (
("full", "checkpoint_folder_full") if stage == "audex_omni" else ("tts", "checkpoint_folder_audiogen")
)
diff --git a/vllm_omni/entrypoints/openai/tts_adapters/base.py b/vllm_omni/entrypoints/openai/tts_adapters/base.py
index ff43b27a5..cb760e865 100644
--- a/vllm_omni/entrypoints/openai/tts_adapters/base.py
+++ b/vllm_omni/entrypoints/openai/tts_adapters/base.py
@@ -39,6 +39,21 @@ DEFAULT_TTS_LANGUAGES = frozenset(
)
+def resolve_stage_model_path(engine_client: Any) -> str | None:
+ """Prefer typed stage model overrides, then legacy and served-model values."""
+ stages = getattr(engine_client, "stage_configs", ()) or ()
+ for stage in stages:
+ model_path = getattr(getattr(stage, "model_config", None), "model", None)
+ if model_path:
+ return str(model_path)
+ for stage in stages:
+ model_path = getattr(getattr(stage, "engine_args", None), "model", None)
+ if model_path:
+ return str(model_path)
+ model_path = getattr(engine_client, "model", None)
+ return str(model_path) if model_path else None
+
+
def conditioning_cache_salt(request: "OpenAICreateSpeechRequest", tts_params: dict | None = None) -> str:
"""Stable hash of the real Stage 0 conditioning for the prefix cache.
diff --git a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py
index 207d8b47b..33f16cf0a 100644
--- a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py
+++ b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py
@@ -9,7 +9,12 @@ import numpy as np
from vllm.inputs import tokens_input
from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
-from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens
+from vllm_omni.entrypoints.openai.tts_adapters.base import (
+ ARTTSAdapter,
+ PreparedRequest,
+ apply_max_new_tokens,
+ resolve_stage_model_path,
+)
if TYPE_CHECKING:
from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
@@ -81,17 +86,7 @@ class HiggsAudioV2Adapter(ARTTSAdapter):
from transformers import AutoProcessor
- model_path = None
- for stage in self.engine_client.stage_configs:
- model_path = getattr(getattr(stage, "engine_args", None), "model", None)
- if model_path:
- break
- if model_path is None:
- # Fallback: the orchestrator stores the served model id on the engine
- # itself (set by AsyncOmniEngine.__init__). Stage-level engine_args
- # may not surface ``model`` when the deploy yaml doesn't set it per
- # stage (the CLI-passed model id is the single source of truth).
- model_path = getattr(self.engine_client, "model", None)
+ model_path = resolve_stage_model_path(self.engine_client)
if model_path is None:
raise RuntimeError("higgs_audio_v2 serving could not resolve the model path from the engine stage configs")
processor = AutoProcessor.from_pretrained(model_path)
diff --git a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py
index 1e64f7a63..d993a7395 100644
--- a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py
+++ b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py
@@ -16,6 +16,7 @@ from vllm_omni.entrypoints.openai.tts_adapters.base import (
PreparedRequest,
apply_max_new_tokens,
conditioning_cache_salt,
+ resolve_stage_model_path,
)
_REF_CODE_CACHE_MAX_ENTRIES = 256
@@ -170,13 +171,7 @@ class HiggsAudioV3Adapter(ARTTSAdapter):
HiggsAudioV3TokenizerAdapter,
)
- model_path = None
- for stage in self.engine_client.stage_configs:
- model_path = getattr(getattr(stage, "engine_args", None), "model", None)
- if model_path:
- break
- if model_path is None:
- model_path = getattr(self.engine_client, "model", None)
+ model_path = resolve_stage_model_path(self.engine_client)
if model_path is None:
raise RuntimeError("higgs_audio_v3 serving could not resolve model path")
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
diff --git a/vllm_omni/entrypoints/openai/video/generation/helpers.py b/vllm_omni/entrypoints/openai/video/generation/helpers.py
index 4473fce07..0ffd1ee83 100644
--- a/vllm_omni/entrypoints/openai/video/generation/helpers.py
+++ b/vllm_omni/entrypoints/openai/video/generation/helpers.py
@@ -58,6 +58,7 @@ from vllm_omni.entrypoints.openai.serving_video import (
ReferenceAudio,
ReferenceImage,
ReferenceVideo,
+ _stage_diffusion_model_class_name,
)
from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER
from vllm_omni.entrypoints.openai.stores import VIDEO_STORE
@@ -136,10 +137,6 @@ def _config_get(config: Any, key: str, default: Any = None) -> Any:
return getattr(config, key, default)
-def _stage_engine_args(stage_cfg: Any) -> Any:
- return _config_get(stage_cfg, "engine_args", {}) or {}
-
-
def _diffusion_model_classes(stage_configs: list[Any] | None) -> list[type]:
if not stage_configs:
return []
@@ -150,7 +147,7 @@ def _diffusion_model_classes(stage_configs: list[Any] | None) -> list[type]:
for stage_cfg in stage_configs:
if get_stage_type(stage_cfg) != "diffusion":
continue
- model_class_name = _config_get(_stage_engine_args(stage_cfg), "model_class_name")
+ model_class_name = _stage_diffusion_model_class_name(stage_cfg)
if not model_class_name:
continue
model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
diff --git a/vllm_omni/entrypoints/openpi/serving.py b/vllm_omni/entrypoints/openpi/serving.py
index 562ed70d4..11fde2bf8 100644
--- a/vllm_omni/entrypoints/openpi/serving.py
+++ b/vllm_omni/entrypoints/openpi/serving.py
@@ -103,6 +103,14 @@ class ServingRealtimeRobotOpenPI:
for stage_config in getattr(engine_client, "stage_configs", []) or []:
if getattr(stage_config, "stage_type", None) != "diffusion":
continue
+ # Typed diffusion stages keep model-owned OpenPI handshake
+ # metadata in diffusion_config.model_config. The out-of-process
+ # head only has this stage view because full od_config lives in
+ # the worker.
+ diffusion_config = getattr(stage_config, "diffusion_config", None)
+ model_config = getattr(diffusion_config, "model_config", None)
+ if model_config is not None:
+ break
engine_args = getattr(stage_config, "engine_args", None)
model_config = getattr(engine_args, "model_config", None)
if model_config is not None:
diff --git a/vllm_omni/entrypoints/serve/profile/utils.py b/vllm_omni/entrypoints/serve/profile/utils.py
index f444a3534..d91ef2cad 100644
--- a/vllm_omni/entrypoints/serve/profile/utils.py
+++ b/vllm_omni/entrypoints/serve/profile/utils.py
@@ -4,18 +4,20 @@
def _should_enable_profiler_endpoints(stage_configs: list | None) -> bool:
- """Check if any stage has profiler_config set in its engine_args."""
+ """Check whether any typed or legacy stage enables a profiler."""
if not stage_configs:
return False
for stage in stage_configs:
- engine_args = stage.get("engine_args") if isinstance(stage, dict) else getattr(stage, "engine_args", None)
- if engine_args is None:
- continue
profiler_config = (
- engine_args.get("profiler_config")
- if isinstance(engine_args, dict)
- else getattr(engine_args, "profiler_config", None)
+ stage.get("profiler_config") if isinstance(stage, dict) else getattr(stage, "profiler_config", None)
)
+ if profiler_config is None:
+ engine_args = stage.get("engine_args") if isinstance(stage, dict) else getattr(stage, "engine_args", None)
+ profiler_config = (
+ engine_args.get("profiler_config")
+ if isinstance(engine_args, dict)
+ else getattr(engine_args, "profiler_config", None)
+ )
if profiler_config is not None:
profiler = (
profiler_config.get("profiler")
diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py
index 281cbc55c..8c2ba0747 100644
--- a/vllm_omni/entrypoints/stage_utils.py
+++ b/vllm_omni/entrypoints/stage_utils.py
@@ -1,7 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
from __future__ import annotations
import logging
import os
+from dataclasses import asdict, is_dataclass
from multiprocessing import shared_memory as _shm
from typing import Any
@@ -274,6 +278,8 @@ def _to_dict(x: Any) -> dict[str, Any]:
try:
if isinstance(x, dict):
return dict(x)
+ if is_dataclass(x) and not isinstance(x, type):
+ return asdict(x)
return _omega_to_dict(x)
except Exception:
try:
diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py
index 39131fd58..c091ae5e8 100644
--- a/vllm_omni/entrypoints/utils.py
+++ b/vllm_omni/entrypoints/utils.py
@@ -4,12 +4,14 @@
import json
import os
import types
-from dataclasses import fields, is_dataclass
+from collections.abc import Mapping
+from dataclasses import dataclass, fields, is_dataclass
from typing import Any, get_args, get_origin
from vllm.logger import init_logger
from vllm.sampling_params import RequestOutputKind, SamplingParams
+from vllm_omni.config.config_factory import with_trust_remote_code_override
from vllm_omni.entrypoints.stage_utils import _to_dict
from vllm_omni.inputs.data import OmniSamplingParams
@@ -18,6 +20,13 @@ logger = init_logger(__name__)
def inject_omni_kv_config(stage: Any, omni_conn_cfg: dict[str, Any], omni_from: str, omni_to: str) -> None:
"""Inject connector configuration into stage engine arguments."""
+ typed_connector_config = getattr(stage, "connector_config", None)
+ if typed_connector_config is not None:
+ omni_conf_dict = dict(typed_connector_config.omni_kv_config or {})
+ omni_conf_dict.update(connector_config=omni_conn_cfg, omni_from_stage=omni_from, omni_to_stage=omni_to)
+ typed_connector_config.omni_kv_config = omni_conf_dict
+ return
+
# Prepare omni_kv_config dict
omni_conf_dict = {}
try:
@@ -79,6 +88,64 @@ def parse_stage_overrides(value: Any) -> dict[str, dict[str, Any]] | None:
return parsed
+@dataclass(frozen=True)
+class StageConfigInputs:
+ """Normalized inputs shared by all stage-config loading paths."""
+
+ model: str
+ kwargs: dict[str, Any]
+ trust_remote_code: bool | None
+ deploy_config_path: str | None
+ stage_overrides: dict[str, dict[str, Any]] | None
+ strategy_config_path: str | None
+
+
+def prepare_stage_config_inputs(
+ model: str,
+ kwargs: Mapping[str, Any],
+ *,
+ trust_remote_code: bool | None,
+ snapshot_model: bool = False,
+) -> StageConfigInputs:
+ """Normalize model/config arguments before resolving stage configs.
+
+ The standard API worker, multi-API parent, and headless entrypoint must
+ apply the same legacy-key filtering, trust-remote-code precedence, and
+ config-file extraction before calling :func:`resolve_omni_config`.
+ """
+ resolved_model = model
+ if snapshot_model:
+ # Keep model materialization at the CLI boundary for the parent-owned
+ # launch path. The API worker's OmniBase has already done this work.
+ from vllm_omni.entrypoints.omni_base import omni_snapshot_download
+
+ resolved_model = omni_snapshot_download(model)
+
+ resolved_kwargs = dict(kwargs)
+ for legacy_arg in ("stage_configs_path", "stage_configs"):
+ if legacy_arg in resolved_kwargs:
+ raise ValueError(f"`{legacy_arg}` is no longer supported; use `deploy_config` instead.")
+ # Frontend process count is not a stage engine override.
+ resolved_kwargs.pop("api_server_count", None)
+ resolved_kwargs.pop("disable_log_stats", None)
+ if resolved_kwargs.get("diffusion_streaming_output") and resolved_kwargs.get("streaming_output") is None:
+ resolved_kwargs["streaming_output"] = True
+ resolved_kwargs.pop("model_tag", None)
+
+ resolved_kwargs = with_trust_remote_code_override(resolved_kwargs, trust_remote_code)
+ deploy_config_path = resolved_kwargs.pop("deploy_config", None)
+ strategy_config_path = resolved_kwargs.pop("strategy_config", None)
+ stage_overrides = parse_stage_overrides(resolved_kwargs.pop("stage_overrides", None))
+ return StageConfigInputs(
+ model=resolved_model,
+ kwargs=resolved_kwargs,
+ trust_remote_code=trust_remote_code,
+ deploy_config_path=deploy_config_path,
+ stage_overrides=stage_overrides,
+ strategy_config_path=strategy_config_path,
+ )
+
+
def get_final_stage_id_for_e2e(
output_modalities: list[str] | None, default_modalities: list[str], stage_list: list
) -> int:
diff --git a/vllm_omni/model_executor/models/minimax_h3/checkpoint.py b/vllm_omni/model_executor/models/minimax_h3/checkpoint.py
index 0143e9675..b39af5cf8 100644
--- a/vllm_omni/model_executor/models/minimax_h3/checkpoint.py
+++ b/vllm_omni/model_executor/models/minimax_h3/checkpoint.py
@@ -1,10 +1,16 @@
# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""MiniMax H3 checkpoint resolution for the AR text-encoder stage."""
+
from pathlib import Path
+from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index
from vllm_omni.model_executor.model_loader.weight_utils import (
download_weights_from_hf_specific,
)
+_FASTH3_MODULAR_MODEL_PREFIX = "FastVideo/"
+
MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS = {
partition: [
f"{subdir}/text_encoder/**",
@@ -16,6 +22,22 @@ MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS = {
MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS["combined"] = MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS["fl2va"]
+def is_minimax_h3_modular(model: str, revision: str | None = None) -> bool:
+ path = Path(model)
+ if path.is_dir():
+ # Native MiniMax-H3 snapshots can advertise the Diffusers modular
+ # class in ``model_index.json`` while retaining the legacy FL2VA /
+ # Ref2VA layout. FastH3 modular releases carry this marker file.
+ return (path / "modular_model_index.json").is_file() or (path / "fastvideo_inference.json").is_file()
+ # The FastVideo release is the modular checkpoint handled by this
+ # resolver. Do not classify the native MiniMax-H3 Hub repository from
+ # its class name alone; its partition files are still required here.
+ if not model.startswith(_FASTH3_MODULAR_MODEL_PREFIX):
+ return False
+ index = get_diffusion_model_index(model, revision=revision) or {}
+ return index.get("_class_name") == "MiniMaxH3ModularPipeline"
+
+
def resolve_minimax_h3_partition(
model: str,
task_type: str | None,
@@ -42,10 +64,13 @@ def resolve_minimax_h3_encoder_model_root(
) -> str:
path = Path(model)
partition = resolve_minimax_h3_partition(model, task_type, auto_partition="fl2va")
+ modular = is_minimax_h3_modular(model, revision)
if path.is_dir():
if path.name == "text_encoder" and (path / "config.json").is_file():
return str(path)
+ if modular:
+ return str(path / "text_encoder")
if path.name in {"FL2VA", "Ref2VA"}:
path = path.parent
subdir = "Ref2VA" if partition == "ref2va" else "FL2VA"
@@ -53,11 +78,13 @@ def resolve_minimax_h3_encoder_model_root(
snapshot = download_weights_from_hf_specific(
model_name_or_path=model,
cache_dir=None,
- allow_patterns=MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS[partition],
+ allow_patterns=["text_encoder/**"] if modular else MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS[partition],
revision=revision,
require_all=True,
)
+ if modular:
+ return str(Path(snapshot) / "text_encoder")
subdir = "Ref2VA" if partition == "ref2va" else "FL2VA"
return str(Path(snapshot) / subdir / "text_encoder")
diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
index 7bc6c21c6..f78b777a5 100644
--- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
+++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
@@ -177,6 +177,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
num_scheduled_tokens_np=np.ones(bsz, dtype=np.int32),
max_num_scheduled_tokens=1,
use_cascade_attn=False,
+ force_uniform_decode=True,
)
n = batch_desc.num_tokens
ids = self.talker_mtp_input_ids.gpu[:n]
diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py
index 8f15aaf50..c0975f651 100644
--- a/vllm_omni/quantization/factory.py
+++ b/vllm_omni/quantization/factory.py
@@ -446,6 +446,33 @@ def _disk_marks_serialized(qc_kwargs: dict[str, Any], quant_config: object) -> b
return False
+_MXFP4_SERIALIZED_FLAGS = {
+ "mxfp4": "is_checkpoint_mxfp4_serialized",
+ "mxfp4_dualscale": "is_checkpoint_serialized",
+}
+
+
+def _normalize_serialized_mxfp4_layer_policy(qc_method: str, qc_kwargs: dict[str, Any]) -> None:
+ """Canonicalize the BF16 layer policy owned by serialized MXFP4 metadata.
+
+ A serialized checkpoint config is a complete storage contract: an omitted
+ layer list means no explicit BF16 layers. Method-only online metadata does
+ not carry that contract and must leave caller-provided layer policy intact.
+ """
+ method = _normalize_quant_method_alias(qc_method)
+ if method is None:
+ return
+ serialized_flag = _MXFP4_SERIALIZED_FLAGS.get(method)
+ if serialized_flag is None or not qc_kwargs.get(serialized_flag):
+ return
+
+ ignored_layers = qc_kwargs.get("ignored_layers")
+ if not ignored_layers:
+ ignored_layers = qc_kwargs.get("modules_to_not_convert")
+ qc_kwargs["ignored_layers"] = ignored_layers or []
+ qc_kwargs.pop("modules_to_not_convert", None)
+
+
def resolve_quant_config_from_disk(
quant_config: QuantizationConfig | None,
disk_qc: dict[str, Any] | str | None,
@@ -460,22 +487,24 @@ def resolve_quant_config_from_disk(
- quant_config is None: auto-detect from disk_qc (full build).
- Methods mismatch: raise ValueError — prevents silent weight corruption.
- Disk marks serialized but quant_config is online: rebuild from disk.
- - ignored_layers differ: rebuild from disk (per-transformer BF16 routing).
+ - Serialized MXFP4 metadata owns the complete BF16 layer policy; omitted
+ ignored_layers means the checkpoint default (no explicit BF16 layers).
+ - Method-only metadata carries no storage policy and preserves caller
+ ignored_layers and runtime fallback policies.
"""
if disk_qc is None:
return quant_config
if isinstance(disk_qc, str):
- if quant_config is None:
- logger.info("Auto-detected quantization from config.json: method=%s", disk_qc)
- return build_quant_config(disk_qc)
- return quant_config
+ qc_method = disk_qc
+ qc_kwargs: dict[str, Any] = {}
+ else:
+ if not isinstance(disk_qc, Mapping) or "quant_method" not in disk_qc:
+ return quant_config
+ qc_method = disk_qc["quant_method"]
+ qc_kwargs = {k: v for k, v in disk_qc.items() if k != "quant_method"}
- if not isinstance(disk_qc, Mapping) or "quant_method" not in disk_qc:
- return quant_config
-
- qc_method: str = disk_qc["quant_method"]
- qc_kwargs: dict[str, Any] = {k: v for k, v in disk_qc.items() if k != "quant_method"}
+ _normalize_serialized_mxfp4_layer_policy(qc_method, qc_kwargs)
if quant_config is None:
logger.info(
@@ -494,6 +523,26 @@ def resolve_quant_config_from_disk(
"Pass a matching --quantization flag or omit it for auto-detection."
)
+ if isinstance(disk_qc, str):
+ return quant_config
+
+ # Serialized checkpoint metadata owns storage and BF16 layer routing. The
+ # active config owns runtime step/layer policies, including explicit empty
+ # lists. Preserve them when each cascade transformer rebuilds from disk.
+ for policy in ("w4a8_fallback_steps", "w4a8_fallback_layers"):
+ if hasattr(quant_config, policy):
+ qc_kwargs[policy] = list(getattr(quant_config, policy))
+ if hasattr(quant_config, "mxfp4_scale_alg"):
+ # Runtime activation policy is independent of how offline W4 was made.
+ qc_kwargs["mxfp4_scale_alg"] = quant_config.mxfp4_scale_alg
+ if hasattr(quant_config, "require_smooth_scale"):
+ disk_requires_smooth = qc_kwargs.get("require_smooth_scale", False)
+ if not isinstance(disk_requires_smooth, bool):
+ raise ValueError("require_smooth_scale must be a boolean")
+ # Either a calibrated checkpoint or an explicit caller can require
+ # Smooth. Per-expert rebuilds must never weaken that requirement.
+ qc_kwargs["require_smooth_scale"] = quant_config.require_smooth_scale or disk_requires_smooth
+
if _disk_marks_serialized(qc_kwargs, quant_config):
logger.info(
"config.json marks checkpoint as serialized; switching to offline %s mode.",
@@ -501,6 +550,14 @@ def resolve_quant_config_from_disk(
)
return build_quant_config(qc_method, **qc_kwargs)
+ if (
+ "require_smooth_scale" in qc_kwargs
+ and hasattr(quant_config, "require_smooth_scale")
+ and qc_kwargs["require_smooth_scale"] != quant_config.require_smooth_scale
+ ):
+ logger.info("config.json Smooth requirement differs from active config; rebuilding quant_config.")
+ return build_quant_config(qc_method, **qc_kwargs)
+
# AutoRound MXFP checkpoints use data_type="mx_fp" instead of
# is_checkpoint_*_serialized; rebuild so the offline MXFP4/MXFP8 path is
# selected according to the checkpoint's bit width.
diff --git a/vllm_omni/quantization/mxfp4_config.py b/vllm_omni/quantization/mxfp4_config.py
index ac93eaa00..43f3efa9d 100644
--- a/vllm_omni/quantization/mxfp4_config.py
+++ b/vllm_omni/quantization/mxfp4_config.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""W4A4 MXFP4 (Microscaling FP4) online/offline quantization for diffusion transformers.
Architecture mirrors mxfp8_config.py:
@@ -41,7 +41,7 @@ Dual-scale (W4A4_MXFP4_DUALSCALE):
and npu_dual_level_quant_matmul. Weight stored in NZ hardware format via npu_format_cast.
Checkpoint tensor shapes for dual-scale:
- weight : (N, K) float8_e4m3fn – FP4 packed
+ weight : (N, K) float8_e4m3fn – numeric FP4 values (packed after loading)
weight_scale : (N, K//32) uint8 – fine scale (float8_e8m0fnu bits)
weight_dual_scale : (N, K//512, 1) float32 – coarse scale (extra dim avoids shape assert)
mul_scale : (K,) float32 – per-input-channel activation pre-scale
@@ -51,9 +51,9 @@ Reference: MindIE-SD W4A4MXFP4QuantLinear / W4A4MXFP4DualQuantLinear (mindiesd/q
from __future__ import annotations
-import re
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
+import regex as re
import torch
from torch.nn import Module
from vllm.logger import init_logger
@@ -83,9 +83,54 @@ from vllm_omni.quantization.mxfp8_config import (
if TYPE_CHECKING:
from vllm.model_executor.models.utils import WeightsMapper
+
logger = init_logger(__name__)
+def _validate_w4a8_fallback_steps(steps: list[int] | None) -> list[int]:
+ if steps is None:
+ return []
+ if not isinstance(steps, list) or any(type(step) is not int or step < 0 for step in steps):
+ raise ValueError("w4a8_fallback_steps must be a list of non-negative integer denoise step indices.")
+ return sorted(set(steps))
+
+
+def _validate_w4a8_fallback_layers(layers: list[str] | None) -> list[str]:
+ """Names are exact runtime Linear paths, not checkpoint aliases or patterns."""
+ if layers is None:
+ return []
+ if not isinstance(layers, list) or any(
+ not isinstance(name, str) or re.fullmatch(r"[A-Za-z_][A-Za-z_0-9]*(?:\.[A-Za-z_0-9]+)*", name) is None
+ for name in layers
+ ):
+ raise ValueError("w4a8_fallback_layers must be a list of non-empty exact runtime Linear paths.")
+ return sorted(set(layers))
+
+
+def _is_w4a8_fallback_step(steps: list[int]) -> bool:
+ """Read the request's zero-based denoise index; never keep state on a Linear."""
+ if not steps:
+ return False
+
+ from vllm_omni.diffusion.forward_context import (
+ get_forward_context,
+ is_forward_context_available,
+ )
+
+ step = get_forward_context().denoise_step_idx if is_forward_context_available() else None
+ if step is None:
+ raise RuntimeError("W4A8 step fallback requires a pipeline that publishes its denoise step in ForwardContext.")
+ return step in steps
+
+
+def _validate_smooth_scale(scale: torch.Tensor) -> None:
+ """Reject malformed calibration before loading or activation dtype conversion."""
+ if scale.ndim != 1 or not scale.is_floating_point():
+ raise ValueError("MXFP4 mul_scale must be a one-dimensional floating-point Smooth tensor.")
+ if not bool(torch.isfinite(scale).all()) or not bool((scale > 0).all()):
+ raise ValueError("MXFP4 mul_scale must contain finite, strictly positive Smooth values.")
+
+
# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
@@ -105,10 +150,24 @@ class DiffusionMXFP4Config(QuantizationConfig):
self,
is_checkpoint_mxfp4_serialized: bool = False,
ignored_layers: list[str] | None = None,
+ w4a8_fallback_steps: list[int] | None = None,
+ require_smooth_scale: bool = False,
+ w4a8_fallback_layers: list[str] | None = None,
+ mxfp4_scale_alg: int = 0,
) -> None:
super().__init__()
self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized
self.ignored_layers = ignored_layers or []
+ self.w4a8_fallback_steps = _validate_w4a8_fallback_steps(w4a8_fallback_steps)
+ self.w4a8_fallback_layers = _validate_w4a8_fallback_layers(w4a8_fallback_layers)
+ if not isinstance(require_smooth_scale, bool):
+ raise ValueError("require_smooth_scale must be a boolean")
+ if require_smooth_scale and not self.is_checkpoint_mxfp4_serialized:
+ raise ValueError("require_smooth_scale requires an offline MXFP4 checkpoint")
+ self.require_smooth_scale = require_smooth_scale
+ if type(mxfp4_scale_alg) is not int or mxfp4_scale_alg not in (0, 2):
+ raise ValueError("mxfp4_scale_alg must be 0 (OCP MX) or 2 (UOS with dst_type_max=7.25).")
+ self.mxfp4_scale_alg = mxfp4_scale_alg
@classmethod
def get_name(cls) -> QuantizationMethods:
@@ -139,6 +198,10 @@ class DiffusionMXFP4Config(QuantizationConfig):
return cls(
is_checkpoint_mxfp4_serialized=is_serialized,
ignored_layers=ignored_layers,
+ w4a8_fallback_steps=config.get("w4a8_fallback_steps"),
+ w4a8_fallback_layers=config.get("w4a8_fallback_layers"),
+ require_smooth_scale=config.get("require_smooth_scale", False),
+ mxfp4_scale_alg=config.get("mxfp4_scale_alg", 0),
)
def get_quant_method(
@@ -155,9 +218,15 @@ class DiffusionMXFP4Config(QuantizationConfig):
return UnquantizedLinearMethod()
if current_omni_platform.is_npu():
if self.is_checkpoint_mxfp4_serialized:
- return NPUMxfp4LinearMethod(self)
- return NPUMxfp4OnlineLinearMethod(self)
+ return NPUMxfp4LinearMethod(self, prefix=prefix)
+ return NPUMxfp4OnlineLinearMethod(self, prefix=prefix)
if current_omni_platform.is_rocm():
+ if self.mxfp4_scale_alg != 0:
+ raise NotImplementedError("MXFP4 UOS quantization is currently only supported on NPU (Ascend).")
+ if self.w4a8_fallback_steps or self.w4a8_fallback_layers:
+ raise NotImplementedError(
+ "MXFP4 W4A8 step/layer fallback is currently only supported on NPU (Ascend)."
+ )
gcn_arch = torch.cuda.get_device_properties(torch.accelerator.current_device_index()).gcnArchName
if "gfx950" not in gcn_arch:
raise NotImplementedError(f"MXFP4 on ROCm requires gfx950 (MI355X). Detected: {gcn_arch}")
@@ -176,6 +245,56 @@ class DiffusionMXFP4Config(QuantizationConfig):
# ---------------------------------------------------------------------------
+def _npu_quantize_mxfp4(x: torch.Tensor, scale_alg: int) -> tuple[torch.Tensor, torch.Tensor]:
+ """Quantize weight or activation using the same explicit FP4 algorithm."""
+ import torch_npu
+
+ # Universal Optimal Scaling (Qmax=7.25), MXAttention Section 4.1:
+ # https://arxiv.org/abs/2607.24377
+ scale_kwargs = {"dst_type_max": 7.25} if scale_alg == 2 else {}
+ return torch_npu.npu_dynamic_mx_quant(
+ x,
+ dst_type=torch_npu.float4_e2m1fn_x2,
+ axis=-1,
+ block_size=32,
+ round_mode="rint",
+ scale_alg=scale_alg,
+ **scale_kwargs,
+ )
+
+
+def _npu_w4a8_matmul(
+ x: torch.Tensor,
+ weight: torch.Tensor,
+ weight_scale: torch.Tensor,
+ bias: torch.Tensor | None,
+ output_dtype: torch.dtype,
+) -> torch.Tensor:
+ """MXFP8 activations with prepared single-scale MXFP4 weights (N, K)."""
+ import torch_npu
+
+ x_q, x_scale = torch_npu.npu_dynamic_mx_quant(
+ x, dst_type=torch_npu.float8_e4m3fn, axis=-1, block_size=32, round_mode="rint", scale_alg=0
+ )
+ # Use the BF16 output/bias pair supported by the mixed FP8/FP4 kernel
+ # in the target A5 runtime. Restore the caller's dtype afterwards.
+ if bias is not None:
+ bias = bias.to(torch.bfloat16).reshape(1, -1)
+ output = torch_npu.npu_quant_matmul(
+ x_q,
+ weight.transpose(0, 1),
+ weight_scale.transpose(0, 1),
+ scale_dtype=torch_npu.float8_e8m0fnu,
+ x2_dtype=torch_npu.float4_e2m1fn_x2,
+ pertoken_scale=x_scale,
+ pertoken_scale_dtype=torch_npu.float8_e8m0fnu,
+ bias=bias,
+ output_dtype=torch.bfloat16,
+ group_sizes=[0, 0, 32],
+ )
+ return output.to(output_dtype)
+
+
class NPUMxfp4LinearMethod(MXFPLinearMethodBase):
"""NPU W4A4 MXFP4 offline linear method for pre-quantized checkpoints.
@@ -187,7 +306,8 @@ class NPUMxfp4LinearMethod(MXFPLinearMethodBase):
NPUMxfp4OnlineLinearMethod normalizes to the same layout so apply() is shared.
"""
- def __init__(self, quant_config: DiffusionMXFP4Config) -> None:
+ def __init__(self, quant_config: DiffusionMXFP4Config, *, prefix: str = "") -> None:
+ self.is_w4a8_fallback_layer = prefix in quant_config.w4a8_fallback_layers
self.quant_config = quant_config
self.out_dtype = torch.get_default_dtype()
@@ -210,11 +330,14 @@ class NPUMxfp4LinearMethod(MXFPLinearMethodBase):
layer.orig_dtype = params_dtype
layer.weight_block_size = None
- # BF16 placeholder; cast to float4_e2m1fn_x2 in process_weights.
layer.register_parameter(
"weight",
ModelWeightParameter(
- data=torch.empty(output_size_per_partition, input_size_per_partition, dtype=params_dtype),
+ data=torch.empty(
+ output_size_per_partition,
+ input_size_per_partition,
+ dtype=params_dtype,
+ ),
input_dim=1,
output_dim=0,
weight_loader=weight_loader,
@@ -234,12 +357,42 @@ class NPUMxfp4LinearMethod(MXFPLinearMethodBase):
),
)
+ # Older single-scale checkpoints have no Smooth tensor. Keep their
+ # identity transform while allowing calibrated checkpoints to load it.
+ def smooth_weight_loader(
+ param: ModelWeightParameter, loaded_weight: torch.Tensor, *args: Any, **kwargs: Any
+ ) -> None:
+ _validate_smooth_scale(loaded_weight)
+ if not callable(weight_loader):
+ raise ValueError("Loading MXFP4 Smooth requires a callable weight_loader")
+ weight_loader(param, loaded_weight, *args, **kwargs)
+
+ layer.register_parameter(
+ "mul_scale",
+ ModelWeightParameter(
+ data=torch.ones(input_size_per_partition, dtype=torch.float32),
+ input_dim=0,
+ output_dim=None,
+ weight_loader=smooth_weight_loader,
+ ),
+ )
+ setattr(layer.mul_scale, "ignore_warning", True)
+ setattr(layer.mul_scale, "is_checkpoint_optional", not self.quant_config.require_smooth_scale)
+ if self.quant_config.is_checkpoint_mxfp4_serialized:
+ setattr(layer.weight_scale, "is_checkpoint_required", True)
+
def process_weights_after_loading(self, layer: Module) -> None:
if getattr(layer, "_already_called_process_weights_after_loading", False):
return
import torch_npu
+ smooth_param = cast(torch.Tensor, layer.mul_scale)
+ _validate_smooth_scale(smooth_param)
+ smooth = smooth_param.data.to(dtype=cast(torch.dtype, layer.orig_dtype)).contiguous()
+ # Finite FP32 values can overflow or underflow in the activation dtype.
+ _validate_smooth_scale(smooth)
+
# NPU: cast to float4_e2m1fn_x2. Weight stays (N, K) — no pre-transpose.
w = layer.weight
if w.dtype != torch_npu.float4_e2m1fn_x2:
@@ -258,15 +411,31 @@ class NPUMxfp4LinearMethod(MXFPLinearMethodBase):
replace_parameter(layer, "weight", w)
replace_parameter(layer, "weight_scale", s)
+ # Keep loader metadata while converting the calibration tensor once.
+ layer.mul_scale.data = smooth.to(device=cast(torch.Tensor, w).device)
layer._already_called_process_weights_after_loading = True
# --- NPU MXFP4 ops — shared with online path via inheritance ---
- def _quantize_activation(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
- import torch_npu
+ def _apply_inner(
+ self,
+ layer: torch.nn.Module,
+ x: torch.Tensor,
+ bias: torch.Tensor | None,
+ ori_dtype: torch.dtype,
+ ) -> torch.Tensor:
+ # Online quantization has no calibrated Smooth tensor. For offline
+ # weights the same pre-scale must precede both A4 and A8 quantization.
+ if (mul_scale := getattr(layer, "mul_scale", None)) is not None:
+ x = x * mul_scale.to(x.dtype)
+ if self.is_w4a8_fallback_layer or _is_w4a8_fallback_step(self.quant_config.w4a8_fallback_steps):
+ return _npu_w4a8_matmul(
+ x, cast(torch.Tensor, layer.weight), cast(torch.Tensor, layer.weight_scale), bias, ori_dtype
+ )
+ return super()._apply_inner(layer, x, bias, ori_dtype)
- # No dst_type: npu_dynamic_mx_quant defaults to float4_e2m1fn_x2.
- return torch_npu.npu_dynamic_mx_quant(x)
+ def _quantize_activation(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ return _npu_quantize_mxfp4(x, self.quant_config.mxfp4_scale_alg)
def _quant_matmul(
self,
@@ -318,8 +487,6 @@ class NPUMxfp4OnlineLinearMethod(_LazyWeightMixin, NPUMxfp4LinearMethod):
if getattr(layer, "_already_called_process_weights_after_loading", False):
return
- import torch_npu
-
if layer.weight.device == torch.device("meta"):
weight = ModelWeightParameter(
data=torch.empty_like(layer.weight, device=layer._load_device),
@@ -331,8 +498,10 @@ class NPUMxfp4OnlineLinearMethod(_LazyWeightMixin, NPUMxfp4LinearMethod):
layer.register_parameter("weight", weight)
initialize_single_dummy_weight(layer.weight)
- # NPU: quantize BF16/FP16 (N, K) → FP4. No dst_type → float4_e2m1fn_x2.
- weight_fp4, weight_scale_raw = torch_npu.npu_dynamic_mx_quant(layer.weight)
+ # Prepare once; offline checkpoints keep their existing weight algorithm.
+ weight_fp4, weight_scale_raw = _npu_quantize_mxfp4(
+ cast(torch.Tensor, layer.weight), self.quant_config.mxfp4_scale_alg
+ )
# Weight stays (N, K) — no pre-transpose for FP4 packed format.
# Scale: (N, S) → (N, S/2, 2). Not pre-transposed; done inline in _quant_matmul.
@@ -497,7 +666,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase):
"""NPU W4A4 MXFP4 dual-scale offline method for pre-quantized checkpoints.
Checkpoint tensors and their canonical post-load shapes:
- weight : (N, K) float8_e4m3fn – FP4 packed (2 values per byte)
+ weight : (N, K) float8_e4m3fn – numeric FP4 values; packed into NZ after loading
weight_scale : (N, K//32) uint8 – fine scale (float8_e8m0fnu bits); reshaped to (N, K//64, 2)
weight_dual_scale : (N, K//512, 1) float32 – coarse scale; transposed to (K//512, N)
mul_scale : (K,) float32 – per-input-channel activation pre-scale (from calibration)
@@ -509,7 +678,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase):
Reference: MindIE-SD W4A4MXFP4DualQuantLinear.
"""
- def __init__(self, quant_config: Any) -> None:
+ def __init__(self, quant_config: DiffusionMXFP4DualScaleMixedConfig) -> None:
self.quant_config = quant_config
self.out_dtype = torch.get_default_dtype()
@@ -532,7 +701,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase):
layer.orig_dtype = params_dtype
layer.weight_block_size = None
- # FP4 packed: 2 values per float8_e4m3fn byte → checkpoint stores as float8_e4m3fn; register as BF16
+ # Checkpoint FP4 values are stored numerically in FP8; BF16 preserves them exactly.
layer.register_parameter(
"weight",
ModelWeightParameter(
@@ -587,6 +756,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase):
),
)
setattr(layer.mul_scale, "ignore_warning", True)
+ setattr(layer.weight_scale, "is_checkpoint_required", True)
def process_weights_after_loading(self, layer: Module) -> None:
if getattr(layer, "_already_called_process_weights_after_loading", False):
@@ -594,7 +764,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase):
import torch_npu
- # float8_e4m3fn (FP4 packed) → float4_e2m1fn_x2 → NZ hardware format (format ID 29).
+ # Numeric FP4 values → packed float4_e2m1fn_x2 → NZ hardware format (format ID 29).
# NZ layout is required by npu_dual_level_quant_matmul for FP4 weight matrices.
w = torch_npu.npu_dtype_cast(layer.weight.data.npu(), torch_npu.float4_e2m1fn_x2)
w = torch_npu.npu_format_cast(w.view(torch.int8), 29, customize_dtype=torch.int8)
@@ -631,7 +801,10 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase):
"""
if ori_dtype not in (torch.bfloat16, torch.float16):
x = x.to(torch.bfloat16)
- x_q, l0_scale, l1_scale = self._quantize_activation(x, layer.mul_scale)
+ # The dual-level activation kernel requires x and Smooth to have the
+ # same dtype. Prepared calibration scales are BF16, including online.
+ smooth_scale = cast(torch.Tensor, layer.mul_scale).to(x.dtype)
+ x_q, l0_scale, l1_scale = self._quantize_activation(x, smooth_scale)
return self._quant_matmul(x_q, l0_scale, l1_scale, layer, bias, ori_dtype)
def _quantize_activation( # type: ignore[override]
@@ -785,11 +958,20 @@ class DiffusionMXFP4DualScaleMixedConfig(QuantizationConfig):
is_checkpoint_serialized: bool = False,
ignored_layers: list[str] | None = None,
num_bf16_fallback_layers: int = 5,
+ w4a8_fallback_steps: list[int] | None = None,
+ w4a8_fallback_layers: list[str] | None = None,
) -> None:
super().__init__()
self.is_checkpoint_serialized = is_checkpoint_serialized
self.ignored_layers = ignored_layers or []
self.num_bf16_fallback_layers = num_bf16_fallback_layers
+ self.w4a8_fallback_steps = _validate_w4a8_fallback_steps(w4a8_fallback_steps)
+ self.w4a8_fallback_layers = _validate_w4a8_fallback_layers(w4a8_fallback_layers)
+ if self.w4a8_fallback_steps or self.w4a8_fallback_layers:
+ raise ValueError(
+ "mxfp4_dualscale does not support W4A8 fallback. Use mxfp4 with single-scale weights; "
+ "changing a DualScale checkpoint's method name does not convert its storage format."
+ )
@classmethod
def get_name(cls) -> QuantizationMethods:
@@ -822,6 +1004,8 @@ class DiffusionMXFP4DualScaleMixedConfig(QuantizationConfig):
is_checkpoint_serialized=is_serialized,
ignored_layers=ignored_layers,
num_bf16_fallback_layers=num_bf16_fallback_layers,
+ w4a8_fallback_steps=config.get("w4a8_fallback_steps"),
+ w4a8_fallback_layers=config.get("w4a8_fallback_layers"),
)
def get_quant_method(
diff --git a/vllm_omni/utils/tracking_parser.py b/vllm_omni/utils/tracking_parser.py
index 64849a5aa..8e5f5832a 100644
--- a/vllm_omni/utils/tracking_parser.py
+++ b/vllm_omni/utils/tracking_parser.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import argparse
from typing import Any
@@ -52,12 +52,28 @@ class TrackingNamespace(argparse.Namespace):
else:
setattr(self.unfiltered_ns, name, value)
+ def __getstate__(self) -> dict[str, Any]:
+ """Preserve wrapper state when API workers use ``spawn``.
+
+ Our __dict__ property exposes the inner namespace, so default state
+ serialization would lose unfiltered_ns and explicit_keys.
+ """
+ return {
+ "unfiltered_ns": self.unfiltered_ns,
+ "explicit_keys": self.explicit_keys,
+ }
+
+ def __setstate__(self, state: dict[str, Any]) -> None:
+ object.__setattr__(self, "unfiltered_ns", state["unfiltered_ns"])
+ object.__setattr__(self, "explicit_keys", state["explicit_keys"])
+
def get_explicit_kwargs_dict(self):
"""Return a dict containing only the explicitly passed key-value pairs."""
return {k: v for k, v in vars(self.unfiltered_ns).items() if k in self.explicit_keys}
def __getattr__(self, name: str) -> Any:
- return getattr(self.unfiltered_ns, name)
+ namespace = object.__getattribute__(self, "unfiltered_ns")
+ return getattr(namespace, name)
@property
def __dict__(self):
diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py
index e43c6f608..4d80d3f8d 100644
--- a/vllm_omni/worker/gpu_generation_model_runner.py
+++ b/vllm_omni/worker/gpu_generation_model_runner.py
@@ -812,7 +812,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
if self.uses_mrope:
positions = self.mrope_positions.gpu[:, :num_tokens_padded]
- elif self.uses_xdrope_dim > 0:
+ elif getattr(self, "uses_xdrope_dim", 0) > 0:
positions = self.xdrope_positions.gpu[:, :num_tokens_padded]
else:
positions = self.positions[:num_tokens_padded]
diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py
index d3ee8de02..37d2ac122 100644
--- a/vllm_omni/worker/gpu_model_runner.py
+++ b/vllm_omni/worker/gpu_model_runner.py
@@ -690,7 +690,7 @@ class OmniGPUModelRunner(GPUModelRunner):
self._init_mrope_positions(req_state)
# Only relevant for models using XD-RoPE (e.g, HunYuan-VL)
- if self.uses_xdrope_dim > 0:
+ if getattr(self, "uses_xdrope_dim", 0) > 0:
self._init_xdrope_positions(req_state)
reqs_to_add.append(self.requests[req_id])
@@ -1164,7 +1164,7 @@ class OmniGPUModelRunner(GPUModelRunner):
if self.uses_mrope:
positions = self.mrope_positions.gpu[:, :num_tokens_padded]
- elif self.uses_xdrope_dim > 0:
+ elif getattr(self, "uses_xdrope_dim", 0) > 0:
positions = self.xdrope_positions.gpu[:, :num_tokens_padded]
else:
positions = self.positions[:num_tokens_padded]
@@ -1701,7 +1701,7 @@ class OmniGPUModelRunner(GPUModelRunner):
if self.uses_mrope:
positions = self.mrope_positions.gpu[:, :num_input_tokens]
- elif self.uses_xdrope_dim > 0:
+ elif getattr(self, "uses_xdrope_dim", 0) > 0:
positions = self.xdrope_positions.gpu[:, :num_input_tokens]
else:
positions = self.positions[:num_input_tokens]