Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
dc072a12bd | ||
|
|
81e5e8e3bb |
No files matched your search
@@ -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.5.0
|
- export CACHE_DIT_VERSION=1.3.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.5.0
|
- export CACHE_DIT_VERSION=1.3.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.5.0
|
- export CACHE_DIT_VERSION=1.3.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.5.0
|
- export CACHE_DIT_VERSION=1.3.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.5.0
|
- export CACHE_DIT_VERSION=1.3.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
|
||||||
|
|||||||
@@ -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.5.0
|
- export CACHE_DIT_VERSION=1.3.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
|
||||||
|
|||||||
+2
-7
@@ -11,15 +11,10 @@ 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. Wheel-only: `--index-url`
|
# 1 GB memory cap and time out the mkdocs build.
|
||||||
# replaces PyPI, so an sdist fallback cannot reach its build deps. The
|
- python -m pip install --index-url https://download.pytorch.org/whl/cpu torch
|
||||||
# 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,11 +167,9 @@ def main() -> None:
|
|||||||
]
|
]
|
||||||
native_cmd = [
|
native_cmd = [
|
||||||
sys.executable,
|
sys.executable,
|
||||||
"examples/offline_inference/text_to_video/text_to_video.py",
|
"examples/offline_inference/text_to_video/text_to_video_lingbot.py",
|
||||||
"--model",
|
"--model",
|
||||||
args.model,
|
args.model,
|
||||||
"--model-class-name",
|
|
||||||
"LingBotVideoPipeline",
|
|
||||||
"--prompt",
|
"--prompt",
|
||||||
args.prompt,
|
args.prompt,
|
||||||
"--output",
|
"--output",
|
||||||
|
|||||||
@@ -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 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.
|
||||||
|
|
||||||
## Objectives
|
## Objectives
|
||||||
|
|
||||||
|
|||||||
+1
-4
@@ -14,10 +14,7 @@ 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: AWQ is not currently listed as a validated vLLM-Omni path. 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).
|
||||||
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?
|
||||||
|
|
||||||
|
|||||||
+26
-50
@@ -1,23 +1,23 @@
|
|||||||
# Production Metrics
|
# Production Metrics
|
||||||
|
|
||||||
vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server.
|
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.
|
||||||
|
|
||||||
```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,66 +25,42 @@ 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.
|
||||||
|
|
||||||
## Diffusion Metrics (`vllm_omni:`)
|
## Cross-Stage Transfer 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:`)
|
||||||
|
|
||||||
@@ -107,9 +83,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` |
|
||||||
|
|
||||||
@@ -118,5 +94,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:`.
|
- 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`.
|
||||||
- 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.
|
||||||
@@ -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 |
|
||||||
| **[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 |
|
| **[GGUF](quantization/gguf.md)** | `quantization="gguf"` | Native GGUF transformer-only weights (Q4, Q8, etc.) | Memory reduction on consumer GPUs |
|
||||||
|
|
||||||
## Supported Models
|
## Supported Models
|
||||||
|
|
||||||
|
|||||||
@@ -12,7 +12,6 @@ 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 |
|
||||||
@@ -36,25 +35,6 @@ 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
|
||||||
@@ -143,8 +123,7 @@ 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`: online quantization method, for example `fp8`; supported
|
- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
|
||||||
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 `modelopt`, `mxfp4`, `mxfp8`,
|
Other vLLM-Omni quantization methods, such as `gguf`, `modelopt`, `mxfp4`,
|
||||||
serialized checkpoints, static FP8 configs, and layer-name skip lists
|
`mxfp8`, 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.
|
||||||
|
|||||||
@@ -10,10 +10,6 @@ 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 |
|
||||||
@@ -25,57 +21,46 @@ checkpoint is not enough by itself to establish end-to-end model support.
|
|||||||
| Intel XPU | ✅ |
|
| Intel XPU | ✅ |
|
||||||
| Ascend NPU | ❌ |
|
| Ascend NPU | ❌ |
|
||||||
|
|
||||||
Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
|
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||||
guide. This table is method-level backend availability. AutoRound is an Intel
|
guide. AutoRound is Intel-supported.
|
||||||
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 | Validation |
|
| Model | Checkpoint | Scope | Scheme | Backend |
|
||||||
|-------|------------|-------|--------|------------|
|
|-------|------------|-------|--------|---------|
|
||||||
| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | Diffusion transformer | W4A16 | Validated |
|
| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
|
||||||
| Qwen-Image | `INC4AI/Qwen-Image-AutoRound-W4A16` | Diffusion transformer | W4A16 | Validated with end-to-end generation and quality comparison |
|
| Qwen-Image | Not listed | Diffusion transformer | W4A16 | Not validated |
|
||||||
| Wan2.2-I2V | `Intel/Wan2.2-I2V-A14B-Diffusers-int4-AutoRound` | Both diffusion transformers | W4A16 | CI-backed |
|
| Wan2.2-I2V | `Intel/Wan2.2-I2V-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-T2V | `Intel/Wan2.2-T2V-A14B-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 |
|
| Wan2.2-TI2V | `Intel/Wan2.2-TI2V-5B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
|
||||||
|
|
||||||
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 | Validation |
|
| Model | Checkpoint | Scope | Scheme | 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-Nano | `Intel/Cosmos3-Nano-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 |
|
| Cosmos3-Super | `Intel/Cosmos3-Super-int4-AutoRound` | World-model transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
|
||||||
|
|
||||||
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 | Validation |
|
| Model | Checkpoint | Scope | Scheme | Backend |
|
||||||
|-------|------------|-------|--------|------------|
|
|-------|------------|-------|--------|---------|
|
||||||
| Qwen2.5-Omni-7B | `Intel/Qwen2.5-Omni-7B-int4-AutoRound` | Thinker language-model stage | W4A16 | Validated outside scheduled CI |
|
| Qwen2.5-Omni-7B | `Intel/Qwen2.5-Omni-7B-int4-AutoRound` | 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-Omni-30B-A3B-Instruct | `Intel/Qwen3-Omni-30B-A3B-Instruct-int4-AutoRound` | Thinker language-model stage | W4A16 | AutoRound |
|
||||||
| 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 | Checkpoint | Scope | Validation |
|
| Model | Scope | Status | Notes |
|
||||||
|-------|------------|-------|------------|
|
|-------|-------|--------|-------|
|
||||||
| GLM-Image | `Intel/GLM-Image-int4-AutoRound` | Diffusion transformer | CI-backed |
|
| GLM-Image | Diffusion transformer | ✅ | `Intel/GLM-Image-int4-AutoRound` |
|
||||||
| BAGEL | Not listed | Checkpoint-defined diffusion or transformer stage | Not validated |
|
| BAGEL | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
|
||||||
|
|
||||||
## Configuration
|
## Configuration
|
||||||
|
|
||||||
@@ -134,14 +119,6 @@ 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
|
||||||
|
|||||||
@@ -9,10 +9,6 @@ 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 |
|
||||||
@@ -25,7 +21,7 @@ not imply that every diffusion transformer has been quality-validated.
|
|||||||
| Intel XPU | ❌ |
|
| Intel XPU | ❌ |
|
||||||
| Ascend NPU | ❌ |
|
| Ascend NPU | ❌ |
|
||||||
|
|
||||||
Legend: `✅` backend available, `❌` unsupported. Non-CUDA platforms raise in
|
Legend: `✅` supported, `❌` 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`).
|
||||||
@@ -34,11 +30,11 @@ Requires the optional `bitsandbytes` package (`pip install bitsandbytes`).
|
|||||||
|
|
||||||
### Diffusion Model (Qwen-Image, Wan2.2)
|
### Diffusion Model (Qwen-Image, Wan2.2)
|
||||||
|
|
||||||
| Model | HF models | Scope | Level | Recommendation |
|
| Model | HF models | CUDA | Mode | Recommendation |
|
||||||
|-------|-----------|-------|-------|----------------|
|
|-------|-----------|:----:|------|----------------|
|
||||||
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Diffusion transformer | Validated | All heavy linear layers; sensitive embedders stay BF16 |
|
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Online W4 weight-only | All heavy linear layers; sensitive embedders stay BF16 |
|
||||||
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Diffusion transformer | Not validated | Compare with BF16 before enabling |
|
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Not validated | Online W4 weight-only | Compare vs BF16 before enabling |
|
||||||
| Wan2.2 | Wan2.2 diffusion pipelines | One or both diffusion transformers | Not validated | Validate quality and cascade behavior before enabling |
|
| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Online W4 weight-only | Validate before enabling in docs |
|
||||||
|
|
||||||
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.
|
||||||
@@ -108,10 +104,6 @@ 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`).
|
||||||
|
|
||||||
|
|||||||
@@ -13,11 +13,6 @@ 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
|
||||||
@@ -63,11 +58,11 @@ stay unquantized unless the model-specific recipe validates otherwise.
|
|||||||
| Intel XPU | ⭕ |
|
| Intel XPU | ⭕ |
|
||||||
| Ascend NPU | ❌ |
|
| Ascend NPU | ❌ |
|
||||||
|
|
||||||
Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
|
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
|
||||||
guide. The optional CUTLASS FP8 runtime override requires CUDA SM89+. ModelOpt
|
The optional CUTLASS FP8 runtime override requires CUDA SM89+. ModelOpt NVFP4
|
||||||
NVFP4 and mixed FP8/NVFP4 diffusion checkpoints are currently validated on
|
and mixed FP8/NVFP4 diffusion checkpoints are currently validated on Blackwell
|
||||||
Blackwell CUDA systems in the recipes below; other CUDA generations require
|
CUDA systems in the recipes below; other CUDA generations require separate
|
||||||
separate backend and quality validation.
|
backend and quality validation.
|
||||||
|
|
||||||
## Model Type Support
|
## Model Type Support
|
||||||
|
|
||||||
@@ -84,9 +79,6 @@ separate 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
|
||||||
@@ -96,8 +88,8 @@ and
|
|||||||
|
|
||||||
| Model | Scope | Status |
|
| Model | Scope | Status |
|
||||||
|-------|-------|--------|
|
|-------|-------|--------|
|
||||||
| Qwen3-Omni | Thinker language-model stage (FP8) | Validated outside scheduled CI |
|
| Qwen3-Omni | Thinker language-model stage | ModelOpt FP8 checkpoint path |
|
||||||
| Qwen3-Omni | Thinker language-model stage (W4A4 NVFP4) | CI-backed for the H100 Marlin fallback; native Blackwell path is recipe-validated |
|
| Qwen3-Omni | Thinker language-model stage (W4A4 NVFP4) | Validated; see [Qwen3-Omni NVFP4 W4A4](#qwen3-omni-nvfp4-w4a4-thinker) below |
|
||||||
| 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
|
||||||
@@ -110,9 +102,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 and validation |
|
| Variant | HF checkpoint | Hardware |
|
||||||
|---------|---------------|-------------------------|
|
|---------|---------------|----------|
|
||||||
| 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 |
|
| W4A4 NVFP4 (full thinker) | `YihongJin/Qwen3-Omni-30B-A3B-Instruct-NVFP4-W4A4-full-thinker-awqclip` | sm_100+ (Blackwell, FlashInfer FP4 GEMM) |
|
||||||
|
|
||||||
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.
|
||||||
@@ -240,6 +232,3 @@ 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.
|
|
||||||
@@ -3,17 +3,12 @@
|
|||||||
## 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. vLLM-Omni consumes its
|
toolkit for producing pre-quantized model checkpoints. In vLLM-Omni, these
|
||||||
output either through a compatible generic Ascend checkpoint path or through
|
checkpoints run through the Ascend/NPU path with `--quantization ascend`.
|
||||||
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 |
|
||||||
@@ -25,28 +20,23 @@ backend alone is an integration point, not model-level validation.
|
|||||||
| Intel XPU | ❌ |
|
| Intel XPU | ❌ |
|
||||||
| Ascend NPU | ✅ |
|
| Ascend NPU | ✅ |
|
||||||
|
|
||||||
Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
|
Legend: `✅` supported, `❌` 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 | Level | Notes |
|
| Model | Base model | Scope | Hardware | Notes |
|
||||||
|-------|------------|-------|-------|-------|
|
|-------|------------|-------|----------|-------|
|
||||||
| 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 |
|
| 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 |
|
||||||
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | DiT or diffusion stage | Not validated | No supported checkpoint workflow is documented |
|
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | DiT or diffusion stage | Ascend NPU | Not validated in this guide |
|
||||||
| 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 |
|
| 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 |
|
||||||
|
|
||||||
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 |
|
||||||
@@ -63,32 +53,25 @@ than passing `--quantization ascend`.
|
|||||||
|
|
||||||
## 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 <compatible-ascend-checkpoint> --quantization ascend
|
python text_to_image.py --model <quantized-model-path> --quantization ascend
|
||||||
```
|
```
|
||||||
|
|
||||||
Online serving:
|
Online serving:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
vllm serve <compatible-ascend-checkpoint> --omni --quantization ascend
|
vllm serve <quantized-model-path> --omni --quantization ascend
|
||||||
```
|
```
|
||||||
|
|
||||||
## Parameters
|
## Parameters
|
||||||
|
|
||||||
| Parameter | Type | Default | Description |
|
| Parameter | Type | Default | Description |
|
||||||
|-----------|------|---------|-------------|
|
|-----------|------|---------|-------------|
|
||||||
| `--quantization` | str | - | Use `ascend` only for a compatible generic Ascend-format checkpoint; omit it for native merged MXFP8/MXFP4 checkpoints |
|
| `--quantization` | str | - | Use `ascend` for msModelSlim-produced 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
|
||||||
@@ -110,6 +93,3 @@ 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.
|
|
||||||
@@ -13,10 +13,6 @@ 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
|
||||||
@@ -38,45 +34,26 @@ Wan diffusion transformers only; the text encoder and VAE stay in BF16.
|
|||||||
|
|
||||||
## Hardware Support
|
## Hardware Support
|
||||||
|
|
||||||
| Device | Online `mxfp4` | Online `mxfp4_dualscale` | Offline `mxfp4_dualscale` |
|
| Device | Support |
|
||||||
|--------|-----------------|---------------------------|----------------------------|
|
|--------|---------|
|
||||||
| 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: `✅` backend available, `❌` unsupported, `⭕` not verified in this
|
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
|
||||||
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 (Ascend) | Offline (Ascend) | Validation and notes |
|
| Model | Online | Offline | Notes |
|
||||||
|-------|-----------------|------------------|----------------------|
|
|-------|--------|---------|-------|
|
||||||
| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | Validated outside scheduled full-model offline NPU CI; both cascade transformers are quantized |
|
| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade (`transformer` + `transformer_2`); both transformers quantized with the same config |
|
||||||
| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | Validated outside scheduled full-model offline NPU CI; same cascade scheme as T2V-A14B |
|
| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade; same scheme as T2V-A14B |
|
||||||
| Wan2.2-TI2V-5B | ❌ | ❌ | Unsupported because W4A4 causes unacceptable accuracy loss |
|
| Wan2.2-TI2V-5B | ❌ | ❌ | Parameter count too small; 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
|
||||||
@@ -106,13 +83,12 @@ 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. On Ascend, vLLM-Omni quantizes BF16
|
Online mode requires no pre-processing. vLLM-Omni quantizes BF16 weights to
|
||||||
weights at load time using `npu_dynamic_mx_quant`. On ROCm gfx950, the same
|
MXFP4 at load time using `npu_dynamic_mx_quant`. A single block scale
|
||||||
`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. The config applies equally to
|
calibration `mul_scale` is available. Applies equally to single-transformer
|
||||||
single-transformer and cascade (A14B) models, but named Wan model validation is
|
and cascade (A14B) models — both transformers in a cascade receive the same
|
||||||
currently limited to the documented Ascend path.
|
quantization config automatically.
|
||||||
|
|
||||||
```python
|
```python
|
||||||
from vllm_omni import Omni
|
from vllm_omni import Omni
|
||||||
@@ -135,10 +111,10 @@ vllm serve <your-model> --omni --quantization mxfp4
|
|||||||
|
|
||||||
### `mxfp4_dualscale` — DualScale Online Mode
|
### `mxfp4_dualscale` — DualScale Online Mode
|
||||||
|
|
||||||
Ascend-only online DualScale mode computes both fine and coarse scales on the
|
Online DualScale mode computes both fine and coarse scales on the fly from BF16
|
||||||
fly from BF16 weights using `npu_dynamic_dual_level_mx_quant`. It applies
|
weights using `npu_dynamic_dual_level_mx_quant`. Applies equally to
|
||||||
equally to single-transformer and cascade (A14B) models. Compared to `mxfp4`
|
single-transformer and cascade (A14B) models. Compared to `mxfp4` online,
|
||||||
online, DualScale provides better quantization accuracy at higher compute cost.
|
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
|
||||||
@@ -365,11 +341,6 @@ 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
|
||||||
|
|||||||
@@ -15,45 +15,30 @@ 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: `✅` backend available, `❌` unsupported, `⭕` not verified in this
|
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
|
||||||
guide. Backend availability does not imply full-model checkpoint validation.
|
|
||||||
|
|
||||||
The native `DiffusionMXFP8Config` currently dispatches only to Intel XPU and
|
**Note**: Intel XPU only supports AutoRound MXFP8 for offline mode. Use AutoRound quantized checkpoints or online mode for XPU.
|
||||||
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 | Scope | Validation |
|
| Model | Mode | Notes |
|
||||||
|-------|------|-------|------------|
|
|-------|------|-------|
|
||||||
| Wan2.2-T2V-A14B | Online + Offline | MoE cascade: `transformer` + `transformer_2` | Validated outside scheduled full-model offline NPU CI |
|
| Wan2.2-T2V-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-I2V-A14B | Online + Offline | MoE cascade; quantizes two transformers (`transformer` + `transformer_2`) |
|
||||||
| Wan2.2-TI2V-5B | Online + Offline | Single diffusion transformer | Validated outside scheduled full-model offline NPU CI |
|
| Wan2.2-TI2V-5B | Online + Offline | Single transformer |
|
||||||
|
|
||||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||||
|
|
||||||
@@ -233,7 +218,3 @@ 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**.
|
|
||||||
@@ -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 ModelOpt,
|
This mode is different from pre-quantized checkpoint formats such as GGUF,
|
||||||
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,22 +15,18 @@ by msModelSlim and the merge tools.
|
|||||||
|
|
||||||
## Hardware Support
|
## Hardware Support
|
||||||
|
|
||||||
| Device | FP8 W8A8 | Int8 W8A8 | BitsAndBytes W4 | MXFP8 W8A8 | MXFP4 W4A4 |
|
| Device | FP8 W8A8 | Int8 W8A8 | 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 (gfx950) | ⭕ | ⭕ | ❌ | ❌ | ✅ |
|
| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ |
|
||||||
| Intel XPU | ⭕ | ⭕ | ❌ | ✅ | ❌ |
|
| Intel XPU | ⭕ | ⭕ | ✅ | ⭕ |
|
||||||
| Ascend NPU | ❌ | ✅ | ❌ | ✅ | ✅ |
|
| Ascend NPU | ❌ | ✅ | ✅ | ✅ |
|
||||||
|
|
||||||
Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
|
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||||
guide. These entries describe method-level backend availability, not validation
|
guide. FP8 on Ampere may use a weight-only path where available. MXFP8 and
|
||||||
for every model. FP8 on Ampere may use a weight-only path where available.
|
MXFP4 are documented for the Ascend NPU path.
|
||||||
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
|
||||||
|
|
||||||
@@ -40,9 +36,8 @@ NVIDIA instead.
|
|||||||
|--------|-------|----------------|--------|
|
|--------|-------|----------------|--------|
|
||||||
| 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 | Validated on Ascend NPU; the ROCm gfx950 backend is integrated but not model-validated; TI2V-5B is unsupported |
|
| MXFP4 W4A4 | [MXFP4](mxfp4.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B | Ascend NPU only; TI2V-5B is not supported |
|
||||||
|
|
||||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||||
|
|
||||||
@@ -64,7 +59,6 @@ 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")
|
||||||
@@ -75,7 +69,6 @@ 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
|
||||||
@@ -92,19 +85,14 @@ 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, BitsAndBytes, MXFP8, MXFP4 | Quantization method: `"fp8"`, `"int8"`, `"bitsandbytes"`, `"mxfp8"`, `"mxfp4"`, or `"mxfp4_dualscale"` |
|
| `method` | FP8, Int8, MXFP8, MXFP4 | Quantization method: `"fp8"`, `"int8"`, `"mxfp8"`, `"mxfp4"`, or `"mxfp4_dualscale"` |
|
||||||
| `ignored_layers` | FP8, Int8, BitsAndBytes, MXFP8, MXFP4 | Layer name patterns to keep in BF16/FP16 |
|
| `ignored_layers` | FP8, Int8, 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
|
||||||
|
|||||||
@@ -16,123 +16,72 @@ 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 |
|
||||||
|
|
||||||
Online quantization starts from a normal BF16/FP16 checkpoint and repeats the
|
## Hardware Support
|
||||||
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.
|
|
||||||
|
|
||||||
## Support Levels
|
| Device | FP8 W8A8 | Int8 W8A8 | BitsAndBytes W4 | ModelOpt | MXFP8 W8A8 | MXFP4 W4A4 | AutoRound | msModelSlim |
|
||||||
|
|--------|----------|-----------|------------------|----------|------------|------------|-----------|-------------|
|
||||||
|
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ❌ |
|
||||||
|
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ❌ |
|
||||||
|
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
|
||||||
|
| AMD ROCm | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ✅ | ⭕ | ❌ |
|
||||||
|
| Intel XPU | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
|
||||||
|
| Ascend NPU | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ |
|
||||||
|
|
||||||
Quantization support is specific to a model, component, checkpoint format, and
|
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||||
hardware backend. The tables in this guide use the following levels:
|
guide. FP8 on Ampere may use a weight-only path where available.
|
||||||
|
|
||||||
| Level | Meaning |
|
## Model Type Support
|
||||||
|-------|---------|
|
|
||||||
| **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. |
|
|
||||||
|
|
||||||
A method being available on a device does not make every model on that device
|
### Diffusion Model (Qwen-Image, Wan2.2)
|
||||||
supported. Likewise, successfully detecting a checkpoint's
|
|
||||||
`quantization_config` establishes loader compatibility, not output quality.
|
|
||||||
|
|
||||||
## Hardware × Quantization Method
|
These models run a diffusion transformer as the primary inference module. The
|
||||||
|
default quantization target is the transformer; tokenizer, scheduler, text
|
||||||
|
encoder, and VAE stay on the base checkpoint unless a method guide says
|
||||||
|
otherwise.
|
||||||
|
|
||||||
This matrix uses concrete inference paths as columns. Producer toolkits are not
|
| Method | Guide | Mode | Example models | Status |
|
||||||
separate methods: NVIDIA ModelOpt produces the ModelOpt FP8/NVFP4 formats,
|
|--------|-------|------|----------------|--------|
|
||||||
Intel AutoRound produces AutoRound checkpoints, and Ascend msModelSlim produces
|
| 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 |
|
||||||
checkpoints consumed by the native MXFP or compatible Ascend paths.
|
| Int8 W8A8 | [Int8](int8.md) | Online or serialized W8A8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
|
||||||
|
| 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 |
|
||||||
|
|
||||||
| 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) |
|
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||||
|----------|-----------------------|------------------|------------------------------------------|--------------------------------|-------------------------------------------------------------------------|--------------------------------------|--------------------------|--------------------------|
|
|
||||||
| 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 | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ |
|
|
||||||
|
|
||||||
Legend: `✅` backend available, `❌` rejected or no implementation, `⭕` not
|
These models combine an AR language model with audio, vision, talker, or TTS
|
||||||
verified in this guide. This matrix does not establish model-level support; use
|
stages. Quantization is scoped to the AR language-model stage when the
|
||||||
the model-and-stage tables below for that.
|
checkpoint contains a supported `quantization_config`; the non-AR stages stay
|
||||||
|
in BF16 unless the model guide explicitly adds support.
|
||||||
|
|
||||||
The Int8 column covers online and serialized Int8 paths. OCP MXFP8 is online on
|
| Method | Guide | Scope | Example models | Status |
|
||||||
Intel XPU and online/offline on Ascend; Intel's separate AutoRound MXFP8
|
|--------|-------|-------|----------------|--------|
|
||||||
checkpoint path is described in the [MXFP8 guide](mxfp8.md). OCP MXFP4 on AMD
|
| ModelOpt | [ModelOpt](modelopt.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path |
|
||||||
means online single-scale `mxfp4` on gfx950, while MXFP4 DualScale is
|
| Int8 | [Int8](int8.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,
|
| BitsAndBytes | [BitsAndBytes](bitsandbytes.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
|
| MXFP8 | [MXFP8](mxfp8.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
|
| MXFP4 | [MXFP4](mxfp4.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||||
may use a weight-only path where available.
|
| AutoRound | [AutoRound](autoround.md) | Thinker or language-model checkpoint config | Qwen2.5-Omni, Qwen3-Omni | Supported through AutoRound checkpoints |
|
||||||
|
| msModelSlim | [msModelSlim](msmodelslim.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||||
|
|
||||||
## Current Model and Stage Support
|
### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
|
||||||
|
|
||||||
The current offline checkpoint guidance is transformer-first. A listed
|
These models split generation across multiple stages. Quantization must be
|
||||||
checkpoint quantizes the stage named in the **Scope** column; it does not imply
|
attached to the intended stage rather than applied globally.
|
||||||
that the model's encoders, VAE, decoder, or other stages are also quantized.
|
|
||||||
|
|
||||||
### Vendor and Runtime Ecosystems
|
| Method | Guide | Scope | Example models | Status |
|
||||||
|
|--------|-------|-------|----------------|--------|
|
||||||
The producer and runtime backend are related, but they are not the same support
|
| FP8 | [FP8](fp8.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
|
| Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
|
||||||
needs a compatible vLLM-Omni loader, model-stage wiring, and validation. A
|
| BitsAndBytes | [BitsAndBytes](bitsandbytes.md) | Stage-specific transformer or DiT module | BAGEL, GLM-Image | Not validated |
|
||||||
runtime may also provide online quantization without an offline producer path.
|
| ModelOpt | [ModelOpt](modelopt.md) | Checkpoint-defined diffusion stage | BAGEL, GLM-Image | Requires model-specific validation |
|
||||||
|
| MXFP8 | [MXFP8](mxfp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated |
|
||||||
| Ecosystem | Current paths | Current vLLM-Omni scope | Model evidence |
|
| MXFP4 | [MXFP4](mxfp4.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated |
|
||||||
|-----------|---------------|-------------------------|----------------|
|
| AutoRound | [AutoRound](autoround.md) | Checkpoint-defined stage | BAGEL, GLM-Image | No validated checkpoint listed |
|
||||||
| 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 |
|
| msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation |
|
||||||
| 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
|
||||||
@@ -143,8 +92,8 @@ mapping, end-to-end output validation, and backend coverage.
|
|||||||
|
|
||||||
### Diffusion Model (Qwen-Image, Wan2.2)
|
### Diffusion Model (Qwen-Image, Wan2.2)
|
||||||
|
|
||||||
The default target is the diffusion transformer. `build_quant_config()` can
|
The default target is the diffusion transformer. Component routing is available
|
||||||
construct a component router:
|
through `build_quant_config()`:
|
||||||
|
|
||||||
```python
|
```python
|
||||||
from vllm_omni.quantization import build_quant_config
|
from vllm_omni.quantization import build_quant_config
|
||||||
@@ -155,19 +104,11 @@ 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, for listed model/method pairs | Primary target for FP8, Int8, BitsAndBytes, ModelOpt, MXFP8, MXFP4, AutoRound, and msModelSlim |
|
| Diffusion transformer | Yes | 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 |
|
| VAE | No | Keep BF16; storage-only paths are method-specific |
|
||||||
| 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)
|
||||||
@@ -180,18 +121,6 @@ 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 |
|
||||||
@@ -212,7 +141,6 @@ 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,7 +8,6 @@ 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 |
|
||||||
@@ -58,31 +57,18 @@ python text_to_video.py \
|
|||||||
|
|
||||||
### LingBot-Video
|
### LingBot-Video
|
||||||
|
|
||||||
The shared runner recognizes the official dense and MoE checkpoint IDs, selects
|
The dedicated LingBot example explicitly selects the video output modality:
|
||||||
`LingBotVideoPipeline`, and builds the canonical video request envelope:
|
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
python text_to_video.py \
|
python examples/offline_inference/text_to_video/text_to_video_lingbot.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. For a local checkpoint path whose name does not contain
|
VAE's `4n+1` grid.
|
||||||
`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, get_output_tensor_range
|
from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
|
||||||
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,17 +80,6 @@ _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,
|
||||||
@@ -121,51 +110,27 @@ _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 "lingbot" in model_lower or "lingbotvideo" in class_lower:
|
if "distilled" in class_lower or "distilled" in model_lower:
|
||||||
return _MODEL_PRESETS["lingbot"]
|
return _MODEL_PRESETS["ltx2_distilled"]
|
||||||
if "ltx" in class_lower or "ltx" in model_lower:
|
if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower:
|
||||||
if "distilled" in class_lower or "distilled" in model_lower:
|
return _MODEL_PRESETS["ltx23"]
|
||||||
return _MODEL_PRESETS["ltx2_distilled"]
|
if "ltx2" in class_lower or "ltx-2" in model_lower or "ltx_2" in model_lower:
|
||||||
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 or "vace" in class_lower:
|
if "vace" in model_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 or "cosmos" in class_lower:
|
if "cosmos" in model_lower:
|
||||||
return _MODEL_PRESETS["cosmos"]
|
return _MODEL_PRESETS["cosmos"]
|
||||||
if "hunyuan" in model_lower or "hunyuan" in class_lower:
|
if "hunyuan" in model_lower:
|
||||||
return _MODEL_PRESETS["hunyuan"]
|
return _MODEL_PRESETS["hunyuan"]
|
||||||
if "helios" in model_lower or "helios" in class_lower:
|
if "helios" in model_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)
|
||||||
@@ -195,7 +160,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, LingBot-Video, and other text-to-video models."
|
"Supports Wan2.2, HunyuanVideo-1.5, Helios, and other text-to-video models."
|
||||||
)
|
)
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--model",
|
"--model",
|
||||||
@@ -443,7 +408,6 @@ 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)
|
||||||
@@ -516,7 +480,6 @@ 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...")
|
||||||
@@ -550,11 +513,12 @@ def main():
|
|||||||
lora_path=lora_path,
|
lora_path=lora_path,
|
||||||
)
|
)
|
||||||
|
|
||||||
negative_prompt = args.negative_prompt
|
prompt_dict = {"prompt": args.prompt}
|
||||||
if negative_prompt is None and all(preset is not _MODEL_PRESETS[name] for name in ("lingbot", "ltx2", "ltx23")):
|
if args.negative_prompt is not None:
|
||||||
|
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.
|
||||||
negative_prompt = ""
|
prompt_dict["negative_prompt"] = ""
|
||||||
prompt_dict = build_text_to_video_prompt(args.prompt, negative_prompt)
|
|
||||||
|
|
||||||
extra_args = {}
|
extra_args = {}
|
||||||
if lora_request:
|
if lora_request:
|
||||||
@@ -666,7 +630,8 @@ 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)
|
||||||
frame_tensor = _normalize_float_tensor(frame_tensor, output_tensor_range)
|
if frame_tensor.is_floating_point():
|
||||||
|
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
|
||||||
@@ -718,7 +683,8 @@ 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)
|
||||||
video_tensor = _normalize_float_tensor(video_tensor, output_tensor_range)
|
if video_tensor.is_floating_point():
|
||||||
|
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
|
||||||
|
|||||||
@@ -0,0 +1,131 @@
|
|||||||
|
# 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()
|
||||||
@@ -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.py`](../../examples/offline_inference/text_to_video/text_to_video.py)
|
- 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 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.py \
|
python examples/offline_inference/text_to_video/text_to_video_lingbot.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.py \
|
|||||||
|
|
||||||
```bash
|
```bash
|
||||||
CUDA_VISIBLE_DEVICES=0 \
|
CUDA_VISIBLE_DEVICES=0 \
|
||||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
python examples/offline_inference/text_to_video/text_to_video_lingbot.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 \
|
||||||
|
|||||||
@@ -82,7 +82,7 @@ example and the DiT sections below.
|
|||||||
- Optional environment variables:
|
- Optional environment variables:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
export CACHE_DIT_VERSION=1.5.0
|
export CACHE_DIT_VERSION=1.3.0
|
||||||
```
|
```
|
||||||
|
|
||||||
HunyuanImage-3.0 sets the diffusion attention backend to `TORCH_SDPA`
|
HunyuanImage-3.0 sets the diffusion attention backend to `TORCH_SDPA`
|
||||||
|
|||||||
@@ -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.5.0
|
cache-dit==1.3.0
|
||||||
tqdm>=4.66.0
|
tqdm>=4.66.0
|
||||||
torchsde>=0.2.6
|
torchsde>=0.2.6
|
||||||
openai-whisper>=20250625
|
openai-whisper>=20250625
|
||||||
@@ -34,5 +34,3 @@ 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
|
|
||||||
@@ -779,6 +779,26 @@ def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
|
|||||||
omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false")
|
omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false")
|
||||||
|
|
||||||
|
|
||||||
|
def test_from_pipeline_config_routes_ltx2_model_extra(tmp_path):
|
||||||
|
deploy_path = tmp_path / "dreamzero_ltx2_extras.yaml"
|
||||||
|
deploy_path.write_text(
|
||||||
|
"\n".join(
|
||||||
|
[
|
||||||
|
"pipeline: dreamzero",
|
||||||
|
"async_chunk: false",
|
||||||
|
"stages:",
|
||||||
|
" - stage_id: 0",
|
||||||
|
" extras:",
|
||||||
|
" ltx2_use_channels_last_3d: true",
|
||||||
|
]
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
|
||||||
|
|
||||||
|
assert stage.diffusion_config.extras["ltx2_use_channels_last_3d"] is True
|
||||||
|
|
||||||
|
|
||||||
def test_structured_diffusion_config_rejects_invalid_compile_granularity():
|
def test_structured_diffusion_config_rejects_invalid_compile_granularity():
|
||||||
with pytest.raises(ValidationError, match="diffusion_compile_granularity"):
|
with pytest.raises(ValidationError, match="diffusion_compile_granularity"):
|
||||||
omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block")
|
omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block")
|
||||||
|
|||||||
@@ -11,6 +11,100 @@ import torch
|
|||||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||||
|
|
||||||
|
|
||||||
|
class TestLTXChannelsLast3D:
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("extras", "expected"),
|
||||||
|
[({}, False), ({"ltx2_use_channels_last_3d": True}, True)],
|
||||||
|
)
|
||||||
|
def test_channels_last_opt_in_is_ltx2_model_extra(self, extras, expected):
|
||||||
|
from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_channels_last_3d
|
||||||
|
|
||||||
|
assert _ltx2_use_channels_last_3d(SimpleNamespace(extras=extras)) is expected
|
||||||
|
|
||||||
|
def test_channels_last_opt_in_rejects_non_boolean_model_extra(self):
|
||||||
|
from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_channels_last_3d
|
||||||
|
|
||||||
|
with pytest.raises(TypeError, match="ltx2_use_channels_last_3d"):
|
||||||
|
_ltx2_use_channels_last_3d(SimpleNamespace(extras={"ltx2_use_channels_last_3d": "true"}))
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("causal", "kernel_size"),
|
||||||
|
[
|
||||||
|
(True, 3),
|
||||||
|
(False, 3),
|
||||||
|
(True, (3, 1, 1)),
|
||||||
|
(True, 1),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_layout_preserving_causal_conv_matches_original(self, causal, kernel_size):
|
||||||
|
from diffusers.models.autoencoders.autoencoder_kl_ltx2 import LTX2VideoCausalConv3d
|
||||||
|
|
||||||
|
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
|
||||||
|
_enable_channels_last_3d,
|
||||||
|
_weight_is_channels_last_3d,
|
||||||
|
)
|
||||||
|
|
||||||
|
torch.manual_seed(0)
|
||||||
|
conv = LTX2VideoCausalConv3d(4, 5, kernel_size).to(dtype=torch.float32).eval()
|
||||||
|
hidden_states = torch.randn(1, 4, 5, 6, 7, dtype=torch.float32)
|
||||||
|
|
||||||
|
# Pointwise Conv3d storage is valid in both contiguous formats because
|
||||||
|
# all spatial/temporal kernel dimensions are singleton.
|
||||||
|
if kernel_size != 1:
|
||||||
|
assert not _weight_is_channels_last_3d(conv)
|
||||||
|
with torch.no_grad():
|
||||||
|
expected = conv(hidden_states.clone(), causal=causal)
|
||||||
|
|
||||||
|
converted, patched = _enable_channels_last_3d(conv)
|
||||||
|
|
||||||
|
assert (converted, patched) == (1, 1)
|
||||||
|
assert _weight_is_channels_last_3d(conv)
|
||||||
|
with torch.no_grad():
|
||||||
|
actual = conv(hidden_states.clone(), causal=causal)
|
||||||
|
|
||||||
|
torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-4)
|
||||||
|
|
||||||
|
def test_temporal_pad_replicates_edges_in_channels_last_3d(self):
|
||||||
|
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
|
||||||
|
_temporal_pad_channels_last_3d,
|
||||||
|
)
|
||||||
|
|
||||||
|
torch.manual_seed(0)
|
||||||
|
hidden_states = torch.randn(1, 4, 3, 2, 2, dtype=torch.float32)
|
||||||
|
|
||||||
|
actual = _temporal_pad_channels_last_3d(hidden_states, left=2, right=1)
|
||||||
|
expected = torch.cat(
|
||||||
|
[
|
||||||
|
hidden_states[:, :, :1].repeat(1, 1, 2, 1, 1),
|
||||||
|
hidden_states,
|
||||||
|
hidden_states[:, :, -1:],
|
||||||
|
],
|
||||||
|
dim=2,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert actual.is_contiguous(memory_format=torch.channels_last_3d)
|
||||||
|
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
||||||
|
|
||||||
|
def test_contiguous_weight_keeps_original_padding_path(self, monkeypatch):
|
||||||
|
from diffusers.models.autoencoders.autoencoder_kl_ltx2 import LTX2VideoCausalConv3d
|
||||||
|
|
||||||
|
import vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 as ltx2_vae
|
||||||
|
|
||||||
|
conv = LTX2VideoCausalConv3d(4, 5, 3).to(dtype=torch.float32).eval()
|
||||||
|
hidden_states = torch.randn(1, 4, 3, 2, 2, dtype=torch.float32)
|
||||||
|
|
||||||
|
def fail_channels_last_pad(*_args, **_kwargs):
|
||||||
|
raise AssertionError("contiguous weights must retain the original padding path")
|
||||||
|
|
||||||
|
monkeypatch.setattr(ltx2_vae, "_temporal_pad_channels_last_3d", fail_channels_last_pad)
|
||||||
|
|
||||||
|
with torch.no_grad():
|
||||||
|
expected = conv(hidden_states, causal=True)
|
||||||
|
actual = ltx2_vae._layout_preserving_causal_conv_forward(conv, hidden_states, causal=True)
|
||||||
|
|
||||||
|
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
||||||
|
|
||||||
|
|
||||||
class TestLTXOutputRank:
|
class TestLTXOutputRank:
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
("distributed_vae_state", "expected_decode_calls"),
|
("distributed_vae_state", "expected_decode_calls"),
|
||||||
|
|||||||
@@ -189,7 +189,7 @@ def _generate_single_stage_video(
|
|||||||
|
|
||||||
|
|
||||||
def _generate_bagel_image(
|
def _generate_bagel_image(
|
||||||
diffusion_quantization_config: str | None = None,
|
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 diffusion_quantization_config:
|
if quantization_config:
|
||||||
omni_kwargs["diffusion_quantization_config"] = diffusion_quantization_config
|
omni_kwargs["quantization_config"] = 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 diffusion-stage FP8 generates a valid image.
|
"""BAGEL with FP8 quantization_config 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(diffusion_quantization_config="fp8")
|
image, _ = _generate_bagel_image(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(diffusion_quantization_config=None)
|
image, _ = _generate_bagel_image(quantization_config=None)
|
||||||
image.save("test_bagel_bf16.png")
|
image.save("test_bagel_bf16.png")
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -95,7 +95,6 @@ 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 public diffusion speech success and extra_params propagation."""
|
"""Test that extra_params are correctly applied to sampling_params_list in diffusion mode."""
|
||||||
# Mock the engine client
|
# Mock the engine client
|
||||||
mock_engine = mocker.MagicMock()
|
mock_engine = mocker.MagicMock()
|
||||||
|
|
||||||
@@ -704,11 +704,7 @@ 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"})
|
||||||
|
|
||||||
response = await server.create_speech(req)
|
await server._create_diffusion_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()
|
||||||
@@ -741,19 +737,6 @@ 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
|
||||||
|
|||||||
@@ -20,6 +20,7 @@ def test_default_stage_config_includes_cache_backend():
|
|||||||
"cache_backend": "cache_dit",
|
"cache_backend": "cache_dit",
|
||||||
"cache_config": '{"Fn_compute_blocks": 2}',
|
"cache_config": '{"Fn_compute_blocks": 2}',
|
||||||
"vae_use_slicing": True,
|
"vae_use_slicing": True,
|
||||||
|
"extras": {"ltx2_use_channels_last_3d": True},
|
||||||
"ulysses_degree": 2,
|
"ulysses_degree": 2,
|
||||||
}
|
}
|
||||||
)[0]
|
)[0]
|
||||||
@@ -29,6 +30,7 @@ def test_default_stage_config_includes_cache_backend():
|
|||||||
assert engine_args["cache_backend"] == "cache_dit"
|
assert engine_args["cache_backend"] == "cache_dit"
|
||||||
assert engine_args["cache_config"]["Fn_compute_blocks"] == 2
|
assert engine_args["cache_config"]["Fn_compute_blocks"] == 2
|
||||||
assert engine_args["vae_use_slicing"] is True
|
assert engine_args["vae_use_slicing"] is True
|
||||||
|
assert engine_args["extras"]["ltx2_use_channels_last_3d"] is True
|
||||||
assert engine_args["parallel_config"].ulysses_degree == 2
|
assert engine_args["parallel_config"].ulysses_degree == 2
|
||||||
assert engine_args["model_stage"] == "diffusion"
|
assert engine_args["model_stage"] == "diffusion"
|
||||||
|
|
||||||
|
|||||||
@@ -1,91 +0,0 @@
|
|||||||
# 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]))
|
|
||||||
@@ -42,10 +42,6 @@ _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,
|
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
@@ -59,10 +55,6 @@ 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:
|
||||||
@@ -164,50 +156,6 @@ 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.
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
@@ -235,18 +183,6 @@ 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):
|
||||||
@@ -303,88 +239,8 @@ 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_diffusion_path_full(self):
|
def test_non_audio_output_type_skipped(self):
|
||||||
stub = _StubModMetrics()
|
# Image / video / text output types are out of scope for this module.
|
||||||
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(
|
||||||
@@ -395,7 +251,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 not any(c[0].startswith("observe_audio") or c[0].startswith("inc_audio") for c in stub.calls)
|
assert stub.calls == []
|
||||||
|
|
||||||
def test_replica_id_none_skipped(self):
|
def test_replica_id_none_skipped(self):
|
||||||
stub = _StubModMetrics()
|
stub = _StubModMetrics()
|
||||||
|
|||||||
@@ -1,24 +1,21 @@
|
|||||||
# 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, pytest.mark.cuda]
|
pytestmark = [pytest.mark.core_model]
|
||||||
|
|
||||||
|
|
||||||
class _F0Predictor(nn.Module):
|
class _F0Predictor(nn.Module):
|
||||||
@@ -145,248 +142,3 @@ 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)
|
|
||||||
@@ -17,7 +17,6 @@ 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,
|
||||||
)
|
)
|
||||||
@@ -135,8 +134,6 @@ 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,13 +13,10 @@ 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
|
||||||
@@ -53,55 +50,22 @@ def _shrink_flux_t5_text_encoder(config: dict) -> dict:
|
|||||||
return config
|
return config
|
||||||
|
|
||||||
|
|
||||||
def _shrink_qwen_text_encoder_config(config: dict, hidden_size: int = 64) -> dict:
|
|
||||||
text_config = config["text_config"]
|
|
||||||
old_head_dim = text_config["hidden_size"] / text_config["num_attention_heads"]
|
|
||||||
text_config["num_hidden_layers"] = 2
|
|
||||||
text_config["intermediate_size"] = 64
|
|
||||||
text_config["hidden_size"] = hidden_size
|
|
||||||
text_config["num_attention_heads"] = 2
|
|
||||||
text_config["num_key_value_heads"] = 2
|
|
||||||
text_config["layer_types"] = text_config["layer_types"][:2]
|
|
||||||
factor = old_head_dim / (hidden_size / 2)
|
|
||||||
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:
|
def tiny_qwen_image_builder() -> str:
|
||||||
|
def shrink_text_encoder(config: dict) -> dict:
|
||||||
|
config["num_hidden_layers"] = 2
|
||||||
|
config["intermediate_size"] = 64
|
||||||
|
config["text_config"]["num_hidden_layers"] = 2
|
||||||
|
config["text_config"]["intermediate_size"] = 64
|
||||||
|
config["text_config"]["layer_types"] = config["text_config"]["layer_types"][:2]
|
||||||
|
config["vision_config"]["depth"] = 2
|
||||||
|
config["vision_config"]["intermediate_size"] = 64
|
||||||
|
config["vision_config"]["fullatt_block_indexes"] = [0, 1]
|
||||||
|
return config
|
||||||
|
|
||||||
return build_tiny_from_configs(
|
return build_tiny_from_configs(
|
||||||
"QwenImagePipeline",
|
"QwenImagePipeline",
|
||||||
"Qwen/Qwen-Image",
|
"Qwen/Qwen-Image",
|
||||||
transform={"text_encoder": _shrink_qwen_text_encoder_config, "transformer": _shrink_qwen_transformer_config},
|
transform={"text_encoder": shrink_text_encoder, "transformer": _shrink_dit_rope_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,14 +74,4 @@ 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,6 +21,8 @@ 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",
|
||||||
|
|||||||
@@ -19,8 +19,6 @@ 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
|
||||||
|
|||||||
@@ -22,8 +22,6 @@ 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
|
||||||
|
|||||||
@@ -22,8 +22,6 @@ 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,8 +33,6 @@ 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
|
||||||
|
|||||||
@@ -1,11 +1,14 @@
|
|||||||
# 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 types import MethodType
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
from diffusers import AutoencoderKLLTX2Video
|
from diffusers import AutoencoderKLLTX2Video
|
||||||
|
from diffusers.models.autoencoders.autoencoder_kl_ltx2 import LTX2VideoCausalConv3d
|
||||||
from diffusers.models.autoencoders.vae import DecoderOutput
|
from diffusers.models.autoencoders.vae import DecoderOutput
|
||||||
|
from torch import nn
|
||||||
from vllm.logger import init_logger
|
from vllm.logger import init_logger
|
||||||
|
|
||||||
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
|
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
|
||||||
@@ -19,6 +22,89 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
|
|||||||
logger = init_logger(__name__)
|
logger = init_logger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
# Adapted from https://github.com/sgl-project/sglang/pull/27431.
|
||||||
|
def _weight_is_channels_last_3d(module: LTX2VideoCausalConv3d) -> bool:
|
||||||
|
weight = module.conv.weight
|
||||||
|
return (
|
||||||
|
weight.dim() == 5
|
||||||
|
and hasattr(torch, "channels_last_3d")
|
||||||
|
and weight.is_contiguous(memory_format=torch.channels_last_3d)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _temporal_pad_channels_last_3d(
|
||||||
|
hidden_states: torch.Tensor,
|
||||||
|
*,
|
||||||
|
left: int,
|
||||||
|
right: int,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
"""Replicate temporal edge frames directly into CL3D storage."""
|
||||||
|
batch, channels, frames, height, width = hidden_states.shape
|
||||||
|
padded = torch.empty(
|
||||||
|
(batch, channels, frames + left + right, height, width),
|
||||||
|
dtype=hidden_states.dtype,
|
||||||
|
device=hidden_states.device,
|
||||||
|
memory_format=torch.channels_last_3d,
|
||||||
|
)
|
||||||
|
padded[:, :, left : left + frames].copy_(hidden_states)
|
||||||
|
if left:
|
||||||
|
padded[:, :, :left].copy_(hidden_states[:, :, :1])
|
||||||
|
if right:
|
||||||
|
padded[:, :, left + frames :].copy_(hidden_states[:, :, -1:])
|
||||||
|
return padded
|
||||||
|
|
||||||
|
|
||||||
|
def _layout_preserving_causal_conv_forward(
|
||||||
|
module: LTX2VideoCausalConv3d,
|
||||||
|
hidden_states: torch.Tensor,
|
||||||
|
causal: bool = True,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
"""Diffusers LTX causal-conv math with a CL3D-preserving pad path."""
|
||||||
|
time_kernel_size = module.kernel_size[0]
|
||||||
|
use_channels_last_pad = hidden_states.dim() == 5 and _weight_is_channels_last_3d(module)
|
||||||
|
|
||||||
|
if use_channels_last_pad:
|
||||||
|
if causal:
|
||||||
|
left, right = time_kernel_size - 1, 0
|
||||||
|
else:
|
||||||
|
left = right = (time_kernel_size - 1) // 2
|
||||||
|
if left == 0 and right == 0:
|
||||||
|
hidden_states = hidden_states.contiguous(memory_format=torch.channels_last_3d)
|
||||||
|
else:
|
||||||
|
hidden_states = _temporal_pad_channels_last_3d(
|
||||||
|
hidden_states,
|
||||||
|
left=left,
|
||||||
|
right=right,
|
||||||
|
)
|
||||||
|
elif causal:
|
||||||
|
pad_left = hidden_states[:, :, :1, :, :].repeat((1, 1, time_kernel_size - 1, 1, 1))
|
||||||
|
hidden_states = torch.concatenate([pad_left, hidden_states], dim=2)
|
||||||
|
else:
|
||||||
|
pad_size = (time_kernel_size - 1) // 2
|
||||||
|
pad_left = hidden_states[:, :, :1, :, :].repeat((1, 1, pad_size, 1, 1))
|
||||||
|
pad_right = hidden_states[:, :, -1:, :, :].repeat((1, 1, pad_size, 1, 1))
|
||||||
|
hidden_states = torch.concatenate([pad_left, hidden_states, pad_right], dim=2)
|
||||||
|
|
||||||
|
return module.conv(hidden_states)
|
||||||
|
|
||||||
|
|
||||||
|
def _enable_channels_last_3d(model: nn.Module) -> tuple[int, int]:
|
||||||
|
"""Convert LTX Conv3d weights and install layout-preserving causal pads."""
|
||||||
|
if not hasattr(torch, "channels_last_3d"):
|
||||||
|
return 0, 0
|
||||||
|
|
||||||
|
converted_convs = 0
|
||||||
|
patched_causal_convs = 0
|
||||||
|
for module in model.modules():
|
||||||
|
if isinstance(module, LTX2VideoCausalConv3d):
|
||||||
|
module.forward = MethodType(_layout_preserving_causal_conv_forward, module)
|
||||||
|
patched_causal_convs += 1
|
||||||
|
if isinstance(module, nn.Conv3d):
|
||||||
|
module.weight.data = module.weight.data.contiguous(memory_format=torch.channels_last_3d)
|
||||||
|
converted_convs += 1
|
||||||
|
return converted_convs, patched_causal_convs
|
||||||
|
|
||||||
|
|
||||||
class LTX2VaeExecutor(DistributedVaeExecutor):
|
class LTX2VaeExecutor(DistributedVaeExecutor):
|
||||||
def _pack_local_tiles_without_meta(self, local_results, global_padding_shape, device, dtype):
|
def _pack_local_tiles_without_meta(self, local_results, global_padding_shape, device, dtype):
|
||||||
tile_tensor = torch.zeros(global_padding_shape, device=device, dtype=dtype)
|
tile_tensor = torch.zeros(global_padding_shape, device=device, dtype=dtype)
|
||||||
@@ -80,6 +166,16 @@ class DistributedAutoencoderKLLTX2Video(AutoencoderKLLTX2Video, DistributedVaeMi
|
|||||||
model.init_distributed()
|
model.init_distributed()
|
||||||
return model
|
return model
|
||||||
|
|
||||||
|
def enable_channels_last_3d(self) -> None:
|
||||||
|
"""Opt into the layout-preserving CL3D path for this video VAE."""
|
||||||
|
converted_convs, patched_causal_convs = _enable_channels_last_3d(self)
|
||||||
|
logger.info(
|
||||||
|
"LTX2 video VAE: converted %d Conv3d weights to channels_last_3d and "
|
||||||
|
"installed %d layout-preserving causal conv paths",
|
||||||
|
converted_convs,
|
||||||
|
patched_causal_convs,
|
||||||
|
)
|
||||||
|
|
||||||
def init_distributed(self):
|
def init_distributed(self):
|
||||||
self.distributed_executor = LTX2VaeExecutor()
|
self.distributed_executor = LTX2VaeExecutor()
|
||||||
|
|
||||||
|
|||||||
@@ -65,10 +65,11 @@ 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"The guardrails are enabled but not functional for {self.__class__}. This is most likely due to a"
|
f"You have disabled the safety checker for {self.__class__}. This is in violation of the "
|
||||||
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."
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -57,6 +57,17 @@ _LTX_COMPONENT_SUBFOLDERS = (
|
|||||||
)
|
)
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
_LTX2_CHANNELS_LAST_3D_EXTRA = "ltx2_use_channels_last_3d"
|
||||||
|
|
||||||
|
|
||||||
|
def _ltx2_use_channels_last_3d(od_config: Any) -> bool:
|
||||||
|
"""Read the LTX2-only channels-last opt-in from model extras."""
|
||||||
|
extras = getattr(od_config, "extras", {}) or {}
|
||||||
|
enabled = extras.get(_LTX2_CHANNELS_LAST_3D_EXTRA, False)
|
||||||
|
if not isinstance(enabled, bool):
|
||||||
|
raise TypeError(f"{_LTX2_CHANNELS_LAST_3D_EXTRA} must be a bool, got {type(enabled)!r}")
|
||||||
|
return enabled
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True)
|
@dataclass(frozen=True)
|
||||||
class LTXComponentProfile:
|
class LTXComponentProfile:
|
||||||
@@ -365,6 +376,8 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
|
|||||||
local_files_only=local_files_only,
|
local_files_only=local_files_only,
|
||||||
dtype=dtype,
|
dtype=dtype,
|
||||||
)
|
)
|
||||||
|
if _ltx2_use_channels_last_3d(od_config):
|
||||||
|
pipeline.vae.enable_channels_last_3d()
|
||||||
pipeline.audio_vae = _load_component(
|
pipeline.audio_vae = _load_component(
|
||||||
AutoencoderKLLTX2Audio,
|
AutoencoderKLLTX2Audio,
|
||||||
model,
|
model,
|
||||||
|
|||||||
@@ -138,7 +138,6 @@ 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
|
||||||
|
|||||||
@@ -1075,6 +1075,7 @@ class AsyncOmniEngine:
|
|||||||
"streaming_output": kwargs.get("diffusion_streaming_output", False),
|
"streaming_output": kwargs.get("diffusion_streaming_output", False),
|
||||||
"enable_ar_profiler": kwargs.get("enable_ar_profiler", False),
|
"enable_ar_profiler": kwargs.get("enable_ar_profiler", False),
|
||||||
"extras": {
|
"extras": {
|
||||||
|
**(kwargs.get("extras") or {}),
|
||||||
"auxiliary_text_encoder": kwargs.get("auxiliary_text_encoder", None),
|
"auxiliary_text_encoder": kwargs.get("auxiliary_text_encoder", None),
|
||||||
"default_llama_model_id": kwargs.get("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"),
|
"default_llama_model_id": kwargs.get("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"),
|
||||||
},
|
},
|
||||||
|
|||||||
@@ -668,7 +668,6 @@ class OmniServeCommand(CLISubcommand):
|
|||||||
action="store_true",
|
action="store_true",
|
||||||
help="Enable VAE tiling for memory optimization (useful for mitigating OOM issues).",
|
help="Enable VAE tiling for memory optimization (useful for mitigating OOM issues).",
|
||||||
)
|
)
|
||||||
|
|
||||||
# Parallel weight loading (faster diffusion startup)
|
# Parallel weight loading (faster diffusion startup)
|
||||||
omni_config_group.add_argument(
|
omni_config_group.add_argument(
|
||||||
"--disable-multithread-weight-load",
|
"--disable-multithread-weight-load",
|
||||||
|
|||||||
@@ -568,7 +568,6 @@ 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,17 +502,11 @@ 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
|
||||||
|
|||||||
@@ -134,15 +134,6 @@ 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)
|
||||||
# ============================================================================
|
# ============================================================================
|
||||||
|
|||||||
+31
-104
@@ -78,35 +78,6 @@ _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
|
||||||
@@ -168,44 +139,6 @@ 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,
|
||||||
@@ -228,49 +161,43 @@ 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)
|
||||||
|
|
||||||
if output_type == "audio":
|
sample_rate = defs.resolve_audio_sample_rate(mm_out)
|
||||||
sample_rate = defs.resolve_audio_sample_rate(mm_out)
|
# `stage_metrics.audio_generated_frames` is the legacy per-chunk
|
||||||
n_frames = int(getattr(stage_metrics, "audio_generated_frames", 0) or 0)
|
# accumulator field on StageRequestStats. No production path currently
|
||||||
if n_frames == 0:
|
# fills it, so the fallback below is the live source — but we leave the
|
||||||
n_frames = count_audio_frames(mm_out)
|
# field lookup in place in case the accumulator gets re-wired upstream.
|
||||||
mod_metrics.inc_audio_frames(stage_label, replica_label, n_frames)
|
n_frames = int(getattr(stage_metrics, "audio_generated_frames", 0) or 0)
|
||||||
duration_s = n_frames / sample_rate if sample_rate > 0 else 0.0
|
if n_frames == 0:
|
||||||
if duration_s > 0:
|
n_frames = count_audio_frames(mm_out)
|
||||||
mod_metrics.observe_audio_duration(stage_label, replica_label, duration_s)
|
mod_metrics.inc_audio_frames(stage_label, replica_label, n_frames)
|
||||||
mod_metrics.observe_audio_rtf(
|
duration_s = n_frames / sample_rate if sample_rate > 0 else 0.0
|
||||||
stage_label,
|
if duration_s > 0:
|
||||||
replica_label,
|
mod_metrics.observe_audio_duration(stage_label, replica_label, duration_s)
|
||||||
defs.compute_audio_rtf(gen_time_s, duration_s),
|
mod_metrics.observe_audio_rtf(
|
||||||
)
|
stage_label,
|
||||||
else:
|
replica_label,
|
||||||
mod_metrics.inc_audio_skipped(stage_label, replica_label, "no_audio_data")
|
defs.compute_audio_rtf(gen_time_s, duration_s),
|
||||||
|
)
|
||||||
dm = getattr(stage_metrics, "diffusion_metrics", None)
|
else:
|
||||||
if dm:
|
# Request completed (finish_reason ∈ {stop, length} — error paths
|
||||||
_key_map = {
|
# don't reach finalize) but no audio samples were produced. Covers
|
||||||
"diffusion_engine_exec_time_s": mod_metrics.observe_diffusion_exec,
|
# silent `return None` skips in the talker→code2wav stage
|
||||||
"preprocess_time_s": mod_metrics.observe_diffusion_preprocess,
|
# processors and the `parsed.append((0,0))` malformed-length path
|
||||||
"postprocess_time_s": mod_metrics.observe_diffusion_postprocess,
|
# in qwen3-tts code2wav. raise-paths surface via the upstream
|
||||||
}
|
# vllm:request_success_total{finished_reason="error"} channel and
|
||||||
for key, observe_fn in _key_map.items():
|
# never reach this branch.
|
||||||
val = dm.get(key)
|
mod_metrics.inc_audio_skipped(stage_label, replica_label, "no_audio_data")
|
||||||
if val is not None:
|
# audio_underrun / continuity are emitted from the streaming path in
|
||||||
observe_fn(stage_label, replica_label, float(val))
|
# observe_audio_streaming_finalize; finalize is too late for the
|
||||||
|
# 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(
|
||||||
|
|||||||
@@ -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, float] = None
|
diffusion_metrics: dict[str, int] = 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: float(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()}
|
{k: int(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,19 +644,8 @@ 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():
|
||||||
if value is None:
|
self.diffusion_metrics[req_id][key] += value
|
||||||
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: tuple[torch.Tensor, torch.Tensor | float] | None = None,
|
rope=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: tuple[torch.Tensor, torch.Tensor | float] | None = None,
|
rope=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,8 +575,6 @@ 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.
|
||||||
|
|
||||||
@@ -590,9 +588,6 @@ 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]
|
||||||
@@ -615,21 +610,13 @@ 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
|
||||||
if text_embed is None:
|
text_embed = self.text_emb_layer(text, seq_len, text_lens=text_lens).to(model_dtype)
|
||||||
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
|
||||||
if rope is None:
|
rope = self.rotary_embed.forward_from_seq_len(seq_len)
|
||||||
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,40 +857,6 @@ 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:
|
||||||
@@ -906,12 +872,24 @@ 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,
|
||||||
@@ -919,10 +897,9 @@ 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,
|
||||||
attn_mask=attn_mask,
|
is_causal=use_causal,
|
||||||
|
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
|
||||||
@@ -933,11 +910,10 @@ 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_uncond,
|
spkr_emb=spkr_embedding_dit if spkr_embedding_dit is not None and self.spkr_emb_adaln else None,
|
||||||
attn_mask=attn_mask,
|
is_causal=use_causal,
|
||||||
|
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 CFMGraphWrapper, HiFTGraphWrapper
|
from .cuda_graph_wrapper import HiFTGraphWrapper
|
||||||
|
|
||||||
logger = init_logger(__name__)
|
logger = init_logger(__name__)
|
||||||
|
|
||||||
@@ -73,7 +73,6 @@ 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
|
||||||
@@ -136,22 +135,6 @@ 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(
|
||||||
@@ -282,19 +265,6 @@ 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,8 +1,6 @@
|
|||||||
# 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
|
||||||
@@ -135,126 +133,3 @@ 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,10 +162,6 @@ 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")
|
||||||
@@ -806,5 +802,4 @@ 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,7 +242,6 @@ 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__()
|
||||||
@@ -316,9 +315,10 @@ 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)
|
||||||
|
|
||||||
# Caller passes a float mask; materialize float form if a bool slips through.
|
# Convert [B, 1, S, S] bool mask to float mask for SDPA
|
||||||
sdpa_mask = attention_mask
|
# (0.0 where attend, -inf where masked)
|
||||||
if sdpa_mask is not None and sdpa_mask.dtype == torch.bool:
|
sdpa_mask = None
|
||||||
|
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,48 +400,24 @@ 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 pre-doubled to full head_dim so the hot path skips a per-call cat."""
|
"""Precompute RoPE cos/sin tensors."""
|
||||||
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 = torch.cat([freqs.cos(), freqs.cos()], dim=-1)
|
cos = freqs.cos()
|
||||||
sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1)
|
sin = freqs.sin()
|
||||||
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). cos/sin are pre-doubled to width D."""
|
"""Apply rotary position embedding. x shape: (B, S, H, D)."""
|
||||||
seq_len = x.shape[1]
|
seq_len = x.shape[1]
|
||||||
cos = cos[:seq_len].unsqueeze(0).unsqueeze(2) # (1, S, 1, D)
|
cos = cos[:seq_len].unsqueeze(0).unsqueeze(2) # (1, S, 1, D/2)
|
||||||
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 * cos + rotated * sin
|
return x * torch.cat([cos, cos], dim=-1) + rotated * torch.cat([sin, sin], dim=-1)
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# 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(),
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
@@ -449,19 +425,6 @@ def _maybe_enable_tf32() -> None:
|
|||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
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.
|
||||||
|
|
||||||
@@ -513,10 +476,11 @@ class _OmniVoiceCUDAGraphForward:
|
|||||||
mask_padded[:, :S] = audio_mask
|
mask_padded[:, :S] = audio_mask
|
||||||
|
|
||||||
if attention_mask is not None:
|
if attention_mask is not None:
|
||||||
# Callers normalize to the additive float form first, so pad with -inf.
|
attn_padded = torch.zeros(
|
||||||
attn_padded = torch.full(
|
two_b,
|
||||||
(two_b, 1, bucket, bucket),
|
1,
|
||||||
float("-inf"),
|
bucket,
|
||||||
|
bucket,
|
||||||
dtype=attention_mask.dtype,
|
dtype=attention_mask.dtype,
|
||||||
device=attention_mask.device,
|
device=attention_mask.device,
|
||||||
)
|
)
|
||||||
@@ -600,8 +564,7 @@ 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)
|
||||||
# Capture with a float mask to match what forward() feeds at replay time.
|
dummy_attn = torch.ones(two_b, 1, bucket, bucket, dtype=torch.bool, device=device)
|
||||||
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))
|
||||||
|
|
||||||
@@ -623,11 +586,6 @@ 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.
|
||||||
@@ -697,10 +655,6 @@ 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)
|
||||||
|
|
||||||
@@ -749,8 +703,6 @@ 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.
|
||||||
|
|
||||||
@@ -758,38 +710,30 @@ 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]
|
||||||
"""
|
"""
|
||||||
# Cached across the denoising loop since text ids don't change.
|
# Text embeddings from first codebook row (all rows identical for text)
|
||||||
if text_embeds is None:
|
text_embeds = self.text_embedding(input_ids[:, 0, :])
|
||||||
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_3d, audio_embeds, text_embeds)
|
return torch.where(audio_mask.unsqueeze(-1), 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]
|
||||||
@@ -799,15 +743,8 @@ class OmniVoiceGenerator(nn.Module):
|
|||||||
self._ensure_rope(seq_len, device)
|
self._ensure_rope(seq_len, device)
|
||||||
|
|
||||||
hidden_states = inputs_embeds
|
hidden_states = inputs_embeds
|
||||||
if cos is None or sin is None:
|
cos = self._rope_cos.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)
|
||||||
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(
|
||||||
@@ -922,37 +859,21 @@ 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)
|
||||||
|
|
||||||
# Single D2H pull for all conditional lengths instead of B per-item .item() syncs.
|
# Compute c_lens for extracting target region from full sequence
|
||||||
c_lens = attention_mask[:B, 0, 0].sum(dim=-1).tolist()
|
c_lens = []
|
||||||
|
for i in range(B):
|
||||||
# Materialize the SDPA float mask once so the captured graph (and eager path) skip per-layer conversion.
|
# Conditional sequence length = number of non-padding positions
|
||||||
sdpa_attn_mask = torch.zeros_like(attention_mask, dtype=torch.float32).masked_fill_(
|
c_len = attention_mask[i, 0, 0].sum().item()
|
||||||
~attention_mask, float("-inf")
|
c_lens.append(int(c_len))
|
||||||
)
|
|
||||||
|
|
||||||
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 use_cuda_graph:
|
if self._cuda_graph_fwd is not None and input_ids.is_cuda:
|
||||||
# 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, attention_mask).to(torch.float32)
|
||||||
batch_logits = self._cuda_graph_fwd(input_ids, audio_mask, sdpa_attn_mask)
|
|
||||||
else:
|
else:
|
||||||
# Eager fallback reuses hoisted constants (text embeds, sdpa mask, cos/sin).
|
inputs_embeds = self._prepare_embeddings(input_ids, audio_mask)
|
||||||
inputs_embeds = self._prepare_embeddings(
|
hidden_states = self._transformer_forward(inputs_embeds, attention_mask)
|
||||||
input_ids, audio_mask, text_embeds=text_embeds_cached, audio_mask_3d=audio_mask_3d
|
batch_logits = self._get_logits(hidden_states).to(torch.float32)
|
||||||
)
|
|
||||||
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):
|
||||||
@@ -963,17 +884,16 @@ 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; upcast only the slices we actually consume.
|
# Extract logits for target region
|
||||||
c_logits = batch_logits[i : i + 1, :, c_len - t_len : c_len, :].to(torch.float32)
|
c_logits = batch_logits[i : i + 1, :, c_len - t_len : c_len, :] # [1, 8, T, 1025]
|
||||||
u_logits = batch_logits[B + i : B + i + 1, :, :t_len, :].to(torch.float32)
|
u_logits = batch_logits[B + i : B + i + 1, :, :t_len, :] # [1, 8, T, 1025]
|
||||||
|
|
||||||
# Classifier-free guidance. Fuse the chain: the two inner
|
# Classifier-free guidance
|
||||||
# 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:
|
||||||
log_probs = F.log_softmax(
|
c_log_probs = F.log_softmax(c_logits, dim=-1)
|
||||||
(1.0 + guidance_scale) * c_logits - guidance_scale * u_logits,
|
u_log_probs = F.log_softmax(u_logits, dim=-1)
|
||||||
|
log_probs = torch.log_softmax(
|
||||||
|
c_log_probs + guidance_scale * (c_log_probs - u_log_probs),
|
||||||
dim=-1,
|
dim=-1,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
@@ -1002,13 +922,15 @@ 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. .flatten() on this non-contiguous view already copies.
|
# Select top-k positions to unmask
|
||||||
_, topk_idx = torch.topk(scores.flatten(), k)
|
_, topk_idx = torch.topk(scores.flatten(), k)
|
||||||
flat_tokens = sample_tokens.flatten()
|
flat_tokens = sample_tokens.flatten().clone()
|
||||||
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))
|
||||||
|
|
||||||
# Mirror update into both cond and uncond input_ids halves for the next step.
|
# Update tokens and batch inputs for next iteration
|
||||||
|
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)
|
||||||
|
|
||||||
|
|||||||
@@ -9,7 +9,6 @@ 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,
|
||||||
)
|
)
|
||||||
@@ -22,7 +21,6 @@ __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",
|
||||||
]
|
]
|
||||||
@@ -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, Literal
|
from typing import Any
|
||||||
|
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
|
|
||||||
@@ -93,7 +93,6 @@ 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(
|
||||||
@@ -194,7 +193,6 @@ _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.
|
||||||
},
|
},
|
||||||
@@ -264,18 +262,6 @@ 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,8 +90,6 @@ 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)
|
||||||
|
|||||||
Reference in new issue
Block a user