Compare commits

...
Author SHA1 Message Date
Hongsheng Liu 00dab2ac95 docs: clarify quantization support matrix
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 15:50:56 +08:00
vraitiandClaude Opus 4.6 c98517ce7f [Metrics] Add Diffusion Metrics (#4755)
Signed-off-by: vraiti <vraiti@redhat.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-08-13 06:59:13 +00:00
NATUREandAlicia d77a35a32b [Bugfix] Scope Bagel FP8 config to diffusion stage (#6085)
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-08-13 11:23:47 +08:00
liwenerandamy-why-3459 877beeaa49 [Perf][MiniCPM-O-4.5]Add CUDA Graph For CFM DiT estimator (#6082)
Signed-off-by: liwener <liwener.c@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-13 09:36:17 +08:00
Hongsheng Liu 8b8d3f56db [Misc] Consolidate LingBot text-to-video runner (#6076)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 08:42:37 +08:00
Hongsheng Liu c57178c0e1 [Dependency] Upgrade Cache-DiT to 1.5.0 (#6065)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 08:41:59 +08:00
mjZhaoElaineandClaude Opus 5 9a59475015 [CI/Build] Fix ReadTheDocs build under RTD's seeded pip 23.1 (#6129)
Signed-off-by: Mengjie Zhao <zmj0129@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
2026-08-13 06:52:01 +08:00
Nick Cao 12f8168179 [Tests] Enable tiny model testing for Qwen-Image Edit and EditPlus (#5656)
Signed-off-by: Nick Cao <ncao@redhat.com>
2026-08-12 16:03:49 -04:00
Zhenghan Song 72a02b492a [Bugfix] Fix diffusion TTS adapter lookup in speech serving (#6121)
Signed-off-by: HaningZS <resossr0v0@gmail.com>
2026-08-12 15:20:12 -04:00
Juan Pablo ZuluagaandYueqian Lin 5358f5db01 [Perf][TTS] omnivoice hot path: D2H batching, mask caching, cached text embeddings (#5174)
Signed-off-by: JuanPZuluaga <juanz9312@gmail.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-08-12 17:15:35 +00:00
MaciejBalaNV e677360168 Add cosmos-guardrail dependency and update error message (#6107)
Signed-off-by: Maciej Bala <mbala@nvidia.com>
2026-08-12 17:04:07 +00:00
Juan Pablo ZuluagaandYueqian Lin 208c6776a8 [Perf][TTS] glm_tts: hoist loop-invariant text embed, RoPE, mask and CFG batch out of Euler loop (#5068)
Signed-off-by: JuanPZuluaga <juanz9312@gmail.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-08-12 17:02:02 +00:00
54 changed files with 1604 additions and 487 deletions

No files matched your search

+5 -5
View File
@@ -345,7 +345,7 @@ steps:
commands: commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0 - export CACHE_DIT_VERSION=1.5.0
- | - |
set +e set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
@@ -361,7 +361,7 @@ steps:
commands: commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0 - export CACHE_DIT_VERSION=1.5.0
- | - |
set +e set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
@@ -377,7 +377,7 @@ steps:
commands: commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0 - export CACHE_DIT_VERSION=1.5.0
- | - |
set +e set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
@@ -393,7 +393,7 @@ steps:
commands: commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0 - export CACHE_DIT_VERSION=1.5.0
- | - |
set +e set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
@@ -409,7 +409,7 @@ steps:
commands: commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0 - export CACHE_DIT_VERSION=1.5.0
- | - |
set +e set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
+1 -1
View File
@@ -115,7 +115,7 @@ steps:
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
commands: commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export CACHE_DIT_VERSION=1.3.0 - export CACHE_DIT_VERSION=1.5.0
- | - |
set +e set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
+7 -2
View File
@@ -11,10 +11,15 @@ build:
post_checkout: post_checkout:
- git fetch --unshallow || true - git fetch --unshallow || true
pre_install: pre_install:
# `pre_install` runs before RTD upgrades pip, and the seeded pip 23.1
# rejects the `typing_extensions` wheel torch needs (fixed in pip 23.2).
- python -m pip install --upgrade pip
# Pre-install CPU torch so the subsequent `pip install .[docs]` does not # Pre-install CPU torch so the subsequent `pip install .[docs]` does not
# pull ~2 GB of CUDA 13 libs (cuBLAS/cuDNN/NCCL/...) that exceed RTD's # pull ~2 GB of CUDA 13 libs (cuBLAS/cuDNN/NCCL/...) that exceed RTD's
# 1 GB memory cap and time out the mkdocs build. # 1 GB memory cap and time out the mkdocs build. Wheel-only: `--index-url`
- python -m pip install --index-url https://download.pytorch.org/whl/cpu torch # replaces PyPI, so an sdist fallback cannot reach its build deps. The
# quotes keep YAML from reading `:all:` as a mapping key.
- "python -m pip install --only-binary=:all: --index-url https://download.pytorch.org/whl/cpu torch"
mkdocs: mkdocs:
configuration: mkdocs.yml configuration: mkdocs.yml
@@ -167,9 +167,11 @@ def main() -> None:
] ]
native_cmd = [ native_cmd = [
sys.executable, sys.executable,
"examples/offline_inference/text_to_video/text_to_video_lingbot.py", "examples/offline_inference/text_to_video/text_to_video.py",
"--model", "--model",
args.model, args.model,
"--model-class-name",
"LingBotVideoPipeline",
"--prompt", "--prompt",
args.prompt, args.prompt,
"--output", "--output",
+1 -1
View File
@@ -1,6 +1,6 @@
# Prometheus Metrics Design # Prometheus Metrics Design
This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage pipelines, the constraints that shaped the design, and how the pipeline-level metrics coexist with upstream vLLM per-engine metrics. This iteration covers the text and audio surface; diffusion / image / video families land in a follow-up. This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage pipelines, the constraints that shaped the design, and how the pipeline-level metrics coexist with upstream vLLM per-engine metrics.
## Objectives ## Objectives
+4 -1
View File
@@ -14,7 +14,10 @@ A: At first, you can check current [issues](https://github.com/vllm-project/vllm
> Q: Does vLLM-Omni support AWQ or any other quantization? > Q: Does vLLM-Omni support AWQ or any other quantization?
A: We plan to introduce GGUF FP8 prequantized models and online FP8 quantization in version 0.16.0. Support for other quantization types will follow in future releases. For details, please see our [Q1 quantization roadmap](https://github.com/vllm-project/vllm-omni/issues/1057). A: AWQ is not currently listed as a validated vLLM-Omni path. Quantization
support depends on the model, pipeline stage, checkpoint format, and hardware
backend. See the current [quantization support matrix](../user_guide/quantization/overview.md)
for online methods, pre-quantized checkpoints, and model-level validation.
> Q: Does vLLM-Omni support multimodal streaming input and output? > Q: Does vLLM-Omni support multimodal streaming input and output?
+50 -26
View File
@@ -1,23 +1,23 @@
# Production Metrics # Production Metrics
vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server. This page covers the text and audio surface; diffusion / image / video metrics are tracked in a follow-up PR. vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server.
```bash ```bash
vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8000 --log-stats vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8000 --log-stats
curl http://localhost:8000/metrics curl http://localhost:8000/metrics
``` ```
**`--log-stats` is required to populate metric data.** Without the flag, the endpoint still returns `200 OK` but the upstream `vllm:*` wrap is not registered at all, and the 15 `vllm:omni_*` families are registered as placeholders with no data written to them. This keeps the runtime cost essentially zero for deployments that don't need monitoring. With the flag, all ~80 families populate. **`--log-stats` is required to populate metric data.** Without the flag, the endpoint still returns `200 OK` but the upstream `vllm:*` wrap is not registered at all, and the 15 `vllm_omni:*` families are registered as placeholders with no data written to them. This keeps the runtime cost essentially zero for deployments that don't need monitoring. With the flag, all ~80 families populate.
## Metric Namespaces ## Metric Namespaces
| Prefix | Source | Present when | | Prefix | Source | Present when |
|--------|--------|--------------| |--------|--------|--------------|
| `vllm:omni_` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent | | `vllm_omni:` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent |
| `vllm:` | Upstream vLLM engine, wrapped by `OmniPrometheusStatLogger` to expose `{stage, replica}` | Pipeline includes an LLM (AR) stage | | `vllm:` | Upstream vLLM engine, wrapped by `OmniPrometheusStatLogger` to expose `{stage, replica}` | Pipeline includes an LLM (AR) stage |
| `http_` / `process_` | Uvicorn / Python runtime | Always | | `http_` / `process_` | Uvicorn / Python runtime | Always |
## Pipeline-Level Metrics (`vllm:omni_`) ## Pipeline-Level Metrics (`vllm_omni:`)
Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the full multi-stage pipeline. Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the full multi-stage pipeline.
@@ -25,42 +25,66 @@ Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the
| Metric | Type | Labels | Description | | Metric | Type | Labels | Description |
|--------|------|--------|-------------| |--------|------|--------|-------------|
| `vllm:omni_num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) | | `vllm_omni:num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) |
| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue | | `vllm_omni:num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue |
| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` | | `vllm_omni:requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
### Latency ### Latency
| Metric | Type | Labels | Description | | Metric | Type | Labels | Description |
|--------|------|--------|-------------| |--------|------|--------|-------------|
| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds | | `vllm_omni:e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds |
## Audio Modality Metrics (`vllm:omni_`) ## Audio Modality Metrics (`vllm_omni:`)
Emitted at request finalize, except for `audio_ttfp_s` (streaming-hook at the first audio packet) and `audio_underrun_s` / `audio_continuity_ok_total` (streaming finalize, after the chunk stream is exhausted). All carry `{model_name, stage, replica}` plus the listed extra label. Emitted at request finalize, except for `audio_ttfp_s` (streaming-hook at the first audio packet) and `audio_underrun_s` / `audio_continuity_ok_total` (streaming finalize, after the chunk stream is exhausted). All carry `{model_name, stage, replica}` plus the listed extra label.
| Metric | Type | Extra label | Description | | Metric | Type | Extra label | Description |
|--------|------|-------------|-------------| |--------|------|-------------|-------------|
| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame | | `vllm_omni:audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) | | `vllm_omni:audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` | | `vllm_omni:audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` |
| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` | | `vllm_omni:audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` |
| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps | | `vllm_omni:audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` | | `vllm_omni:audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio | | `vllm_omni:audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
The continuity math comes from `vllm_omni/benchmarks/audio_continuity.py::compute_continuity_stats` so the server-side observation aligns with the bench-side definition. The continuity math comes from `vllm_omni/benchmarks/audio_continuity.py::compute_continuity_stats` so the server-side observation aligns with the bench-side definition.
## Cross-Stage Transfer Metrics (`vllm:omni_`) ## Diffusion Metrics (`vllm_omni:`)
Per-request timing breakdowns for diffusion (image/video) stages. Emitted at request finalize when `stage_metrics.diffusion_metrics` is present. All carry `{model_name, stage, replica}`.
| Metric | Type | Description |
|--------|------|-------------|
| `vllm_omni:diffusion_exec_s` | Histogram | DiT forward pass execution time per request in seconds |
| `vllm_omni:diffusion_exec_per_step_s` | Histogram | DiT forward pass execution time per denoising step in seconds (`exec_s / num_inference_steps`) |
| `vllm_omni:diffusion_preprocess_s` | Histogram | Diffusion input preprocessing time per request in seconds |
| `vllm_omni:diffusion_postprocess_s` | Histogram | Diffusion output postprocessing (VAE decode) time per request in seconds |
`diffusion_exec_s` uses `SECONDS_BUCKETS` (50 ms → 300 s); the other three use `SECONDS_FAST_BUCKETS` (1 ms → 60 s). `diffusion_exec_per_step_s` is skipped when `num_inference_steps` is not set.
### Example PromQL
```promql
# P99 DiT exec time by stage
histogram_quantile(0.99, rate(vllm_omni:diffusion_exec_s_bucket[5m]))
# Ratio of preprocess to total exec wall time
rate(vllm_omni:diffusion_preprocess_s_sum[5m])
/ rate(vllm_omni:diffusion_exec_s_sum[5m])
```
## Cross-Stage Transfer Metrics (`vllm_omni:`)
Per-physical-transfer histograms tracking the data hop between adjacent stages. Labels `{model_name, from_stage, from_replica, to_stage, to_replica}` let dashboards attribute latency to specific replica edges. `from_replica` / `to_replica` are resolved from the orchestrator's sticky-routing binding (`stage_pool.get_bound_replica_id(request_id)`), so no extra plumbing through `TransferEdgeStats` is needed. Per-physical-transfer histograms tracking the data hop between adjacent stages. Labels `{model_name, from_stage, from_replica, to_stage, to_replica}` let dashboards attribute latency to specific replica edges. `from_replica` / `to_replica` are resolved from the orchestrator's sticky-routing binding (`stage_pool.get_bound_replica_id(request_id)`), so no extra plumbing through `TransferEdgeStats` is needed.
| Metric | Type | Description | | Metric | Type | Description |
|--------|------|-------------| |--------|------|-------------|
| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes | | `vllm_omni:transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) | | `vllm_omni:transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) | | `vllm_omni:transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) | | `vllm_omni:transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
## vLLM Engine Metrics (`vllm:`) ## vLLM Engine Metrics (`vllm:`)
@@ -83,9 +107,9 @@ For the full list of upstream metrics, see [the vLLM docs](https://github.com/vl
| Metric group | Multi-stage LLM (Qwen3-Omni) | | Metric group | Multi-stage LLM (Qwen3-Omni) |
|---|---| |---|---|
| `vllm:omni_` request tracking + latency | With `--log-stats` | | `vllm_omni:` request tracking + latency | With `--log-stats` |
| `vllm:omni_` audio modality | With `--log-stats`, if pipeline has a talker stage | | `vllm_omni:` audio modality | With `--log-stats`, if pipeline has a talker stage |
| `vllm:omni_` transfer | With `--log-stats`, if pipeline has ≥ 2 stages | | `vllm_omni:` transfer | With `--log-stats`, if pipeline has ≥ 2 stages |
| `vllm:` engine metrics (per `(stage, replica)`) | With `--log-stats` | | `vllm:` engine metrics (per `(stage, replica)`) | With `--log-stats` |
| `vllm:` MFU metrics | With `--log-stats --enable-mfu-metrics` | | `vllm:` MFU metrics | With `--log-stats --enable-mfu-metrics` |
@@ -94,5 +118,5 @@ For the full list of upstream metrics, see [the vLLM docs](https://github.com/vl
- All time-bearing metrics use the `_s` suffix (values in seconds). Buckets are `SECONDS_BUCKETS` for e2e / generation-style values and `SECONDS_FAST_BUCKETS` (1 ms → 60 s) for the fine-grained transfer and audio-underrun values. - All time-bearing metrics use the `_s` suffix (values in seconds). Buckets are `SECONDS_BUCKETS` for e2e / generation-style values and `SECONDS_FAST_BUCKETS` (1 ms → 60 s) for the fine-grained transfer and audio-underrun values.
- Counters use the `_total` suffix (auto-appended by `prometheus_client`). - Counters use the `_total` suffix (auto-appended by `prometheus_client`).
- Sizes use the `_bytes` suffix. - Sizes use the `_bytes` suffix.
- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched (see `vllm_omni/patch.py`) to a scoped version that still strips upstream `vllm:*` collectors so multi-engine init within one process does not crash on duplicate registration, but preserves anything prefixed `vllm:omni_` / `vllm_omni`. - All omni-specific families are prefixed `vllm_omni:`. The upstream `unregister_vllm_metrics()` function is monkey-patched (see `vllm_omni/patch.py`) to a scoped version that still strips upstream `vllm:*` collectors so multi-engine init within one process does not crash on duplicate registration, but preserves anything prefixed `vllm_omni:`.
- Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm:omni_audio_ttfp_s` for audio) rather than a single metric with a `modality` label. - Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm_omni:audio_ttfp_s` for audio) rather than a single metric with a `modality` label.
+1 -1
View File
@@ -92,7 +92,7 @@ current support.
|--------|--------------|-------------|----------| |--------|--------------|-------------|----------|
| **[FP8](quantization/fp8.md)** | `quantization="fp8"` | FP8 W8A8 on validated transformer stages | Memory reduction, inference speedup | | **[FP8](quantization/fp8.md)** | `quantization="fp8"` | FP8 W8A8 on validated transformer stages | Memory reduction, inference speedup |
| **[INT8](quantization/int8.md)** | `quantization="int8"` | INT8 W8A8 on validated transformer stages | Memory reduction, broad GPU compatibility | | **[INT8](quantization/int8.md)** | `quantization="int8"` | INT8 W8A8 on validated transformer stages | Memory reduction, broad GPU compatibility |
| **[GGUF](quantization/gguf.md)** | `quantization="gguf"` | Native GGUF transformer-only weights (Q4, Q8, etc.) | Memory reduction on consumer GPUs | | **[BitsAndBytes W4](quantization/bitsandbytes.md)** | `quantization="bitsandbytes"` | Online NF4/FP4 weight-only quantization on validated CUDA transformer stages | Reducing transformer weight memory without preparing a separate checkpoint |
## Supported Models ## Supported Models
@@ -12,6 +12,7 @@ For backend selection and SageAttention usage, see the [Diffusion Attention Back
| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) | | Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
|---|---|---|---|---|---| |---|---|---|---|---|---|
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB | | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
| `robbyant/lingbot-video-dense-1.3b` / `robbyant/lingbot-video-moe-30b-a3b` | 192x320 | 9 | 2 | 3.0 | ~68 GiB (MoE smoke) |
| `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent | | `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
@@ -35,6 +36,25 @@ python text_to_video.py \
--output t2v_out.mp4 --output t2v_out.mp4
``` ```
### LingBot-Video
The shared runner recognizes both official LingBot checkpoint IDs and selects
the native `LingBotVideoPipeline` automatically:
```bash
python text_to_video.py \
--model robbyant/lingbot-video-dense-1.3b \
--prompt "a robotic arm picks up a red block" \
--height 192 --width 320 --num-frames 9 --num-inference-steps 2 \
--guidance-scale 3.0 --flow-shift 3.0 --fps 24 \
--output lingbot_t2v.mp4
```
Use `robbyant/lingbot-video-moe-30b-a3b` for the MoE checkpoint. For local
paths without `lingbot` in the directory name, add
`--model-class-name LingBotVideoPipeline`. Model-only options use
`--extra-body`, for example `'{"batch_cfg": true, "output_type": "np"}'`.
### LTX-2 ### LTX-2
```bash ```bash
@@ -123,7 +143,8 @@ python text_to_video.py \
- `--enable-layerwise-offload`: enable layerwise offloading on DiT modules. - `--enable-layerwise-offload`: enable layerwise offloading on DiT modules.
- `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2). - `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
- `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio. - `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio.
- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF). - `--quantization`: online quantization method, for example `fp8`; supported
pre-quantized checkpoints are normally auto-detected from checkpoint metadata.
- `--flow-shift`: scheduler flow_shift parameter. - `--flow-shift`: scheduler flow_shift parameter.
- `--lora-path`: path to PEFT LoRA adapter folder or checkpoint file. - `--lora-path`: path to PEFT LoRA adapter folder or checkpoint file.
- `--lora-scale`: scale factor for LoRA weights. - `--lora-scale`: scale factor for LoRA weights.
@@ -90,8 +90,8 @@ vllm serve "Qwen/Qwen-Image" \
--quantization-config '{"method": "fp8"}' --quantization-config '{"method": "fp8"}'
``` ```
Other vLLM-Omni quantization methods, such as `gguf`, `modelopt`, `mxfp4`, Other vLLM-Omni quantization methods, such as `modelopt`, `mxfp4`, `mxfp8`,
`mxfp8`, serialized checkpoints, static FP8 configs, and layer-name skip lists serialized checkpoints, static FP8 configs, and layer-name skip lists
such as `ignored_layers`, are intentionally not translated by this compatibility such as `ignored_layers`, are intentionally not translated by this compatibility
shortcut. Use Diffusers-native configuration through `diffusers_load_kwargs` or shortcut. Use Diffusers-native configuration through `diffusers_load_kwargs` or
a native vLLM-Omni pipeline for those cases. a native vLLM-Omni pipeline for those cases.
+44 -21
View File
@@ -10,6 +10,10 @@ checkpoint's `config.json` and auto-detects
AutoRound is static quantization: no `--quantization` flag is needed at AutoRound is static quantization: no `--quantization` flag is needed at
inference time when the checkpoint already contains the quantization config. inference time when the checkpoint already contains the quantization config.
The validation terms on this page follow the
[support levels](overview.md#support-levels) defined in the overview. A listed
checkpoint is not enough by itself to establish end-to-end model support.
## Hardware Support ## Hardware Support
| Device | Support | | Device | Support |
@@ -21,46 +25,57 @@ inference time when the checkpoint already contains the quantization config.
| Intel XPU | ✅ | | Intel XPU | ✅ |
| Ascend NPU | ❌ | | Ascend NPU | ❌ |
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
guide. AutoRound is Intel-supported. guide. This table is method-level backend availability. AutoRound is an Intel
project; vLLM selects a compatible CUDA or XPU compute backend at load time.
## Model Type Support ## Model Type Support
### Diffusion Model (Qwen-Image, Wan2.2) ### Diffusion Model (Qwen-Image, Wan2.2)
| Model | Checkpoint | Scope | Scheme | Backend | | Model | Checkpoint | Scope | Scheme | Validation |
|-------|------------|-------|--------|---------| |-------|------------|-------|--------|------------|
| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend | | FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | Diffusion transformer | W4A16 | Validated |
| Qwen-Image | Not listed | Diffusion transformer | W4A16 | Not validated | | Qwen-Image | `INC4AI/Qwen-Image-AutoRound-W4A16` | Diffusion transformer | W4A16 | Validated with end-to-end generation and quality comparison |
| Wan2.2-I2V | `Intel/Wan2.2-I2V-A14B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend | | Wan2.2-I2V | `Intel/Wan2.2-I2V-A14B-Diffusers-int4-AutoRound` | Both diffusion transformers | W4A16 | CI-backed |
| Wan2.2-T2V | `Intel/Wan2.2-T2V-A14B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend | | Wan2.2-T2V | `Intel/Wan2.2-T2V-A14B-Diffusers-int4-AutoRound` | Both diffusion transformers | W4A16 | CI-backed |
| Wan2.2-TI2V | `Intel/Wan2.2-TI2V-5B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend | | Wan2.2-TI2V | `Intel/Wan2.2-TI2V-5B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | Validated; not in the scheduled Wan AutoRound job |
CUDA execution uses GPTQ-Marlin or another compatible vLLM AutoRound backend;
Intel XPU uses the Intel-supported backend. The text encoder and VAE stay in
BF16 for all entries above.
### World Model (Cosmos3) ### World Model (Cosmos3)
| Model | Checkpoint | Scope | Scheme | Backend | | Model | Checkpoint | Scope | Scheme | Validation |
|-------|------------|-------|--------|---------| |-------|------------|-------|--------|------------|
| Cosmos3-Nano | `Intel/Cosmos3-Nano-int4-AutoRound` | World-model transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend | | Cosmos3-Nano | `Intel/Cosmos3-Nano-int4-AutoRound` | World-model transformer | W4A16 | Integrated; no model-specific end-to-end test is maintained in-tree |
| Cosmos3-Super | `Intel/Cosmos3-Super-int4-AutoRound` | World-model transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend | | Cosmos3-Super | `Intel/Cosmos3-Super-int4-AutoRound` | World-model transformer | W4A16 | Validated with a manual serving run |
The Cosmos VAE and guardrail components stay in BF16.
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS) ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
| Model | Checkpoint | Scope | Scheme | Backend | | Model | Checkpoint | Scope | Scheme | Validation |
|-------|------------|-------|--------|---------| |-------|------------|-------|--------|------------|
| Qwen2.5-Omni-7B | `Intel/Qwen2.5-Omni-7B-int4-AutoRound` | Language-model stage | W4A16 | AutoRound | | Qwen2.5-Omni-7B | `Intel/Qwen2.5-Omni-7B-int4-AutoRound` | Thinker language-model stage | W4A16 | Validated outside scheduled CI |
| Qwen3-Omni-30B-A3B-Instruct | `Intel/Qwen3-Omni-30B-A3B-Instruct-int4-AutoRound` | Thinker language-model stage | W4A16 | AutoRound | | Qwen3-Omni-30B-A3B-Instruct | `Intel/Qwen3-Omni-30B-A3B-Instruct-int4-AutoRound` | Thinker language-model stage | W4A16 | Validated outside scheduled CI |
| Qwen3-TTS | Not listed | TTS language-model stage | W4A16 | Not validated | | Qwen3-TTS | Not listed | TTS language-model stage | W4A16 | Not validated |
AutoRound support is checkpoint-driven. A model is supported when its AutoRound support is checkpoint-driven. A model is supported when its
checkpoint uses a compatible INC/AutoRound config and the target stage maps to checkpoint uses a compatible INC/AutoRound config and the target stage maps to
vLLM-Omni's runtime module names. vLLM-Omni's runtime module names.
For Qwen-Omni, only the thinker language model uses the AutoRound checkpoint.
Audio and vision encoders, the talker, and waveform-decoder stages stay in
BF16.
### Multi-Stage Diffusion Model (BAGEL, GLM-Image) ### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
| Model | Scope | Status | Notes | | Model | Checkpoint | Scope | Validation |
|-------|-------|--------|-------| |-------|------------|-------|------------|
| GLM-Image | Diffusion transformer | ✅ | `Intel/GLM-Image-int4-AutoRound` | | GLM-Image | `Intel/GLM-Image-int4-AutoRound` | Diffusion transformer | CI-backed |
| BAGEL | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint | | BAGEL | Not listed | Checkpoint-defined diffusion or transformer stage | Not validated |
## Configuration ## Configuration
@@ -119,6 +134,14 @@ The checkpoint should contain a config like:
At load time, vLLM-Omni builds an `OmniINCConfig`, remaps checkpoint block names At load time, vLLM-Omni builds an `OmniINCConfig`, remaps checkpoint block names
to runtime module names, and selects the matching vLLM compute backend. to runtime module names, and selects the matching vLLM compute backend.
Checkpoint auto-detection and successful layer construction establish the
**Integrated** level. Promote a model to **Validated** only after the named
checkpoint completes an end-to-end generation or multimodal request and its
output passes a model-appropriate sanity or quality check. Compare against the
BF16 baseline for quality-sensitive image and video paths. Promote it to
**CI-backed** only when that named checkpoint is selected by scheduled hardware
CI.
Example checkpoint creation: Example checkpoint creation:
```bash ```bash
+14 -6
View File
@@ -9,6 +9,10 @@ stay in BF16/FP16. No pre-quantized checkpoint is required.
This is an online (dynamic) path: load the normal HuggingFace checkpoint and This is an online (dynamic) path: load the normal HuggingFace checkpoint and
quantize during model loading via the `bitsandbytes` CUDA kernels. quantize during model loading via the `bitsandbytes` CUDA kernels.
The validation terms on this page follow the
[support levels](overview.md#support-levels). BitsAndBytes kernel support does
not imply that every diffusion transformer has been quality-validated.
## Hardware Support ## Hardware Support
| Device | Support | | Device | Support |
@@ -21,7 +25,7 @@ quantize during model loading via the `bitsandbytes` CUDA kernels.
| Intel XPU | ❌ | | Intel XPU | ❌ |
| Ascend NPU | ❌ | | Ascend NPU | ❌ |
Legend: `✅` supported, `❌` unsupported. Non-CUDA platforms raise in Legend: `✅` backend available, `❌` unsupported. Non-CUDA platforms raise in
`get_quant_method()`; CUDA requires compute capability 7.5+ (`SM 75+`). `get_quant_method()`; CUDA requires compute capability 7.5+ (`SM 75+`).
Requires the optional `bitsandbytes` package (`pip install bitsandbytes`). Requires the optional `bitsandbytes` package (`pip install bitsandbytes`).
@@ -30,11 +34,11 @@ Requires the optional `bitsandbytes` package (`pip install bitsandbytes`).
### Diffusion Model (Qwen-Image, Wan2.2) ### Diffusion Model (Qwen-Image, Wan2.2)
| Model | HF models | CUDA | Mode | Recommendation | | Model | HF models | Scope | Level | Recommendation |
|-------|-----------|:----:|------|----------------| |-------|-----------|-------|-------|----------------|
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Online W4 weight-only | All heavy linear layers; sensitive embedders stay BF16 | | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Diffusion transformer | Validated | All heavy linear layers; sensitive embedders stay BF16 |
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Not validated | Online W4 weight-only | Compare vs BF16 before enabling | | Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Diffusion transformer | Not validated | Compare with BF16 before enabling |
| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Online W4 weight-only | Validate before enabling in docs | | Wan2.2 | Wan2.2 diffusion pipelines | One or both diffusion transformers | Not validated | Validate quality and cascade behavior before enabling |
Other diffusion models may work if their transformer uses supported linear Other diffusion models may work if their transformer uses supported linear
layers, but they are not validated in this guide. layers, but they are not validated in this guide.
@@ -104,6 +108,10 @@ Multi-GPU tensor parallelism (`tensor_parallel_size` > 1) is not validated for
BitsAndBytes in diffusion models; each rank quantizes its own weight shard BitsAndBytes in diffusion models; each rank quantizes its own weight shard
independently. independently.
Z-Image-Turbo is validated outside scheduled full-model BitsAndBytes CI. The
real CUDA kernel smoke tests cover quantization and matrix multiplication, but
do not establish model-level quality for unlisted architectures.
If quality regresses, use `ignored_layers` to keep sensitive projections in BF16 If quality regresses, use `ignored_layers` to keep sensitive projections in BF16
(for example `to_out` or `w2`). (for example `to_out` or `w2`).
+21 -10
View File
@@ -13,6 +13,11 @@ components, such as the tokenizer, scheduler, text encoder, vision/audio
encoder, and VAE, on the base checkpoint unless a model-specific guide says encoder, and VAE, on the base checkpoint unless a model-specific guide says
otherwise. otherwise.
The validation terms on this page follow the
[support levels](overview.md#support-levels) defined in the overview. ModelOpt
format detection and checkpoint adaptation are shared runtime integrations;
each model and stage still requires its own end-to-end validation.
!!! note !!! note
ModelOpt checkpoints are pre-quantized checkpoints. Do not pass ModelOpt checkpoints are pre-quantized checkpoints. Do not pass
`--quantization fp8` for these checkpoints. The checkpoint `--quantization fp8` for these checkpoints. The checkpoint
@@ -58,11 +63,11 @@ stay unquantized unless the model-specific recipe validates otherwise.
| Intel XPU | ⭕ | | Intel XPU | ⭕ |
| Ascend NPU | ❌ | | Ascend NPU | ❌ |
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide. Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
The optional CUTLASS FP8 runtime override requires CUDA SM89+. ModelOpt NVFP4 guide. The optional CUTLASS FP8 runtime override requires CUDA SM89+. ModelOpt
and mixed FP8/NVFP4 diffusion checkpoints are currently validated on Blackwell NVFP4 and mixed FP8/NVFP4 diffusion checkpoints are currently validated on
CUDA systems in the recipes below; other CUDA generations require separate Blackwell CUDA systems in the recipes below; other CUDA generations require
backend and quality validation. separate backend and quality validation.
## Model Type Support ## Model Type Support
@@ -79,6 +84,9 @@ backend and quality validation.
| HunyuanImage-3.0 | `feizhai123/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8` | MoE diffusion transformer | Validated for ModelOpt mixed FP8/NVFP4 checkpoints | | HunyuanImage-3.0 | `feizhai123/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8` | MoE diffusion transformer | Validated for ModelOpt mixed FP8/NVFP4 checkpoints |
| Wan2.2 | Not available | Diffusion transformer | Not validated | | Wan2.2 | Not available | Diffusion transformer | Not validated |
The validated diffusion entries above have named checkpoint and recipe
evidence. They are not all exercised by scheduled full-model CI.
For full serving commands and benchmark context, see For full serving commands and benchmark context, see
[`recipes/Qwen/Qwen-Image.md`](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image.md) [`recipes/Qwen/Qwen-Image.md`](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image.md)
and and
@@ -88,8 +96,8 @@ and
| Model | Scope | Status | | Model | Scope | Status |
|-------|-------|--------| |-------|-------|--------|
| Qwen3-Omni | Thinker language-model stage | ModelOpt FP8 checkpoint path | | Qwen3-Omni | Thinker language-model stage (FP8) | Validated outside scheduled CI |
| Qwen3-Omni | Thinker language-model stage (W4A4 NVFP4) | Validated; see [Qwen3-Omni NVFP4 W4A4](#qwen3-omni-nvfp4-w4a4-thinker) below | | Qwen3-Omni | Thinker language-model stage (W4A4 NVFP4) | CI-backed for the H100 Marlin fallback; native Blackwell path is recipe-validated |
| Qwen3-TTS | TTS language-model stage | Not validated | | Qwen3-TTS | TTS language-model stage | Not validated |
Audio encoder, vision encoder, talker, and code2wav stages stay in BF16 unless Audio encoder, vision encoder, talker, and code2wav stages stay in BF16 unless
@@ -102,9 +110,9 @@ Qwen3-Omni-30B-A3B-Instruct thinker language model. The thinker text body
(attention + MoE experts) is quantized to NVFP4 with FP8 per-tensor input (attention + MoE experts) is quantized to NVFP4 with FP8 per-tensor input
scales; the audio encoder, vision encoder, talker, and code2wav stay in BF16. scales; the audio encoder, vision encoder, talker, and code2wav stay in BF16.
| Variant | HF checkpoint | Hardware | | Variant | HF checkpoint | Hardware and validation |
|---------|---------------|----------| |---------|---------------|-------------------------|
| W4A4 NVFP4 (full thinker) | `YihongJin/Qwen3-Omni-30B-A3B-Instruct-NVFP4-W4A4-full-thinker-awqclip` | sm_100+ (Blackwell, FlashInfer FP4 GEMM) | | W4A4 NVFP4 (full thinker) | `YihongJin/Qwen3-Omni-30B-A3B-Instruct-NVFP4-W4A4-full-thinker-awqclip` | sm_100+ native FlashInfer FP4 GEMM; H100 Marlin fallback is CI-backed |
Calibration uses ModelOpt `mtq.NVFP4_DEFAULT_CFG` with the `awq_clip` algorithm Calibration uses ModelOpt `mtq.NVFP4_DEFAULT_CFG` with the `awq_clip` algorithm
on 1024 ultrachat samples chat-templated through the Qwen3-Omni tokenizer. on 1024 ultrachat samples chat-templated through the Qwen3-Omni tokenizer.
@@ -232,3 +240,6 @@ omni = Omni(
5. Keep `--quantization fp8` for online FP8 from BF16 checkpoints; use this 5. Keep `--quantization fp8` for online FP8 from BF16 checkpoints; use this
ModelOpt path only when the checkpoint already contains ModelOpt quantized ModelOpt path only when the checkpoint already contains ModelOpt quantized
weights and scales. weights and scales.
6. The Qwen3-Omni full-model CI smoke exercises the H100 Marlin fallback. A
scheduled Blackwell test is still required before treating the native
FlashInfer FP4 kernel path as CI-backed.
+31 -11
View File
@@ -3,12 +3,17 @@
## Overview ## Overview
[msModelSlim](https://github.com/Ascend/msmodelslim) is an Ascend compression [msModelSlim](https://github.com/Ascend/msmodelslim) is an Ascend compression
toolkit for producing pre-quantized model checkpoints. In vLLM-Omni, these toolkit for producing pre-quantized model checkpoints. vLLM-Omni consumes its
checkpoints run through the Ascend/NPU path with `--quantization ascend`. output either through a compatible generic Ascend checkpoint path or through
the native MXFP8 and MXFP4 merge workflows documented in this guide.
msModelSlim is static quantization: quantized weights are generated offline msModelSlim is static quantization: quantized weights are generated offline
before vLLM-Omni inference starts. before vLLM-Omni inference starts.
The validation terms on this page follow the
[support levels](overview.md#support-levels). A producer recipe or hardware
backend alone is an integration point, not model-level validation.
## Hardware Support ## Hardware Support
| Device | Support | | Device | Support |
@@ -20,23 +25,28 @@ before vLLM-Omni inference starts.
| Intel XPU | ❌ | | Intel XPU | ❌ |
| Ascend NPU | ✅ | | Ascend NPU | ✅ |
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
guide. guide.
## Model Type Support ## Model Type Support
### Diffusion Model (Qwen-Image, Wan2.2) ### Diffusion Model (Qwen-Image, Wan2.2)
| Model | Base model | Scope | Hardware | Notes | | Model | Base model | Scope | Level | Notes |
|-------|------------|-------|----------|-------| |-------|------------|-------|-------|-------|
| Wan2.2 | Wan2.2 diffusion weights | DiT or diffusion stage | Ascend NPU | Upstream msModelSlim provides a Wan2.2 quantization recipe; vLLM-Omni inference validation is not listed | | Wan2.2 | Wan2.2 diffusion weights | DiT or diffusion stage | Integrated | Upstream msModelSlim provides a producer recipe; use the [MXFP8](mxfp8.md) or [MXFP4](mxfp4.md) workflow for vLLM-Omni's native merged formats |
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | DiT or diffusion stage | Ascend NPU | Not validated in this guide | | Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | DiT or diffusion stage | Not validated | No supported checkpoint workflow is documented |
| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | DiT or diffusion stage | Ascend A2/A3 NPU | Generate quantized weights with the HunyuanImage-3.0 msModelSlim adaptation | | HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | DiT or diffusion stage | Integrated | Generate quantized weights with the HunyuanImage-3.0 msModelSlim adaptation and validate against BF16 |
Public Hugging Face quantized weights are not available yet. Use the Public Hugging Face quantized weights are not available yet. Use the
[HunyuanImage-3.0 msModelSlim adaptation](https://gitcode.com/betta18/msmodelslim/tree/hyimage3_mxfp8) [HunyuanImage-3.0 msModelSlim adaptation](https://gitcode.com/betta18/msmodelslim/tree/hyimage3_mxfp8)
to generate the checkpoint manually. to generate the checkpoint manually.
msModelSlim is also the producer used by the native offline MXFP8 and MXFP4
DualScale workflows. Those workflows preprocess and annotate each transformer
checkpoint and are auto-detected at inference; follow their method pages rather
than passing `--quantization ascend`.
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS) ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
| Model | Scope | Status | Notes | | Model | Scope | Status | Notes |
@@ -53,25 +63,32 @@ to generate the checkpoint manually.
## Configuration ## Configuration
For a compatible generic Ascend checkpoint, select the Ascend quantization
method explicitly:
Offline inference: Offline inference:
```bash ```bash
python text_to_image.py --model <quantized-model-path> --quantization ascend python text_to_image.py --model <compatible-ascend-checkpoint> --quantization ascend
``` ```
Online serving: Online serving:
```bash ```bash
vllm serve <quantized-model-path> --omni --quantization ascend vllm serve <compatible-ascend-checkpoint> --omni --quantization ascend
``` ```
## Parameters ## Parameters
| Parameter | Type | Default | Description | | Parameter | Type | Default | Description |
|-----------|------|---------|-------------| |-----------|------|---------|-------------|
| `--quantization` | str | - | Use `ascend` for msModelSlim-produced checkpoints | | `--quantization` | str | - | Use `ascend` only for a compatible generic Ascend-format checkpoint; omit it for native merged MXFP8/MXFP4 checkpoints |
| `model` | str | - | Path to the quantized checkpoint generated by Ascend tooling | | `model` | str | - | Path to the quantized checkpoint generated by Ascend tooling |
Do not pass `--quantization ascend` for a Wan checkpoint processed by the native
MXFP8 or MXFP4 merge tools. Their transformer metadata selects `mxfp8` or
`mxfp4_dualscale` automatically.
Example msModelSlim command for a Wan2.2 W8A8 checkpoint: Example msModelSlim command for a Wan2.2 W8A8 checkpoint:
```bash ```bash
@@ -93,3 +110,6 @@ For HunyuanImage-3.0, use the Hunyuan-specific adaptation linked above.
tooling; it is not a load-time CUDA quantizer. tooling; it is not a load-time CUDA quantizer.
3. Keep the quantized checkpoint aligned with the same model architecture and 3. Keep the quantized checkpoint aligned with the same model architecture and
stage config used for BF16 inference. stage config used for BF16 inference.
4. No named full-model msModelSlim checkpoint is currently scheduled in NPU
CI. Treat new producer outputs as **Integrated** until an end-to-end BF16
comparison establishes **Validated** support.
+52 -23
View File
@@ -13,6 +13,10 @@ vLLM-Omni provides two quantization methods with different scale structures:
| `mxfp4` | Single-scale (per-32 fine only) | Online only | Quick accuracy baseline; no checkpoint prep needed | | `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_dualscale` | Dual-scale (fine per-32 + coarse per-512 + per-channel `mul_scale`) | Online + Offline | Production; better accuracy; offline recommended |
The validation terms on this page follow the
[support levels](overview.md#support-levels). Current model guidance covers
Wan diffusion transformers only; the text encoder and VAE stay in BF16.
!!! tip "Recommended: `mxfp4_dualscale` offline" !!! tip "Recommended: `mxfp4_dualscale` offline"
For production deployments, use the `mxfp4_dualscale` offline mode with a For production deployments, use the `mxfp4_dualscale` offline mode with a
pre-quantized checkpoint produced by msModelSlim. Offline checkpoints load pre-quantized checkpoint produced by msModelSlim. Offline checkpoints load
@@ -34,26 +38,45 @@ vLLM-Omni provides two quantization methods with different scale structures:
## Hardware Support ## Hardware Support
| Device | Support | | Device | Online `mxfp4` | Online `mxfp4_dualscale` | Offline `mxfp4_dualscale` |
|--------|---------| |--------|-----------------|---------------------------|----------------------------|
| NVIDIA Blackwell GPU (SM 100+) | ⭕ | | NVIDIA Blackwell GPU (SM 100+) | ❌ | ❌ | ❌ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ | | NVIDIA Ada/Hopper GPU (SM 89+) | ❌ | ❌ | ❌ |
| NVIDIA Ampere GPU (SM 80+) | ⭕ | | NVIDIA Ampere GPU (SM 80+) | ❌ | ❌ | ❌ |
| AMD ROCm (gfx950 / MI355X) | ✅ | | AMD ROCm (gfx950 / MI355X) | ✅ | ❌ | ❌ |
| Intel XPU | ⭕ | | Intel XPU | ❌ | ❌ | ❌ |
| Ascend NPU (Atlas 950 A5) | ✅ | | Ascend NPU (Atlas 950 A5) | ✅ | ✅ | ✅ |
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide. Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
guide. Backend availability does not imply full-model checkpoint validation.
!!! note "Why NVIDIA is marked unsupported"
This table describes the methods currently wired into vLLM-Omni, not every
low-precision format NVIDIA hardware could execute. `DiffusionMXFP4Config`
has runtime implementations only for Ascend and ROCm gfx950; on CUDA it
raises `NotImplementedError` because no CUDA MXFP4 linear method or
checkpoint loader is registered.
On NVIDIA, use [ModelOpt](modelopt.md) NVFP4 for the documented 4-bit
checkpoint path. NVFP4 and OCP MXFP4 have different checkpoint metadata,
scale layouts, and kernel contracts, so they are not interchangeable. A
native CUDA MXFP4 path would require a CUDA linear implementation, loader
contract, and model-level output validation.
The ROCm path dispatches online single-scale `mxfp4` to AITER on gfx950. It has
kernel and dispatch coverage, but this guide does not yet claim end-to-end
quality validation for a named Wan checkpoint on ROCm. DualScale and
pre-quantized MXFP4 checkpoints are not supported on ROCm.
## Model Type Support ## Model Type Support
### Diffusion Model (Wan2.2) ### Diffusion Model (Wan2.2)
| Model | Online | Offline | Notes | | Model | Online (Ascend) | Offline (Ascend) | Validation and notes |
|-------|--------|---------|-------| |-------|-----------------|------------------|----------------------|
| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade (`transformer` + `transformer_2`); both transformers quantized with the same config | | Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | Validated outside scheduled full-model offline NPU CI; both cascade transformers are quantized |
| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade; same scheme as T2V-A14B | | Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | Validated outside scheduled full-model offline NPU CI; same cascade scheme as T2V-A14B |
| Wan2.2-TI2V-5B | ❌ | ❌ | Parameter count too small; W4A4 causes unacceptable accuracy loss | | Wan2.2-TI2V-5B | ❌ | ❌ | Unsupported because W4A4 causes unacceptable accuracy loss |
The choice between `mxfp4` and `mxfp4_dualscale` in **online mode** is about The choice between `mxfp4` and `mxfp4_dualscale` in **online mode** is about
quantization quality, not model compatibility — both work on cascade (A14B) and quantization quality, not model compatibility — both work on cascade (A14B) and
@@ -83,12 +106,13 @@ by the merge script); there is no offline `mxfp4` single-scale format.
### `mxfp4` — Single-Scale Online Mode ### `mxfp4` — Single-Scale Online Mode
Online mode requires no pre-processing. vLLM-Omni quantizes BF16 weights to Online mode requires no pre-processing. On Ascend, vLLM-Omni quantizes BF16
MXFP4 at load time using `npu_dynamic_mx_quant`. A single block scale weights at load time using `npu_dynamic_mx_quant`. On ROCm gfx950, the same
`mxfp4` config dispatches to the AITER online kernel. A single block scale
(`float8_e8m0fnu`, one per 32 K elements) is computed on the fly; no (`float8_e8m0fnu`, one per 32 K elements) is computed on the fly; no
calibration `mul_scale` is available. Applies equally to single-transformer calibration `mul_scale` is available. The config applies equally to
and cascade (A14B) models — both transformers in a cascade receive the same single-transformer and cascade (A14B) models, but named Wan model validation is
quantization config automatically. currently limited to the documented Ascend path.
```python ```python
from vllm_omni import Omni from vllm_omni import Omni
@@ -111,10 +135,10 @@ vllm serve <your-model> --omni --quantization mxfp4
### `mxfp4_dualscale` — DualScale Online Mode ### `mxfp4_dualscale` — DualScale Online Mode
Online DualScale mode computes both fine and coarse scales on the fly from BF16 Ascend-only online DualScale mode computes both fine and coarse scales on the
weights using `npu_dynamic_dual_level_mx_quant`. Applies equally to fly from BF16 weights using `npu_dynamic_dual_level_mx_quant`. It applies
single-transformer and cascade (A14B) models. Compared to `mxfp4` online, equally to single-transformer and cascade (A14B) models. Compared to `mxfp4`
DualScale provides better quantization accuracy at higher compute cost. online, DualScale provides better quantization accuracy at higher compute cost.
The default configuration keeps the leading 5 transformer blocks in BF16 The default configuration keeps the leading 5 transformer blocks in BF16
(`num_bf16_fallback_layers=5`). Accuracy evaluation on Wan2.2-A14B shows this (`num_bf16_fallback_layers=5`). Accuracy evaluation on Wan2.2-A14B shows this
@@ -341,6 +365,11 @@ unless they appear in `ignored_layers`.
quantization. Use `ignored_layers` and `num_bf16_fallback_layers` to trade quantization. Use `ignored_layers` and `num_bf16_fallback_layers` to trade
off compression vs. accuracy for precision-sensitive layers. off compression vs. accuracy for precision-sensitive layers.
8. The native methods, merge mappings, and per-transformer cascade metadata
have unit coverage. The current Buildkite configuration does not schedule a
named full-model offline MXFP4 checkpoint on NPU, so this path is
**Validated**, not **CI-backed**.
## Adapting MXFP4 for a New Model ## Adapting MXFP4 for a New Model
This section is aimed at developers who want to add MXFP4 support to a model This section is aimed at developers who want to add MXFP4 support to a model
+30 -11
View File
@@ -15,30 +15,45 @@ This method supports three modes:
| **Offline (Native)** | msModelSlim-exported MXFP8 weights converted to diffusers format via `merge_mxfp8_checkpoint.py` — weights and scales are loaded directly from the preprocessed checkpoint | | **Offline (Native)** | msModelSlim-exported MXFP8 weights converted to diffusers format via `merge_mxfp8_checkpoint.py` — weights and scales are loaded directly from the preprocessed checkpoint |
| **Offline (AutoRound)** | AutoRound MXFP8 checkpoints with `data_type="mx_fp"` — auto-detected from `config.json` | | **Offline (AutoRound)** | AutoRound MXFP8 checkpoints with `data_type="mx_fp"` — auto-detected from `config.json` |
The validation terms on this page follow the
[support levels](overview.md#support-levels). Current model guidance covers
Wan diffusion transformers only; the text encoder and VAE stay in BF16.
## Hardware Support ## Hardware Support
| Device | Online | Offline (Native) | Offline (AutoRound) | | Device | Online | Offline (Native) | Offline (AutoRound) |
|--------|--------|------------------|---------------------| |--------|--------|------------------|---------------------|
| NVIDIA Blackwell GPU (SM 100+) | ⭕ | ⭕ | ⭕ | | NVIDIA Blackwell GPU (SM 100+) | ❌ | ❌ | ⭕ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ | ⭕ | ⭕ | | NVIDIA Ada/Hopper GPU (SM 89+) | ❌ | ❌ | ⭕ |
| NVIDIA Ampere GPU (SM 80+) | ⭕ | ⭕ | ⭕ | | NVIDIA Ampere GPU (SM 80+) | ❌ | ❌ | ⭕ |
| AMD ROCm | ⭕ | ⭕ | ⭕ | | AMD ROCm | ❌ | ❌ | ⭕ |
| Intel XPU | ✅ | ❌ | ✅ | | Intel XPU | ✅ | ❌ | ✅ |
| Ascend NPU (Atlas 950 A5) | ✅ | ✅ | ⭕ | | Ascend NPU (Atlas 950 A5) | ✅ | ✅ | ⭕ |
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide. Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
guide. Backend availability does not imply full-model checkpoint validation.
**Note**: Intel XPU only supports AutoRound MXFP8 for offline mode. Use AutoRound quantized checkpoints or online mode for XPU. The native `DiffusionMXFP8Config` currently dispatches only to Intel XPU and
Ascend NPU implementations. CUDA reaches `NotImplementedError`, so the NVIDIA
rows are marked unsupported even though this is an implementation gap rather
than a general claim about NVIDIA hardware capabilities. On NVIDIA, use
[FP8](fp8.md) for load-time quantization or [ModelOpt](modelopt.md) FP8 for
pre-quantized checkpoints. Adding native CUDA MXFP8 requires a registered CUDA
linear method, a compatible checkpoint contract, and model-level validation.
Native `mxfp8` is available online on Intel XPU and online or offline on
Ascend. Intel XPU only supports AutoRound MXFP8 for offline mode. The AutoRound
column is a separate checkpoint format and compute path.
## Model Type Support ## Model Type Support
### Diffusion Model (Wan2.2) ### Diffusion Model (Wan2.2)
| Model | Mode | Notes | | Model | Mode | Scope | Validation |
|-------|------|-------| |-------|------|-------|------------|
| Wan2.2-T2V-A14B | Online + Offline | MoE cascade; quantizes two transformers (`transformer` + `transformer_2`) | | Wan2.2-T2V-A14B | Online + Offline | MoE cascade: `transformer` + `transformer_2` | Validated outside scheduled full-model offline NPU CI |
| Wan2.2-I2V-A14B | Online + Offline | MoE cascade; quantizes two transformers (`transformer` + `transformer_2`) | | Wan2.2-I2V-A14B | Online + Offline | MoE cascade: `transformer` + `transformer_2` | Validated outside scheduled full-model offline NPU CI |
| Wan2.2-TI2V-5B | Online + Offline | Single transformer | | Wan2.2-TI2V-5B | Online + Offline | Single diffusion transformer | Validated outside scheduled full-model offline NPU CI |
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS) ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
@@ -218,3 +233,7 @@ omni = Omni(model="<autoround-mxfp8-model>")
`--output-dir`), regenerate it with the current script. The old script `--output-dir`), regenerate it with the current script. The old script
wrote a separate `quantization_config.json` that is not read by vLLM-Omni; wrote a separate `quantization_config.json` that is not read by vLLM-Omni;
the current script injects the config directly into `transformer/config.json`. the current script injects the config directly into `transformer/config.json`.
4. The native loader, merge mappings, and Wan cascade config propagation have
unit coverage. The current Buildkite configuration does not schedule a
named full-model offline MXFP8 checkpoint on NPU, so this path is
**Validated**, not **CI-backed**.
+27 -15
View File
@@ -6,7 +6,7 @@ Online quantization means vLLM-Omni computes quantized weights and scales while
loading the model. Use it when you want memory savings without preparing a loading the model. Use it when you want memory savings without preparing a
separate quantized checkpoint. separate quantized checkpoint.
This mode is different from pre-quantized checkpoint formats such as GGUF, This mode is different from pre-quantized checkpoint formats such as ModelOpt,
AutoRound, msModelSlim, or serialized Int8 checkpoints. Those formats are AutoRound, msModelSlim, or serialized Int8 checkpoints. Those formats are
prepared before serving and are documented in their method-specific guides. prepared before serving and are documented in their method-specific guides.
For MXFP8 and MXFP4, use this page for load-time quantization from BF16 For MXFP8 and MXFP4, use this page for load-time quantization from BF16
@@ -15,18 +15,22 @@ by msModelSlim and the merge tools.
## Hardware Support ## Hardware Support
| Device | FP8 W8A8 | Int8 W8A8 | MXFP8 W8A8 | MXFP4 W4A4 | | Device | FP8 W8A8 | Int8 W8A8 | BitsAndBytes W4 | MXFP8 W8A8 | MXFP4 W4A4 |
|--------|----------|-----------|------------|------------| |--------|----------|-----------|------------------|------------|------------|
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ⭕ | ⭕ | | NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ❌ | ❌ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ⭕ | ⭕ | | NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ❌ | ❌ |
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | | NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ✅ | ❌ | ❌ |
| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | | AMD ROCm (gfx950) | ⭕ | ⭕ | ❌ | ❌ | ✅ |
| Intel XPU | ⭕ | ⭕ | ✅ | ⭕ | | Intel XPU | ⭕ | ⭕ | ❌ | ✅ | ❌ |
| Ascend NPU | ❌ | ✅ | ✅ | ✅ | | Ascend NPU | ❌ | ✅ | ❌ | ✅ | ✅ |
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
guide. FP8 on Ampere may use a weight-only path where available. MXFP8 and guide. These entries describe method-level backend availability, not validation
MXFP4 are documented for the Ascend NPU path. for every model. FP8 on Ampere may use a weight-only path where available.
On ROCm gfx950, MXFP4 means online single-scale `mxfp4`; online and offline
`mxfp4_dualscale` are Ascend-only. The native MXFP8/MXFP4 configs do not
currently register CUDA implementations; use FP8 or ModelOpt FP8/NVFP4 on
NVIDIA instead.
## Model Type Support ## Model Type Support
@@ -36,8 +40,9 @@ MXFP4 are documented for the Ascend NPU path.
|--------|-------|----------------|--------| |--------|-------|----------------|--------|
| FP8 W8A8 | [FP8](fp8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models | | FP8 W8A8 | [FP8](fp8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
| Int8 W8A8 | [Int8](int8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image | | Int8 W8A8 | [Int8](int8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
| BitsAndBytes W4 | [BitsAndBytes](bitsandbytes.md) | Z-Image-Turbo; Qwen-Image and Wan2.2 are not validated | Validated for Z-Image-Turbo |
| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B, Wan2.2-TI2V-5B | Validated on Ascend NPU and Intel XPU | | MXFP8 W8A8 | [MXFP8](mxfp8.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B, Wan2.2-TI2V-5B | Validated on Ascend NPU and Intel XPU |
| MXFP4 W4A4 | [MXFP4](mxfp4.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B | Ascend NPU only; TI2V-5B is not supported | | MXFP4 W4A4 | [MXFP4](mxfp4.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B | Validated on Ascend NPU; the ROCm gfx950 backend is integrated but not model-validated; TI2V-5B is unsupported |
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS) ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
@@ -59,6 +64,7 @@ from vllm_omni import Omni
omni_fp8 = Omni(model="<your-model>", quantization="fp8") omni_fp8 = Omni(model="<your-model>", quantization="fp8")
omni_int8 = Omni(model="<your-model>", quantization="int8") omni_int8 = Omni(model="<your-model>", quantization="int8")
omni_bnb = Omni(model="Tongyi-MAI/Z-Image-Turbo", quantization="bitsandbytes")
omni_mxfp8 = Omni(model="<your-model>", quantization="mxfp8") omni_mxfp8 = Omni(model="<your-model>", quantization="mxfp8")
omni_mxfp4 = Omni(model="<your-model>", quantization="mxfp4") omni_mxfp4 = Omni(model="<your-model>", quantization="mxfp4")
omni_mxfp4_dualscale = Omni(model="<your-model>", quantization="mxfp4_dualscale") omni_mxfp4_dualscale = Omni(model="<your-model>", quantization="mxfp4_dualscale")
@@ -69,6 +75,7 @@ CLI:
```bash ```bash
vllm serve <your-model> --omni --quantization fp8 vllm serve <your-model> --omni --quantization fp8
vllm serve <your-model> --omni --quantization int8 vllm serve <your-model> --omni --quantization int8
vllm serve Tongyi-MAI/Z-Image-Turbo --omni --quantization bitsandbytes
vllm serve <your-model> --omni --quantization mxfp8 vllm serve <your-model> --omni --quantization mxfp8
vllm serve <your-model> --omni --quantization mxfp4 vllm serve <your-model> --omni --quantization mxfp4
vllm serve <your-model> --omni --quantization mxfp4_dualscale vllm serve <your-model> --omni --quantization mxfp4_dualscale
@@ -85,14 +92,19 @@ config = build_quant_config({
}) })
``` ```
Component routing is an extension mechanism, not automatic support for every
pipeline component. See [Quantization Scope](overview.md#quantization-scope)
before applying it to an encoder, VAE, or decoder.
## Parameters ## Parameters
| Parameter | Methods | Description | | Parameter | Methods | Description |
|-----------|---------|-------------| |-----------|---------|-------------|
| `method` | FP8, Int8, MXFP8, MXFP4 | Quantization method: `"fp8"`, `"int8"`, `"mxfp8"`, `"mxfp4"`, or `"mxfp4_dualscale"` | | `method` | FP8, Int8, BitsAndBytes, MXFP8, MXFP4 | Quantization method: `"fp8"`, `"int8"`, `"bitsandbytes"`, `"mxfp8"`, `"mxfp4"`, or `"mxfp4_dualscale"` |
| `ignored_layers` | FP8, Int8, MXFP8, MXFP4 | Layer name patterns to keep in BF16/FP16 | | `ignored_layers` | FP8, Int8, BitsAndBytes, MXFP8, MXFP4 | Layer name patterns to keep in BF16/FP16 |
| `activation_scheme` | FP8, Int8 | The runtime value `"dynamic"` selects online activation scaling | | `activation_scheme` | FP8, Int8 | The runtime value `"dynamic"` selects online activation scaling |
| `weight_block_size` | FP8 | Optional block-wise FP8 weight quantization size | | `weight_block_size` | FP8 | Optional block-wise FP8 weight quantization size |
| `quant_type` | BitsAndBytes | Weight format: `"nf4"` (default) or `"fp4"` |
| `num_bf16_fallback_layers` | MXFP4 DualScale | Leading transformer blocks to keep in BF16 for online `mxfp4_dualscale`; defaults to `5` | | `num_bf16_fallback_layers` | MXFP4 DualScale | Leading transformer blocks to keep in BF16 for online `mxfp4_dualscale`; defaults to `5` |
## Validation and Notes ## Validation and Notes
+130 -58
View File
@@ -16,72 +16,123 @@ For the internal architecture and backend extension points, see the
| Runtime attention quantization | [Quantized KV Cache](quantized_kvcache.md) | vLLM-Omni dynamically quantizes eligible diffusion Flash Attention tensors during inference. | FP8 FA | | Runtime attention quantization | [Quantized KV Cache](quantized_kvcache.md) | vLLM-Omni dynamically quantizes eligible diffusion Flash Attention tensors during inference. | FP8 FA |
| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, AutoRound, msModelSlim, serialized Int8, offline MXFP8, offline MXFP4 DualScale | | Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, AutoRound, msModelSlim, serialized Int8, offline MXFP8, offline MXFP4 DualScale |
## Hardware Support Online quantization starts from a normal BF16/FP16 checkpoint and repeats the
conversion on each fresh model load. It is the simplest path for experiments.
Pre-quantized checkpoints store the packed weights and scales produced by a
separate tool; they are easier to reuse across deployments and can carry
calibrated or mixed-precision policies that load-time conversion cannot
reproduce.
| Device | FP8 W8A8 | Int8 W8A8 | BitsAndBytes W4 | ModelOpt | MXFP8 W8A8 | MXFP4 W4A4 | AutoRound | msModelSlim | ## Support Levels
|--------|----------|-----------|------------------|----------|------------|------------|-----------|-------------|
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ❌ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ❌ |
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
| AMD ROCm | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ✅ | ⭕ | ❌ |
| Intel XPU | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
| Ascend NPU | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ |
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this Quantization support is specific to a model, component, checkpoint format, and
guide. FP8 on Ampere may use a weight-only path where available. hardware backend. The tables in this guide use the following levels:
## Model Type Support | Level | Meaning |
|-------|---------|
| **CI-backed** | A named model and quantized checkpoint are exercised by scheduled full-model hardware CI. |
| **Validated** | A named model and checkpoint have end-to-end or quality-validation evidence, but the check is not necessarily scheduled in CI. |
| **Integrated** | The config, loader, quantized layers, or model wiring exists, but this guide has no completed end-to-end validation for the named model and checkpoint. |
| **Not validated** | There is no current support recommendation for this model/component combination. |
| **Unsupported** | The combination is intentionally rejected or known not to work. |
### Diffusion Model (Qwen-Image, Wan2.2) A method being available on a device does not make every model on that device
supported. Likewise, successfully detecting a checkpoint's
`quantization_config` establishes loader compatibility, not output quality.
These models run a diffusion transformer as the primary inference module. The ## Hardware × Quantization Method
default quantization target is the transformer; tokenizer, scheduler, text
encoder, and VAE stay on the base checkpoint unless a method guide says
otherwise.
| Method | Guide | Mode | Example models | Status | This matrix uses concrete inference paths as columns. Producer toolkits are not
|--------|-------|------|----------------|--------| separate methods: NVIDIA ModelOpt produces the ModelOpt FP8/NVFP4 formats,
| FP8 W8A8 | [FP8](fp8.md) | Online W8A8 or checkpoint FP8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models | Intel AutoRound produces AutoRound checkpoints, and Ascend msModelSlim produces
| Int8 W8A8 | [Int8](int8.md) | Online or serialized W8A8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image | checkpoints consumed by the native MXFP or compatible Ascend paths.
| BitsAndBytes W4 | [BitsAndBytes](bitsandbytes.md) | Online W4 weight-only | Z-Image-Turbo; Qwen-Image/Wan2.2 not validated | Validated for Z-Image-Turbo |
| ModelOpt | [ModelOpt](modelopt.md) | Pre-quantized FP8 checkpoints | Qwen-Image, Z-Image, FLUX.2, HunyuanImage-3.0 | Validated for ModelOpt FP8 diffusion checkpoints |
| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Online W8A8 or offline pre-quantized | Wan2.2-T2V-A14B, I2V-A14B, TI2V-5B | Ascend NPU only; validated for Wan2.2 |
| MXFP4 W4A4 | [MXFP4](mxfp4.md) | `mxfp4`: online single-scale only; `mxfp4_dualscale`: online or offline dual-scale (offline recommended) | Wan2.2-T2V-A14B, I2V-A14B | Ascend NPU only; validated for Wan2.2 A14B cascade models; TI2V-5B not supported; offline `mxfp4_dualscale` uses calibrated `mul_scale` for best accuracy |
| AutoRound | [AutoRound](autoround.md) | Pre-quantized W4A16 checkpoints | FLUX.1-dev; Qwen-Image/Wan2.2 not validated | Checkpoint-driven |
| msModelSlim | [msModelSlim](msmodelslim.md) | Pre-quantized Ascend checkpoints | Wan2.2 recipe; HunyuanImage-3.0 inference target | Ascend/NPU path |
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS) | Hardware | [Online FP8](fp8.md) | [Int8](int8.md) | [BitsAndBytes W4](bitsandbytes.md) | [ModelOpt FP8](modelopt.md) | [ModelOpt NVFP4](modelopt.md#supported-modelopt-checkpoint-formats) | [AutoRound W4A16](autoround.md) | [OCP MXFP8](mxfp8.md) | [OCP MXFP4](mxfp4.md) |
|----------|-----------------------|------------------|------------------------------------------|--------------------------------|-------------------------------------------------------------------------|--------------------------------------|--------------------------|--------------------------|
| NVIDIA Blackwell (SM 100+) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
| NVIDIA Hopper (SM 90) | ✅ | ✅ | ✅ | ✅ | ✅ (Marlin fallback) | ✅ | ❌ | ❌ |
| NVIDIA Ada (SM 89) | ✅ | ✅ | ✅ | ✅ | ⭕ | ✅ | ❌ | ❌ |
| NVIDIA Ampere (SM 80+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ❌ | ❌ |
| AMD ROCm (gfx950 for MXFP4) | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ⭕ | ❌ | ✅ (online only) |
| Intel XPU | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
| Ascend NPU | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ |
These models combine an AR language model with audio, vision, talker, or TTS Legend: `✅` backend available, `❌` rejected or no implementation, `⭕` not
stages. Quantization is scoped to the AR language-model stage when the verified in this guide. This matrix does not establish model-level support; use
checkpoint contains a supported `quantization_config`; the non-AR stages stay the model-and-stage tables below for that.
in BF16 unless the model guide explicitly adds support.
| Method | Guide | Scope | Example models | Status | The Int8 column covers online and serialized Int8 paths. OCP MXFP8 is online on
|--------|-------|-------|----------------|--------| Intel XPU and online/offline on Ascend; Intel's separate AutoRound MXFP8
| ModelOpt | [ModelOpt](modelopt.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path | checkpoint path is described in the [MXFP8 guide](mxfp8.md). OCP MXFP4 on AMD
| Int8 | [Int8](int8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | means online single-scale `mxfp4` on gfx950, while MXFP4 DualScale is
| BitsAndBytes | [BitsAndBytes](bitsandbytes.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | Ascend-only. NVIDIA's NVFP4 and OCP MXFP4 have different checkpoint metadata,
| MXFP8 | [MXFP8](mxfp8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | scale layouts, and kernel contracts. Mixed ModelOpt FP8/NVFP4 checkpoints also
| MXFP4 | [MXFP4](mxfp4.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | require a model-specific mixed-precision policy and validation. FP8 on Ampere
| AutoRound | [AutoRound](autoround.md) | Thinker or language-model checkpoint config | Qwen2.5-Omni, Qwen3-Omni | Supported through AutoRound checkpoints | may use a weight-only path where available.
| msModelSlim | [msModelSlim](msmodelslim.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
### Multi-Stage Diffusion Model (BAGEL, GLM-Image) ## Current Model and Stage Support
These models split generation across multiple stages. Quantization must be The current offline checkpoint guidance is transformer-first. A listed
attached to the intended stage rather than applied globally. checkpoint quantizes the stage named in the **Scope** column; it does not imply
that the model's encoders, VAE, decoder, or other stages are also quantized.
| Method | Guide | Scope | Example models | Status | ### Vendor and Runtime Ecosystems
|--------|-------|-------|----------------|--------|
| FP8 | [FP8](fp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation | The producer and runtime backend are related, but they are not the same support
| Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation | claim. A producer can emit several checkpoint formats, and each format still
| BitsAndBytes | [BitsAndBytes](bitsandbytes.md) | Stage-specific transformer or DiT module | BAGEL, GLM-Image | Not validated | needs a compatible vLLM-Omni loader, model-stage wiring, and validation. A
| ModelOpt | [ModelOpt](modelopt.md) | Checkpoint-defined diffusion stage | BAGEL, GLM-Image | Requires model-specific validation | runtime may also provide online quantization without an offline producer path.
| MXFP8 | [MXFP8](mxfp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated |
| MXFP4 | [MXFP4](mxfp4.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated | | Ecosystem | Current paths | Current vLLM-Omni scope | Model evidence |
| AutoRound | [AutoRound](autoround.md) | Checkpoint-defined stage | BAGEL, GLM-Image | No validated checkpoint listed | |-----------|---------------|-------------------------|----------------|
| msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation | | NVIDIA | ModelOpt FP8, NVFP4, and mixed FP8/NVFP4 | Diffusion transformer or Qwen-Omni thinker | Qwen-Image and other named diffusion checkpoints are validated; Qwen3-Omni NVFP4 has H100 CI for the Marlin fallback |
| Intel | AutoRound W4A16 and AutoRound MXFP8 | Diffusion/world-model transformer or Qwen-Omni thinker | Wan A14B and GLM-Image W4A16 are CI-backed; Qwen-Image, Qwen-Omni, and Cosmos3-Super have validation outside scheduled CI |
| AMD ROCm | Online OCP MXFP4 through AITER on gfx950; no pre-quantized MXFP4 or AMD Quark workflow is integrated | Diffusion transformer linear layers | Integrated backend with dispatch and weight-allocation unit coverage; no named-model quality validation or scheduled gfx950 full-model CI |
| Ascend | msModelSlim output through native MXFP8/MXFP4 merge workflows or a compatible generic Ascend format | Wan diffusion transformers | Native Wan MXFP8/MXFP4 paths are validated outside scheduled full-model NPU CI; other model families remain integrated or unvalidated |
### Pre-Quantized Model Matrix
| Model family | Method | Scope | Level | Notes |
|--------------|--------|-------|-------|-------|
| Qwen-Image | [ModelOpt](modelopt.md) FP8 or mixed FP8/NVFP4 | Diffusion transformer | Validated | Named Qwen-Image 2512 checkpoints and recipes are available. Text encoder and VAE stay BF16. |
| Qwen-Image | [AutoRound](autoround.md) W4A16 | Diffusion transformer | Validated | Uses `INC4AI/Qwen-Image-AutoRound-W4A16`; text encoder and VAE stay BF16. |
| Wan2.2 I2V/T2V A14B | [AutoRound](autoround.md) W4A16 | Both diffusion transformers | CI-backed | The high-noise and low-noise DiTs are quantized; text encoder and VAE stay BF16. |
| Wan2.2 TI2V 5B | [AutoRound](autoround.md) W4A16 | Diffusion transformer | Validated | Named checkpoint available; not covered by the scheduled Wan AutoRound job. |
| Wan2.2 | [MXFP8](mxfp8.md) | One or both diffusion transformers | Validated | Native Ascend and AutoRound checkpoint formats are integrated; no scheduled full-model offline NPU job. |
| Wan2.2 I2V/T2V A14B | [MXFP4 DualScale](mxfp4.md) | Both diffusion transformers | Validated | Offline calibrated checkpoints are recommended; no scheduled full-model offline NPU job. |
| Qwen2.5-Omni | [AutoRound](autoround.md) W4A16 | Thinker language model | Validated | Encoders and output/audio stages stay BF16. |
| Qwen3-Omni | [ModelOpt](modelopt.md) FP8 | Thinker language model | Validated | Tested outside scheduled CI. |
| Qwen3-Omni | [ModelOpt](modelopt.md) NVFP4 | Thinker language model | CI-backed | H100 CI exercises the Marlin fallback; the native Blackwell FP4 path has separate recipe validation. |
| Qwen3-Omni | [AutoRound](autoround.md) W4A16 | Thinker language model | Validated | Audio encoder, vision encoder, talker, and Code2Wav stay BF16. |
| GLM-Image | [AutoRound](autoround.md) W4A16 | Diffusion transformer stage | CI-backed | Text-to-image and image-to-image full-model smoke coverage. |
| Cosmos3-Super | [AutoRound](autoround.md) W4A16 | World-model transformer | Validated | Manually validated serving path; VAE and other components stay BF16. |
| Cosmos3-Nano | [AutoRound](autoround.md) W4A16 | World-model transformer | Integrated | Named checkpoint is listed, but no model-specific end-to-end test is maintained in-tree. |
| MiniMax H3 | None | None | Not validated | No offline vendor checkpoint path is documented; online FP8 targets the DiT only. |
The method pages contain the complete checkpoint lists, including FLUX,
Z-Image, and HunyuanImage-3.0 variants.
### Online Quantization
| Model family | Methods | Scope | Level |
|--------------|---------|-------|-------|
| Qwen-Image | FP8, Int8 | Diffusion transformer | Validated |
| Z-Image | Int8, BitsAndBytes W4 | Diffusion transformer | Validated |
| Wan2.2 | MXFP8 | One or both diffusion transformers | Validated on the documented XPU/NPU paths |
| Wan2.2 A14B | MXFP4 | Both diffusion transformers | Validated on Ascend; the ROCm gfx950 online backend is integrated but not model-validated |
| MiniMax H3 | FP8 | DiT and optional reference DiT | Validated |
| Qwen3-Omni | Generic online methods, such as FP8 | Thinker language model only | Integrated; not model-validated as an online recommendation |
See [Online Quantization](online.md) for configuration and hardware details.
### Unlisted and Long-Tail Models
For models not listed above, generic quantized linear layers or compatible
checkpoint metadata are only an integration starting point. The default support
target is the model's primary transformer, DiT, or AR language-model stage.
Text/audio/vision encoders, VAEs, vocoders, decoders, and output stages remain
BF16 until a model-specific integration supplies quant-aware layers, checkpoint
mapping, end-to-end output validation, and backend coverage.
!!! note !!! note
"Online quantization" means vLLM-Omni computes the quantization data while "Online quantization" means vLLM-Omni computes the quantization data while
@@ -92,8 +143,8 @@ attached to the intended stage rather than applied globally.
### Diffusion Model (Qwen-Image, Wan2.2) ### Diffusion Model (Qwen-Image, Wan2.2)
The default target is the diffusion transformer. Component routing is available The default target is the diffusion transformer. `build_quant_config()` can
through `build_quant_config()`: construct a component router:
```python ```python
from vllm_omni.quantization import build_quant_config from vllm_omni.quantization import build_quant_config
@@ -104,11 +155,19 @@ config = build_quant_config({
}) })
``` ```
!!! warning "Routing syntax is not a support claim"
A component dictionary only selects a quantization config by module prefix.
It does not make an encoder or VAE quantizable. The model must construct
that component from quantization-aware layers, and an offline checkpoint
must contain matching packed weights and scales. Current user guidance
keeps diffusion text encoders and VAEs in BF16.
| Component | Default quantized? | Notes | | Component | Default quantized? | Notes |
|-----------|--------------------|-------| |-----------|--------------------|-------|
| Diffusion transformer | Yes | Primary target for FP8, Int8, BitsAndBytes, ModelOpt, MXFP8, MXFP4, AutoRound, and msModelSlim | | Diffusion transformer | Yes, for listed model/method pairs | Primary target for FP8, Int8, BitsAndBytes, ModelOpt, MXFP8, MXFP4, AutoRound, and msModelSlim |
| Second diffusion transformer | Yes, for listed Wan2.2 A14B methods | `transformer` and `transformer_2` read their own offline checkpoint metadata; arbitrary mixed methods are not currently validated |
| Text encoder | No | Keep BF16 unless a method-specific guide documents support | | Text encoder | No | Keep BF16 unless a method-specific guide documents support |
| VAE | No | Keep BF16; storage-only paths are method-specific | | VAE | No | Keep BF16 |
| Scheduler/tokenizer | No | Loaded from the base model repository | | Scheduler/tokenizer | No | Loaded from the base model repository |
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS) ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
@@ -121,6 +180,18 @@ config = build_quant_config({
| Talker or TTS stage | No | BF16 unless model-specific support is documented | | Talker or TTS stage | No | BF16 unless model-specific support is documented |
| Code2Wav | No | BF16 | | Code2Wav | No | BF16 |
An explicit `ComponentQuantizationConfig` can be used by model developers to
route different configs to `audio_tower`, `visual`, and `language_model`, but
the supported Qwen-Omni checkpoints currently quantize only the thinker
language model.
### World Model (Cosmos3)
| Component | Default quantized? | Notes |
|-----------|--------------------|-------|
| World-model transformer | AutoRound checkpoints only | Cosmos3-Super is validated; Cosmos3-Nano is integrated but not end-to-end tested in-tree |
| VAE and guardrail components | No | BF16 |
### Multi-Stage Diffusion Model (BAGEL, GLM-Image) ### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
| Component | Default quantized? | Notes | | Component | Default quantized? | Notes |
@@ -141,6 +212,7 @@ build_quant_config("fp8")
build_quant_config({"method": "fp8", "activation_scheme": "static"}) build_quant_config({"method": "fp8", "activation_scheme": "static"})
build_quant_config("bitsandbytes") build_quant_config("bitsandbytes")
build_quant_config("auto-round", bits=4, group_size=128) build_quant_config("auto-round", bits=4, group_size=128)
# Component routing syntax; model-specific support is still required.
build_quant_config({"transformer": {"method": "fp8"}, "vae": None}) build_quant_config({"transformer": {"method": "fp8"}, "vae": None})
build_quant_config(None) build_quant_config(None)
``` ```
@@ -8,6 +8,7 @@ A unified script for text-to-video generation. Supports multiple models with mod
|---|---|---|---|---|---| |---|---|---|---|---|---|
| `Wan-AI/Wan2.1-VACE-1.3B-diffusers` | 480x832 | 81 | 30 | 5.0 | ~20 GiB (RTX 5090, VAE tiling) | | `Wan-AI/Wan2.1-VACE-1.3B-diffusers` | 480x832 | 81 | 30 | 5.0 | ~20 GiB (RTX 5090, VAE tiling) |
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB | | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
| `robbyant/lingbot-video-dense-1.3b` / `robbyant/lingbot-video-moe-30b-a3b` | 192x320 | 9 | 2 | 3.0 | ~68 GiB (MoE smoke) |
| `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent | | `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
@@ -57,18 +58,31 @@ python text_to_video.py \
### LingBot-Video ### LingBot-Video
The dedicated LingBot example explicitly selects the video output modality: The shared runner recognizes the official dense and MoE checkpoint IDs, selects
`LingBotVideoPipeline`, and builds the canonical video request envelope:
```bash ```bash
python examples/offline_inference/text_to_video/text_to_video_lingbot.py \ python text_to_video.py \
--model robbyant/lingbot-video-dense-1.3b \ --model robbyant/lingbot-video-dense-1.3b \
--prompt "a robotic arm picks up a red block" \
--height 192 --width 320 --num-frames 9 --num-inference-steps 2 \ --height 192 --width 320 --num-frames 9 --num-inference-steps 2 \
--guidance-scale 3.0 --flow-shift 3.0 --fps 24 \
--output lingbot_t2v.mp4 --output lingbot_t2v.mp4
``` ```
Use the same command with `robbyant/lingbot-video-moe-30b-a3b` for the MoE Use the same command with `robbyant/lingbot-video-moe-30b-a3b` for the MoE
checkpoint. Requested frame counts are rounded upward to the causal checkpoint. Requested frame counts are rounded upward to the causal
VAE's `4n+1` grid. VAE's `4n+1` grid. For a local checkpoint path whose name does not contain
`lingbot`, pass `--model-class-name LingBotVideoPipeline`.
LingBot-only options such as `batch_cfg` and `output_type` use the shared
model-extra channel, for example:
```bash
python text_to_video.py \
--model robbyant/lingbot-video-dense-1.3b \
--extra-body '{"batch_cfg": true, "output_type": "np"}'
```
### LTX-2 ### LTX-2
@@ -16,7 +16,7 @@ from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.lora.request import LoRARequest from vllm_omni.lora.request import LoRARequest
from vllm_omni.lora.utils import stable_lora_int_id from vllm_omni.lora.utils import stable_lora_int_id
from vllm_omni.model_extras import get_extra_body_params, get_model_class_name from vllm_omni.model_extras import get_extra_body_params, get_model_class_name, get_output_tensor_range
from vllm_omni.outputs import OmniRequestOutput from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform from vllm_omni.platforms import current_omni_platform
@@ -80,6 +80,17 @@ _MODEL_PRESETS = {
"fps": 16, "fps": 16,
"output": "helios_output.mp4", "output": "helios_output.mp4",
}, },
"lingbot": {
"model_class_name": "LingBotVideoPipeline",
"height": 192,
"width": 320,
"num_frames": 9,
"num_inference_steps": 2,
"guidance_scale": 3.0,
"fps": 24,
"flow_shift": 3.0,
"output": "lingbot_video_output.mp4",
},
"ltx2": { "ltx2": {
"height": 512, "height": 512,
"width": 768, "width": 768,
@@ -110,27 +121,51 @@ _MODEL_PRESETS = {
def _detect_preset(model: str, model_class_name: str | None = None) -> dict: def _detect_preset(model: str, model_class_name: str | None = None) -> dict:
model_lower = model.lower() model_lower = model.lower()
class_lower = (model_class_name or "").lower() class_lower = (model_class_name or "").lower()
if "distilled" in class_lower or "distilled" in model_lower: if "lingbot" in model_lower or "lingbotvideo" in class_lower:
return _MODEL_PRESETS["ltx2_distilled"] return _MODEL_PRESETS["lingbot"]
if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower: if "ltx" in class_lower or "ltx" in model_lower:
return _MODEL_PRESETS["ltx23"] if "distilled" in class_lower or "distilled" in model_lower:
if "ltx2" in class_lower or "ltx-2" in model_lower or "ltx_2" in model_lower: return _MODEL_PRESETS["ltx2_distilled"]
if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower:
return _MODEL_PRESETS["ltx23"]
return _MODEL_PRESETS["ltx2"] return _MODEL_PRESETS["ltx2"]
if "vace" in model_lower: if "vace" in model_lower or "vace" in class_lower:
return _MODEL_PRESETS["vace"] return _MODEL_PRESETS["vace"]
# Edge must be matched before the generic cosmos branch (its "cosmos" substring would # Edge must be matched before the generic cosmos branch (its "cosmos" substring would
# otherwise pick up the Nano/Super 720p / gs 6.0 / flow_shift 10.0 preset). # otherwise pick up the Nano/Super 720p / gs 6.0 / flow_shift 10.0 preset).
if ("cosmos" in model_lower or "cosmos" in class_lower) and ("edge" in model_lower or "edge" in class_lower): if ("cosmos" in model_lower or "cosmos" in class_lower) and ("edge" in model_lower or "edge" in class_lower):
return _MODEL_PRESETS["cosmos3_edge"] return _MODEL_PRESETS["cosmos3_edge"]
if "cosmos" in model_lower: if "cosmos" in model_lower or "cosmos" in class_lower:
return _MODEL_PRESETS["cosmos"] return _MODEL_PRESETS["cosmos"]
if "hunyuan" in model_lower: if "hunyuan" in model_lower or "hunyuan" in class_lower:
return _MODEL_PRESETS["hunyuan"] return _MODEL_PRESETS["hunyuan"]
if "helios" in model_lower: if "helios" in model_lower or "helios" in class_lower:
return _MODEL_PRESETS["helios"] return _MODEL_PRESETS["helios"]
return _MODEL_PRESETS["wan"] return _MODEL_PRESETS["wan"]
def build_text_to_video_prompt(prompt: str, negative_prompt: str | None) -> dict[str, Any]:
"""Build the canonical request envelope for the shared T2V example."""
result: dict[str, Any] = {
"prompt": prompt,
"modalities": ["video"],
}
if negative_prompt is not None:
result["negative_prompt"] = negative_prompt
return result
def _normalize_float_tensor(tensor: torch.Tensor, source_range: str) -> torch.Tensor:
"""Normalize decoded frames according to the pipeline's output contract."""
if not tensor.is_floating_point():
return tensor
if source_range == "negative_one_to_one":
return tensor.clamp(-1, 1) * 0.5 + 0.5
if source_range == "zero_to_one":
return tensor.clamp(0, 1)
raise ValueError(f"Unsupported floating-point tensor range: {source_range!r}")
def parse_profiler_config(value: str) -> dict[str, Any]: def parse_profiler_config(value: str) -> dict[str, Any]:
try: try:
config = json.loads(value) config = json.loads(value)
@@ -160,7 +195,7 @@ def parse_extra_body(value: str) -> dict[str, Any]:
def parse_args() -> argparse.Namespace: def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(
description="Generate a video from a text prompt. " description="Generate a video from a text prompt. "
"Supports Wan2.2, HunyuanVideo-1.5, Helios, and other text-to-video models." "Supports Wan2.2, HunyuanVideo-1.5, Helios, LingBot-Video, and other text-to-video models."
) )
parser.add_argument( parser.add_argument(
"--model", "--model",
@@ -408,6 +443,7 @@ def main():
for key, default_val in preset.items(): for key, default_val in preset.items():
if getattr(args, key.replace("-", "_"), None) is None: if getattr(args, key.replace("-", "_"), None) is None:
setattr(args, key.replace("-", "_"), default_val) setattr(args, key.replace("-", "_"), default_val)
model_class_name = args.model_class_name
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
# Cache-dit config (Wan2.2 only) # Cache-dit config (Wan2.2 only)
@@ -480,6 +516,7 @@ def main():
omni = Omni(**omni_kwargs) omni = Omni(**omni_kwargs)
model_class_name = get_model_class_name(omni) or model_class_name model_class_name = get_model_class_name(omni) or model_class_name
declared_extra_body_params = get_extra_body_params(model_class_name) declared_extra_body_params = get_extra_body_params(model_class_name)
output_tensor_range = get_output_tensor_range(model_class_name)
if profiler_enabled: if profiler_enabled:
print("[Profiler] Starting profiling...") print("[Profiler] Starting profiling...")
@@ -513,12 +550,11 @@ def main():
lora_path=lora_path, lora_path=lora_path,
) )
prompt_dict = {"prompt": args.prompt} negative_prompt = args.negative_prompt
if args.negative_prompt is not None: if negative_prompt is None and all(preset is not _MODEL_PRESETS[name] for name in ("lingbot", "ltx2", "ltx23")):
prompt_dict["negative_prompt"] = args.negative_prompt
elif preset not in (_MODEL_PRESETS["ltx2"], _MODEL_PRESETS["ltx23"]):
# Preserve the historical empty-prompt behavior for non-LTX examples. # Preserve the historical empty-prompt behavior for non-LTX examples.
prompt_dict["negative_prompt"] = "" negative_prompt = ""
prompt_dict = build_text_to_video_prompt(args.prompt, negative_prompt)
extra_args = {} extra_args = {}
if lora_request: if lora_request:
@@ -630,8 +666,7 @@ def main():
frame_tensor = frame_tensor[0] frame_tensor = frame_tensor[0]
if frame_tensor.dim() == 3 and frame_tensor.shape[0] in (3, 4): if frame_tensor.dim() == 3 and frame_tensor.shape[0] in (3, 4):
frame_tensor = frame_tensor.permute(1, 2, 0) frame_tensor = frame_tensor.permute(1, 2, 0)
if frame_tensor.is_floating_point(): frame_tensor = _normalize_float_tensor(frame_tensor, output_tensor_range)
frame_tensor = frame_tensor.clamp(-1, 1) * 0.5 + 0.5
return frame_tensor.float().numpy() return frame_tensor.float().numpy()
if isinstance(frame, np.ndarray): if isinstance(frame, np.ndarray):
frame_array = frame frame_array = frame
@@ -683,8 +718,7 @@ def main():
video_tensor = video_tensor[0] video_tensor = video_tensor[0]
elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4): elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
video_tensor = video_tensor.permute(1, 2, 3, 0) video_tensor = video_tensor.permute(1, 2, 3, 0)
if video_tensor.is_floating_point(): video_tensor = _normalize_float_tensor(video_tensor, output_tensor_range)
video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
video_array = video_tensor.float().numpy() video_array = video_tensor.float().numpy()
elif isinstance(frames, np.ndarray): elif isinstance(frames, np.ndarray):
video_array = frames video_array = frames
@@ -1,131 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import argparse
import time
from pathlib import Path
from typing import Any
import numpy as np
import torch
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Generate text-to-video with a LingBot-Video checkpoint.")
parser.add_argument("--model", default="robbyant/lingbot-video-dense-1.3b")
parser.add_argument("--prompt", default="a robotic arm picks up a red block")
parser.add_argument("--negative-prompt", default=None)
parser.add_argument("--output", default="lingbot_video_output.mp4")
parser.add_argument("--height", type=int, default=192)
parser.add_argument("--width", type=int, default=320)
parser.add_argument("--num-frames", type=int, default=9)
parser.add_argument("--num-inference-steps", type=int, default=2)
parser.add_argument("--guidance-scale", type=float, default=3.0)
parser.add_argument("--flow-shift", type=float, default=3.0)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--fps", type=int, default=24)
parser.add_argument("--output-type", choices=["pt", "np"], default="pt")
parser.add_argument("--batch-cfg", action="store_true")
return parser.parse_args()
def _peak_memory_mb(output: Any) -> float:
result = output[0] if isinstance(output, list) and output else output
return float(getattr(result, "peak_memory_mb", 0.0) or 0.0)
def _extract_video(output: Any) -> Any:
if isinstance(output, list):
output = output[0] if output else None
if isinstance(output, OmniRequestOutput):
if output.final_output_type != "image":
raise ValueError(f"Unexpected output type {output.final_output_type!r}.")
if not output.images:
raise ValueError("No video frames found in OmniRequestOutput.")
return output.images[0] if len(output.images) == 1 else output.images
return output
def _to_frame_array(video: Any) -> np.ndarray:
if isinstance(video, torch.Tensor):
tensor = video.detach().cpu()
if tensor.dim() == 5:
tensor = tensor[0]
if tensor.dim() == 4 and tensor.shape[-1] not in (3, 4):
if tensor.shape[0] in (3, 4):
tensor = tensor.permute(1, 2, 3, 0)
elif tensor.shape[1] in (3, 4):
tensor = tensor.permute(0, 2, 3, 1)
if tensor.is_floating_point():
if float(tensor.min()) < 0.0:
tensor = tensor.clamp(-1, 1) * 0.5 + 0.5
else:
tensor = tensor.clamp(0, 1)
return tensor.float().numpy()
if isinstance(video, np.ndarray):
array = video[0] if video.ndim == 5 else video
if np.issubdtype(array.dtype, np.integer):
return array.astype(np.float32) / 255.0
return array.astype(np.float32, copy=False)
if isinstance(video, list):
return np.stack([_to_frame_array(frame) for frame in video], axis=0)
return np.asarray(video, dtype=np.float32)
def main() -> None:
args = parse_args()
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
omni = Omni(
model=args.model,
model_class_name="LingBotVideoPipeline",
flow_shift=args.flow_shift,
parallel_config=DiffusionParallelConfig(),
)
prompt: dict[str, Any] = {
"prompt": args.prompt,
"modalities": ["video"],
"num_frames": args.num_frames,
}
if args.negative_prompt is not None:
prompt["negative_prompt"] = args.negative_prompt
sampling_params = OmniDiffusionSamplingParams(
height=args.height,
width=args.width,
num_frames=args.num_frames,
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
generator=generator,
output_type=args.output_type,
extra_args={"batch_cfg": args.batch_cfg, "shift": args.flow_shift},
)
start = time.perf_counter()
output = omni.generate(prompt, sampling_params)
elapsed = time.perf_counter() - start
print(f"Total generation time: {elapsed:.4f} seconds ({elapsed * 1000:.2f} ms)")
peak_mb = _peak_memory_mb(output)
if peak_mb:
print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
video = _to_frame_array(_extract_video(output))
output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
from diffusers.utils import export_to_video
export_to_video(list(video), str(output_path), fps=args.fps)
print(f"Saved generated video to {output_path}")
omni.close()
if __name__ == "__main__":
main()
+3 -3
View File
@@ -27,7 +27,7 @@ parallel, cache, quantized, or expert-kernel backends.
- Dense checkpoint: <https://huggingface.co/robbyant/lingbot-video-dense-1.3b> - Dense checkpoint: <https://huggingface.co/robbyant/lingbot-video-dense-1.3b>
- MoE checkpoint: <https://huggingface.co/robbyant/lingbot-video-moe-30b-a3b> - MoE checkpoint: <https://huggingface.co/robbyant/lingbot-video-moe-30b-a3b>
- Upstream project: <https://github.com/Robbyant/lingbot-video> - Upstream project: <https://github.com/Robbyant/lingbot-video>
- Related offline example: [`examples/offline_inference/text_to_video/text_to_video_lingbot.py`](../../examples/offline_inference/text_to_video/text_to_video_lingbot.py) - Related offline example: [`examples/offline_inference/text_to_video/text_to_video.py`](../../examples/offline_inference/text_to_video/text_to_video.py)
- Related online video API docs: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md) - Related online video API docs: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md)
## Hardware Support ## Hardware Support
@@ -52,7 +52,7 @@ The MoE path is validated with BF16 expert weights.
```bash ```bash
CUDA_VISIBLE_DEVICES=0 \ CUDA_VISIBLE_DEVICES=0 \
python examples/offline_inference/text_to_video/text_to_video_lingbot.py \ python examples/offline_inference/text_to_video/text_to_video.py \
--model robbyant/lingbot-video-dense-1.3b \ --model robbyant/lingbot-video-dense-1.3b \
--prompt "a robotic arm picks up a red block" \ --prompt "a robotic arm picks up a red block" \
--output lingbot_t2v.mp4 \ --output lingbot_t2v.mp4 \
@@ -70,7 +70,7 @@ python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
```bash ```bash
CUDA_VISIBLE_DEVICES=0 \ CUDA_VISIBLE_DEVICES=0 \
python examples/offline_inference/text_to_video/text_to_video_lingbot.py \ python examples/offline_inference/text_to_video/text_to_video.py \
--model robbyant/lingbot-video-moe-30b-a3b \ --model robbyant/lingbot-video-moe-30b-a3b \
--prompt "a robotic arm picks up a red block" \ --prompt "a robotic arm picks up a red block" \
--output lingbot_moe_t2v.mp4 \ --output lingbot_moe_t2v.mp4 \
+1 -1
View File
@@ -82,7 +82,7 @@ example and the DiT sections below.
- Optional environment variables: - Optional environment variables:
```bash ```bash
export CACHE_DIT_VERSION=1.3.0 export CACHE_DIT_VERSION=1.5.0
``` ```
HunyuanImage-3.0 sets the diffusion attention backend to `TORCH_SDPA` HunyuanImage-3.0 sets the diffusion attention backend to `TORCH_SDPA`
+3 -1
View File
@@ -16,7 +16,7 @@ diffusers==0.38.0
safetensors>=0.8.0 safetensors>=0.8.0
accelerate==1.12.0 accelerate==1.12.0
soundfile>=0.13.1 soundfile>=0.13.1
cache-dit==1.3.0 cache-dit==1.5.0
tqdm>=4.66.0 tqdm>=4.66.0
torchsde>=0.2.6 torchsde>=0.2.6
openai-whisper>=20250625 openai-whisper>=20250625
@@ -34,3 +34,5 @@ pydantic-settings>=2.0.0
# base image does not ship it, so declare it explicitly. The loader also guards # base image does not ship it, so declare it explicitly. The loader also guards
# this import so a missing gguf no longer breaks all diffusion model loading. # this import so a missing gguf no longer breaks all diffusion model loading.
gguf>=0.10.0 gguf>=0.10.0
# Required for running the Cosmos3 model with guardrails enabled (default behavior).
cosmos-guardrail>=0.3.1
@@ -189,7 +189,7 @@ def _generate_single_stage_video(
def _generate_bagel_image( def _generate_bagel_image(
quantization_config: str | None = None, diffusion_quantization_config: str | None = None,
num_inference_steps: int = 15, num_inference_steps: int = 15,
) -> tuple[Any, float]: ) -> tuple[Any, float]:
"""Generate an image with BAGEL (multi-stage: LLM + Diffusion). """Generate an image with BAGEL (multi-stage: LLM + Diffusion).
@@ -202,8 +202,8 @@ def _generate_bagel_image(
"stage_configs_path": config_path, "stage_configs_path": config_path,
"stage_init_timeout": 300, "stage_init_timeout": 300,
} }
if quantization_config: if diffusion_quantization_config:
omni_kwargs["quantization_config"] = quantization_config omni_kwargs["diffusion_quantization_config"] = diffusion_quantization_config
model_name = omni_kwargs.pop("model") model_name = omni_kwargs.pop("model")
with OmniRunner(model_name, **omni_kwargs) as runner: with OmniRunner(model_name, **omni_kwargs) as runner:
@@ -380,21 +380,21 @@ def test_single_stage_ltx2_fp8_uses_less_memory():
@hardware_test(res={"cuda": "H100"}) @hardware_test(res={"cuda": "H100"})
def test_bagel_fp8_generates_image(): def test_bagel_fp8_generates_image():
"""BAGEL with FP8 quantization_config generates a valid image. """BAGEL with diffusion-stage FP8 generates a valid image.
FP8 should only apply to the diffusion stage (Stage-1), not the FP8 should only apply to the diffusion stage (Stage-1), not the
LLM stage (Stage-0). We verify this by checking: LLM stage (Stage-0). We verify this by checking:
1. Image is generated successfully 1. Image is generated successfully
2. LLM stage finish_reason is 'stop' (not 'length' from garbled output) 2. LLM stage finish_reason is 'stop' (not 'length' from garbled output)
""" """
image, _ = _generate_bagel_image(quantization_config="fp8") image, _ = _generate_bagel_image(diffusion_quantization_config="fp8")
image.save("test_bagel_fp8.png") image.save("test_bagel_fp8.png")
@hardware_test(res={"cuda": "H100"}) @hardware_test(res={"cuda": "H100"})
def test_bagel_bf16_generates_image(): def test_bagel_bf16_generates_image():
"""BAGEL without quantization generates a valid image (baseline).""" """BAGEL without quantization generates a valid image (baseline)."""
image, _ = _generate_bagel_image(quantization_config=None) image, _ = _generate_bagel_image(diffusion_quantization_config=None)
image.save("test_bagel_bf16.png") image.save("test_bagel_bf16.png")
@@ -95,6 +95,7 @@ def test_formatter_preserves_single_video_audio_actions_and_metadata(
"resolution": 512, "resolution": 512,
"width": None, "width": None,
"height": None, "height": None,
"num_inference_steps": 1,
} }
@@ -680,7 +680,7 @@ class TestSpeechAPI:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_create_diffusion_speech_extra_params(self, mocker: MockerFixture): async def test_create_diffusion_speech_extra_params(self, mocker: MockerFixture):
"""Test that extra_params are correctly applied to sampling_params_list in diffusion mode.""" """Test public diffusion speech success and extra_params propagation."""
# Mock the engine client # Mock the engine client
mock_engine = mocker.MagicMock() mock_engine = mocker.MagicMock()
@@ -704,7 +704,11 @@ class TestSpeechAPI:
req = OpenAICreateSpeechRequest(input="Hello", extra_params={"new_arg": 123, "existing_arg": "new_value"}) req = OpenAICreateSpeechRequest(input="Hello", extra_params={"new_arg": 123, "existing_arg": "new_value"})
await server._create_diffusion_speech(req) response = await server.create_speech(req)
assert response.status_code == 200
assert response.media_type == "audio/wav"
assert response.body == b"dummy"
# Verify generate was called # Verify generate was called
mock_engine.generate.assert_called_once() mock_engine.generate.assert_called_once()
@@ -737,6 +741,19 @@ class TestTTSMethods:
yield server yield server
server.shutdown() server.shutdown()
def test_diffusion_audio_encode_speed_skips_ar_adapter(self, mocker: MockerFixture):
"""Pure-diffusion speech keeps generic speed handling outside AR adapters."""
server = OmniOpenAIServingSpeech.for_diffusion(
diffusion_engine=mocker.MagicMock(),
model_name="test-model",
)
resolve_adapter = mocker.patch.object(serving_speech_module, "resolve_adapter")
speed = server._audio_encode_speed(OpenAICreateSpeechRequest(input="Hello", speed=1.25))
assert speed == 1.25
resolve_adapter.assert_not_called()
def test_is_tts_detection_no_stage(self, speech_server): def test_is_tts_detection_no_stage(self, speech_server):
"""Test TTS model detection when no TTS stage exists.""" """Test TTS model detection when no TTS stage exists."""
# Fixture creates server with stage_configs = [] -> _is_tts should be False # Fixture creates server with stage_configs = [] -> _is_tts should be False
@@ -0,0 +1,91 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import pytest
import torch
from examples.offline_inference.text_to_video.text_to_video import (
_MODEL_PRESETS,
_detect_preset,
_normalize_float_tensor,
build_text_to_video_prompt,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.mark.parametrize(
("negative_prompt", "expected_negative_prompt"),
[(None, None), ("", ""), ("blurry", "blurry")],
)
def test_text_to_video_builds_canonical_prompt(
negative_prompt: str | None,
expected_negative_prompt: str | None,
) -> None:
result = build_text_to_video_prompt("a robot waves", negative_prompt)
assert result["prompt"] == "a robot waves"
assert result["modalities"] == ["video"]
if expected_negative_prompt is None:
assert "negative_prompt" not in result
else:
assert result["negative_prompt"] == expected_negative_prompt
@pytest.mark.parametrize(
("model", "model_class_name", "preset_name"),
[
("robbyant/lingbot-video-dense-1.3b", None, "lingbot"),
("/models/custom-checkpoint", "LingBotVideoPipeline", "lingbot"),
("Lightricks/LTX-2", "LTX2DistilledPipeline", "ltx2_distilled"),
("BestWishYsh/Helios-Distilled", None, "helios"),
("/models/cosmos-edge", "Cosmos3EdgeVFMTransformer", "cosmos3_edge"),
("/models/vace", "WanVACEPipeline", "vace"),
],
)
def test_detect_preset_is_scoped_to_model_family(
model: str,
model_class_name: str | None,
preset_name: str,
) -> None:
assert _detect_preset(model, model_class_name) is _MODEL_PRESETS[preset_name]
def test_lingbot_preset_matches_lingbot_defaults() -> None:
assert _MODEL_PRESETS["lingbot"] == {
"model_class_name": "LingBotVideoPipeline",
"height": 192,
"width": 320,
"num_frames": 9,
"num_inference_steps": 2,
"guidance_scale": 3.0,
"fps": 24,
"flow_shift": 3.0,
"output": "lingbot_video_output.mp4",
}
def test_normalize_float_tensor_preserves_zero_to_one_frames() -> None:
frames = torch.tensor([0.0, 0.25, 1.0, 1.5])
assert torch.equal(_normalize_float_tensor(frames, "zero_to_one"), torch.tensor([0.0, 0.25, 1.0, 1.0]))
def test_normalize_float_tensor_maps_negative_one_to_one_frames() -> None:
frames = torch.tensor([-1.5, -1.0, 0.0, 1.0, 1.5])
assert torch.equal(
_normalize_float_tensor(frames, "negative_one_to_one"),
torch.tensor([0.0, 0.0, 0.5, 1.0, 1.0]),
)
def test_normalize_float_tensor_uses_one_contract_for_ambiguous_frames() -> None:
frames = [torch.tensor([0.0, 0.5, 1.0]), torch.tensor([-1.0, 0.0, 1.0])]
normalized = [_normalize_float_tensor(frame, "negative_one_to_one") for frame in frames]
assert torch.equal(normalized[0], torch.tensor([0.5, 0.75, 1.0]))
assert torch.equal(normalized[1], torch.tensor([0.0, 0.5, 1.0]))
+147 -3
View File
@@ -42,6 +42,10 @@ _EXPECTED_FAMILIES = [
defs.AUDIO_UNDERRUN_S, defs.AUDIO_UNDERRUN_S,
defs.AUDIO_CONTINUITY_OK_METRIC, defs.AUDIO_CONTINUITY_OK_METRIC,
defs.AUDIO_SKIPPED_REQUESTS_METRIC, defs.AUDIO_SKIPPED_REQUESTS_METRIC,
defs.DIFFUSION_EXEC_S,
defs.DIFFUSION_EXEC_PER_STEP_S,
defs.DIFFUSION_PREPROCESS_S,
defs.DIFFUSION_POSTPROCESS_S,
] ]
@@ -55,6 +59,10 @@ class TestRegistration:
mod.observe_audio_underrun("s", "r", 0.01) mod.observe_audio_underrun("s", "r", 0.01)
mod.inc_audio_continuity_ok("s", "r", 100) mod.inc_audio_continuity_ok("s", "r", 100)
mod.inc_audio_skipped("s", "r", "malformed_codec") mod.inc_audio_skipped("s", "r", "malformed_codec")
mod.observe_diffusion_exec("s", "r", 0.5)
mod.observe_diffusion_exec_per_step("s", "r", 0.01)
mod.observe_diffusion_preprocess("s", "r", 0.01)
mod.observe_diffusion_postprocess("s", "r", 0.2)
out = generate_latest(REGISTRY).decode() out = generate_latest(REGISTRY).decode()
for name in _EXPECTED_FAMILIES: for name in _EXPECTED_FAMILIES:
@@ -156,6 +164,50 @@ class TestAudio:
assert unknown == 1.0 assert unknown == 1.0
# ---------------------------------------------------------------------------
# Diffusion observe API
# ---------------------------------------------------------------------------
class TestDiffusion:
def test_diffusion_exec_observed(self, mod: OmniModalityMetrics) -> None:
stage, replica = "dit_exec", "0"
mod.observe_diffusion_exec(stage, replica, 1.25)
out = generate_latest(REGISTRY).decode()
prefix = f'{defs.DIFFUSION_EXEC_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
assert _sample_value(out, prefix) == 1.25
def test_diffusion_preprocess_observed(self, mod: OmniModalityMetrics) -> None:
stage, replica = "dit_pre", "0"
mod.observe_diffusion_preprocess(stage, replica, 0.03)
out = generate_latest(REGISTRY).decode()
prefix = f'{defs.DIFFUSION_PREPROCESS_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
assert _sample_value(out, prefix) == pytest.approx(0.03)
def test_diffusion_postprocess_observed(self, mod: OmniModalityMetrics) -> None:
stage, replica = "dit_post", "0"
mod.observe_diffusion_postprocess(stage, replica, 0.8)
out = generate_latest(REGISTRY).decode()
prefix = f'{defs.DIFFUSION_POSTPROCESS_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
assert _sample_value(out, prefix) == 0.8
def test_diffusion_exec_per_step_observed(self, mod: OmniModalityMetrics) -> None:
stage, replica = "dit_per_step", "0"
mod.observe_diffusion_exec_per_step(stage, replica, 0.078)
out = generate_latest(REGISTRY).decode()
prefix = f'{defs.DIFFUSION_EXEC_PER_STEP_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
assert _sample_value(out, prefix) == pytest.approx(0.078)
def test_log_stats_false_skips(self) -> None:
silent = OmniModalityMetrics(model_name="silent-model", log_stats=False)
silent.observe_diffusion_exec("s", "r", 1.0)
silent.observe_diffusion_exec_per_step("s", "r", 0.02)
silent.observe_diffusion_preprocess("s", "r", 0.01)
silent.observe_diffusion_postprocess("s", "r", 0.5)
out = generate_latest(REGISTRY).decode()
assert 'model_name="silent-model"' not in out
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Stub-driven routing tests for the finalize / streaming helpers. # Stub-driven routing tests for the finalize / streaming helpers.
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -183,6 +235,18 @@ class _StubModMetrics:
def inc_audio_skipped(self, s, r, reason): def inc_audio_skipped(self, s, r, reason):
self.calls.append(("inc_audio_skipped", s, r, reason)) self.calls.append(("inc_audio_skipped", s, r, reason))
def observe_diffusion_exec(self, s, r, v):
self.calls.append(("observe_diffusion_exec", s, r, v))
def observe_diffusion_exec_per_step(self, s, r, v):
self.calls.append(("observe_diffusion_exec_per_step", s, r, v))
def observe_diffusion_preprocess(self, s, r, v):
self.calls.append(("observe_diffusion_preprocess", s, r, v))
def observe_diffusion_postprocess(self, s, r, v):
self.calls.append(("observe_diffusion_postprocess", s, r, v))
class _Bag: class _Bag:
def __init__(self, **kw): def __init__(self, **kw):
@@ -239,8 +303,88 @@ class TestObserveModalityAtFinalize:
) )
assert ("observe_audio_duration", "1", "0", 1.0) in stub.calls assert ("observe_audio_duration", "1", "0", 1.0) in stub.calls
def test_non_audio_output_type_skipped(self): def test_diffusion_path_full(self):
# Image / video / text output types are out of scope for this module. stub = _StubModMetrics()
stage_metrics = _Bag(
stage_gen_time_ms=1000.0,
diffusion_metrics={
"preprocess_time_s": 0.05,
"diffusion_engine_exec_time_s": 1.2,
"postprocess_time_s": 0.3,
"image_num": 1,
"resolution": 1024,
"num_inference_steps": 50,
},
)
observe_modality_at_finalize(
stub,
output_type="image",
stage_id=2,
replica_id=0,
stage_metrics=stage_metrics,
engine_outputs=_Bag(),
)
assert ("observe_diffusion_preprocess", "2", "0", 0.05) in stub.calls
assert ("observe_diffusion_exec", "2", "0", 1.2) in stub.calls
assert ("observe_diffusion_postprocess", "2", "0", 0.3) in stub.calls
assert ("observe_diffusion_exec_per_step", "2", "0", pytest.approx(0.024)) in stub.calls
def test_diffusion_per_step_skipped_without_num_inference_steps(self):
stub = _StubModMetrics()
stage_metrics = _Bag(
stage_gen_time_ms=1000.0,
diffusion_metrics={
"diffusion_engine_exec_time_s": 1.2,
"preprocess_time_s": 0.05,
"postprocess_time_s": 0.3,
},
)
observe_modality_at_finalize(
stub,
output_type="image",
stage_id=2,
replica_id=0,
stage_metrics=stage_metrics,
engine_outputs=_Bag(),
)
assert ("observe_diffusion_exec", "2", "0", 1.2) in stub.calls
assert not any(c[0] == "observe_diffusion_exec_per_step" for c in stub.calls)
def test_diffusion_metrics_none_skipped(self):
stub = _StubModMetrics()
observe_modality_at_finalize(
stub,
output_type="image",
stage_id=2,
replica_id=0,
stage_metrics=_Bag(stage_gen_time_ms=500.0, diffusion_metrics=None),
engine_outputs=_Bag(),
)
assert not any(c[0].startswith("observe_diffusion") for c in stub.calls)
def test_diffusion_zero_values_observed(self):
stub = _StubModMetrics()
stage_metrics = _Bag(
stage_gen_time_ms=500.0,
diffusion_metrics={
"preprocess_time_s": 0.0,
"diffusion_engine_exec_time_s": 0.0,
"postprocess_time_s": 0.0,
},
)
observe_modality_at_finalize(
stub,
output_type="image",
stage_id=2,
replica_id=0,
stage_metrics=stage_metrics,
engine_outputs=_Bag(),
)
assert ("observe_diffusion_preprocess", "2", "0", 0.0) in stub.calls
assert ("observe_diffusion_exec", "2", "0", 0.0) in stub.calls
assert ("observe_diffusion_postprocess", "2", "0", 0.0) in stub.calls
def test_non_audio_output_type_skips_audio_but_still_checks_diffusion(self):
stub = _StubModMetrics() stub = _StubModMetrics()
for output_type in ("text", "image", "video"): for output_type in ("text", "image", "video"):
observe_modality_at_finalize( observe_modality_at_finalize(
@@ -251,7 +395,7 @@ class TestObserveModalityAtFinalize:
stage_metrics=_Bag(stage_gen_time_ms=100.0), stage_metrics=_Bag(stage_gen_time_ms=100.0),
engine_outputs=_Bag(), engine_outputs=_Bag(),
) )
assert stub.calls == [] assert not any(c[0].startswith("observe_audio") or c[0].startswith("inc_audio") for c in stub.calls)
def test_replica_id_none_skipped(self): def test_replica_id_none_skipped(self):
stub = _StubModMetrics() stub = _StubModMetrics()
@@ -1,21 +1,24 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from functools import lru_cache
from types import SimpleNamespace from types import SimpleNamespace
from unittest.mock import Mock from unittest.mock import Mock
import pytest import pytest
import torch import torch
import torch.nn as nn import torch.nn as nn
from vllm.platforms import current_platform
from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
HiFTGenerator, HiFTGenerator,
) )
from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import ( from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import (
CFMGraphWrapper,
HiFTGraphWrapper, HiFTGraphWrapper,
) )
pytestmark = [pytest.mark.core_model] pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
class _F0Predictor(nn.Module): class _F0Predictor(nn.Module):
@@ -142,3 +145,248 @@ def test_lazy_capture_limit_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch)
wrapper._capture.assert_not_called() wrapper._capture.assert_not_called()
wrapper.decode_fn.assert_called_once_with(speech_feat, cache_source) wrapper.decode_fn.assert_called_once_with(speech_feat, cache_source)
assert result is wrapper.decode_fn.return_value assert result is wrapper.decode_fn.return_value
# ---------------------------------------------------------------------------
# CFMGraphWrapper tests
# ---------------------------------------------------------------------------
class _MiniDiT(nn.Module):
"""Minimal DiT-like module with blocks_forward_chunk for CFM graph testing.
Mimics the upstream cosyvoice2 DiT's cache semantics:
- CausalConv1d.forward_chunk: cat([cnn_cache, x], dim=time) when cache is not None
- Attention.forward_chunk: cat([k, k_cache], dim=seq) when att_cache is not None
"""
def __init__(self, in_dim: int = 16, hidden: int = 8, depth: int = 2) -> None:
super().__init__()
self.in_proj = nn.Linear(in_dim, hidden)
self.blocks = nn.ModuleList([nn.LayerNorm(hidden) for _ in range(depth)])
self.final_layer = nn.Linear(hidden, in_dim)
def blocks_forward_chunk(
self,
x: torch.Tensor,
t: torch.Tensor,
mask: torch.Tensor | None,
cnn_cache: torch.Tensor | None = None,
att_cache: torch.Tensor | None = None,
cnn_cache_buffer: torch.Tensor | None = None,
att_cache_buffer: torch.Tensor | None = None,
) -> torch.Tensor:
x = x.transpose(1, 2)
x = self.in_proj(x)
for b_idx in range(len(self.blocks)):
# Simulate CausalConv1d: cnn_cache contributes to the first frames.
# Upstream: if cnn_cache[b_idx] is None, creates zeros internally (same as zeros).
cnn_b = cnn_cache[b_idx] if cnn_cache is not None else None
if cnn_b is not None:
x[:, : cnn_b.shape[2], :] += cnn_b.transpose(1, 2)
# Simulate Attention: att_cache contributes a bias when non-empty.
# Upstream: if att_cache[b_idx] is None, skips cat entirely (different path).
att_b = att_cache[b_idx] if att_cache is not None else None
if att_b is not None and att_b.shape[3] > 0:
x += att_b.sum(dim=(1, 2), keepdim=False).unsqueeze(1)
x = self.blocks[b_idx](x)
x = x + t
cnn_cache_buffer[b_idx] = x[:, -2:, :].transpose(1, 2).contiguous()
dt = x.shape[1]
att_cache_buffer[b_idx][:, :, :dt, :] = x.unsqueeze(1)
x = self.final_layer(x)
x = x.transpose(1, 2)
return x
def _cfm_inputs(
batch_size: int, chunk_size: int, old_att_len: int, *, device: str = "cuda"
) -> tuple[torch.Tensor, ...]:
depth = 2
hidden = 8
estimator_input = torch.randn(batch_size, 16, chunk_size, device=device)
time_emb = torch.randn(batch_size, 1, hidden, device=device)
cnn_cache = torch.randn(depth, batch_size, hidden, 2, device=device)
att_cache = torch.randn(depth, batch_size, 1, old_att_len, hidden, device=device)
cnn_out = torch.empty(depth, batch_size, hidden, 2, device=device)
att_out = torch.empty(depth, batch_size, 1, old_att_len + chunk_size, hidden, device=device)
return estimator_input, time_emb, cnn_cache, att_cache, cnn_out, att_out
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
def test_cfm_graph_replay_matches_eager_for_uncached_and_cached_shapes(
monkeypatch: pytest.MonkeyPatch,
) -> None:
pool = torch.cuda.graph_pool_handle()
monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
torch.manual_seed(0)
estimator = _MiniDiT().eval().cuda()
wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=32)
with torch.inference_mode():
for _, chunk_size, old_att_len in ((2, 10, 0), (2, 10, 5)):
inputs = _cfm_inputs(2, chunk_size, old_att_len)
eager_inputs = tuple(v.clone() for v in inputs)
with torch.no_grad():
eager_result = estimator.blocks_forward_chunk(
eager_inputs[0],
eager_inputs[1],
None,
eager_inputs[2],
eager_inputs[3],
eager_inputs[4],
eager_inputs[5],
)
wrapper.replay(*inputs)
replay_inputs = tuple(v.clone() for v in inputs)
graph_result, graph_cnn, graph_att = wrapper.replay(*replay_inputs)
torch.testing.assert_close(graph_result, eager_result, rtol=1e-4, atol=1e-5)
torch.testing.assert_close(graph_cnn, eager_inputs[4], rtol=1e-4, atol=1e-5)
torch.testing.assert_close(graph_att, eager_inputs[5], rtol=1e-4, atol=1e-5)
wrapper._capture_graph.cache_clear()
def _cfm_mock_wrapper(monkeypatch: pytest.MonkeyPatch, *, max_graphs: int = 1) -> CFMGraphWrapper:
monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
wrapper = object.__new__(CFMGraphWrapper)
wrapper.max_graphs = max_graphs
wrapper.graph_fn = Mock(return_value=torch.tensor([42.0]))
wrapper.device = torch.device("cuda")
@lru_cache(maxsize=max_graphs)
def _capture_graph(key: tuple):
return wrapper._do_capture(key)
wrapper._capture_graph = _capture_graph
wrapper._do_capture = Mock(return_value=None)
return wrapper
def test_cfm_unseen_shape_is_lazily_captured(monkeypatch: pytest.MonkeyPatch) -> None:
wrapper = _cfm_mock_wrapper(monkeypatch)
wrapper._capture_graph = Mock(return_value=None)
inputs = _cfm_inputs(2, 10, 0)
wrapper.replay(*inputs)
wrapper._capture_graph.assert_called_once()
wrapper.graph_fn.assert_called_once()
def test_cfm_capture_failure_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) -> None:
wrapper = _cfm_mock_wrapper(monkeypatch)
wrapper._capture_graph = lru_cache(maxsize=1)(lambda key: None)
inputs = _cfm_inputs(2, 10, 0)
result = wrapper.replay(*inputs)
wrapper.graph_fn.assert_called_once()
assert result[0] is wrapper.graph_fn.return_value
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
def test_cfm_lru_eviction(monkeypatch: pytest.MonkeyPatch) -> None:
"""Verify that LRU cache evicts oldest entries when full."""
pool = torch.cuda.graph_pool_handle()
monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
torch.manual_seed(0)
estimator = _MiniDiT().eval().cuda()
wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=2)
with torch.inference_mode():
inputs_a = _cfm_inputs(2, 10, 0)
wrapper.replay(*inputs_a)
assert wrapper._capture_graph.cache_info().currsize == 1
inputs_b = _cfm_inputs(2, 12, 0)
wrapper.replay(*inputs_b)
assert wrapper._capture_graph.cache_info().currsize == 2
wrapper.replay(*inputs_a)
assert wrapper._capture_graph.cache_info().currsize == 2
inputs_c = _cfm_inputs(2, 14, 0)
wrapper.replay(*inputs_c)
assert wrapper._capture_graph.cache_info().currsize == 2
wrapper.replay(*inputs_b)
hits_before = wrapper._capture_graph.cache_info().hits
wrapper.replay(*inputs_b)
hits_after = wrapper._capture_graph.cache_info().hits
assert hits_after == hits_before + 1
wrapper._capture_graph.cache_clear()
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
def test_cfm_none_cache_parity_between_graph_and_eager(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Verify that None→zeros (graph path) matches [None]*depth (eager path).
setup_batch calls _decode_cfm with cnn_cache=None on every request's
prompt pass. The graph path replaces None with zeros; the eager path
passes [None]*depth. Both must produce identical outputs.
"""
torch.accelerator.synchronize()
pool = torch.cuda.graph_pool_handle()
monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
torch.manual_seed(0)
estimator = _MiniDiT().eval().cuda()
wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=32)
depth = len(estimator.blocks)
hidden = 8
batch_size = 2
chunk_size = 20
estimator_input = torch.randn(batch_size, 16, chunk_size, device="cuda")
time_emb = torch.randn(batch_size, 1, hidden, device="cuda")
cnn_out = torch.empty(depth, batch_size, hidden, 2, device="cuda")
att_out = torch.empty(depth, batch_size, 1, chunk_size, hidden, device="cuda")
# Eager path: cnn_cache=[None]*depth, att_cache=[None]*depth
eager_input = estimator_input.clone()
eager_time = time_emb.clone()
eager_cnn_out = torch.empty_like(cnn_out)
eager_att_out = torch.empty_like(att_out)
with torch.no_grad():
eager_result = estimator.blocks_forward_chunk(
eager_input,
eager_time,
None,
[None] * depth,
[None] * depth,
eager_cnn_out,
eager_att_out,
)
# Graph path: cnn_cache=zeros, att_cache=zero-length (simulating _estimator_step's None→zeros)
zero_cnn = torch.zeros_like(cnn_out)
zero_att = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:])
wrapper.replay(estimator_input, time_emb, zero_cnn, zero_att, cnn_out, att_out)
graph_input = estimator_input.clone()
graph_time = time_emb.clone()
graph_cnn_out = torch.empty_like(cnn_out)
graph_att_out = torch.empty_like(att_out)
graph_zero_cnn = torch.zeros_like(cnn_out)
graph_zero_att = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:])
graph_result, graph_cnn, graph_att = wrapper.replay(
graph_input,
graph_time,
graph_zero_cnn,
graph_zero_att,
graph_cnn_out,
graph_att_out,
)
torch.testing.assert_close(graph_result, eager_result, rtol=1e-4, atol=1e-5)
torch.testing.assert_close(graph_cnn, eager_cnn_out, rtol=1e-4, atol=1e-5)
torch.testing.assert_close(graph_att, eager_att_out, rtol=1e-4, atol=1e-5)
+3
View File
@@ -17,6 +17,7 @@ from vllm_omni.model_extras import (
build_x_to_text_prompt, build_x_to_text_prompt,
get_extra_body_params, get_extra_body_params,
get_extra_output_params, get_extra_output_params,
get_output_tensor_range,
get_x_to_text_model_family, get_x_to_text_model_family,
should_init_extra_args_for_non_diffusion_stages, should_init_extra_args_for_non_diffusion_stages,
) )
@@ -134,6 +135,8 @@ def test_lingbot_extra_registry_declares_request_params() -> None:
} }
) )
assert get_extra_output_params("LingBotVideoPipeline") == frozenset() assert get_extra_output_params("LingBotVideoPipeline") == frozenset()
assert get_output_tensor_range("LingBotVideoPipeline") == "zero_to_one"
assert get_output_tensor_range("WanPipeline") == "negative_one_to_one"
assert should_init_extra_args_for_non_diffusion_stages("LingBotVideoPipeline") is False assert should_init_extra_args_for_non_diffusion_stages("LingBotVideoPipeline") is False
@@ -13,10 +13,13 @@ def _shrink_dit_rope_config(
target_head_dim: int = 32, target_head_dim: int = 32,
target_heads: int = 4, target_heads: int = 4,
default_axes_dims_rope: list[int] | None = None, default_axes_dims_rope: list[int] | None = None,
joint_attention_dim: int | None = None,
) -> dict: ) -> dict:
config["num_layers"] = num_layers config["num_layers"] = num_layers
if num_single_layers is not None: if num_single_layers is not None:
config["num_single_layers"] = num_single_layers config["num_single_layers"] = num_single_layers
if joint_attention_dim is not None:
config["joint_attention_dim"] = joint_attention_dim
axes_dims_rope = config.get("axes_dims_rope", default_axes_dims_rope) axes_dims_rope = config.get("axes_dims_rope", default_axes_dims_rope)
factor = config["attention_head_dim"] / target_head_dim factor = config["attention_head_dim"] / target_head_dim
config["attention_head_dim"] = target_head_dim config["attention_head_dim"] = target_head_dim
@@ -50,22 +53,55 @@ def _shrink_flux_t5_text_encoder(config: dict) -> dict:
return config return config
def tiny_qwen_image_builder() -> str: def _shrink_qwen_text_encoder_config(config: dict, hidden_size: int = 64) -> dict:
def shrink_text_encoder(config: dict) -> dict: text_config = config["text_config"]
config["num_hidden_layers"] = 2 old_head_dim = text_config["hidden_size"] / text_config["num_attention_heads"]
config["intermediate_size"] = 64 text_config["num_hidden_layers"] = 2
config["text_config"]["num_hidden_layers"] = 2 text_config["intermediate_size"] = 64
config["text_config"]["intermediate_size"] = 64 text_config["hidden_size"] = hidden_size
config["text_config"]["layer_types"] = config["text_config"]["layer_types"][:2] text_config["num_attention_heads"] = 2
config["vision_config"]["depth"] = 2 text_config["num_key_value_heads"] = 2
config["vision_config"]["intermediate_size"] = 64 text_config["layer_types"] = text_config["layer_types"][:2]
config["vision_config"]["fullatt_block_indexes"] = [0, 1] factor = old_head_dim / (hidden_size / 2)
return config mrope_section = text_config["rope_scaling"]["mrope_section"]
text_config["rope_scaling"]["mrope_section"] = [round(d / factor) for d in mrope_section]
config["num_hidden_layers"] = 2
config["intermediate_size"] = 64
config["hidden_size"] = hidden_size
config["num_attention_heads"] = 2
config["num_key_value_heads"] = 2
config["vision_config"]["depth"] = 2
config["vision_config"]["intermediate_size"] = 64
config["vision_config"]["fullatt_block_indexes"] = [0, 1]
config["vision_config"]["out_hidden_size"] = hidden_size
return config
def _shrink_qwen_transformer_config(config: dict, joint_attention_dim: int = 64) -> dict:
return _shrink_dit_rope_config(config, joint_attention_dim=joint_attention_dim)
def tiny_qwen_image_builder() -> str:
return build_tiny_from_configs( return build_tiny_from_configs(
"QwenImagePipeline", "QwenImagePipeline",
"Qwen/Qwen-Image", "Qwen/Qwen-Image",
transform={"text_encoder": shrink_text_encoder, "transformer": _shrink_dit_rope_config}, transform={"text_encoder": _shrink_qwen_text_encoder_config, "transformer": _shrink_qwen_transformer_config},
)
def tiny_qwen_image_edit_builder() -> str:
return build_tiny_from_configs(
"QwenImageEditPipeline",
"Qwen/Qwen-Image-Edit",
transform={"text_encoder": _shrink_qwen_text_encoder_config, "transformer": _shrink_qwen_transformer_config},
)
def tiny_qwen_image_edit_plus_builder() -> str:
return build_tiny_from_configs(
"QwenImageEditPlusPipeline",
"Qwen/Qwen-Image-Edit-2511",
transform={"text_encoder": _shrink_qwen_text_encoder_config, "transformer": _shrink_qwen_transformer_config},
) )
@@ -74,4 +74,14 @@ DIFFUSION_TEST_SETTINGS = {
builder=diff_model_builders.tiny_flux2_builder, builder=diff_model_builders.tiny_flux2_builder,
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE], supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
), ),
"QwenImageEditPipeline": DiffusionModelTestOpts(
model="Qwen/Qwen-Image-Edit",
builder=diff_model_builders.tiny_qwen_image_edit_builder,
supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
),
"QwenImageEditPlusPipeline": DiffusionModelTestOpts(
model="Qwen/Qwen-Image-Edit-2511",
builder=diff_model_builders.tiny_qwen_image_edit_plus_builder,
supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
),
} }
@@ -21,8 +21,6 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
# if you are adding a new model and see the tests below fail, please follow the pattern # if you are adding a new model and see the tests below fail, please follow the pattern
# for adding a new tiny model builder and corresponding entry in DIFFUSION_TEST_SETTINGS. # for adding a new tiny model builder and corresponding entry in DIFFUSION_TEST_SETTINGS.
EXCLUDED_MODELS = [ EXCLUDED_MODELS = [
"QwenImageEditPipeline",
"QwenImageEditPlusPipeline",
"QwenImageLayeredPipeline", "QwenImageLayeredPipeline",
"GlmImagePipeline", "GlmImagePipeline",
"ZImagePipeline", "ZImagePipeline",
+2
View File
@@ -19,6 +19,8 @@ connectors:
# A future initial_codec_chunk_frames shape will use lazy capture unless # A future initial_codec_chunk_frames shape will use lazy capture unless
# HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it. # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it.
enable_hift_graph: true enable_hift_graph: true
enable_cfm_graph: true
cfm_max_graphs: 32
stages: stages:
- stage_id: 0 - stage_id: 0
+2
View File
@@ -22,6 +22,8 @@ connectors:
# A future initial_codec_chunk_frames shape will use lazy capture unless # A future initial_codec_chunk_frames shape will use lazy capture unless
# HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it. # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it.
enable_hift_graph: true enable_hift_graph: true
enable_cfm_graph: true
cfm_max_graphs: 32
stages: stages:
- stage_id: 0 - stage_id: 0
+2
View File
@@ -22,6 +22,8 @@ connectors:
# A future initial_codec_chunk_frames shape will use lazy capture unless # A future initial_codec_chunk_frames shape will use lazy capture unless
# HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it. # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it.
enable_hift_graph: true enable_hift_graph: true
enable_cfm_graph: true
cfm_max_graphs: 32
stages: stages:
- stage_id: 0 - stage_id: 0
@@ -33,6 +33,8 @@ connectors:
# A future initial_codec_chunk_frames shape will use lazy capture unless # A future initial_codec_chunk_frames shape will use lazy capture unless
# HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it. # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it.
enable_hift_graph: true enable_hift_graph: true
enable_cfm_graph: true
cfm_max_graphs: 32
stages: stages:
- stage_id: 0 - stage_id: 0
@@ -65,11 +65,10 @@ except ImportError:
# ``model_config["guardrails"]`` is False. # ``model_config["guardrails"]`` is False.
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
raise ValueError( raise ValueError(
f"You have disabled the safety checker for {self.__class__}. This is in violation of the " f"The guardrails are enabled but not functional for {self.__class__}. This is most likely due to a"
f"missing cosmos-guardrail package. Please install cosmos-guardrail to be compliant with the "
"[NVIDIA Open Model License Agreement]" "[NVIDIA Open Model License Agreement]"
"(https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license). " "(https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license)."
f"Please ensure that you are compliant with the license agreement."
f"Please install cosmos-guardrail package to enable safety checks."
) )
+1
View File
@@ -138,6 +138,7 @@ def format_diffusion_outputs(
), ),
"width": request.sampling_params.width, "width": request.sampling_params.width,
"height": request.sampling_params.height, "height": request.sampling_params.height,
"num_inference_steps": request.sampling_params.num_inference_steps,
} }
# Only the primary payload determines the response type. Some image models # Only the primary payload determines the response type. Some image models
+1
View File
@@ -568,6 +568,7 @@ class OmniBase(PDDisaggregationMixin):
msg_id = id(result) msg_id = id(result)
consumed = self._consumed_metric_message_ids(req_id) consumed = self._consumed_metric_message_ids(req_id)
if msg_id not in consumed: if msg_id not in consumed:
metrics.accumulate_diffusion_metrics(stage_meta.stage_type, req_id, engine_outputs)
metrics.on_stage_metrics(stage_id, req_id, _m, output_type) metrics.on_stage_metrics(stage_id, req_id, _m, output_type)
consumed.add(msg_id) consumed.add(msg_id)
@@ -502,11 +502,17 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
def _get_tts_adapter(self): def _get_tts_adapter(self):
"""Return the per-model serving adapter for the current ``_tts_model_type``. """Return the per-model serving adapter for the current ``_tts_model_type``.
Pure-diffusion speech uses its dedicated request path and does not
resolve adapters for AR-stage TTS models.
Resolved lazily (rebuilt if ``_tts_model_type`` changed since the cached Resolved lazily (rebuilt if ``_tts_model_type`` changed since the cached
instance was built) so callers that set ``_tts_model_type`` after instance was built) so callers that set ``_tts_model_type`` after
construction still dispatch to the matching adapter. In production construction still dispatch to the matching adapter. In production
``_tts_model_type`` is fixed at init, so the cached instance is reused. ``_tts_model_type`` is fixed at init, so the cached instance is reused.
""" """
if self._diffusion_mode:
return None
adapter_cls = resolve_adapter(self._tts_model_type) adapter_cls = resolve_adapter(self._tts_model_type)
if adapter_cls is None: if adapter_cls is None:
self._adapter = None self._adapter = None
+9
View File
@@ -134,6 +134,15 @@ AUDIO_CONTINUITY_OK_METRIC = METRIC_PREFIX + AUDIO_CONTINUITY_OK
AUDIO_SKIPPED_REQUESTS_METRIC = METRIC_PREFIX + AUDIO_SKIPPED_REQUESTS AUDIO_SKIPPED_REQUESTS_METRIC = METRIC_PREFIX + AUDIO_SKIPPED_REQUESTS
# ============================================================================
# Diffusion family (per-stage + per-replica diffusion timing breakdowns)
# ============================================================================
DIFFUSION_EXEC_S = METRIC_PREFIX + "diffusion_exec_s"
DIFFUSION_EXEC_PER_STEP_S = METRIC_PREFIX + "diffusion_exec_per_step_s"
DIFFUSION_PREPROCESS_S = METRIC_PREFIX + "diffusion_preprocess_s"
DIFFUSION_POSTPROCESS_S = METRIC_PREFIX + "diffusion_postprocess_s"
# ============================================================================ # ============================================================================
# Cross-stage Transfer family (per-physical-hop TX/RX/in-flight timings) # Cross-stage Transfer family (per-physical-hop TX/RX/in-flight timings)
# ============================================================================ # ============================================================================
+104 -31
View File
@@ -78,6 +78,35 @@ _audio_skipped_family = Counter(
) )
# ----------------------------------------------------------------------------
# Diffusion family
# ----------------------------------------------------------------------------
_diffusion_exec_family = Histogram(
defs.DIFFUSION_EXEC_S,
"DiT forward pass execution time per request in seconds.",
labelnames=_stage_labels,
buckets=defs.SECONDS_BUCKETS,
)
_diffusion_exec_per_step_family = Histogram(
defs.DIFFUSION_EXEC_PER_STEP_S,
"DiT forward pass execution time per denoising step in seconds.",
labelnames=_stage_labels,
buckets=defs.SECONDS_FAST_BUCKETS,
)
_diffusion_preprocess_family = Histogram(
defs.DIFFUSION_PREPROCESS_S,
"Diffusion input preprocessing time per request in seconds.",
labelnames=_stage_labels,
buckets=defs.SECONDS_FAST_BUCKETS,
)
_diffusion_postprocess_family = Histogram(
defs.DIFFUSION_POSTPROCESS_S,
"Diffusion output postprocessing (VAE decode) time per request in seconds.",
labelnames=_stage_labels,
buckets=defs.SECONDS_FAST_BUCKETS,
)
class OmniModalityMetrics: class OmniModalityMetrics:
"""Per-modality observe API. Stage/replica are passed at observe time """Per-modality observe API. Stage/replica are passed at observe time
because a single OmniModalityMetrics instance per pipeline serves all because a single OmniModalityMetrics instance per pipeline serves all
@@ -139,6 +168,44 @@ class OmniModalityMetrics:
reason=reason or "unknown", reason=reason or "unknown",
).inc() ).inc()
# ---- Diffusion --------------------------------------------------------
def observe_diffusion_exec(self, stage: str, replica: str, seconds: float) -> None:
if not self._log_stats:
return
_diffusion_exec_family.labels(
model_name=self._model_name,
stage=stage,
replica=replica,
).observe(seconds)
def observe_diffusion_exec_per_step(self, stage: str, replica: str, seconds: float) -> None:
if not self._log_stats:
return
_diffusion_exec_per_step_family.labels(
model_name=self._model_name,
stage=stage,
replica=replica,
).observe(seconds)
def observe_diffusion_preprocess(self, stage: str, replica: str, seconds: float) -> None:
if not self._log_stats:
return
_diffusion_preprocess_family.labels(
model_name=self._model_name,
stage=stage,
replica=replica,
).observe(seconds)
def observe_diffusion_postprocess(self, stage: str, replica: str, seconds: float) -> None:
if not self._log_stats:
return
_diffusion_postprocess_family.labels(
model_name=self._model_name,
stage=stage,
replica=replica,
).observe(seconds)
def observe_modality_at_finalize( def observe_modality_at_finalize(
mod_metrics: OmniModalityMetrics, mod_metrics: OmniModalityMetrics,
@@ -161,43 +228,49 @@ def observe_modality_at_finalize(
""" """
if replica_id is None or stage_metrics is None or output_type is None: if replica_id is None or stage_metrics is None or output_type is None:
return return
if output_type != "audio":
return
stage_label = str(stage_id) stage_label = str(stage_id)
replica_label = str(replica_id) replica_label = str(replica_id)
gen_time_s = float(getattr(stage_metrics, "stage_gen_time_ms", 0.0)) / 1000.0 gen_time_s = float(getattr(stage_metrics, "stage_gen_time_ms", 0.0)) / 1000.0
mm_out = extract_mm_output(engine_outputs) mm_out = extract_mm_output(engine_outputs)
sample_rate = defs.resolve_audio_sample_rate(mm_out) if output_type == "audio":
# `stage_metrics.audio_generated_frames` is the legacy per-chunk sample_rate = defs.resolve_audio_sample_rate(mm_out)
# accumulator field on StageRequestStats. No production path currently n_frames = int(getattr(stage_metrics, "audio_generated_frames", 0) or 0)
# fills it, so the fallback below is the live source — but we leave the if n_frames == 0:
# field lookup in place in case the accumulator gets re-wired upstream. n_frames = count_audio_frames(mm_out)
n_frames = int(getattr(stage_metrics, "audio_generated_frames", 0) or 0) mod_metrics.inc_audio_frames(stage_label, replica_label, n_frames)
if n_frames == 0: duration_s = n_frames / sample_rate if sample_rate > 0 else 0.0
n_frames = count_audio_frames(mm_out) if duration_s > 0:
mod_metrics.inc_audio_frames(stage_label, replica_label, n_frames) mod_metrics.observe_audio_duration(stage_label, replica_label, duration_s)
duration_s = n_frames / sample_rate if sample_rate > 0 else 0.0 mod_metrics.observe_audio_rtf(
if duration_s > 0: stage_label,
mod_metrics.observe_audio_duration(stage_label, replica_label, duration_s) replica_label,
mod_metrics.observe_audio_rtf( defs.compute_audio_rtf(gen_time_s, duration_s),
stage_label, )
replica_label, else:
defs.compute_audio_rtf(gen_time_s, duration_s), mod_metrics.inc_audio_skipped(stage_label, replica_label, "no_audio_data")
)
else: dm = getattr(stage_metrics, "diffusion_metrics", None)
# Request completed (finish_reason ∈ {stop, length} — error paths if dm:
# don't reach finalize) but no audio samples were produced. Covers _key_map = {
# silent `return None` skips in the talker→code2wav stage "diffusion_engine_exec_time_s": mod_metrics.observe_diffusion_exec,
# processors and the `parsed.append((0,0))` malformed-length path "preprocess_time_s": mod_metrics.observe_diffusion_preprocess,
# in qwen3-tts code2wav. raise-paths surface via the upstream "postprocess_time_s": mod_metrics.observe_diffusion_postprocess,
# vllm:request_success_total{finished_reason="error"} channel and }
# never reach this branch. for key, observe_fn in _key_map.items():
mod_metrics.inc_audio_skipped(stage_label, replica_label, "no_audio_data") val = dm.get(key)
# audio_underrun / continuity are emitted from the streaming path in if val is not None:
# observe_audio_streaming_finalize; finalize is too late for the observe_fn(stage_label, replica_label, float(val))
# per-chunk timeline they need.
exec_time = dm.get("diffusion_engine_exec_time_s")
num_steps = dm.get("num_inference_steps")
if exec_time is not None and num_steps and num_steps > 0:
mod_metrics.observe_diffusion_exec_per_step(
stage_label,
replica_label,
float(exec_time) / int(num_steps),
)
def observe_audio_first_packet( def observe_audio_first_packet(
+14 -3
View File
@@ -44,7 +44,7 @@ class StageRequestStats:
final_output_type: str | None = None final_output_type: str | None = None
request_id: str | None = None request_id: str | None = None
postprocess_time_ms: float = 0.0 postprocess_time_ms: float = 0.0
diffusion_metrics: dict[str, int] = None diffusion_metrics: dict[str, float] = None
audio_generated_frames: int = 0 audio_generated_frames: int = 0
audio_sample_rate: int = 0 audio_sample_rate: int = 0
audio_duration_s: float = 0.0 audio_duration_s: float = 0.0
@@ -579,7 +579,7 @@ class OrchestratorAggregator:
if final_output_type is not None: if final_output_type is not None:
stats.final_output_type = final_output_type stats.final_output_type = final_output_type
stats.diffusion_metrics = ( stats.diffusion_metrics = (
{k: int(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()} {k: float(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()}
if req_id in self.diffusion_metrics if req_id in self.diffusion_metrics
else None else None
) )
@@ -644,8 +644,19 @@ class OrchestratorAggregator:
if isinstance(diffusion_metrics, list): if isinstance(diffusion_metrics, list):
diffusion_metrics = diffusion_metrics[0] diffusion_metrics = diffusion_metrics[0]
if diffusion_metrics: if diffusion_metrics:
_MS_TO_S = {
"diffusion_engine_exec_time_ms": "diffusion_engine_exec_time_s",
"preprocess_time_ms": "preprocess_time_s",
"postprocess_time_ms": "postprocess_time_s",
"diffusion_engine_total_time_ms": "diffusion_engine_total_time_s",
}
for key, value in diffusion_metrics.items(): for key, value in diffusion_metrics.items():
self.diffusion_metrics[req_id][key] += value if value is None:
continue
if key in _MS_TO_S:
self.diffusion_metrics[req_id][_MS_TO_S[key]] = float(value) / 1000.0
else:
self.diffusion_metrics[req_id][key] = value
def on_forward( def on_forward(
self, self,
@@ -332,7 +332,7 @@ class DiTAttention(nn.Module):
self, self,
x: torch.Tensor, x: torch.Tensor,
padding_mask: torch.Tensor | None = None, padding_mask: torch.Tensor | None = None,
rope=None, rope: tuple[torch.Tensor, torch.Tensor | float] | None = None,
attn_mask: torch.Tensor | None = None, attn_mask: torch.Tensor | None = None,
) -> torch.Tensor: ) -> torch.Tensor:
batch_size, seq_len = x.shape[0], x.shape[1] batch_size, seq_len = x.shape[0], x.shape[1]
@@ -406,7 +406,7 @@ class DiTBlock(nn.Module):
x: torch.Tensor, x: torch.Tensor,
t: torch.Tensor, t: torch.Tensor,
padding_mask: torch.Tensor | None = None, padding_mask: torch.Tensor | None = None,
rope=None, rope: tuple[torch.Tensor, torch.Tensor | float] | None = None,
attn_mask: torch.Tensor | None = None, attn_mask: torch.Tensor | None = None,
) -> torch.Tensor: ) -> torch.Tensor:
# Pre-norm & modulation for attention # Pre-norm & modulation for attention
@@ -575,6 +575,8 @@ class GLMTTSDiT(nn.Module):
is_causal: bool = False, is_causal: bool = False,
attn_mask: torch.Tensor | None = None, attn_mask: torch.Tensor | None = None,
block_pattern: list[int] | None = None, block_pattern: list[int] | None = None,
text_embed: torch.Tensor | None = None,
rope: tuple[torch.Tensor, torch.Tensor | float] | None = None,
) -> torch.Tensor: ) -> torch.Tensor:
"""Forward pass. """Forward pass.
@@ -588,6 +590,9 @@ class GLMTTSDiT(nn.Module):
is_causal: Use block-causal attention is_causal: Use block-causal attention
attn_mask: Custom attention mask (overrides is_causal) attn_mask: Custom attention mask (overrides is_causal)
block_pattern: Token-level block sizes for block-causal mask block_pattern: Token-level block sizes for block-causal mask
text_embed: Precomputed text embedding (must already be in model
dtype); skips the internal embedding + ConvNeXt pass
rope: Precomputed rotary embedding tuple for ``seq_len``
Returns: Returns:
Predicted mel-spectrogram [B, T, mel_dim] Predicted mel-spectrogram [B, T, mel_dim]
@@ -610,13 +615,21 @@ class GLMTTSDiT(nn.Module):
time_emb = torch.cat([time_emb, spkr_emb], dim=-1) time_emb = torch.cat([time_emb, spkr_emb], dim=-1)
# Text embedding - ensure cast to model dtype # Text embedding - ensure cast to model dtype
text_embed = self.text_emb_layer(text, seq_len, text_lens=text_lens).to(model_dtype) if text_embed is None:
text_embed = self.text_emb_layer(text, seq_len, text_lens=text_lens).to(model_dtype)
else:
# A precomputed text_embed replaces `text` entirely, so its length must
# match middle_point; a mismatch would apply the wrong rotary/mask silently.
assert text_embed.shape[1] == seq_len, (
f"precomputed text_embed len {text_embed.shape[1]} != seq_len {seq_len}"
)
# Input projection # Input projection
x = self.emb_concator(middle_point, condition, text_embed, drop_audio_cond=False) x = self.emb_concator(middle_point, condition, text_embed, drop_audio_cond=False)
# Rotary embeddings # Rotary embeddings
rope = self.rotary_embed.forward_from_seq_len(seq_len) if rope is None:
rope = self.rotary_embed.forward_from_seq_len(seq_len)
# Build block-causal mask when is_causal and no explicit attn_mask # Build block-causal mask when is_causal and no explicit attn_mask
if attn_mask is None and is_causal and block_pattern is not None: if attn_mask is None and is_causal and block_pattern is not None:
@@ -857,6 +857,40 @@ class GLMTTSDiTForGeneration(nn.Module):
if self.t_scheduler == "cosine": if self.t_scheduler == "cosine":
t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) t_span = 1 - torch.cos(t_span * 0.5 * torch.pi)
# Hoist loop-invariant DiT inputs (text embedding, RoPE table,
# block-causal mask, CFG double-batch conditioning) out of the Euler
# loop; they depend only on the conditioning and seq_len, not on the
# per-step state x or the timestep.
seq_len = int(mel_cond.shape[1])
# Must match the dtype forward() casts with internally (both are
# proj_out.weight.dtype); keep them in sync if that derivation changes.
dit_dtype = self.dit.proj_out.weight.dtype
# `device` is speech_tokens.device; in this call path the sampled state x
# lives on the same device, so the hoisted mask matches the old in-forward one.
attn_mask = (
self.dit._create_block_causal_mask(block_pattern, seq_len, device, dit_dtype)
if block_pattern is not None
else None
)
rope = self.dit.rotary_embed.forward_from_seq_len(seq_len)
if self.inference_cfg_rate > 0:
condition_b = torch.cat([condition_dit, torch.zeros_like(condition_dit)], dim=0)
if self.speech_token_cfg:
text_b = torch.cat([speech_tokens, torch.zeros_like(speech_tokens)], dim=0)
else:
text_b = speech_tokens.expand(2, -1)
padding_b = padding_mask.expand(2, -1)
if spkr_embedding_dit is not None and self.spkr_emb_adaln:
spkr_b = torch.cat([spkr_embedding_dit, torch.zeros_like(spkr_embedding_dit)], dim=0)
else:
spkr_b = None
text_lens_b = text_lens.expand(2) if text_lens is not None else None
text_embed_b = self.dit.text_emb_layer(text_b, seq_len, text_lens=text_lens_b).to(dit_dtype)
else:
spkr_uncond = spkr_embedding_dit if (spkr_embedding_dit is not None and self.spkr_emb_adaln) else None
text_embed_u = self.dit.text_emb_layer(speech_tokens, seq_len, text_lens=text_lens).to(dit_dtype)
for step in range(n_timesteps): for step in range(n_timesteps):
cache_step = step + 1 cache_step = step + 1
if last_step_cache is not None and cache_step in last_step_cache: if last_step_cache is not None and cache_step in last_step_cache:
@@ -872,24 +906,12 @@ class GLMTTSDiTForGeneration(nn.Module):
t_current = t_span[step] t_current = t_span[step]
dt = t_span[step + 1] - t_current dt = t_span[step + 1] - t_current
use_causal = block_pattern is not None
x_dit = x.to(mel_cond.dtype) x_dit = x.to(mel_cond.dtype)
if self.inference_cfg_rate > 0: if self.inference_cfg_rate > 0:
# Batch conditional + unconditional into a single forward # Batch conditional + unconditional into a single forward
middle_b = x_dit.expand(2, -1, -1) middle_b = x_dit.expand(2, -1, -1)
condition_b = torch.cat([condition_dit, torch.zeros_like(condition_dit)], dim=0)
if self.speech_token_cfg:
text_b = torch.cat([speech_tokens, torch.zeros_like(speech_tokens)], dim=0)
else:
text_b = speech_tokens.expand(2, -1)
time_b = t_current.unsqueeze(0).expand(2) time_b = t_current.unsqueeze(0).expand(2)
padding_b = padding_mask.expand(2, -1)
if spkr_embedding_dit is not None and self.spkr_emb_adaln:
spkr_b = torch.cat([spkr_embedding_dit, torch.zeros_like(spkr_embedding_dit)], dim=0)
else:
spkr_b = None
text_lens_b = text_lens.expand(2) if text_lens is not None else None
pred_b = self.dit( pred_b = self.dit(
middle_point=middle_b, middle_point=middle_b,
condition=condition_b, condition=condition_b,
@@ -897,9 +919,10 @@ class GLMTTSDiTForGeneration(nn.Module):
time_step=time_b, time_step=time_b,
padding_mask=padding_b, padding_mask=padding_b,
spkr_emb=spkr_b, spkr_emb=spkr_b,
is_causal=use_causal, attn_mask=attn_mask,
block_pattern=block_pattern,
text_lens=text_lens_b, text_lens=text_lens_b,
text_embed=text_embed_b,
rope=rope,
).to(torch.float32) ).to(torch.float32)
dphi_dt, cfg_dphi_dt = pred_b.chunk(2, dim=0) dphi_dt, cfg_dphi_dt = pred_b.chunk(2, dim=0)
dphi_dt = (1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt dphi_dt = (1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt
@@ -910,10 +933,11 @@ class GLMTTSDiTForGeneration(nn.Module):
text=speech_tokens, text=speech_tokens,
time_step=t_current.unsqueeze(0), time_step=t_current.unsqueeze(0),
padding_mask=padding_mask, padding_mask=padding_mask,
spkr_emb=spkr_embedding_dit if spkr_embedding_dit is not None and self.spkr_emb_adaln else None, spkr_emb=spkr_uncond,
is_causal=use_causal, attn_mask=attn_mask,
block_pattern=block_pattern,
text_lens=text_lens, text_lens=text_lens,
text_embed=text_embed_u,
rope=rope,
).to(torch.float32) ).to(torch.float32)
x = x + dt * dphi_dt x = x + dt * dphi_dt
@@ -14,7 +14,7 @@ import torch.nn as nn
import torch.nn.functional as F import torch.nn.functional as F
from vllm.logger import init_logger from vllm.logger import init_logger
from .cuda_graph_wrapper import HiFTGraphWrapper from .cuda_graph_wrapper import CFMGraphWrapper, HiFTGraphWrapper
logger = init_logger(__name__) logger = init_logger(__name__)
@@ -73,6 +73,7 @@ class BatchedToken2Wav(nn.Module):
*, *,
connector_config: Mapping[str, int] | None = None, connector_config: Mapping[str, int] | None = None,
hift_graph_config: Mapping[str, Any] | None = None, hift_graph_config: Mapping[str, Any] | None = None,
cfm_graph_config: Mapping[str, Any] | None = None,
): ):
super().__init__() super().__init__()
self._token2wav = token2wav self._token2wav = token2wav
@@ -135,6 +136,22 @@ class BatchedToken2Wav(nn.Module):
with torch.inference_mode(), _autocast_disabled(hift_parameter.device): with torch.inference_mode(), _autocast_disabled(hift_parameter.device):
self.hift_graph_wrapper.capture() self.hift_graph_wrapper.capture()
logger.info("HiFT CUDA Graph captured successfully") logger.info("HiFT CUDA Graph captured successfully")
self._cfm_graph_wrapper: CFMGraphWrapper | None = None
cfm_graph_cfg = dict(cfm_graph_config or {})
if bool(cfm_graph_cfg.get("enabled", False)):
flow_parameter = next(self.flow.parameters(), None)
if flow_parameter is not None and flow_parameter.device.type == "cuda":
estimator = self.flow.decoder.estimator
self._cfm_graph_wrapper = CFMGraphWrapper(
graph_fn=estimator.blocks_forward_chunk,
max_graphs=int(cfm_graph_cfg.get("max_graphs", 32)),
)
logger.info("CFM CUDA Graph enabled (max_graphs=%d)", int(cfm_graph_cfg.get("max_graphs", 32)))
else:
logger.info(
"CFM CUDA Graph is disabled on device type %s",
flow_parameter.device.type if flow_parameter is not None else "unknown",
)
self._prompt_features: dict[tuple[str, str], PromptFeatures] = {} self._prompt_features: dict[tuple[str, str], PromptFeatures] = {}
def _hift_inference( def _hift_inference(
@@ -265,6 +282,19 @@ class BatchedToken2Wav(nn.Module):
speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width) speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width)
estimator_input = torch.cat((x, mu, speaker_features, cond), dim=1) estimator_input = torch.cat((x, mu, speaker_features, cond), dim=1)
cnn_out, att_out = self._estimator_buffers(estimator, estimator_input, att_cache) cnn_out, att_out = self._estimator_buffers(estimator, estimator_input, att_cache)
if self._cfm_graph_wrapper is not None:
if cnn_cache is None:
cnn_cache = torch.zeros_like(cnn_out)
if att_cache is None:
att_cache = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:])
return self._cfm_graph_wrapper.replay(
estimator_input,
time_embedding,
cnn_cache,
att_cache,
cnn_out,
att_out,
)
old_cnn: Any = cnn_cache if cnn_cache is not None else [None] * len(estimator.blocks) old_cnn: Any = cnn_cache if cnn_cache is not None else [None] * len(estimator.blocks)
old_att: Any = att_cache if att_cache is not None else [None] * len(estimator.blocks) old_att: Any = att_cache if att_cache is not None else [None] * len(estimator.blocks)
result = estimator.blocks_forward_chunk( result = estimator.blocks_forward_chunk(
@@ -1,6 +1,8 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from functools import lru_cache
import torch import torch
from torch.cuda import CUDAGraph from torch.cuda import CUDAGraph
from vllm.logger import init_logger from vllm.logger import init_logger
@@ -133,3 +135,126 @@ class HiFTGraphWrapper:
cache_source = static_cache_source_output[:batch_size].clone() cache_source = static_cache_source_output[:batch_size].clone()
speech = self.finalize_fn(static_magnitude_output[:batch_size], static_phase_output[:batch_size]).clone() speech = self.finalize_fn(static_magnitude_output[:batch_size], static_phase_output[:batch_size]).clone()
return speech, cache_source return speech, cache_source
def _tensor_signature(value: torch.Tensor) -> tuple:
return tuple(value.shape), str(value.dtype), str(value.device)
_DTYPE_MAP = {
"torch.float32": torch.float32,
"torch.float16": torch.float16,
"torch.bfloat16": torch.bfloat16,
"torch.float64": torch.float64,
}
def _tensors_from_key(key: tuple) -> tuple[torch.Tensor, ...]:
"""Rebuild zero tensors from a cache key (shape, dtype, device tuples)."""
tensors = []
for shape, dtype_str, device_str in key[1:]:
dtype = _DTYPE_MAP.get(dtype_str, torch.float32)
device = torch.device(device_str)
tensors.append(torch.zeros(shape, dtype=dtype, device=device))
return tuple(tensors)
class CFMGraphWrapper:
"""Per-shape CUDA graph capture/replay for the CFM DiT estimator.
Captures one blocks_forward_chunk call (in_proj -> DiT blocks -> final_layer)
as the graph target. The 10-step Euler loop stays in Python, replaying
the graph 10 times per decode.
Uses functools.lru_cache for automatic LRU eviction. Cache misses
trigger capture; capture failures fall back to eager. Outputs are
cloned after replay to prevent streaming cache corruption.
"""
def __init__(self, graph_fn, *, max_graphs: int = 32) -> None:
self.graph_fn = graph_fn
self.max_graphs = int(max_graphs)
self.device = next(graph_fn.__self__.parameters()).device
@lru_cache(maxsize=self.max_graphs)
def _capture_graph(key: tuple):
return self._capture(key)
self._capture_graph = _capture_graph
def _call_graph_fn(self, args: tuple[torch.Tensor, ...]) -> torch.Tensor:
return self.graph_fn(args[0], args[1], None, args[2], args[3], args[4], args[5])
def _capture(self, key: tuple) -> tuple | None:
"""Capture a CUDA graph for the given key. Returns None on failure."""
static_inputs = _tensors_from_key(key)
current_stream = torch.cuda.current_stream(self.device)
warmup_stream = torch.cuda.Stream(device=self.device)
warmup_stream.wait_stream(current_stream)
with torch.cuda.stream(warmup_stream), torch.no_grad():
for _ in range(3):
self._call_graph_fn(static_inputs)
current_stream.wait_stream(warmup_stream)
try:
graph = CUDAGraph()
with torch.no_grad(), torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()):
static_output = self._call_graph_fn(static_inputs)
except Exception:
logger.warning(
"CFM graph capture failed for shape=%s; using eager",
key,
exc_info=True,
)
return None
logger.info(
"Captured CFM CUDA Graph for shape %s (cache=%d/%d, hits=%d, misses=%d)",
key,
self._capture_graph.cache_info().currsize,
self.max_graphs,
self._capture_graph.cache_info().hits,
self._capture_graph.cache_info().misses,
)
return (static_inputs, static_output, graph)
def replay(
self,
estimator_input: torch.Tensor,
time_emb: torch.Tensor,
cnn_cache: torch.Tensor,
att_cache: torch.Tensor,
cnn_out: torch.Tensor,
att_out: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
inputs = (estimator_input, time_emb, cnn_cache, att_cache, cnn_out, att_out)
if torch.cuda.is_current_stream_capturing() or estimator_input.device.type != "cuda":
with torch.no_grad():
result = self._call_graph_fn(inputs)
return result, inputs[4], inputs[5]
key = ("estimator_step",) + tuple(_tensor_signature(v) for v in inputs)
entry = self._capture_graph(key)
if entry is None:
logger.debug(
"CFM graph eager fallback for shape=%s (hits=%d, misses=%d)",
key,
self._capture_graph.cache_info().hits,
self._capture_graph.cache_info().misses,
)
with torch.no_grad():
result = self._call_graph_fn(inputs)
return result, inputs[4], inputs[5]
static_inputs, static_output, graph = entry
for static, current in zip(static_inputs, inputs, strict=True):
static.copy_(current)
graph.replay()
return (
static_output.detach().clone(),
static_inputs[4].detach().clone(),
static_inputs[5].detach().clone(),
)
@@ -162,6 +162,10 @@ class MiniCPMO45Code2Wav(nn.Module):
"enabled": bool(extra.get("enable_hift_graph", False)), "enabled": bool(extra.get("enable_hift_graph", False)),
"capture_batch_sizes": capture_batch_sizes, "capture_batch_sizes": capture_batch_sizes,
} }
self._cfm_graph_config = {
"enabled": bool(extra.get("enable_cfm_graph", False)),
"max_graphs": int(extra.get("cfm_max_graphs", 32)),
}
self._min_batch_size = int(extra.get("code2wav_min_batch_size", 1)) self._min_batch_size = int(extra.get("code2wav_min_batch_size", 1))
if self._min_batch_size < 1: if self._min_batch_size < 1:
raise ValueError("MiniCPM-o Code2Wav code2wav_min_batch_size must be >= 1") raise ValueError("MiniCPM-o Code2Wav code2wav_min_batch_size must be >= 1")
@@ -802,4 +806,5 @@ class MiniCPMO45Code2Wav(nn.Module):
trt_stepper=trt_stepper, trt_stepper=trt_stepper,
connector_config=self._connector_config, connector_config=self._connector_config,
hift_graph_config=self._hift_graph_config, hift_graph_config=self._hift_graph_config,
cfm_graph_config=self._cfm_graph_config,
) )
@@ -242,6 +242,7 @@ def _gumbel_sample(logits: torch.Tensor, temperature: float, generator: torch.Ge
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Subclass keeps .weight name + ctor shape so the state_dict loader stays unchanged.
class OmniVoiceRMSNorm(nn.Module): class OmniVoiceRMSNorm(nn.Module):
def __init__(self, hidden_size: int, eps: float = 1e-6): def __init__(self, hidden_size: int, eps: float = 1e-6):
super().__init__() super().__init__()
@@ -315,10 +316,9 @@ class OmniVoiceAttention(nn.Module):
k = k.permute(0, 2, 1, 3) k = k.permute(0, 2, 1, 3)
v = v.permute(0, 2, 1, 3) v = v.permute(0, 2, 1, 3)
# Convert [B, 1, S, S] bool mask to float mask for SDPA # Caller passes a float mask; materialize float form if a bool slips through.
# (0.0 where attend, -inf where masked) sdpa_mask = attention_mask
sdpa_mask = None if sdpa_mask is not None and sdpa_mask.dtype == torch.bool:
if attention_mask is not None:
sdpa_mask = torch.zeros_like(attention_mask, dtype=q.dtype).masked_fill_(~attention_mask, float("-inf")) sdpa_mask = torch.zeros_like(attention_mask, dtype=q.dtype).masked_fill_(~attention_mask, float("-inf"))
out = F.scaled_dot_product_attention( out = F.scaled_dot_product_attention(
@@ -400,24 +400,48 @@ def _precompute_rope(
theta: float = 1000000.0, theta: float = 1000000.0,
device: torch.device | None = None, device: torch.device | None = None,
) -> tuple[torch.Tensor, torch.Tensor]: ) -> tuple[torch.Tensor, torch.Tensor]:
"""Precompute RoPE cos/sin tensors.""" """Precompute RoPE cos/sin pre-doubled to full head_dim so the hot path skips a per-call cat."""
inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim)) inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim))
t = torch.arange(max_seq_len, device=device, dtype=torch.float32) t = torch.arange(max_seq_len, device=device, dtype=torch.float32)
freqs = torch.outer(t, inv_freq) freqs = torch.outer(t, inv_freq)
cos = freqs.cos() cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1)
sin = freqs.sin() sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1)
return cos, sin return cos, sin
def _apply_rotary_pos_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: def _apply_rotary_pos_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
"""Apply rotary position embedding. x shape: (B, S, H, D).""" """Apply rotary position embedding. x shape: (B, S, H, D). cos/sin are pre-doubled to width D."""
seq_len = x.shape[1] seq_len = x.shape[1]
cos = cos[:seq_len].unsqueeze(0).unsqueeze(2) # (1, S, 1, D/2) cos = cos[:seq_len].unsqueeze(0).unsqueeze(2) # (1, S, 1, D)
sin = sin[:seq_len].unsqueeze(0).unsqueeze(2) sin = sin[:seq_len].unsqueeze(0).unsqueeze(2)
x1 = x[..., : x.shape[-1] // 2] x1 = x[..., : x.shape[-1] // 2]
x2 = x[..., x.shape[-1] // 2 :] x2 = x[..., x.shape[-1] // 2 :]
rotated = torch.cat([-x2, x1], dim=-1) rotated = torch.cat([-x2, x1], dim=-1)
return x * torch.cat([cos, cos], dim=-1) + rotated * torch.cat([sin, sin], dim=-1) return x * cos + rotated * sin
# ---------------------------------------------------------------------------
# TF32 opt-in (process-wide; default off)
# ---------------------------------------------------------------------------
_TF32_ENABLED = False
def _maybe_enable_tf32() -> None:
"""Enable TF32 matmuls process-wide (idempotent). Not bit-identical; opt-in via config.enable_tf32."""
global _TF32_ENABLED
if _TF32_ENABLED or not torch.cuda.is_available():
return
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.set_float32_matmul_precision("high")
_TF32_ENABLED = True
logger.info(
"OmniVoice TF32 enabled process-wide: matmul.allow_tf32=%s cudnn.allow_tf32=%s float32_matmul_precision=%s",
torch.backends.cuda.matmul.allow_tf32,
torch.backends.cudnn.allow_tf32,
torch.get_float32_matmul_precision(),
)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -425,6 +449,19 @@ def _apply_rotary_pos_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
def _additive_float_mask(mask: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor:
"""Convert a boolean attention mask to its additive float form.
``True`` (attend) maps to ``0.0`` and ``False`` (masked) to ``-inf``. A bool
mask must never be copied straight into a float buffer: the implicit cast
maps True/False to 1.0/0.0, which leaves masked positions at 0.0 and so
silently *unmasks* them.
"""
if mask.dtype != torch.bool:
return mask
return torch.zeros_like(mask, dtype=dtype).masked_fill_(~mask, float("-inf"))
class _OmniVoiceCUDAGraphForward: class _OmniVoiceCUDAGraphForward:
"""Pre-captures CUDA graphs for predefined sequence-length buckets. """Pre-captures CUDA graphs for predefined sequence-length buckets.
@@ -476,11 +513,10 @@ class _OmniVoiceCUDAGraphForward:
mask_padded[:, :S] = audio_mask mask_padded[:, :S] = audio_mask
if attention_mask is not None: if attention_mask is not None:
attn_padded = torch.zeros( # Callers normalize to the additive float form first, so pad with -inf.
two_b, attn_padded = torch.full(
1, (two_b, 1, bucket, bucket),
bucket, float("-inf"),
bucket,
dtype=attention_mask.dtype, dtype=attention_mask.dtype,
device=attention_mask.device, device=attention_mask.device,
) )
@@ -564,7 +600,8 @@ class _OmniVoiceCUDAGraphForward:
key = (two_b, bucket) key = (two_b, bucket)
dummy_ids = torch.zeros(two_b, num_cb, bucket, dtype=torch.long, device=device) dummy_ids = torch.zeros(two_b, num_cb, bucket, dtype=torch.long, device=device)
dummy_mask = torch.zeros(two_b, bucket, dtype=torch.bool, device=device) dummy_mask = torch.zeros(two_b, bucket, dtype=torch.bool, device=device)
dummy_attn = torch.ones(two_b, 1, bucket, bucket, dtype=torch.bool, device=device) # Capture with a float mask to match what forward() feeds at replay time.
dummy_attn = torch.zeros(two_b, 1, bucket, bucket, dtype=torch.float32, device=device)
self._graphs[key] = self._capture_for_key(key, dummy_ids, dummy_mask, dummy_attn) self._graphs[key] = self._capture_for_key(key, dummy_ids, dummy_mask, dummy_attn)
logger.info("OmniVoice CUDA Graph warmup complete (%d graphs)", len(self._graphs)) logger.info("OmniVoice CUDA Graph warmup complete (%d graphs)", len(self._graphs))
@@ -586,6 +623,11 @@ class _OmniVoiceCUDAGraphForward:
two_b = input_ids.shape[0] two_b = input_ids.shape[0]
bucket = self._find_bucket(seq_len) if two_b == 2 else None bucket = self._find_bucket(seq_len) if two_b == 2 else None
# Graphs are captured with (and their static buffers hold) the additive
# float mask, so normalize here, before padding or any copy_ into them.
if attention_mask is not None:
attention_mask = _additive_float_mask(attention_mask)
if bucket is None: if bucket is None:
# Lazy capture: oversized sequence or non-unit batch (no pre-warmed bucket). # Lazy capture: oversized sequence or non-unit batch (no pre-warmed bucket).
# Lock prevents concurrent threads from double-capturing the same key. # Lock prevents concurrent threads from double-capturing the same key.
@@ -655,6 +697,10 @@ class OmniVoiceGenerator(nn.Module):
super().__init__() super().__init__()
self.config = config self.config = config
# Opt-in TF32; must run before any CUDA-graph capture so captured kernels honour it.
if getattr(config, "enable_tf32", False):
_maybe_enable_tf32()
# Text embedding (shared with LLM) # Text embedding (shared with LLM)
self.text_embedding = nn.Embedding(config.llm_vocab_size, config.llm_hidden_size) self.text_embedding = nn.Embedding(config.llm_vocab_size, config.llm_hidden_size)
@@ -703,6 +749,8 @@ class OmniVoiceGenerator(nn.Module):
self, self,
input_ids: torch.Tensor, input_ids: torch.Tensor,
audio_mask: torch.Tensor, audio_mask: torch.Tensor,
text_embeds: torch.Tensor | None = None,
audio_mask_3d: torch.Tensor | None = None,
) -> torch.Tensor: ) -> torch.Tensor:
"""Prepare mixed text+audio embeddings. """Prepare mixed text+audio embeddings.
@@ -710,30 +758,38 @@ class OmniVoiceGenerator(nn.Module):
input_ids: [B, 8, S] - text tokens replicated across codebooks, input_ids: [B, 8, S] - text tokens replicated across codebooks,
audio positions have per-codebook token IDs audio positions have per-codebook token IDs
audio_mask: [B, S] - True for audio positions, False for text audio_mask: [B, S] - True for audio positions, False for text
text_embeds: optional cached [B, S, H] text-position embeddings
audio_mask_3d: optional cached [B, S, 1] audio_mask.unsqueeze(-1)
Returns: Returns:
embeddings: [B, S, hidden_size] embeddings: [B, S, hidden_size]
""" """
# Text embeddings from first codebook row (all rows identical for text) # Cached across the denoising loop since text ids don't change.
text_embeds = self.text_embedding(input_ids[:, 0, :]) if text_embeds is None:
text_embeds = self.text_embedding(input_ids[:, 0, :])
if audio_mask_3d is None:
audio_mask_3d = audio_mask.unsqueeze(-1)
# Audio embeddings: offset per codebook, then sum across codebooks # Audio embeddings: offset per codebook, then sum across codebooks
shifted_ids = (input_ids * audio_mask.unsqueeze(1)) + self.codebook_layer_offsets.view(1, -1, 1) shifted_ids = (input_ids * audio_mask.unsqueeze(1)) + self.codebook_layer_offsets.view(1, -1, 1)
audio_embeds = self.audio_embeddings(shifted_ids).sum(dim=1) audio_embeds = self.audio_embeddings(shifted_ids).sum(dim=1)
# Merge: audio where audio_mask=True, text elsewhere # Merge: audio where audio_mask=True, text elsewhere
return torch.where(audio_mask.unsqueeze(-1), audio_embeds, text_embeds) return torch.where(audio_mask_3d, audio_embeds, text_embeds)
def _transformer_forward( def _transformer_forward(
self, self,
inputs_embeds: torch.Tensor, inputs_embeds: torch.Tensor,
attention_mask: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None,
cos: torch.Tensor | None = None,
sin: torch.Tensor | None = None,
) -> torch.Tensor: ) -> torch.Tensor:
"""Run through transformer layers. """Run through transformer layers.
Args: Args:
inputs_embeds: [B, S, hidden_size] inputs_embeds: [B, S, hidden_size]
attention_mask: [B, 1, S, S] or None attention_mask: [B, 1, S, S] or None
cos, sin: optional precomputed RoPE tables in target dtype
Returns: Returns:
hidden_states: [B, S, hidden_size] hidden_states: [B, S, hidden_size]
@@ -743,8 +799,15 @@ class OmniVoiceGenerator(nn.Module):
self._ensure_rope(seq_len, device) self._ensure_rope(seq_len, device)
hidden_states = inputs_embeds hidden_states = inputs_embeds
cos = self._rope_cos.to(device=device, dtype=hidden_states.dtype) if cos is None or sin is None:
sin = self._rope_sin.to(device=device, dtype=hidden_states.dtype) cos = self._rope_cos.to(device=device, dtype=hidden_states.dtype)
sin = self._rope_sin.to(device=device, dtype=hidden_states.dtype)
# Safety: convert bool mask if caller hasn't (e.g. external paths beyond forward()).
if attention_mask is not None and attention_mask.dtype == torch.bool:
attention_mask = torch.zeros_like(attention_mask, dtype=hidden_states.dtype).masked_fill_(
~attention_mask, float("-inf")
)
for layer in self.layers: for layer in self.layers:
hidden_states = layer( hidden_states = layer(
@@ -859,21 +922,37 @@ class OmniVoiceGenerator(nn.Module):
layer_ids = torch.arange(num_codebooks, device=device).view(1, -1, 1) layer_ids = torch.arange(num_codebooks, device=device).view(1, -1, 1)
# Compute c_lens for extracting target region from full sequence # Single D2H pull for all conditional lengths instead of B per-item .item() syncs.
c_lens = [] c_lens = attention_mask[:B, 0, 0].sum(dim=-1).tolist()
for i in range(B):
# Conditional sequence length = number of non-padding positions # Materialize the SDPA float mask once so the captured graph (and eager path) skip per-layer conversion.
c_len = attention_mask[i, 0, 0].sum().item() sdpa_attn_mask = torch.zeros_like(attention_mask, dtype=torch.float32).masked_fill_(
c_lens.append(int(c_len)) ~attention_mask, float("-inf")
)
use_cuda_graph = self._cuda_graph_fwd is not None and input_ids.is_cuda
if not use_cuda_graph:
# Eager-path-only constants (the cuda-graph captures its own).
text_embeds_cached = self.text_embedding(input_ids[:, 0, :])
audio_mask_3d = audio_mask.unsqueeze(-1)
self._ensure_rope(input_ids.shape[-1], device)
target_dtype = text_embeds_cached.dtype
cos = self._rope_cos.to(device=device, dtype=target_dtype)
sin = self._rope_sin.to(device=device, dtype=target_dtype)
# Main iterative loop # Main iterative loop
for step in range(num_step): for step in range(num_step):
if self._cuda_graph_fwd is not None and input_ids.is_cuda: if use_cuda_graph:
batch_logits = self._cuda_graph_fwd(input_ids, audio_mask, attention_mask).to(torch.float32) # Float mask skips per-layer conversion; fp32 cast deferred to the per-item slices below.
batch_logits = self._cuda_graph_fwd(input_ids, audio_mask, sdpa_attn_mask)
else: else:
inputs_embeds = self._prepare_embeddings(input_ids, audio_mask) # Eager fallback reuses hoisted constants (text embeds, sdpa mask, cos/sin).
hidden_states = self._transformer_forward(inputs_embeds, attention_mask) inputs_embeds = self._prepare_embeddings(
batch_logits = self._get_logits(hidden_states).to(torch.float32) input_ids, audio_mask, text_embeds=text_embeds_cached, audio_mask_3d=audio_mask_3d
)
hidden_states = self._transformer_forward(inputs_embeds, sdpa_attn_mask, cos=cos, sin=sin)
# fp32 cast deferred to the per-item slices below.
batch_logits = self._get_logits(hidden_states)
# batch_logits: [2*B, 8, S, 1025] # batch_logits: [2*B, 8, S, 1025]
for i in range(B): for i in range(B):
@@ -884,16 +963,17 @@ class OmniVoiceGenerator(nn.Module):
c_len = c_lens[i] c_len = c_lens[i]
t_len = target_lens[i] t_len = target_lens[i]
# Extract logits for target region # Extract logits for target region; upcast only the slices we actually consume.
c_logits = batch_logits[i : i + 1, :, c_len - t_len : c_len, :] # [1, 8, T, 1025] c_logits = batch_logits[i : i + 1, :, c_len - t_len : c_len, :].to(torch.float32)
u_logits = batch_logits[B + i : B + i + 1, :, :t_len, :] # [1, 8, T, 1025] u_logits = batch_logits[B + i : B + i + 1, :, :t_len, :].to(torch.float32)
# Classifier-free guidance # Classifier-free guidance. Fuse the chain: the two inner
# log_softmax normalizers are per-position scalars that the final
# shift-invariant log_softmax cancels, so guide on the raw logits
# with a single softmax: log_softmax((1+s)*c - s*u). Exact.
if guidance_scale != 0: if guidance_scale != 0:
c_log_probs = F.log_softmax(c_logits, dim=-1) log_probs = F.log_softmax(
u_log_probs = F.log_softmax(u_logits, dim=-1) (1.0 + guidance_scale) * c_logits - guidance_scale * u_logits,
log_probs = torch.log_softmax(
c_log_probs + guidance_scale * (c_log_probs - u_log_probs),
dim=-1, dim=-1,
) )
else: else:
@@ -922,15 +1002,13 @@ class OmniVoiceGenerator(nn.Module):
sample_tokens = tokens[i : i + 1, :, :t_len] sample_tokens = tokens[i : i + 1, :, :t_len]
scores.masked_fill_(sample_tokens != mask_id, -float("inf")) scores.masked_fill_(sample_tokens != mask_id, -float("inf"))
# Select top-k positions to unmask # Select top-k positions to unmask. .flatten() on this non-contiguous view already copies.
_, topk_idx = torch.topk(scores.flatten(), k) _, topk_idx = torch.topk(scores.flatten(), k)
flat_tokens = sample_tokens.flatten().clone() flat_tokens = sample_tokens.flatten()
flat_tokens[topk_idx] = pred_tokens.flatten()[topk_idx] flat_tokens[topk_idx] = pred_tokens.flatten()[topk_idx]
sample_tokens.copy_(flat_tokens.view_as(sample_tokens)) sample_tokens.copy_(flat_tokens.view_as(sample_tokens))
# Update tokens and batch inputs for next iteration # Mirror update into both cond and uncond input_ids halves for the next step.
tokens[i : i + 1, :, :t_len] = sample_tokens
input_ids = input_ids.clone()
input_ids[i, :, c_len - t_len : c_len] = sample_tokens.squeeze(0) input_ids[i, :, c_len - t_len : c_len] = sample_tokens.squeeze(0)
input_ids[B + i, :, :t_len] = sample_tokens.squeeze(0) input_ids[B + i, :, :t_len] = sample_tokens.squeeze(0)
+2
View File
@@ -9,6 +9,7 @@ from vllm_omni.model_extras.registry import (
get_extra_body_params, get_extra_body_params,
get_extra_output_params, get_extra_output_params,
get_model_class_name, get_model_class_name,
get_output_tensor_range,
get_x_to_text_model_family, get_x_to_text_model_family,
should_init_extra_args_for_non_diffusion_stages, should_init_extra_args_for_non_diffusion_stages,
) )
@@ -21,6 +22,7 @@ __all__ = [
"get_extra_body_params", "get_extra_body_params",
"get_extra_output_params", "get_extra_output_params",
"get_model_class_name", "get_model_class_name",
"get_output_tensor_range",
"get_x_to_text_model_family", "get_x_to_text_model_family",
"should_init_extra_args_for_non_diffusion_stages", "should_init_extra_args_for_non_diffusion_stages",
] ]
+15 -1
View File
@@ -4,7 +4,7 @@
from __future__ import annotations from __future__ import annotations
from collections.abc import Callable, Mapping from collections.abc import Callable, Mapping
from typing import Any from typing import Any, Literal
from PIL import Image from PIL import Image
@@ -93,6 +93,7 @@ ImageToVideoPromptBuilder = Callable[
dict[str, Any], dict[str, Any],
] ]
XToTextPromptBuilder = Callable[[str, str, bool], tuple[dict[str, Any], list[int] | None]] XToTextPromptBuilder = Callable[[str, str, bool], tuple[dict[str, Any], list[int] | None]]
OutputTensorRange = Literal["negative_one_to_one", "zero_to_one"]
def default_x_to_text_prompt( def default_x_to_text_prompt(
@@ -193,6 +194,7 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
}, },
"LingBotVideoPipeline": { "LingBotVideoPipeline": {
"extra_body_params": LINGBOT_VIDEO_EXTRA_BODY_PARAMS, "extra_body_params": LINGBOT_VIDEO_EXTRA_BODY_PARAMS,
"output_tensor_range": "zero_to_one",
# Shared T2I/I2V envelopes select the output modality. LingBot's # Shared T2I/I2V envelopes select the output modality. LingBot's
# pipeline owns model-specific validation and normalization. # pipeline owns model-specific validation and normalization.
}, },
@@ -262,6 +264,18 @@ def get_extra_output_params(model_class_name: str | None) -> frozenset[str]:
return spec.get("extra_output_params", frozenset()) if spec is not None else frozenset() return spec.get("extra_output_params", frozenset()) if spec is not None else frozenset()
def get_output_tensor_range(model_class_name: str | None) -> OutputTensorRange:
"""Return the declared range for floating-point tensor outputs.
The default preserves the shared examples' historical handling. Pipelines
that already return normalized tensors declare ``zero_to_one`` explicitly.
"""
spec = _get_spec(model_class_name)
if spec is None:
return "negative_one_to_one"
return spec.get("output_tensor_range", "negative_one_to_one")
def should_init_extra_args_for_non_diffusion_stages(model_class_name: str | None) -> bool: def should_init_extra_args_for_non_diffusion_stages(model_class_name: str | None) -> bool:
spec = _get_spec(model_class_name) spec = _get_spec(model_class_name)
return bool(spec and spec.get("init_extra_args_for_non_diffusion_stages", False)) return bool(spec and spec.get("init_extra_args_for_non_diffusion_stages", False))
@@ -90,6 +90,8 @@ class OmniVoiceConfig(PretrainedConfig):
"cuda_graph_capture_sizes", "cuda_graph_capture_sizes",
[128, 192, 256, 320, 384, 448, 512, 640, 768, 1024], [128, 192, 256, 320, 384, 448, 512, 640, 768, 1024],
) )
# TF32 matmuls: not bit-identical; opt-in (matches vLLM default-off), set OMNIVOICE_TF32=1 to enable.
self.enable_tf32 = getattr(self, "enable_tf32", os.environ.get("OMNIVOICE_TF32", "0") != "0")
AutoConfig.register("omnivoice", OmniVoiceConfig) AutoConfig.register("omnivoice", OmniVoiceConfig)