Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
089851e55a | ||
|
|
4b66f62e10 | ||
|
|
f41451f64c | ||
|
|
bbe6ccc512 | ||
|
|
bbb709cfb5 | ||
|
|
3466ff9b98 | ||
|
|
e16473c6b5 | ||
|
|
14d67bcd60 | ||
|
|
5234e3599c | ||
|
|
3aefb31727 | ||
|
|
c8210b8b5c | ||
|
|
879d19c4c2 | ||
|
|
af1825e07d |
No files matched your search
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.26.0
|
||||
ARG VLLM_ASCEND_TAG=v0.27.0
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
# WORKDIR /vllm-workspace/vllm
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.26.0-a3
|
||||
ARG VLLM_ASCEND_TAG=v0.27.0-a3
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
# WORKDIR /vllm-workspace/vllm
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.26.0
|
||||
ARG VLLM_ASCEND_TAG=v0.27.0
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
ARG APP_DIR=/vllm-workspace/vllm-omni
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.26.0-a3
|
||||
ARG VLLM_ASCEND_TAG=v0.27.0-a3
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
ARG APP_DIR=/vllm-workspace/vllm-omni
|
||||
|
||||
+4
-3
@@ -71,7 +71,7 @@ nav:
|
||||
- Hybrid Sharded Data Parallel: user_guide/diffusion/parallelism/hsdp.md
|
||||
- Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md
|
||||
- Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md
|
||||
- VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md
|
||||
- VAE Parallel: user_guide/diffusion/parallelism/vae_parallel.md
|
||||
- Attention Backends: user_guide/diffusion/attention_backends.md
|
||||
- Regional Compilation: user_guide/diffusion/regional_compilation.md
|
||||
- Frame Interpolation: user_guide/diffusion/frame_interpolation.md
|
||||
@@ -123,7 +123,7 @@ nav:
|
||||
- Pipeline Parallel: design/feature/pipeline_parallel.md
|
||||
- Sequence Parallel: design/feature/sequence_parallel.md
|
||||
- Tensor Parallel: design/feature/tensor_parallel.md
|
||||
- VAE Patch Parallelism: design/feature/vae_parallel.md
|
||||
- VAE Parallel: design/feature/vae_parallel.md
|
||||
- Attention Optimization:
|
||||
- Skip-Softmax: design/feature/skip_softmax.md
|
||||
- Cache-DiT: design/feature/cache_dit.md
|
||||
@@ -163,7 +163,8 @@ nav:
|
||||
- CLI Reference: cli
|
||||
- Community:
|
||||
- Governance: community/governance.md
|
||||
- community/*
|
||||
- Contact Us: community/contact_us.md
|
||||
- Community Meetings: community/meetings.md
|
||||
- Slack: https://slack.vllm.ai
|
||||
- Blog: https://blog.vllm.ai
|
||||
- Forum: https://discuss.vllm.ai
|
||||
@@ -1,5 +1,9 @@
|
||||
# Contact Us
|
||||
|
||||
- For technical questions and feature requests, please use GitHub [Issues](https://github.com/vllm-project/vllm-omni/issues)
|
||||
- For coordinating contributions and development and discussing with other users and developers, please join `sig-omni` channel in our [Slack](https://slack.vllm.ai/) or use the [vLLM Forum](https://discuss.vllm.ai/)
|
||||
- For security disclosures, please use GitHub's [Security Advisories](https://github.com/vllm-project/vllm-omni/security/advisories) feature
|
||||
- For bugs and feature requests, search or open a GitHub [issue](https://github.com/vllm-project/vllm-omni/issues).
|
||||
- For technical and usage questions, join the `sig-omni` channel in [vLLM Slack](https://slack.vllm.ai/) or use the [vLLM Forum](https://discuss.vllm.ai/).
|
||||
- For coordinating contributions and development, use `sig-omni`, the forum, or join a [community meeting](meetings.md).
|
||||
|
||||
## Security Disclosures
|
||||
|
||||
Do not share vulnerability details in public GitHub issues, pull requests, discussions, or Slack channels. To arrange a private report, send a direct message to one of the [lead maintainers](governance.md#lead-maintainers) in [vLLM Slack](https://slack.vllm.ai/). If you cannot send a direct message, you may post a detail-free request in `sig-omni` asking a lead maintainer to contact you privately; do not include any vulnerability details in that request.
|
||||
@@ -52,31 +52,11 @@ Committers have write access and merge rights. They typically have deep expertis
|
||||
- [@ZJY0516](https://github.com/ZJY0516): Diffusion attention backend, kernel fusion, and CustomOp
|
||||
- [@yuanheng-zhao](https://github.com/yuanheng-zhao): Diffusion cache, offload, and Omni Support
|
||||
|
||||
## Diffusion Workload Division
|
||||
## Reviewer Routing
|
||||
|
||||
This section breaks down diffusion-related responsibilities. If you have PRs touching these areas, please ping the listed owners for review.
|
||||
For path-based reviewer routing, consult [`.github/CODEOWNERS`](https://github.com/vllm-project/vllm-omni/blob/main/.github/CODEOWNERS). Detailed subsystem boundaries and safe-change guidance live in the [module design documents](../design/index.md#module-design-documents).
|
||||
|
||||
### Runtime-related
|
||||
|
||||
- Diffusion models: @RuixiangMa, @wtomin, @xuechendi
|
||||
|
||||
### Optimization
|
||||
|
||||
- Parallel: @RuixiangMa, @wtomin
|
||||
- Attention backend: @ZJY0516
|
||||
- Cache: @yuanheng-zhao, @RuixiangMa, @SamitHuang
|
||||
- Offload: @yuanheng-zhao
|
||||
- Intel GPU support: @xuechendi
|
||||
- Quantization: @lishunyang12, @david6666666
|
||||
- Kernel fusion / communication-computation: @ZJY0516
|
||||
|
||||
### Docs & Test
|
||||
|
||||
- Diffusion docs: @RuixiangMa, @wtomin
|
||||
|
||||
## Meetings
|
||||
|
||||
Committers hold **bi-weekly meetings** to discuss future directions and collaborations of the project.
|
||||
`CODEOWNERS` is operational routing metadata, not a governance or approval policy. Being listed does not grant committer status or merge rights, and it does not mean every listed reviewer must approve a change. The maintainer roster above is authoritative for project roles and merge rights.
|
||||
|
||||
## Committer Nomination Process
|
||||
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
# Community Meetings
|
||||
|
||||
vLLM-Omni hosts developer-facing meetings for Chinese- and English-language audiences. Schedules and access details can change, so use these maintained documents as the source of truth:
|
||||
|
||||
- [Chinese-language meeting document](https://docs.google.com/document/d/1pdUBiS_7mdOUNDtdwy-9OUf7jsMWN324BIbps5olDME/edit?tab=t.0)
|
||||
- [English-language meeting document](https://docs.google.com/document/d/19AlD08T8uJ7wqfs8VtYCi0iGMTwlTYVVlAxFZ3S7OIg/edit?tab=t.0)
|
||||
|
||||
Each document contains the current schedule, access details, agendas, and past meeting notes. For broader vLLM community events, see the [vLLM events calendar](https://vllm.ai/events/).
|
||||
@@ -1 +0,0 @@
|
||||
# Meetups
|
||||
@@ -1,12 +0,0 @@
|
||||
# Volunteers for Bugfix and CI
|
||||
|
||||
We encourage you to check current docs and [issues](https://github.com/vllm-project/vllm-omni/issues) to find possible solutions for your questions. If non of these can solve it, please propose an issue to describe your questions about bug or CI problems for developing.
|
||||
|
||||
If you have urgent need for locating and solving bugfix or CI problems, please find community volunteers below.
|
||||
|
||||
| Dec 4-Dec 12 | Dec 15-Dec 19 | Dec 22-Dec 26 | Dec 29- Jan 2, 2026| Jan 5-Jan 9 | Jan 12-Jan 16 |
|
||||
|----------|----------|----------|----------|----------|----------|
|
||||
| <a href="https://github.com/congw729">Conw729</a> | <a href="https://github.com/yinpeiqi">yinpeiqi</a> | <a href="https://github.com/tzhouam">tzhouam</a> | <a href="https://github.com/SamitHuang">SamitHuang</a> | <a href="https://github.com/gcanlin">gcanlin</a> | <a href="https://github.com/natureofnature">natureofnature</a> |
|
||||
| <a href="https://github.com/david6666666">david6666666</a> | <a href="https://github.com/R2-Y">R2-Y</a> | <a href="https://github.com/hsliuustc0106">hsliuustc0106</a> | <a href="https://github.com/Gaohan123">Gaohan123</a> | <a href="https://github.com/ZJY0516">ZJY0516</a> | <a href="https://github.com/qibaoyuan">qibaoyuan</a> |
|
||||
|
||||
We kindly welcome more contributors to fix bugs and contribute new features!
|
||||
@@ -3,7 +3,7 @@
|
||||
Thank you for your interest in contributing to vLLM-Omni! This document provides guidelines and instructions for contributing.
|
||||
|
||||
!!! note
|
||||
We host weekly developer-facing online meetings to discuss milestones and updates **every Tuesday at 19:30 PDT**. Meeting link as well as the past meeting notes can be found [here](https://tinyurl.com/vllm-omni-meeting).
|
||||
vLLM-Omni hosts developer-facing meetings for Chinese- and English-language audiences. See [Community Meetings](../community/meetings.md) for current schedules, access details, agendas, and past notes.
|
||||
|
||||
## Getting Started
|
||||
|
||||
@@ -122,7 +122,7 @@ Repository skills guide the agent, but they do not replace contributor judgment,
|
||||
If you encounter a bug or have a feature request, please search existing issues first to see if it has already been reported. If not, please file a new issue, providing as much relevant information as possible.
|
||||
|
||||
!!! important
|
||||
If you discover a security vulnerability, please report it by creating a GitHub issue with the `security` label.
|
||||
Do not report suspected security vulnerabilities through a public issue, pull request, discussion, or Slack channel. Follow the [security disclosure instructions](../community/contact_us.md#security-disclosures) to arrange a private report.
|
||||
|
||||
## Pull Requests & Code Reviews
|
||||
|
||||
|
||||
@@ -398,7 +398,7 @@ the source for configuration, compatibility, and implementation details.
|
||||
* **Composable parallelism:** Diffusion stages can combine [CFG-Parallel](feature/cfg_parallel.md),
|
||||
[Expert Parallel](feature/expert_parallel.md), [HSDP](feature/hsdp.md),
|
||||
[Pipeline Parallel](feature/pipeline_parallel.md), [Sequence Parallel](feature/sequence_parallel.md),
|
||||
[Tensor Parallel](feature/tensor_parallel.md), and [VAE Patch Parallelism](feature/vae_parallel.md)
|
||||
[Tensor Parallel](feature/tensor_parallel.md), and [VAE Parallel](feature/vae_parallel.md)
|
||||
according to the pipeline and hardware topology.
|
||||
* **Attention and cache acceleration:** [Skip-Softmax](feature/skip_softmax.md),
|
||||
[Cache-DiT](feature/cache_dit.md), and [TeaCache](feature/teacache.md)
|
||||
|
||||
@@ -1,42 +1,57 @@
|
||||
# VAE Patch Parallelism
|
||||
# VAE Parallel
|
||||
|
||||
This document describes how to add **VAE Patch Parallelism** support to a diffusion model.
|
||||
We use **Qwen-Image** as the reference implementation for decode parallel, and **Wan2.2** for encode parallel.
|
||||
This document defines the VAE parallelism strategy contract and describes how to
|
||||
add a strategy to a diffusion model. Tile/patch parallelism is the portable
|
||||
baseline. Spatial-shard parallelism is a model-specific decode strategy that can
|
||||
be selected explicitly or through a request-aware policy.
|
||||
|
||||
We use **Qwen-Image** as the reference implementation for tile-parallel decode,
|
||||
**Wan2.2** for tile-parallel encode, and **Wan2.2** as the first spatial-shard
|
||||
decode implementation. Other VAE families can add spatial strategies later
|
||||
without changing the public mode meanings defined here.
|
||||
|
||||
---
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Overview](#overview)
|
||||
- [Step-by-Step Implementation (Decode)](#step-by-step-implementation-decode)
|
||||
- [Encode Parallel Implementation](#encode-parallel-implementation)
|
||||
- [Strategy Selection Contract](#strategy-selection-contract)
|
||||
- [Tile/Patch Implementation (Decode)](#tilepatch-implementation-decode)
|
||||
- [Tile/Patch Encode Implementation](#tilepatch-encode-implementation)
|
||||
- [Testing](#testing)
|
||||
- [Reference Implementations](#reference-implementations)
|
||||
- [Spatially-Sharded Decode](#spatially-sharded-decode)
|
||||
- [Adding Another VAE Strategy](#adding-another-vae-strategy)
|
||||
- [Summary](#summary)
|
||||
|
||||
---
|
||||
|
||||
## Overview
|
||||
|
||||
### What is Vae Patch parallel?
|
||||
### What is VAE parallelism?
|
||||
|
||||
**VAE Patch Parallelism** is an acceleration technique for both **encoding** and **decoding**. Instead of processing the entire tensor at once, the tensor is:
|
||||
VAE parallelism distributes VAE **encoding** or **decoding** across multiple
|
||||
ranks. The implementation currently has two strategy families:
|
||||
|
||||
+ Split into multiple spatial tiles
|
||||
| Strategy | Unit of work | Communication | Current scope |
|
||||
|----------|--------------|---------------|---------------|
|
||||
| Tile/patch | Independent overlapping tiles | Gather, stitch, and optionally broadcast | Encode and decode across supported distributed VAEs |
|
||||
| Spatial shard | One feature map sharded along a spatial axis | Per-layer halo exchange followed by output gather | Wan decode |
|
||||
|
||||
+ Distributed across multiple ranks
|
||||
Both strategies preserve the model-specific receptive field and reconstruct the
|
||||
same logical output as single-rank execution. They differ in how work is divided:
|
||||
tile parallelism assigns complete overlapping tiles to ranks, while spatial
|
||||
sharding keeps one logical feature map partitioned throughout the decoder.
|
||||
|
||||
+ Encoded/Decoded in parallel
|
||||
This approach can:
|
||||
|
||||
+ Merged to reconstruct the final output
|
||||
+ Distribute computation across multiple devices
|
||||
+ Reduce peak memory usage per device
|
||||
+ Accelerate encoding or decoding latency
|
||||
|
||||
This approach:
|
||||
|
||||
+ Distributes computation across multiple devices
|
||||
|
||||
+ Reduces peak memory usage per device
|
||||
|
||||
+ Accelerates encoding/decoding latency
|
||||
The actual benefit depends on shape, topology, communication cost, and VAE
|
||||
architecture. A strategy must not become the global default without parity and
|
||||
performance evidence on its supported models.
|
||||
|
||||
### When to Use Encode vs Decode Parallel
|
||||
|
||||
@@ -45,8 +60,16 @@ This approach:
|
||||
| **Decode Parallel** | Text-to-Image, Text-to-Video | Latent → Image/Video |
|
||||
| **Encode Parallel** | Image-to-Video (I2V) | Image → Latent (for conditioning) |
|
||||
|
||||
### Architecture
|
||||
We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE encoding/decoding.
|
||||
### Architecture and ownership
|
||||
|
||||
`DiffusionParallelConfig.vae_parallel_mode` carries the public strategy request.
|
||||
The shared registry forwards that value to the VAE, but each distributed VAE
|
||||
adapter owns its supported modes, request-shape selection, topology validation,
|
||||
and fallback behavior. Model-specific policy must not be added to the shared
|
||||
runner or registry.
|
||||
|
||||
**DistributedVaeExecutor** is the model-agnostic core for tile-parallel VAE
|
||||
encoding and decoding.
|
||||
|
||||
The executor is model-agnostic and accepts three function parameters:
|
||||
|
||||
@@ -76,6 +99,11 @@ This design separates:
|
||||
|
||||
+ Model-specific tiling and merging logic
|
||||
|
||||
Spatial-shard implementations use the same configured VAE process group but do
|
||||
not use the tile task executor. They keep the feature map sharded across decoder
|
||||
layers and own their halo-exchange and final-gather contracts inside the model's
|
||||
distributed autoencoder implementation.
|
||||
|
||||
#### Why split / exec / merge is necessary?
|
||||
|
||||
The latent tensor cannot be arbitrarily partitioned.
|
||||
@@ -92,7 +120,34 @@ Therefore:
|
||||
|
||||
+ Merge must perform blending to avoid seams
|
||||
|
||||
## Step-by-Step Implementation (Decode)
|
||||
## Strategy Selection Contract
|
||||
|
||||
`vae_parallel_mode="tile"` is the compatibility-preserving global default.
|
||||
Other values are stable strategy requests, but support remains model-specific.
|
||||
|
||||
| Mode | Meaning | Current implementation | Unsupported or ineligible request |
|
||||
|------|---------|------------------------|-----------------------------------|
|
||||
| `tile` | Use the existing tile/patch executor | All distributed VAE adapters | Not applicable |
|
||||
| `auto` | Let the VAE adapter select a validated strategy for each request | Wan decode selects height or width spatial sharding | Wan uses tile when ineligible; unsupported VAE families reject or document their fallback |
|
||||
| `spatial_shard_height` | Force spatial-shard decode along height | Wan decode | Wan warns and uses tile when ineligible; unsupported VAE families reject |
|
||||
| `spatial_shard_width` | Force spatial-shard decode along width | Wan decode | Wan warns and uses tile when ineligible; unsupported VAE families reject |
|
||||
|
||||
The shared configuration layer validates the mode name. It does not decide
|
||||
whether a model supports that mode. A VAE adapter adding a mode must:
|
||||
|
||||
1. Keep `tile` as an exact feature-off path.
|
||||
2. Select a strategy per request rather than from stale process-global state.
|
||||
3. Validate input rank/shape and process-group membership before collectives.
|
||||
4. Reject or explicitly fall back from unsupported combinations; never silently
|
||||
reinterpret one forced strategy as another.
|
||||
5. Preserve direct encode/decode behavior when Diffusers does not select a tiled
|
||||
path.
|
||||
|
||||
`auto` is a policy hook, not a universal heuristic. Each VAE family may select
|
||||
from only the strategies it has validated. Adding another family must not change
|
||||
Wan's selector or the meanings of the explicit modes.
|
||||
|
||||
## Tile/Patch Implementation (Decode)
|
||||
|
||||
### Step 1: Implement DistributedAutoencoderKLQwenImage
|
||||
`QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
|
||||
@@ -251,7 +306,7 @@ class YourModelPipeline(nn.Module):
|
||||
+ ).to(self.device)
|
||||
```
|
||||
|
||||
## Encode Parallel Implementation
|
||||
## Tile/Patch Encode Implementation
|
||||
|
||||
For models that require VAE encoding (e.g., Image-to-Video), you can also parallelize the encode operation. We use **Wan2.2** as the reference implementation.
|
||||
|
||||
@@ -449,18 +504,24 @@ Complete examples in the codebase:
|
||||
|
||||
---
|
||||
|
||||
## Spatially-Sharded Decode (Wan)
|
||||
## Spatially-Sharded Decode
|
||||
|
||||
The tile-parallel executor above assigns independent spatial **tiles** to ranks. The Wan VAE additionally supports a **spatially-sharded decode** backend, selected through `DiffusionParallelConfig.vae_parallel_mode` (`"spatial_shard_height"` or `"spatial_shard_width"`; default is `"tile"`).
|
||||
The tile-parallel executor above assigns independent spatial **tiles** to ranks.
|
||||
A spatial-shard backend instead partitions one global decoder feature map along
|
||||
height or width. Wan is the first supported implementation.
|
||||
|
||||
`DiffusionParallelConfig.vae_parallel_mode="tile"` remains the default. For Wan,
|
||||
opt into request-shape selection with `"auto"`, or force an axis with
|
||||
`"spatial_shard_height"` or `"spatial_shard_width"`.
|
||||
|
||||
### How it differs from tile parallel
|
||||
|
||||
| Aspect | Tile parallel (`"tile"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
|
||||
|--------|--------------------------|-----------------------------------------------|
|
||||
| Unit of work | Independent overlapping tiles | A single global feature map sharded along H or W |
|
||||
| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Per-conv **halo exchange** of boundary rows/cols (P2P) |
|
||||
| Output assembly | Blend overlapping tiles | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) |
|
||||
| Scope | Decode + encode | Decode only |
|
||||
| Aspect | Tile parallel (`"tile"`) | Automatic (`"auto"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
|
||||
|--------|--------------------------|----------------------|-----------------------------------------------|
|
||||
| Unit of work | Independent overlapping tiles | Chooses tile or a spatial shard per request | A single global feature map sharded along H or W |
|
||||
| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Depends on the selected path | Per-conv **halo exchange** of boundary rows/cols (P2P) |
|
||||
| Output assembly | Blend overlapping tiles | Depends on the selected path | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) |
|
||||
| Scope | Decode + encode | Wan decode selection; encode remains tiled | Wan decode only |
|
||||
|
||||
Spatial-shard decode swaps the decoder's spatial convolutions/padding for halo-exchanging variants (`WanDistConv2d`, `WanDistCausalConv3d`, `WanDistZeroPad2d`) so each rank only holds a shard of the activations but still sees the correct receptive field at shard boundaries. Implementation lives in `vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py`.
|
||||
|
||||
@@ -476,21 +537,65 @@ serve.py (--vae-parallel-mode) / OmniEngineArgs
|
||||
-> DistributedAutoencoderKLWan.tiled_decode dispatch
|
||||
```
|
||||
|
||||
`DistributedAutoencoderKLWan._spatial_shard_decode_enabled()` gates the path: it requires distributed decode to be enabled, a 5D latent, and `vae_patch_parallel_size == DiT group size`. Otherwise it logs a warning and falls back to tile-parallel decode.
|
||||
`DistributedAutoencoderKLWan._spatial_shard_decode_enabled()` gates the path: it
|
||||
requires distributed decode to be enabled, a 5D latent, and
|
||||
`vae_patch_parallel_size == DiT group size`. The method is reached only after
|
||||
Diffusers has selected tiled decode for the request. In `auto` mode the longer
|
||||
axis that can cover every rank is selected, with width winning ties. If neither
|
||||
axis can cover the group, auto mode uses tile decode. This policy is cached only
|
||||
through normal kernel shape caches; it is reevaluated for every request.
|
||||
|
||||
### Notes
|
||||
|
||||
- The decoder is patched **in place** the first time spatial-shard decode runs and is bound to a single split dimension for the lifetime of the VAE instance.
|
||||
- The decoder is patched **in place** the first time spatial-shard decode runs. Its wrappers consult a context-local split dimension and retain an exact direct fallback, allowing later requests to select tile, height, or width safely.
|
||||
- Numerical correctness vs. single-GPU decode is covered by `tests/diffusion/distributed/test_wan_spatial_shard.py::test_spatial_shard_decode_matches_reference` (multi-GPU, nightly `full_model` + `distributed_cuda`).
|
||||
|
||||
## Adding Another VAE Strategy
|
||||
|
||||
Add future strategies through the distributed autoencoder that owns the model
|
||||
semantics. Do not add model-name dispatch to the shared runner, registry, or
|
||||
`DistributedVaeExecutor`.
|
||||
|
||||
For a new VAE family:
|
||||
|
||||
1. **Declare capability**: document which modes, operations, shapes, dtypes, and
|
||||
process-group layouts the adapter supports.
|
||||
2. **Implement model-local selection**: keep `auto` policy and explicit-mode
|
||||
dispatch in the distributed autoencoder. The selector may use request shape,
|
||||
but must be deterministic and side-effect free.
|
||||
3. **Preserve the baseline**: retain exact tile/direct fallbacks and state-dict
|
||||
compatibility. Structural wrappers must not register duplicate parameters or
|
||||
alter checkpoint names.
|
||||
4. **Define communication**: derive shards from the configured VAE group, keep
|
||||
collectives symmetric, specify padding/halo ownership, and return the same
|
||||
output placement as the tile path.
|
||||
5. **Bound request state**: make the selected axis and scratch buffers
|
||||
request-scoped, release temporary buffers on success and failure, and support
|
||||
alternating modes in one long-lived worker.
|
||||
6. **Validate independently**: compare each explicit strategy against the
|
||||
single-rank reference, then test `auto`, tile fallback, unsupported topology,
|
||||
and a mixed request sequence. Record per-rank peak memory and decode latency
|
||||
before documenting a preferred mode.
|
||||
|
||||
Promote a selector into shared infrastructure only after at least two VAE
|
||||
families use the same inputs, invariants, fallback behavior, and lifecycle. Until
|
||||
then, shared code owns transport and configuration while each model adapter owns
|
||||
policy.
|
||||
|
||||
---
|
||||
|
||||
## Summary
|
||||
|
||||
Adding VAE Patch Parallel support to diffusion model:
|
||||
Adding VAE parallel support to a diffusion model:
|
||||
|
||||
1. **Implement Distributed VAE** - Inherit from base VAE class and `DistributedVaeMixin`
|
||||
2. **Decode Parallel** - Refactor `tiled_decode` into `tile_split`/`tile_exec`/`tile_merge`
|
||||
3. **Encode Parallel** (optional) - Implement `encode_tile_split`/`encode_tile_exec`/`encode_tile_merge` for I2V models
|
||||
4. **Change VAE model in pipeline** - Use the distributed version
|
||||
5. **Test** - Verify numerical consistency with `vae_patch_parallel_size=1` vs `N`
|
||||
1. **Implement the distributed VAE** - Inherit from the base VAE class and
|
||||
`DistributedVaeMixin`.
|
||||
2. **Keep tile as the baseline** - For tile decode, refactor `tiled_decode` into
|
||||
`tile_split`/`tile_exec`/`tile_merge`.
|
||||
3. **Add encode parallelism when needed** - Implement
|
||||
`encode_tile_split`/`encode_tile_exec`/`encode_tile_merge` for I2V models.
|
||||
4. **Add model-specific strategies locally** - Own selectors, topology checks,
|
||||
communication, and fallback in the distributed autoencoder.
|
||||
5. **Wire the pipeline** - Replace the pipeline VAE with its distributed version.
|
||||
6. **Test strategy and lifecycle parity** - Compare one rank with every supported
|
||||
strategy, including feature-off, fallback, and mixed-request sequences.
|
||||
@@ -48,7 +48,7 @@ implementation contract; it is not, by itself, a general support claim.
|
||||
- [Pipeline Parallel](feature/pipeline_parallel.md)
|
||||
- [Sequence Parallel](feature/sequence_parallel.md)
|
||||
- [Tensor Parallel](feature/tensor_parallel.md)
|
||||
- [VAE Patch Parallelism](feature/vae_parallel.md)
|
||||
- [VAE Parallel](feature/vae_parallel.md)
|
||||
|
||||
#### Attention optimization
|
||||
|
||||
|
||||
+3
-6
@@ -10,15 +10,12 @@ A: Refer to [GPU memory calculation and configuration](../configuration/gpu_memo
|
||||
|
||||
> Q: I encounter some bugs or CI problems, which is urgent. How can I solve it?
|
||||
|
||||
A: At first, you can check current [issues](https://github.com/vllm-project/vllm-omni/issues) to find possible solutions. If none of these satisfy your demand and it is urgent, please find these [volunteers](https://docs.vllm.ai/projects/vllm-omni/en/latest/community/volunteers/) for help.
|
||||
A: Search the current [issues](https://github.com/vllm-project/vllm-omni/issues) first. If no existing issue matches a bug or CI failure, open a [bug report](https://github.com/vllm-project/vllm-omni/issues/new?template=400-bug-report.yml) with a minimal reproducer, environment details, and relevant logs. For technical or usage questions, use the `sig-omni` channel in [vLLM Slack](https://slack.vllm.ai/) or the [vLLM Forum](https://discuss.vllm.ai/).
|
||||
|
||||
> Q: Does vLLM-Omni support AWQ or any other quantization?
|
||||
|
||||
A: AWQ is not currently listed as a validated vLLM-Omni path. Quantization
|
||||
support depends on the model, pipeline stage, checkpoint format, and hardware
|
||||
backend. See the current [quantization support matrix](../user_guide/quantization/overview.md)
|
||||
for online methods, pre-quantized checkpoints, and model-level validation.
|
||||
A: AWQ is available through the inherited vLLM quantization registry, and vLLM-Omni supports several additional methods. Actual support and validation vary by model, component, hardware, and checkpoint. See the [quantization overview](../user_guide/quantization/overview.md) for validated combinations and method-specific guides.
|
||||
|
||||
> Q: Does vLLM-Omni support multimodal streaming input and output?
|
||||
|
||||
A: Not yet. We already put it on the [Roadmap](https://github.com/vllm-project/vllm-omni/issues/165). Please stay tuned!
|
||||
A: Yes, for specific models and serving paths. Current capabilities include [streaming video input](../serving/video_stream_api.md), [streaming diffusion output](../user_guide/diffusion/execution_modes.md#streaming-output), and experimental full-duplex real-time audio described in the [project overview](../README.md). Check the model and API documentation because input and output support is model-specific.
|
||||
@@ -10,7 +10,7 @@ This guide covers the parallelism methods in vLLM-Omni for speeding up diffusion
|
||||
| **[Sequence Parallelism](sequence_parallel.md)** | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos |
|
||||
| **[CFG-Parallel](cfg_parallel.md)** | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation |
|
||||
| **[Pipeline Parallelism](pipeline_parallel.md)** | Splits the denoising transformer block-wise across sequential GPU stages to reduce per-GPU model memory |
|
||||
| **[VAE Parallelism](vae_parallelism.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory |
|
||||
| **[VAE Parallel](vae_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory |
|
||||
| **[HSDP](hsdp.md)** | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs |
|
||||
| **[Expert Parallelism](expert_parallel.md)** | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0) |
|
||||
|
||||
|
||||
+8
-8
@@ -1,4 +1,4 @@
|
||||
# VAE Parallelism Guide
|
||||
# VAE Parallel Guide
|
||||
|
||||
|
||||
## Table of Content
|
||||
@@ -113,7 +113,7 @@ In `DiffusionParallelConfig`:
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `vae_patch_parallel_size` | int | 1 | Number of GPUs for VAE patch/tile parallelism. Set to 2 or higher to enable. Should typically match `tensor_parallel_size` as they share the same process group. |
|
||||
| `vae_parallel_mode` | str | `"tile"` | VAE parallel decode strategy: `"tile"` (default tile/patch parallel decode), `"spatial_shard_height"`, or `"spatial_shard_width"` (spatially-sharded decode, Wan only). See [Spatially-Sharded Decode](#spatially-sharded-decode-wan). |
|
||||
| `vae_parallel_mode` | str | `"tile"` | VAE parallel decode strategy: `"tile"`, `"auto"` (Wan request-shape selection), `"spatial_shard_height"`, or `"spatial_shard_width"`. Spatial sharding is currently Wan-only. See [Spatially-Sharded Decode](#spatially-sharded-decode-wan). |
|
||||
|
||||
Additional requirements:
|
||||
|
||||
@@ -128,7 +128,7 @@ Additional requirements:
|
||||
|
||||
## Spatially-Sharded Decode (Wan)
|
||||
|
||||
The default `vae_parallel_mode="tile"` distributes whole tiles across ranks. For the **Wan** VAE there is an alternative decode strategy, **spatially-sharded decode**, selected via `vae_parallel_mode="spatial_shard_height"` or `vae_parallel_mode="spatial_shard_width"`.
|
||||
The default `vae_parallel_mode="tile"` keeps patch/tile parallel decode unchanged. For Wan, set `vae_parallel_mode="auto"` to select spatial-shard decode only after a request is already large enough to enter Diffusers' tiled-decode path and the VAE uses the full DiT decode group. Auto mode splits the longer eligible spatial axis to reduce halo traffic. Direct-decode requests and partial groups fall back to tile decode. Use `"spatial_shard_height"` or `"spatial_shard_width"` to force an axis.
|
||||
|
||||
Instead of assigning independent tiles to ranks, spatial-shard decode shards the decoder feature maps along the height (`spatial_shard_height`) or width (`spatial_shard_width`) dimension and exchanges halo rows/columns between neighboring ranks around the spatial convolutions. This keeps the receptive field correct across shard boundaries, so the result matches the single-GPU decode within numerical tolerance.
|
||||
|
||||
@@ -141,7 +141,7 @@ omni = Omni(
|
||||
parallel_config=DiffusionParallelConfig(
|
||||
tensor_parallel_size=2,
|
||||
vae_patch_parallel_size=2, # must match the DiT group size
|
||||
vae_parallel_mode="spatial_shard_width", # or "spatial_shard_height"
|
||||
vae_parallel_mode="auto", # opt in to Wan request-shape selection
|
||||
),
|
||||
)
|
||||
```
|
||||
@@ -152,14 +152,14 @@ Or from the CLI / serving entrypoint:
|
||||
vllm serve Wan-AI/Wan2.1-T2V-1.3B-Diffusers --omni \
|
||||
--tensor-parallel-size 2 \
|
||||
--vae-patch-parallel-size 2 \
|
||||
--vae-parallel-mode spatial_shard_width
|
||||
--vae-parallel-mode auto
|
||||
```
|
||||
|
||||
**Constraints and behavior:**
|
||||
|
||||
- Spatial-shard decode is **decode-only** and currently implemented for the **Wan** VAE. Other models ignore `spatial_shard_*` modes.
|
||||
- It requires `vae_patch_parallel_size` to **match the DiT process group size**. If it does not, the VAE logs a warning and **falls back to tile-parallel decode** at runtime.
|
||||
- `spatial_shard_height` and `spatial_shard_width` are mutually exclusive for a given VAE instance (the decoder is patched in place for a single split dimension).
|
||||
- Spatial-shard decode is **decode-only** and currently implemented for the **Wan** VAE. Other VAE families retain their existing mode handling and may reject explicit Wan spatial modes.
|
||||
- It requires `vae_patch_parallel_size` to **match the DiT process group size**. `auto` quietly falls back to tile decode for a partial group; an explicitly forced spatial mode logs a warning.
|
||||
- The decoder's context-aware wrappers retain the direct path, so one process can alternate small/large and portrait/landscape requests safely.
|
||||
|
||||
For end-to-end latency/throughput, launch serving with the desired `vae_parallel_mode` and use the existing diffusion serving benchmark:
|
||||
|
||||
@@ -61,7 +61,7 @@ Memory optimization methods help reduce GPU memory usage, enabling inference on
|
||||
|--------|-------------|----------|
|
||||
| **[CPU Offload](diffusion/cpu_offload.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs |
|
||||
| **[Quantization](quantization/overview.md)** | Reduces transformer stages from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss |
|
||||
| **[VAE Parallelism](diffusion/parallelism/vae_parallelism.md)** | Distributes VAE decode work across GPUs | High-resolution generation with reduced VAE memory peak |
|
||||
| **[VAE Parallel](diffusion/parallelism/vae_parallel.md)** | Distributes VAE decode work across GPUs | High-resolution generation with reduced VAE memory peak |
|
||||
|
||||
### Extensions
|
||||
|
||||
@@ -92,7 +92,7 @@ current support.
|
||||
|--------|--------------|-------------|----------|
|
||||
| **[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 |
|
||||
| **[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
|
||||
|
||||
|
||||
@@ -143,8 +143,7 @@ python text_to_video.py \
|
||||
- `--enable-layerwise-offload`: enable layerwise offloading on DiT modules.
|
||||
- `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
|
||||
- `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio.
|
||||
- `--quantization`: online quantization method, for example `fp8`; supported
|
||||
pre-quantized checkpoints are normally auto-detected from checkpoint metadata.
|
||||
- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
|
||||
- `--flow-shift`: scheduler flow_shift parameter.
|
||||
- `--lora-path`: path to PEFT LoRA adapter folder or checkpoint file.
|
||||
- `--lora-scale`: scale factor for LoRA weights.
|
||||
|
||||
@@ -90,8 +90,8 @@ vllm serve "Qwen/Qwen-Image" \
|
||||
--quantization-config '{"method": "fp8"}'
|
||||
```
|
||||
|
||||
Other vLLM-Omni quantization methods, such as `modelopt`, `mxfp4`, `mxfp8`,
|
||||
serialized checkpoints, static FP8 configs, and layer-name skip lists
|
||||
Other vLLM-Omni quantization methods, such as `gguf`, `modelopt`, `mxfp4`,
|
||||
`mxfp8`, serialized checkpoints, static FP8 configs, and layer-name skip lists
|
||||
such as `ignored_layers`, are intentionally not translated by this compatibility
|
||||
shortcut. Use Diffusers-native configuration through `diffusers_load_kwargs` or
|
||||
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
|
||||
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
|
||||
|
||||
| Device | Support |
|
||||
@@ -25,57 +21,46 @@ checkpoint is not enough by itself to establish end-to-end model support.
|
||||
| Intel XPU | ✅ |
|
||||
| Ascend NPU | ❌ |
|
||||
|
||||
Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
|
||||
guide. This table is method-level backend availability. AutoRound is an Intel
|
||||
project; vLLM selects a compatible CUDA or XPU compute backend at load time.
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide. AutoRound is Intel-supported.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
### Diffusion Model (Qwen-Image, Wan2.2)
|
||||
|
||||
| Model | Checkpoint | Scope | Scheme | Validation |
|
||||
|-------|------------|-------|--------|------------|
|
||||
| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | Diffusion transformer | W4A16 | Validated |
|
||||
| Qwen-Image | `INC4AI/Qwen-Image-AutoRound-W4A16` | Diffusion transformer | W4A16 | Validated with end-to-end generation and quality comparison |
|
||||
| Wan2.2-I2V | `Intel/Wan2.2-I2V-A14B-Diffusers-int4-AutoRound` | Both diffusion transformers | W4A16 | CI-backed |
|
||||
| Wan2.2-T2V | `Intel/Wan2.2-T2V-A14B-Diffusers-int4-AutoRound` | Both diffusion transformers | W4A16 | CI-backed |
|
||||
| Wan2.2-TI2V | `Intel/Wan2.2-TI2V-5B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | Validated; not in the scheduled Wan AutoRound job |
|
||||
|
||||
CUDA execution uses GPTQ-Marlin or another compatible vLLM AutoRound backend;
|
||||
Intel XPU uses the Intel-supported backend. The text encoder and VAE stay in
|
||||
BF16 for all entries above.
|
||||
| Model | Checkpoint | Scope | Scheme | Backend |
|
||||
|-------|------------|-------|--------|---------|
|
||||
| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
|
||||
| Qwen-Image | Not listed | Diffusion transformer | W4A16 | Not validated |
|
||||
| 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` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
|
||||
| Wan2.2-TI2V | `Intel/Wan2.2-TI2V-5B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
|
||||
|
||||
### World Model (Cosmos3)
|
||||
|
||||
| Model | Checkpoint | Scope | Scheme | Validation |
|
||||
|-------|------------|-------|--------|------------|
|
||||
| Cosmos3-Nano | `Intel/Cosmos3-Nano-int4-AutoRound` | World-model transformer | W4A16 | Integrated; no model-specific end-to-end test is maintained in-tree |
|
||||
| Cosmos3-Super | `Intel/Cosmos3-Super-int4-AutoRound` | World-model transformer | W4A16 | Validated with a manual serving run |
|
||||
|
||||
The Cosmos VAE and guardrail components stay in BF16.
|
||||
| Model | Checkpoint | Scope | Scheme | Backend |
|
||||
|-------|------------|-------|--------|---------|
|
||||
| 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 | GPTQ-Marlin or Intel-supported AutoRound backend |
|
||||
|
||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||
|
||||
| Model | Checkpoint | Scope | Scheme | Validation |
|
||||
|-------|------------|-------|--------|------------|
|
||||
| Qwen2.5-Omni-7B | `Intel/Qwen2.5-Omni-7B-int4-AutoRound` | Thinker language-model stage | W4A16 | Validated outside scheduled CI |
|
||||
| Qwen3-Omni-30B-A3B-Instruct | `Intel/Qwen3-Omni-30B-A3B-Instruct-int4-AutoRound` | Thinker language-model stage | W4A16 | Validated outside scheduled CI |
|
||||
| Model | Checkpoint | Scope | Scheme | Backend |
|
||||
|-------|------------|-------|--------|---------|
|
||||
| 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 | AutoRound |
|
||||
| Qwen3-TTS | Not listed | TTS language-model stage | W4A16 | Not validated |
|
||||
|
||||
AutoRound support is checkpoint-driven. A model is supported when its
|
||||
checkpoint uses a compatible INC/AutoRound config and the target stage maps to
|
||||
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)
|
||||
|
||||
| Model | Checkpoint | Scope | Validation |
|
||||
|-------|------------|-------|------------|
|
||||
| GLM-Image | `Intel/GLM-Image-int4-AutoRound` | Diffusion transformer | CI-backed |
|
||||
| BAGEL | Not listed | Checkpoint-defined diffusion or transformer stage | Not validated |
|
||||
| Model | Scope | Status | Notes |
|
||||
|-------|-------|--------|-------|
|
||||
| GLM-Image | Diffusion transformer | ✅ | `Intel/GLM-Image-int4-AutoRound` |
|
||||
| BAGEL | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
|
||||
|
||||
## Configuration
|
||||
|
||||
@@ -134,14 +119,6 @@ The checkpoint should contain a config like:
|
||||
At load time, vLLM-Omni builds an `OmniINCConfig`, remaps checkpoint block names
|
||||
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:
|
||||
|
||||
```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
|
||||
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
|
||||
|
||||
| Device | Support |
|
||||
@@ -25,7 +21,7 @@ not imply that every diffusion transformer has been quality-validated.
|
||||
| Intel XPU | ❌ |
|
||||
| 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+`).
|
||||
|
||||
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)
|
||||
|
||||
| Model | HF models | Scope | Level | Recommendation |
|
||||
|-------|-----------|-------|-------|----------------|
|
||||
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Diffusion transformer | Validated | All heavy linear layers; sensitive embedders stay BF16 |
|
||||
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Diffusion transformer | Not validated | Compare with BF16 before enabling |
|
||||
| Wan2.2 | Wan2.2 diffusion pipelines | One or both diffusion transformers | Not validated | Validate quality and cascade behavior before enabling |
|
||||
| Model | HF models | CUDA | Mode | Recommendation |
|
||||
|-------|-----------|:----:|------|----------------|
|
||||
| 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` | Not validated | Online W4 weight-only | Compare vs BF16 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
|
||||
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
|
||||
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
|
||||
(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
|
||||
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
|
||||
ModelOpt checkpoints are pre-quantized checkpoints. Do not pass
|
||||
`--quantization fp8` for these checkpoints. The checkpoint
|
||||
@@ -63,11 +58,11 @@ stay unquantized unless the model-specific recipe validates otherwise.
|
||||
| Intel XPU | ⭕ |
|
||||
| Ascend NPU | ❌ |
|
||||
|
||||
Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
|
||||
guide. The optional CUTLASS FP8 runtime override requires CUDA SM89+. ModelOpt
|
||||
NVFP4 and mixed FP8/NVFP4 diffusion checkpoints are currently validated on
|
||||
Blackwell CUDA systems in the recipes below; other CUDA generations require
|
||||
separate backend and quality validation.
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
|
||||
The optional CUTLASS FP8 runtime override requires CUDA SM89+. ModelOpt NVFP4
|
||||
and mixed FP8/NVFP4 diffusion checkpoints are currently validated on Blackwell
|
||||
CUDA systems in the recipes below; other CUDA generations require separate
|
||||
backend and quality validation.
|
||||
|
||||
## 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 |
|
||||
| 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
|
||||
[`recipes/Qwen/Qwen-Image.md`](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image.md)
|
||||
and
|
||||
@@ -96,8 +88,8 @@ and
|
||||
|
||||
| Model | Scope | Status |
|
||||
|-------|-------|--------|
|
||||
| Qwen3-Omni | Thinker language-model stage (FP8) | Validated outside scheduled CI |
|
||||
| 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 | ModelOpt FP8 checkpoint path |
|
||||
| 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 |
|
||||
|
||||
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
|
||||
scales; the audio encoder, vision encoder, talker, and code2wav stay in BF16.
|
||||
|
||||
| Variant | HF checkpoint | Hardware and validation |
|
||||
|---------|---------------|-------------------------|
|
||||
| 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 |
|
||||
| Variant | HF checkpoint | Hardware |
|
||||
|---------|---------------|----------|
|
||||
| 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
|
||||
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
|
||||
ModelOpt path only when the checkpoint already contains ModelOpt quantized
|
||||
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
|
||||
|
||||
[msModelSlim](https://github.com/Ascend/msmodelslim) is an Ascend compression
|
||||
toolkit for producing pre-quantized model checkpoints. vLLM-Omni consumes its
|
||||
output either through a compatible generic Ascend checkpoint path or through
|
||||
the native MXFP8 and MXFP4 merge workflows documented in this guide.
|
||||
toolkit for producing pre-quantized model checkpoints. In vLLM-Omni, these
|
||||
checkpoints run through the Ascend/NPU path with `--quantization ascend`.
|
||||
|
||||
msModelSlim is static quantization: quantized weights are generated offline
|
||||
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
|
||||
|
||||
| Device | Support |
|
||||
@@ -25,28 +20,23 @@ backend alone is an integration point, not model-level validation.
|
||||
| Intel XPU | ❌ |
|
||||
| Ascend NPU | ✅ |
|
||||
|
||||
Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
### Diffusion Model (Qwen-Image, Wan2.2)
|
||||
|
||||
| Model | Base model | Scope | Level | 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 |
|
||||
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | DiT or diffusion stage | Not validated | No supported checkpoint workflow is documented |
|
||||
| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | DiT or diffusion stage | Integrated | Generate quantized weights with the HunyuanImage-3.0 msModelSlim adaptation and validate against BF16 |
|
||||
| Model | Base model | Scope | Hardware | Notes |
|
||||
|-------|------------|-------|----------|-------|
|
||||
| Wan2.2 | Wan2.2 diffusion weights | DiT or diffusion stage | Ascend NPU | Upstream msModelSlim provides a Wan2.2 quantization recipe; vLLM-Omni inference validation is not listed |
|
||||
| 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 | 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
|
||||
[HunyuanImage-3.0 msModelSlim adaptation](https://gitcode.com/betta18/msmodelslim/tree/hyimage3_mxfp8)
|
||||
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)
|
||||
|
||||
| Model | Scope | Status | Notes |
|
||||
@@ -63,32 +53,25 @@ than passing `--quantization ascend`.
|
||||
|
||||
## Configuration
|
||||
|
||||
For a compatible generic Ascend checkpoint, select the Ascend quantization
|
||||
method explicitly:
|
||||
|
||||
Offline inference:
|
||||
|
||||
```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:
|
||||
|
||||
```bash
|
||||
vllm serve <compatible-ascend-checkpoint> --omni --quantization ascend
|
||||
vllm serve <quantized-model-path> --omni --quantization ascend
|
||||
```
|
||||
|
||||
## Parameters
|
||||
|
||||
| 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 |
|
||||
|
||||
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:
|
||||
|
||||
```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.
|
||||
3. Keep the quantized checkpoint aligned with the same model architecture and
|
||||
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_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"
|
||||
For production deployments, use the `mxfp4_dualscale` offline mode with a
|
||||
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
|
||||
|
||||
| Device | Online `mxfp4` | Online `mxfp4_dualscale` | Offline `mxfp4_dualscale` |
|
||||
|--------|-----------------|---------------------------|----------------------------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ❌ | ❌ | ❌ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ❌ | ❌ | ❌ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ❌ | ❌ | ❌ |
|
||||
| AMD ROCm (gfx950 / MI355X) | ✅ | ❌ | ❌ |
|
||||
| Intel XPU | ❌ | ❌ | ❌ |
|
||||
| Ascend NPU (Atlas 950 A5) | ✅ | ✅ | ✅ |
|
||||
| Device | Support |
|
||||
|--------|---------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ⭕ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ⭕ |
|
||||
| AMD ROCm (gfx950 / MI355X) | ✅ |
|
||||
| Intel XPU | ⭕ |
|
||||
| Ascend NPU (Atlas 950 A5) | ✅ |
|
||||
|
||||
Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
|
||||
guide. Backend availability does not imply full-model checkpoint validation.
|
||||
|
||||
!!! note "Why NVIDIA is marked unsupported"
|
||||
This table describes the methods currently wired into vLLM-Omni, not every
|
||||
low-precision format NVIDIA hardware could execute. `DiffusionMXFP4Config`
|
||||
has runtime implementations only for Ascend and ROCm gfx950; on CUDA it
|
||||
raises `NotImplementedError` because no CUDA MXFP4 linear method or
|
||||
checkpoint loader is registered.
|
||||
|
||||
On NVIDIA, use [ModelOpt](modelopt.md) NVFP4 for the documented 4-bit
|
||||
checkpoint path. NVFP4 and OCP MXFP4 have different checkpoint metadata,
|
||||
scale layouts, and kernel contracts, so they are not interchangeable. A
|
||||
native CUDA MXFP4 path would require a CUDA linear implementation, loader
|
||||
contract, and model-level output validation.
|
||||
|
||||
The ROCm path dispatches online single-scale `mxfp4` to AITER on gfx950. It has
|
||||
kernel and dispatch coverage, but this guide does not yet claim end-to-end
|
||||
quality validation for a named Wan checkpoint on ROCm. DualScale and
|
||||
pre-quantized MXFP4 checkpoints are not supported on ROCm.
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
### Diffusion Model (Wan2.2)
|
||||
|
||||
| Model | Online (Ascend) | Offline (Ascend) | Validation and 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-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | Validated outside scheduled full-model offline NPU CI; same cascade scheme as T2V-A14B |
|
||||
| Wan2.2-TI2V-5B | ❌ | ❌ | Unsupported because W4A4 causes unacceptable accuracy loss |
|
||||
| Model | Online | Offline | Notes |
|
||||
|-------|--------|---------|-------|
|
||||
| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade (`transformer` + `transformer_2`); both transformers quantized with the same config |
|
||||
| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade; same scheme as T2V-A14B |
|
||||
| Wan2.2-TI2V-5B | ❌ | ❌ | Parameter count too small; W4A4 causes unacceptable accuracy loss |
|
||||
|
||||
The choice between `mxfp4` and `mxfp4_dualscale` in **online mode** is about
|
||||
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
|
||||
|
||||
Online mode requires no pre-processing. On Ascend, vLLM-Omni quantizes BF16
|
||||
weights at load time using `npu_dynamic_mx_quant`. On ROCm gfx950, the same
|
||||
`mxfp4` config dispatches to the AITER online kernel. A single block scale
|
||||
Online mode requires no pre-processing. vLLM-Omni quantizes BF16 weights to
|
||||
MXFP4 at load time using `npu_dynamic_mx_quant`. A single block scale
|
||||
(`float8_e8m0fnu`, one per 32 K elements) is computed on the fly; no
|
||||
calibration `mul_scale` is available. The config applies equally to
|
||||
single-transformer and cascade (A14B) models, but named Wan model validation is
|
||||
currently limited to the documented Ascend path.
|
||||
calibration `mul_scale` is available. Applies equally to single-transformer
|
||||
and cascade (A14B) models — both transformers in a cascade receive the same
|
||||
quantization config automatically.
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
@@ -135,10 +111,10 @@ vllm serve <your-model> --omni --quantization mxfp4
|
||||
|
||||
### `mxfp4_dualscale` — DualScale Online Mode
|
||||
|
||||
Ascend-only online DualScale mode computes both fine and coarse scales on the
|
||||
fly from BF16 weights using `npu_dynamic_dual_level_mx_quant`. It applies
|
||||
equally to single-transformer and cascade (A14B) models. Compared to `mxfp4`
|
||||
online, DualScale provides better quantization accuracy at higher compute cost.
|
||||
Online DualScale mode computes both fine and coarse scales on the fly from BF16
|
||||
weights using `npu_dynamic_dual_level_mx_quant`. Applies equally to
|
||||
single-transformer and cascade (A14B) models. Compared to `mxfp4` online,
|
||||
DualScale provides better quantization accuracy at higher compute cost.
|
||||
|
||||
The default configuration keeps the leading 5 transformer blocks in BF16
|
||||
(`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
|
||||
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
|
||||
|
||||
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 (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
|
||||
|
||||
| Device | Online | Offline (Native) | Offline (AutoRound) |
|
||||
|--------|--------|------------------|---------------------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ❌ | ❌ | ⭕ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ❌ | ❌ | ⭕ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ❌ | ❌ | ⭕ |
|
||||
| AMD ROCm | ❌ | ❌ | ⭕ |
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ⭕ | ⭕ | ⭕ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ | ⭕ | ⭕ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ⭕ | ⭕ | ⭕ |
|
||||
| AMD ROCm | ⭕ | ⭕ | ⭕ |
|
||||
| Intel XPU | ✅ | ❌ | ✅ |
|
||||
| Ascend NPU (Atlas 950 A5) | ✅ | ✅ | ⭕ |
|
||||
|
||||
Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
|
||||
guide. Backend availability does not imply full-model checkpoint validation.
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
|
||||
|
||||
The native `DiffusionMXFP8Config` currently dispatches only to Intel XPU and
|
||||
Ascend NPU implementations. CUDA reaches `NotImplementedError`, so the NVIDIA
|
||||
rows are marked unsupported even though this is an implementation gap rather
|
||||
than a general claim about NVIDIA hardware capabilities. On NVIDIA, use
|
||||
[FP8](fp8.md) for load-time quantization or [ModelOpt](modelopt.md) FP8 for
|
||||
pre-quantized checkpoints. Adding native CUDA MXFP8 requires a registered CUDA
|
||||
linear method, a compatible checkpoint contract, and model-level validation.
|
||||
|
||||
Native `mxfp8` is available online on Intel XPU and online or offline on
|
||||
Ascend. Intel XPU only supports AutoRound MXFP8 for offline mode. The AutoRound
|
||||
column is a separate checkpoint format and compute path.
|
||||
**Note**: Intel XPU only supports AutoRound MXFP8 for offline mode. Use AutoRound quantized checkpoints or online mode for XPU.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
### Diffusion Model (Wan2.2)
|
||||
|
||||
| Model | Mode | Scope | Validation |
|
||||
|-------|------|-------|------------|
|
||||
| Wan2.2-T2V-A14B | Online + Offline | MoE cascade: `transformer` + `transformer_2` | Validated outside scheduled full-model offline NPU CI |
|
||||
| Wan2.2-I2V-A14B | Online + Offline | MoE cascade: `transformer` + `transformer_2` | Validated outside scheduled full-model offline NPU CI |
|
||||
| Wan2.2-TI2V-5B | Online + Offline | Single diffusion transformer | Validated outside scheduled full-model offline NPU CI |
|
||||
| Model | Mode | Notes |
|
||||
|-------|------|-------|
|
||||
| Wan2.2-T2V-A14B | Online + Offline | MoE cascade; quantizes two transformers (`transformer` + `transformer_2`) |
|
||||
| Wan2.2-I2V-A14B | Online + Offline | MoE cascade; quantizes two transformers (`transformer` + `transformer_2`) |
|
||||
| Wan2.2-TI2V-5B | Online + Offline | Single transformer |
|
||||
|
||||
### 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
|
||||
wrote a separate `quantization_config.json` that is not read by vLLM-Omni;
|
||||
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
|
||||
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
|
||||
prepared before serving and are documented in their method-specific guides.
|
||||
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
|
||||
|
||||
| Device | FP8 W8A8 | Int8 W8A8 | BitsAndBytes W4 | MXFP8 W8A8 | MXFP4 W4A4 |
|
||||
|--------|----------|-----------|------------------|------------|------------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ❌ | ❌ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ❌ | ❌ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ✅ | ❌ | ❌ |
|
||||
| AMD ROCm (gfx950) | ⭕ | ⭕ | ❌ | ❌ | ✅ |
|
||||
| Intel XPU | ⭕ | ⭕ | ❌ | ✅ | ❌ |
|
||||
| Ascend NPU | ❌ | ✅ | ❌ | ✅ | ✅ |
|
||||
| Device | FP8 W8A8 | Int8 W8A8 | MXFP8 W8A8 | MXFP4 W4A4 |
|
||||
|--------|----------|-----------|------------|------------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ⭕ | ⭕ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ⭕ | ⭕ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ |
|
||||
| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ |
|
||||
| Intel XPU | ⭕ | ⭕ | ✅ | ⭕ |
|
||||
| Ascend NPU | ❌ | ✅ | ✅ | ✅ |
|
||||
|
||||
Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
|
||||
guide. These entries describe method-level backend availability, not validation
|
||||
for every model. FP8 on Ampere may use a weight-only path where available.
|
||||
On ROCm gfx950, MXFP4 means online single-scale `mxfp4`; online and offline
|
||||
`mxfp4_dualscale` are Ascend-only. The native MXFP8/MXFP4 configs do not
|
||||
currently register CUDA implementations; use FP8 or ModelOpt FP8/NVFP4 on
|
||||
NVIDIA instead.
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide. FP8 on Ampere may use a weight-only path where available. MXFP8 and
|
||||
MXFP4 are documented for the Ascend NPU path.
|
||||
|
||||
## 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 |
|
||||
| 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 |
|
||||
| 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)
|
||||
|
||||
@@ -64,7 +59,6 @@ from vllm_omni import Omni
|
||||
|
||||
omni_fp8 = Omni(model="<your-model>", quantization="fp8")
|
||||
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_mxfp4 = Omni(model="<your-model>", quantization="mxfp4")
|
||||
omni_mxfp4_dualscale = Omni(model="<your-model>", quantization="mxfp4_dualscale")
|
||||
@@ -75,7 +69,6 @@ CLI:
|
||||
```bash
|
||||
vllm serve <your-model> --omni --quantization fp8
|
||||
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 mxfp4
|
||||
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
|
||||
|
||||
| Parameter | Methods | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `method` | FP8, Int8, BitsAndBytes, MXFP8, MXFP4 | Quantization method: `"fp8"`, `"int8"`, `"bitsandbytes"`, `"mxfp8"`, `"mxfp4"`, or `"mxfp4_dualscale"` |
|
||||
| `ignored_layers` | FP8, Int8, BitsAndBytes, MXFP8, MXFP4 | Layer name patterns to keep in BF16/FP16 |
|
||||
| `method` | FP8, Int8, MXFP8, MXFP4 | Quantization method: `"fp8"`, `"int8"`, `"mxfp8"`, `"mxfp4"`, or `"mxfp4_dualscale"` |
|
||||
| `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 |
|
||||
| `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` |
|
||||
|
||||
## 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 |
|
||||
| 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
|
||||
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.
|
||||
## Hardware Support
|
||||
|
||||
## 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
|
||||
hardware backend. The tables in this guide use the following levels:
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide. FP8 on Ampere may use a weight-only path where available.
|
||||
|
||||
| Level | Meaning |
|
||||
|-------|---------|
|
||||
| **CI-backed** | A named model and quantized checkpoint are exercised by scheduled full-model hardware CI. |
|
||||
| **Validated** | A named model and checkpoint have end-to-end or quality-validation evidence, but the check is not necessarily scheduled in CI. |
|
||||
| **Integrated** | The config, loader, quantized layers, or model wiring exists, but this guide has no completed end-to-end validation for the named model and checkpoint. |
|
||||
| **Not validated** | There is no current support recommendation for this model/component combination. |
|
||||
| **Unsupported** | The combination is intentionally rejected or known not to work. |
|
||||
## Model Type Support
|
||||
|
||||
A method being available on a device does not make every model on that device
|
||||
supported. Likewise, successfully detecting a checkpoint's
|
||||
`quantization_config` establishes loader compatibility, not output quality.
|
||||
### Diffusion Model (Qwen-Image, Wan2.2)
|
||||
|
||||
## 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
|
||||
separate methods: NVIDIA ModelOpt produces the ModelOpt FP8/NVFP4 formats,
|
||||
Intel AutoRound produces AutoRound checkpoints, and Ascend msModelSlim produces
|
||||
checkpoints consumed by the native MXFP or compatible Ascend paths.
|
||||
| Method | Guide | Mode | Example models | Status |
|
||||
|--------|-------|------|----------------|--------|
|
||||
| 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 |
|
||||
| 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) |
|
||||
|----------|-----------------------|------------------|------------------------------------------|--------------------------------|-------------------------------------------------------------------------|--------------------------------------|--------------------------|--------------------------|
|
||||
| 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 | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ |
|
||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||
|
||||
Legend: `✅` backend available, `❌` rejected or no implementation, `⭕` not
|
||||
verified in this guide. This matrix does not establish model-level support; use
|
||||
the model-and-stage tables below for that.
|
||||
These models combine an AR language model with audio, vision, talker, or TTS
|
||||
stages. Quantization is scoped to the AR language-model stage when the
|
||||
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
|
||||
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
|
||||
means online single-scale `mxfp4` on gfx950, while MXFP4 DualScale is
|
||||
Ascend-only. NVIDIA's NVFP4 and OCP MXFP4 have different checkpoint metadata,
|
||||
scale layouts, and kernel contracts. Mixed ModelOpt FP8/NVFP4 checkpoints also
|
||||
require a model-specific mixed-precision policy and validation. FP8 on Ampere
|
||||
may use a weight-only path where available.
|
||||
| Method | Guide | Scope | Example models | Status |
|
||||
|--------|-------|-------|----------------|--------|
|
||||
| ModelOpt | [ModelOpt](modelopt.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path |
|
||||
| Int8 | [Int8](int8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
| BitsAndBytes | [BitsAndBytes](bitsandbytes.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
| MXFP8 | [MXFP8](mxfp8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
| MXFP4 | [MXFP4](mxfp4.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
| 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
|
||||
checkpoint quantizes the stage named in the **Scope** column; it does not imply
|
||||
that the model's encoders, VAE, decoder, or other stages are also quantized.
|
||||
These models split generation across multiple stages. Quantization must be
|
||||
attached to the intended stage rather than applied globally.
|
||||
|
||||
### Vendor and Runtime Ecosystems
|
||||
|
||||
The producer and runtime backend are related, but they are not the same support
|
||||
claim. A producer can emit several checkpoint formats, and each format still
|
||||
needs a compatible vLLM-Omni loader, model-stage wiring, and validation. A
|
||||
runtime may also provide online quantization without an offline producer path.
|
||||
|
||||
| Ecosystem | Current paths | Current vLLM-Omni scope | Model evidence |
|
||||
|-----------|---------------|-------------------------|----------------|
|
||||
| NVIDIA | ModelOpt FP8, NVFP4, and mixed FP8/NVFP4 | Diffusion transformer or Qwen-Omni thinker | Qwen-Image and other named diffusion checkpoints are validated; Qwen3-Omni NVFP4 has H100 CI for the Marlin fallback |
|
||||
| Intel | AutoRound W4A16 and AutoRound MXFP8 | Diffusion/world-model transformer or Qwen-Omni thinker | Wan A14B and GLM-Image W4A16 are CI-backed; Qwen-Image, Qwen-Omni, and Cosmos3-Super have validation outside scheduled CI |
|
||||
| AMD ROCm | Online OCP MXFP4 through AITER on gfx950; no pre-quantized MXFP4 or AMD Quark workflow is integrated | Diffusion transformer linear layers | Integrated backend with dispatch and weight-allocation unit coverage; no named-model quality validation or scheduled gfx950 full-model CI |
|
||||
| Ascend | msModelSlim output through native MXFP8/MXFP4 merge workflows or a compatible generic Ascend format | Wan diffusion transformers | Native Wan MXFP8/MXFP4 paths are validated outside scheduled full-model NPU CI; other model families remain integrated or unvalidated |
|
||||
|
||||
### Pre-Quantized Model Matrix
|
||||
|
||||
| Model family | Method | Scope | Level | Notes |
|
||||
|--------------|--------|-------|-------|-------|
|
||||
| Qwen-Image | [ModelOpt](modelopt.md) FP8 or mixed FP8/NVFP4 | Diffusion transformer | Validated | Named Qwen-Image 2512 checkpoints and recipes are available. Text encoder and VAE stay BF16. |
|
||||
| Qwen-Image | [AutoRound](autoround.md) W4A16 | Diffusion transformer | Validated | Uses `INC4AI/Qwen-Image-AutoRound-W4A16`; text encoder and VAE stay BF16. |
|
||||
| Wan2.2 I2V/T2V A14B | [AutoRound](autoround.md) W4A16 | Both diffusion transformers | CI-backed | The high-noise and low-noise DiTs are quantized; text encoder and VAE stay BF16. |
|
||||
| Wan2.2 TI2V 5B | [AutoRound](autoround.md) W4A16 | Diffusion transformer | Validated | Named checkpoint available; not covered by the scheduled Wan AutoRound job. |
|
||||
| Wan2.2 | [MXFP8](mxfp8.md) | One or both diffusion transformers | Validated | Native Ascend and AutoRound checkpoint formats are integrated; no scheduled full-model offline NPU job. |
|
||||
| Wan2.2 I2V/T2V A14B | [MXFP4 DualScale](mxfp4.md) | Both diffusion transformers | Validated | Offline calibrated checkpoints are recommended; no scheduled full-model offline NPU job. |
|
||||
| Qwen2.5-Omni | [AutoRound](autoround.md) W4A16 | Thinker language model | Validated | Encoders and output/audio stages stay BF16. |
|
||||
| Qwen3-Omni | [ModelOpt](modelopt.md) FP8 | Thinker language model | Validated | Tested outside scheduled CI. |
|
||||
| Qwen3-Omni | [ModelOpt](modelopt.md) NVFP4 | Thinker language model | CI-backed | H100 CI exercises the Marlin fallback; the native Blackwell FP4 path has separate recipe validation. |
|
||||
| Qwen3-Omni | [AutoRound](autoround.md) W4A16 | Thinker language model | Validated | Audio encoder, vision encoder, talker, and Code2Wav stay BF16. |
|
||||
| GLM-Image | [AutoRound](autoround.md) W4A16 | Diffusion transformer stage | CI-backed | Text-to-image and image-to-image full-model smoke coverage. |
|
||||
| Cosmos3-Super | [AutoRound](autoround.md) W4A16 | World-model transformer | Validated | Manually validated serving path; VAE and other components stay BF16. |
|
||||
| Cosmos3-Nano | [AutoRound](autoround.md) W4A16 | World-model transformer | Integrated | Named checkpoint is listed, but no model-specific end-to-end test is maintained in-tree. |
|
||||
| MiniMax H3 | None | None | Not validated | No offline vendor checkpoint path is documented; online FP8 targets the DiT only. |
|
||||
|
||||
The method pages contain the complete checkpoint lists, including FLUX,
|
||||
Z-Image, and HunyuanImage-3.0 variants.
|
||||
|
||||
### Online Quantization
|
||||
|
||||
| Model family | Methods | Scope | Level |
|
||||
|--------------|---------|-------|-------|
|
||||
| Qwen-Image | FP8, Int8 | Diffusion transformer | Validated |
|
||||
| Z-Image | Int8, BitsAndBytes W4 | Diffusion transformer | Validated |
|
||||
| Wan2.2 | MXFP8 | One or both diffusion transformers | Validated on the documented XPU/NPU paths |
|
||||
| Wan2.2 A14B | MXFP4 | Both diffusion transformers | Validated on Ascend; the ROCm gfx950 online backend is integrated but not model-validated |
|
||||
| MiniMax H3 | FP8 | DiT and optional reference DiT | Validated |
|
||||
| Qwen3-Omni | Generic online methods, such as FP8 | Thinker language model only | Integrated; not model-validated as an online recommendation |
|
||||
|
||||
See [Online Quantization](online.md) for configuration and hardware details.
|
||||
|
||||
### Unlisted and Long-Tail Models
|
||||
|
||||
For models not listed above, generic quantized linear layers or compatible
|
||||
checkpoint metadata are only an integration starting point. The default support
|
||||
target is the model's primary transformer, DiT, or AR language-model stage.
|
||||
Text/audio/vision encoders, VAEs, vocoders, decoders, and output stages remain
|
||||
BF16 until a model-specific integration supplies quant-aware layers, checkpoint
|
||||
mapping, end-to-end output validation, and backend coverage.
|
||||
| Method | Guide | Scope | Example models | Status |
|
||||
|--------|-------|-------|----------------|--------|
|
||||
| FP8 | [FP8](fp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
|
||||
| Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
|
||||
| BitsAndBytes | [BitsAndBytes](bitsandbytes.md) | Stage-specific transformer or DiT module | BAGEL, GLM-Image | Not validated |
|
||||
| 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 |
|
||||
| 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 |
|
||||
| msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation |
|
||||
|
||||
!!! note
|
||||
"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)
|
||||
|
||||
The default target is the diffusion transformer. `build_quant_config()` can
|
||||
construct a component router:
|
||||
The default target is the diffusion transformer. Component routing is available
|
||||
through `build_quant_config()`:
|
||||
|
||||
```python
|
||||
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 |
|
||||
|-----------|--------------------|-------|
|
||||
| Diffusion transformer | Yes, for listed model/method pairs | Primary target for FP8, Int8, BitsAndBytes, ModelOpt, MXFP8, MXFP4, AutoRound, and msModelSlim |
|
||||
| Second diffusion transformer | Yes, for listed Wan2.2 A14B methods | `transformer` and `transformer_2` read their own offline checkpoint metadata; arbitrary mixed methods are not currently validated |
|
||||
| Diffusion transformer | Yes | Primary target for FP8, Int8, BitsAndBytes, ModelOpt, MXFP8, MXFP4, AutoRound, and msModelSlim |
|
||||
| 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 |
|
||||
|
||||
### 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 |
|
||||
| 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)
|
||||
|
||||
| Component | Default quantized? | Notes |
|
||||
@@ -212,7 +141,6 @@ build_quant_config("fp8")
|
||||
build_quant_config({"method": "fp8", "activation_scheme": "static"})
|
||||
build_quant_config("bitsandbytes")
|
||||
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(None)
|
||||
```
|
||||
|
||||
@@ -59,7 +59,7 @@ def main(args):
|
||||
|
||||
omni = Omni(
|
||||
model=args.model_name,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
deploy_config=args.deploy_config,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
batch_timeout=args.batch_timeout,
|
||||
@@ -167,10 +167,10 @@ def parse_args():
|
||||
help="Sampling rate for audio loading (default: 16000).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to stage configs YAML file.",
|
||||
help="Path to a deploy config YAML file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
|
||||
@@ -876,7 +876,7 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
|
||||
parser.add_argument("--task", type=str, required=True, choices=TASK_CHOICES)
|
||||
parser.add_argument("--model", type=str, required=True, help="HF repo id or local model directory.")
|
||||
parser.add_argument(
|
||||
"--stage-config-path",
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=str(repo_root / "vllm_omni/deploy/dynin_omni.yaml"),
|
||||
help="Path to deploy config yaml.",
|
||||
@@ -1394,8 +1394,8 @@ def main() -> None:
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
stage_config_path = str(Path(args.stage_config_path).expanduser())
|
||||
omni = Omni(model=model_source, stage_configs_path=stage_config_path, dtype=args.dtype)
|
||||
deploy_config_path = str(Path(args.deploy_config).expanduser())
|
||||
omni = Omni(model=model_source, deploy_config=deploy_config_path, dtype=args.dtype)
|
||||
sampling_params_list = [
|
||||
SamplingParams(max_tokens=int(args.max_tokens_per_stage), temperature=0.0, top_p=1.0, detokenize=False)
|
||||
for _ in range(omni.num_stages)
|
||||
|
||||
@@ -73,7 +73,6 @@ def parse_args():
|
||||
)
|
||||
parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.")
|
||||
parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.")
|
||||
parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom legacy stage config YAML path.")
|
||||
parser.add_argument(
|
||||
"--stream",
|
||||
action="store_true",
|
||||
@@ -146,13 +145,7 @@ def main():
|
||||
else:
|
||||
bot_task = args.bot_task
|
||||
|
||||
if args.deploy_config is not None and args.stage_configs_path is not None:
|
||||
raise ValueError("--deploy-config and --stage-configs-path are mutually exclusive.")
|
||||
|
||||
deploy_config = args.deploy_config
|
||||
stage_configs_path = args.stage_configs_path
|
||||
if deploy_config is None and stage_configs_path is None:
|
||||
deploy_config = _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality]
|
||||
deploy_config = args.deploy_config or _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality]
|
||||
|
||||
omni_kwargs = {
|
||||
"model": args.model,
|
||||
@@ -168,10 +161,7 @@ def main():
|
||||
|
||||
if additional_config is not None:
|
||||
omni_kwargs["additional_config"] = additional_config
|
||||
if deploy_config is not None:
|
||||
omni_kwargs["deploy_config"] = deploy_config
|
||||
else:
|
||||
omni_kwargs["stage_configs_path"] = stage_configs_path
|
||||
|
||||
omni = Omni(**omni_kwargs)
|
||||
|
||||
@@ -272,10 +262,7 @@ def main():
|
||||
print(f" Modality: {args.modality}")
|
||||
print(f" Prompt task: {task}")
|
||||
print(f" Bot task: {bot_task}")
|
||||
if deploy_config is not None:
|
||||
print(f" Deploy config: {deploy_config}")
|
||||
else:
|
||||
print(f" Stage config: {stage_configs_path}")
|
||||
print(f" Num stages: {omni.num_stages}")
|
||||
if args.modality in ("text2img", "img2img"):
|
||||
print(f" Inference steps: {args.steps}")
|
||||
|
||||
@@ -452,12 +452,6 @@ def parse_args():
|
||||
default=None,
|
||||
help="Path to a .txt file with one prompt per line (preferred).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a stage configs file (deprecated; prefer --deploy-config).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
|
||||
@@ -484,10 +484,10 @@ def parse_args():
|
||||
help="Path to a .txt file with one prompt per line (preferred).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a stage configs file.",
|
||||
help="Path to a deploy config YAML.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video-path",
|
||||
|
||||
@@ -232,15 +232,8 @@ def main(args):
|
||||
raise ValueError(f"Unknown query type: {args.query_type}")
|
||||
|
||||
# Initialize vLLM-Omni with Step-Audio2
|
||||
# Resolve deploy config path. Keep explicit stage_configs_path for legacy
|
||||
# custom configs, but use bundled deploy configs by default.
|
||||
deploy_config_path = None
|
||||
stage_config_path = None
|
||||
if args.stage_configs_path and args.deploy_config:
|
||||
raise ValueError("--stage-configs-path and --deploy-config are mutually exclusive")
|
||||
if args.stage_configs_path:
|
||||
stage_config_path = args.stage_configs_path
|
||||
elif args.deploy_config:
|
||||
# Resolve an explicit deploy config or use the bundled task default.
|
||||
if args.deploy_config:
|
||||
deploy_config_path = args.deploy_config
|
||||
else:
|
||||
configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/deploy"
|
||||
@@ -262,9 +255,6 @@ def main(args):
|
||||
"ray_address": args.ray_address,
|
||||
"trust_remote_code": True,
|
||||
}
|
||||
if stage_config_path is not None:
|
||||
omni_kwargs["stage_configs_path"] = stage_config_path
|
||||
else:
|
||||
omni_kwargs["deploy_config"] = deploy_config_path
|
||||
|
||||
omni_llm = Omni(**omni_kwargs)
|
||||
@@ -418,12 +408,6 @@ def parse_args():
|
||||
" - HuggingFace ID: stepfun-ai/Step-Audio-2-mini"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a legacy stage config YAML file (default: use bundled deploy config)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
|
||||
@@ -92,12 +92,6 @@ def parse_args() -> argparse.Namespace:
|
||||
"meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, "
|
||||
"stabilityai/stable-diffusion-3.5-medium, Tongyi-MAI/Z-Image-Turbo and etc.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="[Deprecated] Path to a legacy stage_args-format YAML. Prefer --deploy-config.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
@@ -475,8 +469,6 @@ def main():
|
||||
omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size
|
||||
if args.enforce_eager is not None:
|
||||
omni_kwargs["enforce_eager"] = args.enforce_eager
|
||||
if args.stage_configs_path:
|
||||
omni_kwargs["stage_configs_path"] = args.stage_configs_path
|
||||
if args.deploy_config:
|
||||
omni_kwargs["deploy_config"] = args.deploy_config
|
||||
if use_nextstep:
|
||||
@@ -515,8 +507,8 @@ def main():
|
||||
print(f" Image size: {args.width}x{args.height}")
|
||||
if args.lora_path:
|
||||
print(f" LoRA: scale={args.lora_scale}")
|
||||
if args.stage_configs_path:
|
||||
print(f" stage-configs-path: {args.stage_configs_path}")
|
||||
if args.deploy_config:
|
||||
print(f" deploy-config: {args.deploy_config}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
# Build LoRA request when --lora-path is set
|
||||
|
||||
@@ -38,7 +38,7 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_MODEL = "fishaudio/s2-pro"
|
||||
DEFAULT_STAGE_CONFIG = os.path.join(
|
||||
DEFAULT_DEPLOY_CONFIG = os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
"..",
|
||||
"..",
|
||||
@@ -152,10 +152,7 @@ def _build_omni_kwargs(args, model_name: str) -> dict:
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
if args.deploy_config:
|
||||
omni_kwargs["deploy_config"] = args.deploy_config
|
||||
else:
|
||||
omni_kwargs["stage_configs_path"] = args.stage_configs_path or DEFAULT_STAGE_CONFIG
|
||||
omni_kwargs["deploy_config"] = args.deploy_config or DEFAULT_DEPLOY_CONFIG
|
||||
return omni_kwargs
|
||||
|
||||
|
||||
@@ -278,12 +275,6 @@ def parse_args():
|
||||
default=None,
|
||||
help="Deploy config YAML path (under vllm_omni/deploy/).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to stage configs YAML (legacy, prefer --deploy-config).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
default="output_audio",
|
||||
|
||||
@@ -145,7 +145,7 @@ def main(args):
|
||||
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
stage_configs_path=deploy_config_path,
|
||||
deploy_config=deploy_config_path,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
|
||||
@@ -490,10 +490,10 @@ def parse_args():
|
||||
help="Path to a .txt file with one prompt per line (preferred).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a stage configs file.",
|
||||
help="Path to a deploy config YAML.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-path",
|
||||
|
||||
@@ -73,6 +73,8 @@ plugins:
|
||||
- awesome-nav
|
||||
- redirects:
|
||||
redirect_maps:
|
||||
community/meetups.md: community/meetings.md
|
||||
community/volunteers.md: community/contact_us.md
|
||||
user_guide/diffusion/request_batching.md: user_guide/diffusion/execution_modes.md
|
||||
user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md
|
||||
user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md
|
||||
|
||||
@@ -6,7 +6,7 @@ Unit tests for StageConfigFactory and related classes.
|
||||
|
||||
import importlib
|
||||
import warnings
|
||||
from dataclasses import dataclass, fields
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
@@ -34,9 +34,8 @@ from vllm_omni.config.stage_config import (
|
||||
load_deploy_config,
|
||||
merge_pipeline_deploy,
|
||||
pipeline_cfg_resolver,
|
||||
strip_parent_engine_args,
|
||||
)
|
||||
from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_blacklist_keys
|
||||
from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -475,7 +474,7 @@ class TestStageConfigFactory:
|
||||
cli_overrides = {
|
||||
"gpu_memory_utilization": 0.9,
|
||||
"model": "some_model", # Internal
|
||||
"stage_configs_path": "/path", # Internal
|
||||
"deploy_config": "/path", # Internal
|
||||
"batch_timeout": 10, # Internal
|
||||
}
|
||||
|
||||
@@ -483,7 +482,7 @@ class TestStageConfigFactory:
|
||||
|
||||
assert overrides["gpu_memory_utilization"] == 0.9
|
||||
assert "model" not in overrides
|
||||
assert "stage_configs_path" not in overrides
|
||||
assert "deploy_config" not in overrides
|
||||
assert "batch_timeout" not in overrides
|
||||
|
||||
def test_per_stage_override_excludes_internal_keys(self):
|
||||
@@ -525,49 +524,6 @@ class TestStageResolutionHelpers:
|
||||
assert "model" not in overrides
|
||||
assert "parallel_config" not in overrides
|
||||
|
||||
def test_strip_parent_engine_args_reports_only_surprising_parent_overrides(self):
|
||||
parent_fields = {f.name: f for f in fields(EngineArgs)}
|
||||
filtered, overridden = strip_parent_engine_args(
|
||||
{
|
||||
"model": "some/model",
|
||||
"stage_configs_path": "/tmp/stages.yaml",
|
||||
"tensor_parallel_size": 4,
|
||||
"worker_extension_cls": "some.Extension",
|
||||
"custom_pipeline_args": {"pipeline_class": "demo.Pipeline"},
|
||||
},
|
||||
parent_fields=parent_fields,
|
||||
keep_keys={"worker_extension_cls"},
|
||||
strip_keys={"stage_configs_path"},
|
||||
no_warn_keys={"model"},
|
||||
)
|
||||
|
||||
assert filtered == {
|
||||
"worker_extension_cls": "some.Extension",
|
||||
"custom_pipeline_args": {"pipeline_class": "demo.Pipeline"},
|
||||
}
|
||||
assert overridden == ["tensor_parallel_size"]
|
||||
|
||||
def test_strip_parent_engine_args_keeps_allowed_media_access_controls(self):
|
||||
parent_fields = {f.name: f for f in fields(EngineArgs)}
|
||||
filtered, overridden = strip_parent_engine_args(
|
||||
{
|
||||
"model": "some/model",
|
||||
"stage_configs_path": "/tmp/stages.yaml",
|
||||
"allowed_local_media_path": "/data/qwentts",
|
||||
"allowed_media_domains": ["example.com"],
|
||||
},
|
||||
parent_fields=parent_fields,
|
||||
keep_keys={"allowed_local_media_path", "allowed_media_domains"},
|
||||
strip_keys={"stage_configs_path"},
|
||||
no_warn_keys={"model"},
|
||||
)
|
||||
|
||||
assert filtered == {
|
||||
"allowed_local_media_path": "/data/qwentts",
|
||||
"allowed_media_domains": ["example.com"],
|
||||
}
|
||||
assert overridden == []
|
||||
|
||||
|
||||
class TestPipelineDiscovery:
|
||||
"""Tests for the central pipeline registry (``OMNI_PIPELINES``)."""
|
||||
|
||||
@@ -582,8 +582,12 @@ def test_diffusion_parallel_config_keeps_current_diffusion_parallel_surface():
|
||||
assert cfg.data_parallel_size == 3
|
||||
assert cfg.cfg_parallel_size == 3
|
||||
assert cfg.mask_sp_padding is True
|
||||
assert cfg.vae_parallel_mode == "tile"
|
||||
assert cfg.world_size == 72
|
||||
|
||||
auto_cfg = OmniStageDiffusionParallelConfig(vae_parallel_mode="auto")
|
||||
assert auto_cfg.vae_parallel_mode == "auto"
|
||||
|
||||
|
||||
def test_parallel_config_derived_fields_are_not_init_inputs():
|
||||
with pytest.raises(ValidationError):
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import os
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
@@ -10,6 +11,7 @@ from tests.helpers.mark import hardware_test
|
||||
from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
|
||||
DistributedAutoencoderKLWan,
|
||||
select_auto_spatial_shard_split_dim,
|
||||
)
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
@@ -17,6 +19,23 @@ from vllm_omni.platforms import current_omni_platform
|
||||
# the end uses hardware_test (H100 x2) plus full_model / diffusion / parallel.
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _active_spatial_shard(split_dim: str):
|
||||
token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
|
||||
wan_spatial_shard.SpatialShardContext(
|
||||
input_extent=2,
|
||||
local_input_extent=2,
|
||||
split_dim=split_dim,
|
||||
rank=0,
|
||||
world_size=3,
|
||||
)
|
||||
)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_split_for_parallel_decode_pads_uneven_height():
|
||||
@@ -332,7 +351,7 @@ def test_attention_wrapper_rejects_spatial_extent_change(monkeypatch: pytest.Mon
|
||||
lambda x, *, group, dim, dst=None: torch.cat([x, x], dim=dim),
|
||||
)
|
||||
module = ShrinkingWanAttentionBlock()
|
||||
wan_spatial_shard._patch_attention_block(module, group=object(), split_dim="width")
|
||||
wan_spatial_shard._patch_attention_block(module, group=object())
|
||||
x = torch.zeros((1, 1, 1, 1, 2), dtype=torch.float32)
|
||||
token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
|
||||
wan_spatial_shard.SpatialShardContext(
|
||||
@@ -377,11 +396,13 @@ def test_dist_zero_pad_only_applies_global_height_edges(monkeypatch: pytest.Monk
|
||||
|
||||
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3))
|
||||
mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object())
|
||||
with _active_spatial_shard("height"):
|
||||
mid = mid_rank_pad(x)
|
||||
assert mid.shape == (1, 1, 2, 3)
|
||||
|
||||
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3))
|
||||
last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object())
|
||||
with _active_spatial_shard("height"):
|
||||
last = last_rank_pad(x)
|
||||
assert last.shape == (1, 1, 3, 3)
|
||||
|
||||
@@ -392,12 +413,14 @@ def test_dist_zero_pad_only_applies_global_width_edges(monkeypatch: pytest.Monke
|
||||
x = torch.ones((1, 1, 2, 2))
|
||||
|
||||
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3))
|
||||
mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width")
|
||||
mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object())
|
||||
with _active_spatial_shard("width"):
|
||||
mid = mid_rank_pad(x)
|
||||
assert mid.shape == (1, 1, 2, 2)
|
||||
|
||||
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3))
|
||||
last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width")
|
||||
last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object())
|
||||
with _active_spatial_shard("width"):
|
||||
last = last_rank_pad(x)
|
||||
assert last.shape == (1, 1, 2, 3)
|
||||
|
||||
@@ -421,6 +444,111 @@ def test_spatial_shard_height_gate_falls_back_for_partial_group(monkeypatch: pyt
|
||||
assert vae._spatial_shard_decode_enabled(z) is False
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.parametrize(
|
||||
"z_shape,world_size,expected",
|
||||
[
|
||||
((1, 16, 1, 60, 104), 4, "width"),
|
||||
((1, 16, 5, 60, 104), 4, "width"),
|
||||
((1, 16, 5, 104, 60), 4, "height"),
|
||||
((1, 16, 5, 2, 2), 4, None),
|
||||
((1, 16, 5, 64, 64), 4, "width"),
|
||||
((1, 16, 64, 64), 4, None),
|
||||
((1, 16, 5, 64, 64), 1, None),
|
||||
],
|
||||
)
|
||||
def test_auto_spatial_shard_selector_is_shape_aware(z_shape, world_size, expected):
|
||||
assert select_auto_spatial_shard_split_dim(z_shape, world_size) == expected
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_auto_spatial_shard_gate_requires_full_group(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size",
|
||||
lambda group=None: 4,
|
||||
)
|
||||
|
||||
vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
|
||||
vae.use_tiling = True
|
||||
vae.distributed_executor = SimpleNamespace(group=object(), parallel_size=4, parallel_mode="auto")
|
||||
vae.is_distributed_enabled = lambda: True
|
||||
large = torch.zeros((1, 16, 5, 60, 104))
|
||||
unshardable = torch.zeros((1, 16, 1, 2, 2))
|
||||
unsupported = torch.zeros((1, 16, 60, 104))
|
||||
|
||||
assert vae._spatial_shard_decode_split_dim(large, 4) == "width"
|
||||
assert vae._spatial_shard_decode_enabled(large) is True
|
||||
assert vae._spatial_shard_decode_enabled(unshardable) is False
|
||||
assert vae._spatial_shard_decode_enabled(unsupported) is False
|
||||
|
||||
vae.distributed_executor.parallel_size = 2
|
||||
assert vae._spatial_shard_decode_enabled(large) is False
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_set_parallel_size_prepares_spatial_wrappers_only_for_a_full_group(monkeypatch: pytest.MonkeyPatch):
|
||||
vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
|
||||
executor = SimpleNamespace(group=object(), parallel_size=1, parallel_mode="tile")
|
||||
|
||||
def set_parallel_size(parallel_size: int, mode: str = "tile") -> None:
|
||||
executor.parallel_size = parallel_size
|
||||
executor.parallel_mode = mode
|
||||
|
||||
executor.set_parallel_size = set_parallel_size
|
||||
object.__setattr__(vae, "distributed_executor", executor)
|
||||
installed = []
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size",
|
||||
lambda group=None: 4,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
wan_spatial_shard,
|
||||
"install_wan_spatial_shard_decode",
|
||||
lambda candidate, group, split_dim: installed.append((candidate, group, split_dim)),
|
||||
)
|
||||
|
||||
vae.set_parallel_size(2, mode="auto")
|
||||
assert installed == []
|
||||
|
||||
vae.set_parallel_size(4, mode="auto")
|
||||
assert installed == [(vae, vae.distributed_executor.group, "height")]
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
@torch.no_grad()
|
||||
def test_installed_spatial_decoder_keeps_direct_fallback_and_can_change_axis():
|
||||
vae = DistributedAutoencoderKLWan(
|
||||
base_dim=8,
|
||||
decoder_base_dim=8,
|
||||
z_dim=4,
|
||||
dim_mult=[1, 2],
|
||||
num_res_blocks=1,
|
||||
attn_scales=[],
|
||||
temperal_downsample=[False],
|
||||
latents_mean=[0.0] * 4,
|
||||
latents_std=[1.0] * 4,
|
||||
scale_factor_temporal=1,
|
||||
scale_factor_spatial=2,
|
||||
).eval()
|
||||
latents = torch.randn(1, 4, 1, 4, 5)
|
||||
parameter_names = {name for name, _ in vae.named_parameters()}
|
||||
state_dict = {name: value.clone() for name, value in vae.state_dict().items()}
|
||||
reference = vae.decode(latents, return_dict=False)[0]
|
||||
|
||||
wan_spatial_shard.install_wan_spatial_shard_decode(vae, group=object(), split_dim="height")
|
||||
assert any(getattr(module, "_vllm_omni_dynamic_spatial_shard_conv", False) for module in vae.decoder.modules())
|
||||
wan_spatial_shard.install_wan_spatial_shard_decode(vae, group=object(), split_dim="width")
|
||||
|
||||
assert {name for name, _ in vae.named_parameters()} == parameter_names
|
||||
vae.load_state_dict(state_dict, strict=True)
|
||||
assert torch.equal(vae.decode(latents, return_dict=False)[0], reference)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_spatial_shard_width_gate_selects_width(monkeypatch: pytest.MonkeyPatch):
|
||||
@@ -566,3 +694,99 @@ def test_spatial_shard_decode_matches_reference(split_dim: str):
|
||||
assert mean_abs_diff <= _SPATIAL_SHARD_TOLERANCE, (
|
||||
f"spatial_shard_{split_dim} mean_abs_diff {mean_abs_diff} exceeds {_SPATIAL_SHARD_TOLERANCE}"
|
||||
)
|
||||
|
||||
|
||||
def _auto_spatial_shard_decode_worker(rank: int, return_dict, master_port: str) -> None:
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
destroy_model_parallel,
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
|
||||
os.environ["MASTER_ADDR"] = "127.0.0.1"
|
||||
os.environ["MASTER_PORT"] = master_port
|
||||
device = current_omni_platform.get_torch_device(rank)
|
||||
current_omni_platform.set_device(device)
|
||||
backend = current_omni_platform.dist_backend
|
||||
init_distributed_environment(world_size=2, rank=rank, local_rank=rank, backend=backend)
|
||||
initialize_model_parallel(sequence_parallel_size=2, ulysses_degree=2, backend=backend)
|
||||
|
||||
try:
|
||||
vae = DistributedAutoencoderKLWan.from_pretrained(
|
||||
_SPATIAL_SHARD_MODEL,
|
||||
subfolder=_SPATIAL_SHARD_SUBFOLDER,
|
||||
torch_dtype=torch.float32,
|
||||
).to(device=device, dtype=torch.float32)
|
||||
vae.eval()
|
||||
generator = torch.Generator(device=device).manual_seed(1)
|
||||
small = torch.randn((1, vae.config.z_dim, 1, 60, 104), generator=generator, device=device)
|
||||
landscape = torch.randn((1, vae.config.z_dim, 2, 60, 104), generator=generator, device=device)
|
||||
portrait = torch.randn((1, vae.config.z_dim, 2, 104, 60), generator=generator, device=device)
|
||||
|
||||
with torch.inference_mode():
|
||||
vae.use_tiling = False
|
||||
vae.set_parallel_size(1, mode="tile")
|
||||
small_reference = vae.decode(small, return_dict=False)[0]
|
||||
|
||||
vae.use_tiling = True
|
||||
vae.set_parallel_size(2, mode="tile")
|
||||
tile_before_install = vae.decode(small, return_dict=False)[0]
|
||||
|
||||
vae.set_parallel_size(2, mode="auto")
|
||||
assert vae._spatial_shard_decode_split_dim(small, 2) == "width"
|
||||
small_sharded = vae.decode(small, return_dict=False)[0]
|
||||
assert vae._spatial_shard_decode_split_dim(landscape, 2) == "width"
|
||||
vae.decode(landscape, return_dict=False)
|
||||
|
||||
vae.set_parallel_size(2, mode="tile")
|
||||
tile_after_install = vae.decode(small, return_dict=False)[0]
|
||||
|
||||
vae.use_tiling = False
|
||||
vae.set_parallel_size(1, mode="tile")
|
||||
portrait_reference = vae.decode(portrait, return_dict=False)[0]
|
||||
|
||||
vae.use_tiling = True
|
||||
vae.set_parallel_size(2, mode="auto")
|
||||
assert vae._spatial_shard_decode_split_dim(portrait, 2) == "height"
|
||||
portrait_sharded = vae.decode(portrait, return_dict=False)[0]
|
||||
|
||||
assert all(
|
||||
buffer is None
|
||||
for module in vae.decoder.modules()
|
||||
for name, buffer in module._buffers.items()
|
||||
if name in {"_halo_recv_top_buf", "_halo_recv_bottom_buf"}
|
||||
)
|
||||
|
||||
if rank == 0:
|
||||
small_diff = (small_sharded.float() - small_reference.float()).abs()
|
||||
portrait_diff = (portrait_sharded.float() - portrait_reference.float()).abs()
|
||||
return_dict["tile_fallback_exact"] = torch.equal(tile_after_install, tile_before_install)
|
||||
return_dict["small_max_abs_diff"] = small_diff.max().item()
|
||||
return_dict["small_mean_abs_diff"] = small_diff.mean().item()
|
||||
return_dict["portrait_max_abs_diff"] = portrait_diff.max().item()
|
||||
return_dict["portrait_mean_abs_diff"] = portrait_diff.mean().item()
|
||||
finally:
|
||||
destroy_model_parallel()
|
||||
if dist.is_initialized():
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
@pytest.mark.full_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parallel
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
||||
def test_auto_spatial_shard_switches_shape_paths_safely():
|
||||
manager = mp.get_context("spawn").Manager()
|
||||
return_dict = manager.dict()
|
||||
mp.spawn(
|
||||
_auto_spatial_shard_decode_worker,
|
||||
args=(return_dict, "29502"),
|
||||
nprocs=2,
|
||||
join=True,
|
||||
)
|
||||
|
||||
assert return_dict.get("tile_fallback_exact") is True
|
||||
assert return_dict["small_max_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
|
||||
assert return_dict["small_mean_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
|
||||
assert return_dict["portrait_max_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
|
||||
assert return_dict["portrait_mean_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
|
||||
@@ -199,7 +199,7 @@ def _generate_bagel_image(
|
||||
config_path = get_deploy_config_path("ci/bagel.yaml")
|
||||
omni_kwargs: dict[str, Any] = {
|
||||
"model": "ByteDance-Seed/BAGEL-7B-MoT",
|
||||
"stage_configs_path": config_path,
|
||||
"deploy_config": config_path,
|
||||
"stage_init_timeout": 300,
|
||||
}
|
||||
if diffusion_quantization_config:
|
||||
|
||||
@@ -314,7 +314,7 @@ def test_bagel_text2img_mooncake_connector(run_level):
|
||||
temp_config_file = _resolve_deploy_config(temp_config_file, run_level)
|
||||
with OmniRunner(
|
||||
"ByteDance-Seed/BAGEL-7B-MoT",
|
||||
stage_configs_path=temp_config_file,
|
||||
deploy_config=temp_config_file,
|
||||
stage_init_timeout=300,
|
||||
) as runner:
|
||||
generated_image = _generate_bagel_image(runner.omni)
|
||||
|
||||
@@ -251,7 +251,7 @@ def test_bagel_img2img_shared_memory_connector(run_level):
|
||||
config_path = _resolve_deploy_config(BAGEL_CI_DEPLOY, run_level)
|
||||
with OmniRunner(
|
||||
"ByteDance-Seed/BAGEL-7B-MoT",
|
||||
stage_configs_path=config_path,
|
||||
deploy_config=config_path,
|
||||
) as runner:
|
||||
generated_image = _generate_bagel_img2img(runner.omni, input_image)
|
||||
if run_level == "advanced_model":
|
||||
@@ -266,7 +266,7 @@ def test_bagel_text2img_shared_memory_connector(run_level):
|
||||
config_path = _resolve_deploy_config(BAGEL_CI_DEPLOY, run_level)
|
||||
with OmniRunner(
|
||||
"ByteDance-Seed/BAGEL-7B-MoT",
|
||||
stage_configs_path=config_path,
|
||||
deploy_config=config_path,
|
||||
) as runner:
|
||||
generated_image = _generate_bagel_text2img(runner.omni)
|
||||
if run_level == "advanced_model":
|
||||
|
||||
@@ -17,7 +17,7 @@ from vllm_omni.platforms import current_omni_platform
|
||||
# Override in CI if needed: AUDIOX_TEST_MODEL=<model-or-local-path>
|
||||
models = [os.environ.get("AUDIOX_TEST_MODEL", "zhangj1an/audiox_random")]
|
||||
|
||||
# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
|
||||
# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
|
||||
_OMNI_RUNNER_PARAMS = [(m, None, {"model_class_name": "AudioXPipeline"}) for m in models]
|
||||
|
||||
pytestmark = [
|
||||
|
||||
@@ -43,7 +43,7 @@ MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
|
||||
BAGEL_STAGE_CONFIG = get_deploy_config_path("ci/bagel.yaml")
|
||||
DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>"
|
||||
|
||||
# (model, stage_configs_path) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
|
||||
# (model, deploy_config_path) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
|
||||
_OMNI_RUNNER_PARAM = (MODEL, BAGEL_STAGE_CONFIG)
|
||||
|
||||
|
||||
|
||||
@@ -126,7 +126,7 @@ def _build_reference_inputs(prompt_audio: tuple[np.ndarray, int]) -> list[dict[s
|
||||
]
|
||||
|
||||
|
||||
# (model, stage_configs_path, extra_omni_kwargs) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
|
||||
# (model, deploy_config_path, extra_omni_kwargs) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
|
||||
_cosy_deployment = get_deploy_config_path("cosyvoice3.yaml")
|
||||
_cosy_model_path = str(_resolve_model_dir())
|
||||
_OMNI_RUNNER_PARAMS = [
|
||||
|
||||
@@ -17,7 +17,7 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
MODEL = "baidu/ERNIE-Image"
|
||||
|
||||
# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
|
||||
# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
|
||||
_OMNI_RUNNER_PARAM = (
|
||||
MODEL,
|
||||
None,
|
||||
|
||||
@@ -96,7 +96,7 @@ def test_offline_voice_clone_zh(async_chunk: bool) -> None:
|
||||
|
||||
with OmniRunner(
|
||||
MODEL,
|
||||
stage_configs_path=_get_deploy_config(async_chunk=async_chunk),
|
||||
deploy_config=_get_deploy_config(async_chunk=async_chunk),
|
||||
stage_init_timeout=600,
|
||||
trust_remote_code=True,
|
||||
) as omni_runner:
|
||||
|
||||
@@ -52,7 +52,7 @@ pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
|
||||
def omni() -> Generator[Omni, None, None]:
|
||||
with OmniRunner(
|
||||
MODEL_NAME,
|
||||
stage_configs_path=str(STAGE_CONFIG_PATH),
|
||||
deploy_config=str(STAGE_CONFIG_PATH),
|
||||
trust_remote_code=True,
|
||||
) as runner:
|
||||
yield runner.omni
|
||||
|
||||
@@ -84,7 +84,6 @@ def longcat_omni_runner() -> Generator[OmniRunner, None, None]:
|
||||
with OmniRunner(
|
||||
model,
|
||||
seed=42,
|
||||
stage_configs_path=None,
|
||||
model_class_name="LongCatVideoAvatarPipeline",
|
||||
additional_config={
|
||||
"model_type": "avatar-v1.5",
|
||||
|
||||
@@ -38,7 +38,7 @@ _AR_PATCH_SIZE = 16
|
||||
MODEL_PATH = "bytedance-research/MammothModa2-Preview"
|
||||
T2I_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")
|
||||
|
||||
_OMNI_RUNNER_PARAM = (MODEL_PATH, None, {"deploy_config": T2I_DEPLOY_CONFIG})
|
||||
_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_DEPLOY_CONFIG)
|
||||
|
||||
# Golden pixel reference file. Set UPDATE_GOLDEN=1 to regenerate.
|
||||
_GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json"
|
||||
|
||||
@@ -17,7 +17,7 @@ models = ["openbmb/MiniCPM-o-4_5"]
|
||||
_CI_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml")
|
||||
|
||||
|
||||
test_params = [(model, None, {"deploy_config": _CI_DEPLOY, "trust_remote_code": True}) for model in models]
|
||||
test_params = [(model, _CI_DEPLOY, {"trust_remote_code": True}) for model in models]
|
||||
|
||||
|
||||
def get_question(prompt_type: str = "text") -> str:
|
||||
|
||||
@@ -25,7 +25,7 @@ from vllm_omni import Omni
|
||||
MODEL_NAME = "OpenMOSS-Team/MOSS-TTS-Nano"
|
||||
STAGE_CONFIG = get_deploy_config_path("moss_tts_nano.yaml")
|
||||
|
||||
# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
|
||||
# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
|
||||
_OMNI_RUNNER_PARAM = (
|
||||
MODEL_NAME,
|
||||
STAGE_CONFIG,
|
||||
|
||||
@@ -20,13 +20,11 @@ from tests.helpers.stage_config import get_deploy_config_path
|
||||
MODEL = "k2-fsa/OmniVoice"
|
||||
DEPLOY_CONFIG = get_deploy_config_path("omnivoice.yaml")
|
||||
|
||||
# OmniRunner tuple: model, legacy stage config path, extra Omni kwargs.
|
||||
# The migrated test passes deploy_config through extra Omni kwargs.
|
||||
# OmniRunner tuple: model, deploy config path, extra Omni kwargs.
|
||||
_OMNI_RUNNER_PARAM = (
|
||||
MODEL,
|
||||
None,
|
||||
DEPLOY_CONFIG,
|
||||
{
|
||||
"deploy_config": DEPLOY_CONFIG,
|
||||
"trust_remote_code": True,
|
||||
"log_stats": True,
|
||||
},
|
||||
|
||||
@@ -131,7 +131,7 @@ def test_sensenova_u1_lora_scale_and_deactivation(tmp_path) -> None:
|
||||
"""Validate LoRA effect, scale sensitivity, bounded perturbation, and clean deactivation."""
|
||||
lora_request = _make_file_lora_request(tmp_path / "sensenova_lora")
|
||||
|
||||
with OmniRunner(MODEL, stage_configs_path=None) as runner:
|
||||
with OmniRunner(MODEL) as runner:
|
||||
omni = runner.omni
|
||||
baseline = _generate(omni)
|
||||
img_1x = _generate(omni, lora_request, lora_scale=1.0)
|
||||
|
||||
@@ -164,10 +164,7 @@ def _generate_sensenova_u1_img2img(
|
||||
def test_sensenova_u1_img2img(run_level):
|
||||
"""Test SenseNova-U1 img2img (single-stage diffusion, no deploy YAML)."""
|
||||
input_image = _load_input_image()
|
||||
with OmniRunner(
|
||||
"SenseNova/SenseNova-U1-8B-MoT",
|
||||
stage_configs_path=None,
|
||||
) as runner:
|
||||
with OmniRunner("SenseNova/SenseNova-U1-8B-MoT") as runner:
|
||||
generated_image = _generate_sensenova_u1_img2img(runner.omni, input_image)
|
||||
if run_level == "full_model" or run_level == "advanced_model":
|
||||
_validate_pixels(generated_image)
|
||||
@@ -182,7 +179,6 @@ def test_sensenova_u1_img2img_cache_dit():
|
||||
input_image = _load_input_image()
|
||||
with OmniRunner(
|
||||
"SenseNova/SenseNova-U1-8B-MoT",
|
||||
stage_configs_path=None,
|
||||
cache_backend="cache_dit",
|
||||
) as runner:
|
||||
_generate_sensenova_u1_img2img(runner.omni, input_image)
|
||||
@@ -151,10 +151,7 @@ def _generate_sensenova_u1_image(
|
||||
@hardware_test(res={"cuda": "H100"})
|
||||
def test_sensenova_u1_text2img(run_level):
|
||||
"""Test SenseNova-U1 text2img (single-stage diffusion, no deploy YAML)."""
|
||||
with OmniRunner(
|
||||
"SenseNova/SenseNova-U1-8B-MoT",
|
||||
stage_configs_path=None,
|
||||
) as runner:
|
||||
with OmniRunner("SenseNova/SenseNova-U1-8B-MoT") as runner:
|
||||
generated_image = _generate_sensenova_u1_image(runner.omni)
|
||||
if run_level == "full_model" or run_level == "advanced_model":
|
||||
_validate_pixels(generated_image)
|
||||
@@ -168,7 +165,6 @@ def test_sensenova_u1_text2img_cache_dit():
|
||||
"""Test SenseNova-U1 text2img with Cache-DiT enabled."""
|
||||
with OmniRunner(
|
||||
"SenseNova/SenseNova-U1-8B-MoT",
|
||||
stage_configs_path=None,
|
||||
cache_backend="cache_dit",
|
||||
) as runner:
|
||||
_generate_sensenova_u1_image(runner.omni)
|
||||
@@ -159,7 +159,7 @@ def test_soulxsinger_multistage_from_audio(
|
||||
|
||||
with OmniRunner(
|
||||
model,
|
||||
stage_configs_path=get_deploy_config_path(deploy_yaml),
|
||||
deploy_config=get_deploy_config_path(deploy_yaml),
|
||||
async_chunk=False,
|
||||
) as runner:
|
||||
sampling = OmniDiffusionSamplingParams(
|
||||
@@ -212,7 +212,7 @@ def test_soulxsinger_svs_precomputed(soulx_weights: tuple[Path, Path, Path]) ->
|
||||
|
||||
with OmniRunner(
|
||||
str(base_dir),
|
||||
stage_configs_path=get_deploy_config_path("soulxsinger_svs.yaml"),
|
||||
deploy_config=get_deploy_config_path("soulxsinger_svs.yaml"),
|
||||
async_chunk=False,
|
||||
) as runner:
|
||||
sampling = OmniDiffusionSamplingParams(
|
||||
|
||||
@@ -19,7 +19,7 @@ from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
MODEL = "stepfun-ai/Step-Audio-2-mini"
|
||||
DEPLOY_CONFIG = get_deploy_config_path("step_audio2_ci.yaml")
|
||||
TEST_PARAMS = [(MODEL, None, {"deploy_config": DEPLOY_CONFIG, "trust_remote_code": True})]
|
||||
TEST_PARAMS = [(MODEL, DEPLOY_CONFIG, {"trust_remote_code": True})]
|
||||
|
||||
pytestmark = [pytest.mark.slow, pytest.mark.tts]
|
||||
|
||||
|
||||
@@ -102,7 +102,7 @@ def test_voxtral_tts_offline_streaming():
|
||||
async def _run():
|
||||
async_omni = AsyncOmni(
|
||||
model=MODEL,
|
||||
stage_configs_path=STAGE_CONFIG,
|
||||
deploy_config=STAGE_CONFIG,
|
||||
stage_init_timeout=300,
|
||||
enforce_eager=True,
|
||||
)
|
||||
|
||||
@@ -4,7 +4,6 @@ invariant to the specific attributes of vLLM config except in cases where we
|
||||
explicitly patch values that differ from vLLM.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
@@ -20,7 +19,6 @@ from vllm.engine.arg_utils import EngineArgs
|
||||
|
||||
from vllm_omni.config.model import OmniModelConfig
|
||||
from vllm_omni.engine.arg_utils import OmniEngineArgs
|
||||
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
|
||||
from vllm_omni.engine.stage_init_utils import build_engine_args_dict
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -149,19 +147,6 @@ def test_qwen3_tts_startup_task_type_is_validated():
|
||||
)
|
||||
|
||||
|
||||
def test_from_cli_args_picks_up_stage_configs_path():
|
||||
"""from_cli_args should pick up stage_configs_path from namespace."""
|
||||
ns = argparse.Namespace(
|
||||
model="facebook/opt-125m",
|
||||
stage_configs_path="/some/path.yaml",
|
||||
custom_pipeline_args=None,
|
||||
)
|
||||
|
||||
args = OmniEngineArgs.from_cli_args(ns)
|
||||
assert args.stage_configs_path == "/some/path.yaml"
|
||||
assert args.custom_pipeline_args is None
|
||||
|
||||
|
||||
def test_qwen3_tts_code2wav_injects_max_position_embeddings(monkeypatch):
|
||||
"""Ensure Code2Wav mirrors stage max_model_len into nested HF overrides.
|
||||
|
||||
@@ -418,130 +403,6 @@ def test_non_override_ar_stage_inputs_embeds_size_matches_hidden_size():
|
||||
assert omni_config.get_inputs_embeds_size() == omni_config.hf_text_config.hidden_size
|
||||
|
||||
|
||||
def test_stage_configs_path_field():
|
||||
"""OmniEngineArgs with stage_configs_path should construct without error."""
|
||||
args = OmniEngineArgs(stage_configs_path="/some/path.yaml")
|
||||
assert args.stage_configs_path == "/some/path.yaml"
|
||||
|
||||
|
||||
def test_strip_single_engine_args():
|
||||
"""_strip_single_engine_args should remove EngineArgs fields but keep omni fields."""
|
||||
kwargs = {
|
||||
# Parent EngineArgs fields — stripped unless explicitly allowlisted
|
||||
"compilation_config": '{"cudagraph_mode": "FULL_AND_PIECEWISE"}',
|
||||
"tensor_parallel_size": 4,
|
||||
"gpu_memory_utilization": 0.9,
|
||||
"model": "some/model",
|
||||
# Parent field that should be kept (allowlisted)
|
||||
"worker_extension_cls": "some.Extension",
|
||||
# OmniEngineArgs-only / non-engine fields — should pass through
|
||||
"stage_configs_path": "/path/to/yaml",
|
||||
"custom_pipeline_args": {"pipeline_class": "my.Pipeline"},
|
||||
"mode": "text-to-image",
|
||||
"lora_path": "/some/lora",
|
||||
}
|
||||
|
||||
filtered = AsyncOmniEngine._strip_single_engine_args(kwargs)
|
||||
|
||||
# Stripped — parent EngineArgs fields
|
||||
assert "compilation_config" not in filtered
|
||||
assert filtered["tensor_parallel_size"] == 4
|
||||
assert "gpu_memory_utilization" not in filtered
|
||||
assert "model" not in filtered
|
||||
|
||||
# Stripped — orchestrator-level OmniEngineArgs field
|
||||
assert "stage_configs_path" not in filtered
|
||||
|
||||
# Kept
|
||||
assert filtered["worker_extension_cls"] == "some.Extension"
|
||||
assert filtered["custom_pipeline_args"] == {"pipeline_class": "my.Pipeline"}
|
||||
assert filtered["mode"] == "text-to-image"
|
||||
assert filtered["lora_path"] == "/some/lora"
|
||||
|
||||
|
||||
def test_strip_single_engine_args_model_does_not_trigger_warning(mocker):
|
||||
"""model is always in kwargs (callers set it via from_cli_args/asdict),
|
||||
so it should not cause the override warning by itself or appear in it."""
|
||||
mock_warn = mocker.patch("vllm_omni.engine.async_omni_engine.logger.warning")
|
||||
|
||||
# Typical caller kwargs: model is always present, no other parent
|
||||
# EngineArgs fields are explicitly overridden.
|
||||
AsyncOmniEngine._strip_single_engine_args(
|
||||
{
|
||||
"model": "some/model",
|
||||
"custom_pipeline_args": {"pipeline_class": "my.Pipeline"},
|
||||
}
|
||||
)
|
||||
mock_warn.assert_not_called()
|
||||
|
||||
# When there *are* genuinely surprising overrides alongside model,
|
||||
# the warning should mention them but not model. Keep-listed fields such as
|
||||
# tensor_parallel_size are intentionally passed through and should not warn.
|
||||
AsyncOmniEngine._strip_single_engine_args(
|
||||
{
|
||||
"model": "some/model",
|
||||
"compilation_config": '{"cudagraph_mode": "FULL_AND_PIECEWISE"}',
|
||||
"tensor_parallel_size": 4,
|
||||
"custom_pipeline_args": {"pipeline_class": "my.Pipeline"},
|
||||
}
|
||||
)
|
||||
mock_warn.assert_called_once()
|
||||
warned_args = mock_warn.call_args[0][-1] # the formatted arg list
|
||||
assert "compilation_config" in warned_args
|
||||
assert "tensor_parallel_size" not in warned_args
|
||||
assert "model" not in warned_args
|
||||
|
||||
|
||||
def test_new_format_deploy_yaml_does_not_strip_interleave_cli_overrides(tmp_path):
|
||||
"""New-format deploy YAMLs must not strip interleave_mm_strings / media_io_kwargs.
|
||||
|
||||
Pytest historically passed ``--stage-configs-path`` for deploy YAMLs; stripping
|
||||
those CLI flags dropped MiniCPM interleaved AV packing and cut Daily-Omni
|
||||
accuracy from ~78% to ~70%.
|
||||
"""
|
||||
from vllm_omni.entrypoints.utils import is_new_format_deploy_config
|
||||
|
||||
yaml_path = tmp_path / "minicpm_deploy.yaml"
|
||||
yaml_path.write_text(
|
||||
"""
|
||||
pipeline: minicpmo_4_5
|
||||
async_chunk: true
|
||||
stages:
|
||||
- stage_id: 0
|
||||
devices: "0"
|
||||
""".strip(),
|
||||
encoding="utf-8",
|
||||
)
|
||||
path = str(yaml_path)
|
||||
assert is_new_format_deploy_config(path)
|
||||
|
||||
kwargs = {
|
||||
"stage_configs_path": path,
|
||||
"trust_remote_code": True,
|
||||
"interleave_mm_strings": True,
|
||||
"media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}},
|
||||
"gpu_memory_utilization": 0.55,
|
||||
}
|
||||
# Mirror AsyncOmniEngine._resolve_stage_configs strip decision.
|
||||
if kwargs.get("stage_configs_path") is not None and not is_new_format_deploy_config(kwargs["stage_configs_path"]):
|
||||
base = AsyncOmniEngine._strip_single_engine_args(dict(kwargs))
|
||||
else:
|
||||
base = dict(kwargs)
|
||||
assert base.get("interleave_mm_strings") is True
|
||||
assert base.get("media_io_kwargs") == {"video": {"fps": 1, "num_frames": 128}}
|
||||
# Parent EngineArgs fields also remain available for deploy merge.
|
||||
assert base.get("gpu_memory_utilization") == 0.55
|
||||
|
||||
legacy = tmp_path / "legacy.yaml"
|
||||
legacy.write_text("stage_args:\n - stage_id: 0\n", encoding="utf-8")
|
||||
assert not is_new_format_deploy_config(str(legacy))
|
||||
legacy_kwargs = dict(kwargs, stage_configs_path=str(legacy))
|
||||
stripped = AsyncOmniEngine._strip_single_engine_args(legacy_kwargs)
|
||||
assert "interleave_mm_strings" not in stripped
|
||||
assert "media_io_kwargs" not in stripped
|
||||
assert "gpu_memory_utilization" not in stripped
|
||||
|
||||
|
||||
# For https://github.com/vllm-project/vllm-omni/issues/3293
|
||||
def test_tensor_parallel_size_none_is_handled():
|
||||
"""Ensure the tensor parallel size of None isn't forwarded."""
|
||||
|
||||
@@ -68,7 +68,7 @@ async def test_abort():
|
||||
with ExitStack() as after:
|
||||
engine = AsyncOmni(
|
||||
model=model,
|
||||
stage_configs_path=stage_config,
|
||||
deploy_config=stage_config,
|
||||
)
|
||||
after.callback(engine.shutdown)
|
||||
|
||||
|
||||
@@ -428,7 +428,7 @@ class TestSingleStageModeDetection:
|
||||
assert engine.single_stage_mode is False
|
||||
assert engine._single_stage_id_filter is None
|
||||
|
||||
def test_stage_configs_path_loads_duplex_runtime_config(self, mocker: MockerFixture):
|
||||
def test_deploy_config_loads_duplex_runtime_config(self, mocker: MockerFixture):
|
||||
duplex_session = DuplexSessionRuntimeConfig(max_sessions=2)
|
||||
get_pipeline_config = mocker.patch(
|
||||
"vllm_omni.engine.async_omni_engine.StageConfigFactory.get_pipeline_config",
|
||||
@@ -441,7 +441,7 @@ class TestSingleStageModeDetection:
|
||||
|
||||
engine = self._make_engine_no_thread(
|
||||
mocker,
|
||||
stage_configs_path="/fake/duplex.yaml",
|
||||
deploy_config="/fake/duplex.yaml",
|
||||
)
|
||||
|
||||
get_pipeline_config.assert_called_once_with(
|
||||
|
||||
@@ -364,7 +364,7 @@ async def test_omni_generate_request_id():
|
||||
Text modality only; asserts caller-visible ids are preserved across the
|
||||
multi-stage orchestrator path on H100.
|
||||
"""
|
||||
engine = AsyncOmni(model=OMNI_MODEL, stage_configs_path=OMNI_STAGE_CONFIG)
|
||||
engine = AsyncOmni(model=OMNI_MODEL, deploy_config=OMNI_STAGE_CONFIG)
|
||||
try:
|
||||
for req_id in _OMNI_REQ_IDS:
|
||||
async for output in engine.generate(
|
||||
|
||||
@@ -478,12 +478,11 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo
|
||||
)
|
||||
|
||||
engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
|
||||
engine._strip_single_engine_args = lambda kwargs: kwargs
|
||||
|
||||
_, stage_configs = engine._resolve_stage_configs(
|
||||
"dummy-model",
|
||||
{
|
||||
"stage_configs_path": "dummy.yaml",
|
||||
"deploy_config": "dummy.yaml",
|
||||
"additional_config": {"torchair_graph_config": {"enabled": True}},
|
||||
},
|
||||
trust_remote_code=False,
|
||||
@@ -516,12 +515,11 @@ def test_resolve_stage_configs_injects_quantization_config_into_diffusion_stage(
|
||||
)
|
||||
|
||||
engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
|
||||
engine._strip_single_engine_args = lambda kwargs: kwargs
|
||||
|
||||
_, stage_configs = engine._resolve_stage_configs(
|
||||
"dummy-model",
|
||||
{
|
||||
"stage_configs_path": "dummy.yaml",
|
||||
"deploy_config": "dummy.yaml",
|
||||
"quantization_config": {"method": "bitsandbytes"},
|
||||
},
|
||||
trust_remote_code=False,
|
||||
|
||||
@@ -76,6 +76,18 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None:
|
||||
assert args.cfg_parallel_size == 4
|
||||
|
||||
|
||||
def test_serve_parser_defaults_vae_parallel_mode_to_tile_and_accepts_auto() -> None:
|
||||
parser = TrackingArgumentParser()
|
||||
subparsers = parser.add_subparsers(dest="subcommand")
|
||||
OmniServeCommand().subparser_init(subparsers)
|
||||
|
||||
default_args = parser.parse_args(["serve", "fake-model", "--omni"])
|
||||
auto_args = parser.parse_args(["serve", "fake-model", "--omni", "--vae-parallel-mode", "auto"])
|
||||
|
||||
assert default_args.vae_parallel_mode == "tile"
|
||||
assert auto_args.vae_parallel_mode == "auto"
|
||||
|
||||
|
||||
def _make_headless_args(**kwargs) -> TrackingNamespace:
|
||||
defaults = {
|
||||
"model": "fake-model",
|
||||
@@ -193,7 +205,7 @@ def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture)
|
||||
with pytest.raises(ValueError, match="No stage config found for stage_id=0"):
|
||||
run_headless(args)
|
||||
|
||||
assert captured["args"][1] is None
|
||||
assert len(captured["args"]) == 2
|
||||
assert captured["deploy_config_path"] == "/tmp/deploy.yaml"
|
||||
assert captured["stage_overrides"] == {"0": {"devices": "0,1"}, "1": {"devices": "2"}}
|
||||
assert captured["strategy_config_path"] == "/tmp/strategy.yaml"
|
||||
|
||||
@@ -381,7 +381,6 @@ class TestLoadAndResolveStageConfigs:
|
||||
|
||||
config_path, stage_configs, _ = load_and_resolve_stage_configs(
|
||||
model="black-forest-labs/FLUX.2-klein-4B",
|
||||
stage_configs_path=None,
|
||||
kwargs=kwargs,
|
||||
trust_remote_code=False,
|
||||
default_stage_cfg_factory=lambda: AsyncOmniEngine._create_default_diffusion_stage_cfg(kwargs),
|
||||
@@ -390,9 +389,7 @@ class TestLoadAndResolveStageConfigs:
|
||||
assert len(stage_configs) == 1
|
||||
assert "dtype" in stage_configs[0]["engine_args"]
|
||||
|
||||
def test_stage_configs_path_promotes_new_deploy_yaml_without_expanding_replicas(
|
||||
self, tmp_path, mocker: MockerFixture
|
||||
):
|
||||
def test_deploy_config_preserves_cli_overrides_and_replicas(self, tmp_path, mocker: MockerFixture):
|
||||
deploy_path = tmp_path / "qwen3_multi.yaml"
|
||||
deploy_path.write_text(
|
||||
'stages:\n - stage_id: 0\n devices: "0"\n - stage_id: 1\n devices: "1,2,3"\n num_replicas: 3\n',
|
||||
@@ -413,18 +410,23 @@ class TestLoadAndResolveStageConfigs:
|
||||
"vllm_omni.entrypoints.utils.load_stage_configs_from_model",
|
||||
return_value=(returned_stage_configs, None),
|
||||
)
|
||||
cli_overrides = {
|
||||
"interleave_mm_strings": True,
|
||||
"media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}},
|
||||
"gpu_memory_utilization": 0.55,
|
||||
}
|
||||
|
||||
config_path, stage_configs, _ = load_and_resolve_stage_configs(
|
||||
model="dummy-model",
|
||||
stage_configs_path=str(deploy_path),
|
||||
kwargs={},
|
||||
kwargs=cli_overrides,
|
||||
trust_remote_code=True,
|
||||
deploy_config_path=str(deploy_path),
|
||||
)
|
||||
|
||||
load_stage_configs.assert_called_once_with(
|
||||
"dummy-model",
|
||||
trust_remote_code=True,
|
||||
base_engine_args={},
|
||||
base_engine_args=cli_overrides,
|
||||
deploy_config_path=str(deploy_path),
|
||||
stage_overrides=None,
|
||||
strategy_config_path=None,
|
||||
|
||||
@@ -2635,7 +2635,7 @@ class OmniRunner:
|
||||
# affects the test runner wrapper.
|
||||
init_timeout: int = 1800,
|
||||
log_stats: bool = False,
|
||||
stage_configs_path: str | None = None,
|
||||
deploy_config: str | None = None,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
startup_t0 = time.perf_counter()
|
||||
@@ -2653,7 +2653,7 @@ class OmniRunner:
|
||||
stage_init_timeout=stage_init_timeout,
|
||||
batch_timeout=batch_timeout,
|
||||
init_timeout=init_timeout,
|
||||
stage_configs_path=stage_configs_path,
|
||||
deploy_config=deploy_config,
|
||||
**kwargs,
|
||||
)
|
||||
startup_s = time.perf_counter() - startup_t0
|
||||
@@ -3179,12 +3179,7 @@ def iter_omni_server(
|
||||
print("OmniServer stopping...")
|
||||
else:
|
||||
if stage_config_path is not None:
|
||||
from vllm_omni.entrypoints.utils import is_new_format_deploy_config
|
||||
|
||||
if is_new_format_deploy_config(stage_config_path):
|
||||
server_args += ["--deploy-config", stage_config_path]
|
||||
else:
|
||||
server_args += ["--stage-configs-path", stage_config_path]
|
||||
|
||||
with (
|
||||
OmniServer(
|
||||
@@ -3237,7 +3232,7 @@ def iter_omni_runner(
|
||||
model = model_prefix + model
|
||||
if run_level == "core_model" and request.node.get_closest_marker("diffusion"):
|
||||
model = resolve_tiny_model_path(model)
|
||||
with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, **extra_omni_kwargs) as runner:
|
||||
with OmniRunner(model, seed=42, deploy_config=stage_config_path, **extra_omni_kwargs) as runner:
|
||||
print("OmniRunner started successfully")
|
||||
yield runner
|
||||
print("OmniRunner stopping...")
|
||||
|
||||
@@ -147,7 +147,7 @@ def _run_once(model: str, yaml_path: str, *, async_mode: bool) -> list[dict]:
|
||||
# (see tests/e2e/accuracy/qwen3_omni). Confirm the exact knob at bench time.
|
||||
omni = Omni(
|
||||
model,
|
||||
stage_configs_path=yaml_path,
|
||||
deploy_config=yaml_path,
|
||||
stage_init_timeout=1800,
|
||||
init_timeout=3600,
|
||||
async_chunk=async_mode,
|
||||
|
||||
@@ -53,7 +53,6 @@ _NON_STAGE_ENGINE_CLI_FIELDS = frozenset(
|
||||
"model",
|
||||
"omni",
|
||||
"output_modalities",
|
||||
"stage_configs_path",
|
||||
"stage_id",
|
||||
"tokenizer",
|
||||
}
|
||||
@@ -474,9 +473,10 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
|
||||
raise ValueError("allgather_degree > 1 is mutually exclusive with ulysses_degree/ring_degree > 1")
|
||||
if self.ulysses_mode not in {"strict", "advanced_uaa"}:
|
||||
raise ValueError("ulysses_mode must be 'strict' or 'advanced_uaa'")
|
||||
if self.vae_parallel_mode not in {"tile", "spatial_shard_height", "spatial_shard_width"}:
|
||||
if self.vae_parallel_mode not in {"tile", "auto", "spatial_shard_height", "spatial_shard_width"}:
|
||||
raise ValueError(
|
||||
"vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
|
||||
"vae_parallel_mode must be one of "
|
||||
"{'tile', 'auto', 'spatial_shard_height', 'spatial_shard_width'}, "
|
||||
f"but got {self.vae_parallel_mode!r}."
|
||||
)
|
||||
|
||||
|
||||
@@ -4,7 +4,6 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import functools
|
||||
import re
|
||||
import warnings
|
||||
@@ -56,7 +55,7 @@ def build_stage_runtime_overrides(
|
||||
``internal_keys`` defaults to the union of
|
||||
``arg_utils.internal_blacklist_keys()`` and ``arg_utils.SHARED_FIELDS``
|
||||
so that neither orchestrator-only fields nor shared-pipeline fields
|
||||
(``model`` / ``stage_configs_path`` / ``log_stats`` / ``stage_id``) leak
|
||||
(``model`` / ``log_stats`` / ``stage_id``) leak
|
||||
into a stage's per-stage runtime overrides — the orchestrator sets those
|
||||
uniformly for every stage, they are not per-stage knobs. Callers can
|
||||
pass an explicit set for tests or specialized flows.
|
||||
@@ -90,46 +89,6 @@ def build_stage_runtime_overrides(
|
||||
return result
|
||||
|
||||
|
||||
def strip_parent_engine_args(
|
||||
kwargs: dict[str, Any],
|
||||
*,
|
||||
parent_fields: dict[str, dataclasses.Field],
|
||||
keep_keys: set[str] | frozenset[str] = frozenset(),
|
||||
strip_keys: set[str] | frozenset[str] = frozenset(),
|
||||
no_warn_keys: set[str] | frozenset[str] = frozenset(),
|
||||
) -> tuple[dict[str, Any], list[str]]:
|
||||
"""Strip parent ``EngineArgs`` fields before merging into stage YAML."""
|
||||
overridden: list[str] = []
|
||||
result: dict[str, Any] = {}
|
||||
|
||||
for key, value in kwargs.items():
|
||||
if key in strip_keys:
|
||||
continue
|
||||
|
||||
if key not in parent_fields or key in keep_keys:
|
||||
result[key] = value
|
||||
continue
|
||||
|
||||
field_def = parent_fields[key]
|
||||
if field_def.default is not dataclasses.MISSING:
|
||||
default = field_def.default
|
||||
elif field_def.default_factory is not dataclasses.MISSING:
|
||||
default = field_def.default_factory()
|
||||
else:
|
||||
default = dataclasses.MISSING
|
||||
|
||||
if default is dataclasses.MISSING or value is None:
|
||||
continue
|
||||
|
||||
if dataclasses.is_dataclass(default) and not isinstance(default, type):
|
||||
default = asdict(default)
|
||||
|
||||
if value != default and key not in no_warn_keys:
|
||||
overridden.append(key)
|
||||
|
||||
return result, sorted(overridden)
|
||||
|
||||
|
||||
def _apply_diffusion_parallel_runtime_overrides(
|
||||
engine_args: dict[str, Any],
|
||||
runtime_overrides: dict[str, Any],
|
||||
|
||||
@@ -31,12 +31,12 @@ stages:
|
||||
output_connectors:
|
||||
to_stage_1: shm
|
||||
default_sampling_params:
|
||||
temperature: 1.0
|
||||
top_p: 0.95
|
||||
top_k: 50
|
||||
temperature: 1.7
|
||||
top_p: 0.8
|
||||
top_k: 25
|
||||
max_tokens: 4096
|
||||
seed: 42
|
||||
repetition_penalty: 1.1
|
||||
repetition_penalty: 1.0
|
||||
|
||||
- stage_id: 1
|
||||
max_num_seqs: 64
|
||||
|
||||
@@ -212,6 +212,9 @@ class DiffusionParallelConfig:
|
||||
|
||||
- "tile": Patch/tile parallel decode (default). Each rank decodes a subset
|
||||
of spatial tiles and the results are stitched on rank 0.
|
||||
- "auto": Select the longer spatial-shard axis for validated Wan tiled
|
||||
decodes on a full DiT-group topology; otherwise use patch/tile parallel
|
||||
decode.
|
||||
- "spatial_shard_height": Spatially-sharded decode that splits decoder
|
||||
feature maps along height and exchanges halo rows around spatial
|
||||
convolutions.
|
||||
@@ -251,8 +254,9 @@ class DiffusionParallelConfig:
|
||||
assert self.allgather_degree > 0, "AllGather degree must be > 0"
|
||||
assert self.cfg_parallel_size > 0, "CFG parallel size must be > 0"
|
||||
assert self.vae_patch_parallel_size > 0, "VAE patch parallel size must be > 0"
|
||||
assert self.vae_parallel_mode in {"tile", "spatial_shard_height", "spatial_shard_width"}, (
|
||||
"vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
|
||||
assert self.vae_parallel_mode in {"tile", "auto", "spatial_shard_height", "spatial_shard_width"}, (
|
||||
"vae_parallel_mode must be one of "
|
||||
"{'tile', 'auto', 'spatial_shard_height', 'spatial_shard_width'}, "
|
||||
f"but got {self.vae_parallel_mode!r}."
|
||||
)
|
||||
if self.allgather_degree > 1:
|
||||
|
||||
@@ -23,6 +23,30 @@ from vllm_omni.platforms import current_omni_platform
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def select_auto_spatial_shard_split_dim(
|
||||
z_shape: tuple[int, ...],
|
||||
world_size: int,
|
||||
) -> str | None:
|
||||
"""Select a profitable Wan spatial-shard axis for one latent shape.
|
||||
|
||||
This selector runs only after Diffusers has already chosen ``tiled_decode``
|
||||
for a spatially large request. Splitting the longer eligible axis reduces
|
||||
halo traffic for rectangular requests and improves per-rank balance.
|
||||
"""
|
||||
if len(z_shape) != 5 or world_size <= 1:
|
||||
return None
|
||||
|
||||
_, _, _, height, width = z_shape
|
||||
candidates = [
|
||||
(extent, split_dim) for extent, split_dim in ((height, "height"), (width, "width")) if extent >= world_size
|
||||
]
|
||||
if not candidates:
|
||||
return None
|
||||
# Prefer width for square latents: it was modestly faster in the existing
|
||||
# four-GPU Wan benchmark and keeps the choice deterministic.
|
||||
return max(candidates, key=lambda candidate: (candidate[0], candidate[1] == "width"))[1]
|
||||
|
||||
|
||||
class OmniAutoencoderKLWan(AutoencoderKLWan):
|
||||
def _execution_context(self):
|
||||
try:
|
||||
@@ -56,6 +80,19 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
|
||||
model.init_distributed()
|
||||
return model
|
||||
|
||||
def set_parallel_size(self, parallel_size: int, mode: str = "tile") -> None:
|
||||
super().set_parallel_size(parallel_size, mode=mode)
|
||||
if mode not in {"auto", "spatial_shard_height", "spatial_shard_width"}:
|
||||
return
|
||||
|
||||
executor = self.distributed_executor
|
||||
world_size = dist.get_world_size(group=executor.group)
|
||||
if parallel_size <= 1 or parallel_size != world_size:
|
||||
return
|
||||
|
||||
split_dim = "width" if mode == "spatial_shard_width" else "height"
|
||||
wan_spatial_shard.install_wan_spatial_shard_decode(self, executor.group, split_dim=split_dim)
|
||||
|
||||
def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
|
||||
_, _, num_frames, height, width = z.shape
|
||||
sample_height = height * self.spatial_compression_ratio
|
||||
@@ -270,19 +307,25 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
|
||||
dec = torch.clamp(dec, min=-1.0, max=1.0)
|
||||
return dec
|
||||
|
||||
def _spatial_shard_decode_split_dim(self) -> str | None:
|
||||
def _spatial_shard_decode_split_dim(
|
||||
self,
|
||||
z: torch.Tensor | None = None,
|
||||
world_size: int | None = None,
|
||||
) -> str | None:
|
||||
mode = self.distributed_executor.parallel_mode
|
||||
if mode == "spatial_shard_width":
|
||||
return "width"
|
||||
if mode == "spatial_shard_height":
|
||||
return "height"
|
||||
if mode == "auto" and z is not None and world_size is not None:
|
||||
return select_auto_spatial_shard_split_dim(tuple(z.shape), world_size)
|
||||
return None
|
||||
|
||||
def _spatial_shard_decode_enabled(self, z: torch.Tensor) -> bool:
|
||||
split_dim = self._spatial_shard_decode_split_dim()
|
||||
if split_dim is None:
|
||||
if self.distributed_executor.parallel_mode == "tile":
|
||||
return False
|
||||
if z.ndim != 5:
|
||||
if self.distributed_executor.parallel_mode != "tile":
|
||||
logger.warning("Wan VAE spatial sharded decode expects 5D latent input; falling back to tiled decode.")
|
||||
return False
|
||||
if not self.is_distributed_enabled():
|
||||
@@ -292,23 +335,32 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
|
||||
world_size = dist.get_world_size(group=group)
|
||||
requested_size = int(self.distributed_executor.parallel_size)
|
||||
if requested_size != world_size:
|
||||
if self.distributed_executor.parallel_mode == "auto":
|
||||
logger.debug(
|
||||
"Wan VAE auto decode selected tile mode for a partial DiT group (requested=%s dit_group=%s)",
|
||||
requested_size,
|
||||
world_size,
|
||||
)
|
||||
elif self._spatial_shard_decode_split_dim() is not None:
|
||||
logger.warning(
|
||||
"Wan VAE spatial sharded decode currently requires vae_patch_parallel_size "
|
||||
"to match the DIT group size; falling back to tiled decode. "
|
||||
"requested=%s dit_group=%s split_dim=%s",
|
||||
requested_size,
|
||||
world_size,
|
||||
split_dim,
|
||||
self._spatial_shard_decode_split_dim(),
|
||||
)
|
||||
return False
|
||||
return True
|
||||
return self._spatial_shard_decode_split_dim(z, world_size) is not None
|
||||
|
||||
def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
|
||||
if not self.is_distributed_enabled():
|
||||
return super().tiled_decode(z, return_dict=return_dict)
|
||||
|
||||
if self._spatial_shard_decode_enabled(z):
|
||||
split_dim = self._spatial_shard_decode_split_dim()
|
||||
world_size = dist.get_world_size(group=self.distributed_executor.group)
|
||||
split_dim = self._spatial_shard_decode_split_dim(z, world_size)
|
||||
assert split_dim is not None
|
||||
logger.debug("Decode running with Wan VAE spatial_shard_%s mode", split_dim)
|
||||
return wan_spatial_shard.spatial_shard_decode(
|
||||
self,
|
||||
|
||||
@@ -52,6 +52,11 @@ _SPATIAL_SHARD_CONTEXT: ContextVar[SpatialShardContext | None] = ContextVar(
|
||||
)
|
||||
|
||||
|
||||
def _active_spatial_shard_split_dim() -> str | None:
|
||||
context = _SPATIAL_SHARD_CONTEXT.get()
|
||||
return context.split_dim if context is not None else None
|
||||
|
||||
|
||||
def _spatial_dim(split_dim: str) -> int:
|
||||
if split_dim == "height":
|
||||
return -2
|
||||
@@ -332,21 +337,27 @@ class WanDistZeroPad2d(nn.Module):
|
||||
padding: tuple[int, int, int, int],
|
||||
group: dist.ProcessGroup,
|
||||
*,
|
||||
split_dim: str = "height",
|
||||
split_padding: tuple[int, int] | None = None,
|
||||
defer_split_after_padding: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.padding = padding
|
||||
self.split_dim = split_dim
|
||||
default_split_padding = (padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1])
|
||||
self.split_padding = split_padding or default_split_padding
|
||||
self.defer_split_after_padding = defer_split_after_padding
|
||||
self.group = group
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
split_dim = _active_spatial_shard_split_dim()
|
||||
if split_dim is None:
|
||||
return F.pad(x, self.padding)
|
||||
|
||||
rank, world_size = _rank_world(self.group)
|
||||
left, right, top, bottom = self.padding
|
||||
if self.defer_split_after_padding:
|
||||
if split_dim == "height":
|
||||
bottom = 0
|
||||
else:
|
||||
right = 0
|
||||
if world_size > 1:
|
||||
if self.split_dim == "height":
|
||||
if split_dim == "height":
|
||||
top = top if rank == 0 else 0
|
||||
bottom = bottom if rank == world_size - 1 else 0
|
||||
else:
|
||||
@@ -360,8 +371,7 @@ class WanDistConv2d(nn.Conv2d):
|
||||
self,
|
||||
source: nn.Conv2d,
|
||||
group: dist.ProcessGroup,
|
||||
split_dim: str = "height",
|
||||
split_padding: tuple[int, int] | None = None,
|
||||
deferred_padding: tuple[int, int, int, int] | None = None,
|
||||
):
|
||||
super().__init__(
|
||||
source.in_channels,
|
||||
@@ -378,75 +388,117 @@ class WanDistConv2d(nn.Conv2d):
|
||||
)
|
||||
self.load_state_dict(source.state_dict())
|
||||
self.group = group
|
||||
self.split_dim = split_dim
|
||||
self.split_tensor_dim = _spatial_dim(split_dim)
|
||||
kernel_extent = self.kernel_size[-2] if split_dim == "height" else self.kernel_size[-1]
|
||||
self.halo_size = (kernel_extent - 1) // 2
|
||||
pad_h = source.padding[-2] if isinstance(source.padding, tuple) else source.padding
|
||||
pad_w = source.padding[-1] if isinstance(source.padding, tuple) else source.padding
|
||||
if split_dim == "height":
|
||||
if split_padding is None:
|
||||
split_padding = (pad_h, pad_h)
|
||||
self._non_split_padding = (pad_w, pad_w, 0, 0)
|
||||
self.kernel_extent = self.kernel_size[-2]
|
||||
self.stride_extent = self.stride[-2]
|
||||
else:
|
||||
if split_padding is None:
|
||||
split_padding = (pad_w, pad_w)
|
||||
self._non_split_padding = (0, 0, pad_h, pad_h)
|
||||
self.kernel_extent = self.kernel_size[-1]
|
||||
self.stride_extent = self.stride[-1]
|
||||
self.split_pad_left, self.split_pad_right = split_padding
|
||||
self._halo_recv_top_buf: torch.Tensor | None = None
|
||||
self._halo_recv_bottom_buf: torch.Tensor | None = None
|
||||
self._trim_cache: dict[int, tuple[int, int, int]] = {}
|
||||
self._direct_padding = source.padding
|
||||
self._direct_padding_mode = source.padding_mode
|
||||
self._direct_reversed_padding = source._reversed_padding_repeated_twice
|
||||
self._deferred_padding = deferred_padding
|
||||
source_padding = source.padding if isinstance(source.padding, tuple) else (source.padding, source.padding)
|
||||
if len(source_padding) != 2 or not all(isinstance(value, int) for value in source_padding):
|
||||
raise ValueError(f"Wan spatial-shard Conv2d requires integer padding, got {source.padding!r}.")
|
||||
self._spatial_pad_h = int(source_padding[-2])
|
||||
self._spatial_pad_w = int(source_padding[-1])
|
||||
self.register_buffer("_halo_recv_top_buf", None, persistent=False)
|
||||
self.register_buffer("_halo_recv_bottom_buf", None, persistent=False)
|
||||
self._trim_cache: dict[tuple[str, int], tuple[int, int, int]] = {}
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = F.pad(x, self._non_split_padding)
|
||||
split_dim = _active_spatial_shard_split_dim()
|
||||
if split_dim is None:
|
||||
if self._direct_padding_mode != "zeros":
|
||||
x = F.pad(x, self._direct_reversed_padding, mode=self._direct_padding_mode)
|
||||
padding: str | tuple[int, int] = (0, 0)
|
||||
else:
|
||||
padding = self._direct_padding
|
||||
return F.conv2d(x, self.weight, self.bias, self.stride, padding, self.dilation, self.groups)
|
||||
|
||||
split_tensor_dim = _spatial_dim(split_dim)
|
||||
if split_dim == "height":
|
||||
kernel_extent = self.kernel_size[-2]
|
||||
stride_extent = self.stride[-2]
|
||||
non_split_padding = (self._spatial_pad_w, self._spatial_pad_w, 0, 0)
|
||||
default_split_padding = (self._spatial_pad_h, self._spatial_pad_h)
|
||||
deferred_split_padding = (
|
||||
(self._deferred_padding[2], self._deferred_padding[3]) if self._deferred_padding is not None else None
|
||||
)
|
||||
else:
|
||||
kernel_extent = self.kernel_size[-1]
|
||||
stride_extent = self.stride[-1]
|
||||
non_split_padding = (0, 0, self._spatial_pad_h, self._spatial_pad_h)
|
||||
default_split_padding = (self._spatial_pad_w, self._spatial_pad_w)
|
||||
deferred_split_padding = (
|
||||
(self._deferred_padding[0], self._deferred_padding[1]) if self._deferred_padding is not None else None
|
||||
)
|
||||
split_pad_left, split_pad_right = deferred_split_padding or default_split_padding
|
||||
halo_size = (kernel_extent - 1) // 2
|
||||
|
||||
x = F.pad(x, non_split_padding)
|
||||
x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(
|
||||
x,
|
||||
group=self.group,
|
||||
halo_size=self.halo_size,
|
||||
split_dim=self.split_dim,
|
||||
halo_size=halo_size,
|
||||
split_dim=split_dim,
|
||||
recv_top_buf=self._halo_recv_top_buf,
|
||||
recv_bottom_buf=self._halo_recv_bottom_buf,
|
||||
)
|
||||
shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim])
|
||||
shift, start, upper_bound = self._get_trim_params(
|
||||
x.shape[split_tensor_dim],
|
||||
split_dim=split_dim,
|
||||
halo_size=halo_size,
|
||||
split_pad_left=split_pad_left,
|
||||
split_pad_right=split_pad_right,
|
||||
kernel_extent=kernel_extent,
|
||||
stride_extent=stride_extent,
|
||||
)
|
||||
if shift:
|
||||
x_padded = _narrow_along_dim(
|
||||
x_padded,
|
||||
self.split_tensor_dim,
|
||||
split_tensor_dim,
|
||||
shift,
|
||||
x_padded.shape[self.split_tensor_dim] - shift,
|
||||
x_padded.shape[split_tensor_dim] - shift,
|
||||
)
|
||||
out = super().forward(x_padded)
|
||||
out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim)
|
||||
return _zero_invalid_extent(out, split_dim=self.split_dim)
|
||||
out = _trim_local_conv_output(out, halo_size, start, upper_bound, split_dim=split_dim)
|
||||
return _zero_invalid_extent(out, split_dim=split_dim)
|
||||
|
||||
def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
|
||||
trim_params = self._trim_cache.get(local_extent)
|
||||
def _get_trim_params(
|
||||
self,
|
||||
local_extent: int,
|
||||
*,
|
||||
split_dim: str,
|
||||
halo_size: int,
|
||||
split_pad_left: int,
|
||||
split_pad_right: int,
|
||||
kernel_extent: int,
|
||||
stride_extent: int,
|
||||
) -> tuple[int, int, int]:
|
||||
cache_key = (split_dim, local_extent)
|
||||
trim_params = self._trim_cache.get(cache_key)
|
||||
if trim_params is None:
|
||||
rank, world_size = _rank_world(self.group)
|
||||
trim_params = _compute_conv_trim_params(
|
||||
local_extent=local_extent,
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
halo_size=self.halo_size,
|
||||
pad_before=self.split_pad_left,
|
||||
pad_after=self.split_pad_right,
|
||||
kernel_extent=self.kernel_extent,
|
||||
stride_extent=self.stride_extent,
|
||||
halo_size=halo_size,
|
||||
pad_before=split_pad_left,
|
||||
pad_after=split_pad_right,
|
||||
kernel_extent=kernel_extent,
|
||||
stride_extent=stride_extent,
|
||||
)
|
||||
self._trim_cache[local_extent] = trim_params
|
||||
self._trim_cache[cache_key] = trim_params
|
||||
return trim_params
|
||||
|
||||
|
||||
class WanDistCausalConv3d(nn.Conv3d):
|
||||
# Cross-PR compatibility contract for the lossless Wan data-movement
|
||||
# installer: it may fuse the no-context direct path, but must leave this
|
||||
# wrapper in control whenever a spatial-shard context is active.
|
||||
_vllm_omni_dynamic_spatial_shard_conv = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
source: nn.Conv3d,
|
||||
group: dist.ProcessGroup,
|
||||
split_dim: str = "height",
|
||||
):
|
||||
super().__init__(
|
||||
source.in_channels,
|
||||
@@ -463,88 +515,98 @@ class WanDistCausalConv3d(nn.Conv3d):
|
||||
)
|
||||
self.load_state_dict(source.state_dict())
|
||||
self.group = group
|
||||
self.split_dim = split_dim
|
||||
self.split_tensor_dim = _spatial_dim(split_dim)
|
||||
source_padding = getattr(source, "_padding", None)
|
||||
if source_padding is None:
|
||||
p_t, p_h, p_w = source.padding
|
||||
source_padding = (p_w, p_w, p_h, p_h, 2 * p_t, 0)
|
||||
self._source_padding = tuple(source_padding)
|
||||
if split_dim == "height":
|
||||
self.split_pad_left = int(self._source_padding[2])
|
||||
self.split_pad_right = int(self._source_padding[3])
|
||||
self.kernel_extent = self.kernel_size[-2]
|
||||
self.stride_extent = self.stride[-2]
|
||||
padding = (
|
||||
self._source_padding[0],
|
||||
self._source_padding[1],
|
||||
0,
|
||||
0,
|
||||
self._source_padding[4],
|
||||
self._source_padding[5],
|
||||
)
|
||||
else:
|
||||
self.split_pad_left = int(self._source_padding[0])
|
||||
self.split_pad_right = int(self._source_padding[1])
|
||||
self.kernel_extent = self.kernel_size[-1]
|
||||
self.stride_extent = self.stride[-1]
|
||||
padding = (
|
||||
0,
|
||||
0,
|
||||
self._source_padding[2],
|
||||
self._source_padding[3],
|
||||
self._source_padding[4],
|
||||
self._source_padding[5],
|
||||
)
|
||||
self.halo_size = (self.kernel_extent - 1) // 2
|
||||
self._padding = padding if self.halo_size > 0 else self._source_padding
|
||||
self._halo_recv_top_buf: torch.Tensor | None = None
|
||||
self._halo_recv_bottom_buf: torch.Tensor | None = None
|
||||
self._trim_cache: dict[int, tuple[int, int, int]] = {}
|
||||
self.register_buffer("_halo_recv_top_buf", None, persistent=False)
|
||||
self.register_buffer("_halo_recv_bottom_buf", None, persistent=False)
|
||||
self._trim_cache: dict[tuple[str, int], tuple[int, int, int]] = {}
|
||||
|
||||
def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
|
||||
padding = list(self._padding)
|
||||
split_dim = _active_spatial_shard_split_dim()
|
||||
padding = list(self._source_padding)
|
||||
if cache_x is not None and padding[4] > 0:
|
||||
cache_x = cache_x.to(x.device)
|
||||
x = torch.cat([cache_x, x], dim=2)
|
||||
padding[4] -= cache_x.shape[2]
|
||||
|
||||
if split_dim is None:
|
||||
return super().forward(F.pad(x, padding))
|
||||
|
||||
split_tensor_dim = _spatial_dim(split_dim)
|
||||
if split_dim == "height":
|
||||
split_pad_left = int(padding[2])
|
||||
split_pad_right = int(padding[3])
|
||||
kernel_extent = self.kernel_size[-2]
|
||||
stride_extent = self.stride[-2]
|
||||
padding[2] = 0
|
||||
padding[3] = 0
|
||||
else:
|
||||
split_pad_left = int(padding[0])
|
||||
split_pad_right = int(padding[1])
|
||||
kernel_extent = self.kernel_size[-1]
|
||||
stride_extent = self.stride[-1]
|
||||
padding[0] = 0
|
||||
padding[1] = 0
|
||||
halo_size = (kernel_extent - 1) // 2
|
||||
|
||||
x = F.pad(x, padding)
|
||||
x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(
|
||||
x,
|
||||
group=self.group,
|
||||
halo_size=self.halo_size,
|
||||
split_dim=self.split_dim,
|
||||
halo_size=halo_size,
|
||||
split_dim=split_dim,
|
||||
recv_top_buf=self._halo_recv_top_buf,
|
||||
recv_bottom_buf=self._halo_recv_bottom_buf,
|
||||
)
|
||||
shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim])
|
||||
shift, start, upper_bound = self._get_trim_params(
|
||||
x.shape[split_tensor_dim],
|
||||
split_dim=split_dim,
|
||||
halo_size=halo_size,
|
||||
split_pad_left=split_pad_left,
|
||||
split_pad_right=split_pad_right,
|
||||
kernel_extent=kernel_extent,
|
||||
stride_extent=stride_extent,
|
||||
)
|
||||
if shift:
|
||||
x_padded = _narrow_along_dim(
|
||||
x_padded,
|
||||
self.split_tensor_dim,
|
||||
split_tensor_dim,
|
||||
shift,
|
||||
x_padded.shape[self.split_tensor_dim] - shift,
|
||||
x_padded.shape[split_tensor_dim] - shift,
|
||||
)
|
||||
out = super().forward(x_padded)
|
||||
out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim)
|
||||
return _zero_invalid_extent(out, split_dim=self.split_dim)
|
||||
out = _trim_local_conv_output(out, halo_size, start, upper_bound, split_dim=split_dim)
|
||||
return _zero_invalid_extent(out, split_dim=split_dim)
|
||||
|
||||
def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
|
||||
trim_params = self._trim_cache.get(local_extent)
|
||||
def _get_trim_params(
|
||||
self,
|
||||
local_extent: int,
|
||||
*,
|
||||
split_dim: str,
|
||||
halo_size: int,
|
||||
split_pad_left: int,
|
||||
split_pad_right: int,
|
||||
kernel_extent: int,
|
||||
stride_extent: int,
|
||||
) -> tuple[int, int, int]:
|
||||
cache_key = (split_dim, local_extent)
|
||||
trim_params = self._trim_cache.get(cache_key)
|
||||
if trim_params is None:
|
||||
rank, world_size = _rank_world(self.group)
|
||||
trim_params = _compute_conv_trim_params(
|
||||
local_extent=local_extent,
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
halo_size=self.halo_size,
|
||||
pad_before=self.split_pad_left,
|
||||
pad_after=self.split_pad_right,
|
||||
kernel_extent=self.kernel_extent,
|
||||
stride_extent=self.stride_extent,
|
||||
halo_size=halo_size,
|
||||
pad_before=split_pad_left,
|
||||
pad_after=split_pad_right,
|
||||
kernel_extent=kernel_extent,
|
||||
stride_extent=stride_extent,
|
||||
)
|
||||
self._trim_cache[local_extent] = trim_params
|
||||
self._trim_cache[cache_key] = trim_params
|
||||
return trim_params
|
||||
|
||||
|
||||
@@ -591,12 +653,15 @@ def _trim_local_conv_output(
|
||||
return out
|
||||
|
||||
|
||||
def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup, split_dim: str) -> None:
|
||||
def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup) -> None:
|
||||
if getattr(module, "_vllm_omni_spatial_shard_attention", False):
|
||||
return
|
||||
orig_forward = module.forward
|
||||
|
||||
def _forward(self: nn.Module, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
|
||||
split_dim = _active_spatial_shard_split_dim()
|
||||
if split_dim is None:
|
||||
return orig_forward(x, *args, **kwargs)
|
||||
_, world_size = _rank_world(group)
|
||||
if world_size <= 1:
|
||||
return orig_forward(x, *args, **kwargs)
|
||||
@@ -630,7 +695,6 @@ def _replace_child(
|
||||
name: str,
|
||||
child: nn.Module,
|
||||
group: dist.ProcessGroup,
|
||||
split_dim: str,
|
||||
) -> None:
|
||||
if child.__class__.__name__ == "WanCausalConv3d":
|
||||
setattr(
|
||||
@@ -639,51 +703,39 @@ def _replace_child(
|
||||
WanDistCausalConv3d(
|
||||
child,
|
||||
group,
|
||||
split_dim=split_dim,
|
||||
),
|
||||
)
|
||||
return
|
||||
if isinstance(child, nn.ZeroPad2d):
|
||||
padding = tuple(int(p) for p in child.padding)
|
||||
module_padding = padding
|
||||
if parent.__class__.__name__ == "Sequential":
|
||||
# Let the following WanDistConv2d account for global after-edge
|
||||
# padding; this module handles the non-split dimension and before edge.
|
||||
if split_dim == "height":
|
||||
module_padding = (padding[0], padding[1], padding[2], 0)
|
||||
else:
|
||||
module_padding = (padding[0], 0, padding[2], padding[3])
|
||||
setattr(
|
||||
parent,
|
||||
name,
|
||||
WanDistZeroPad2d(
|
||||
module_padding,
|
||||
padding,
|
||||
group,
|
||||
split_dim=split_dim,
|
||||
split_padding=(padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1]),
|
||||
defer_split_after_padding=parent.__class__.__name__ == "Sequential",
|
||||
),
|
||||
)
|
||||
return
|
||||
if isinstance(child, nn.Conv2d):
|
||||
split_padding = None
|
||||
deferred_padding = None
|
||||
if name == "1" and parent.__class__.__name__ == "Sequential":
|
||||
# WanResample downsample uses ZeroPad2d((0, 1, 0, 1)) before a
|
||||
# stride-2 conv with padding=0. Only the last rank should see the
|
||||
# bottom/right global padding, which is approximated by split_padding.
|
||||
prev = getattr(parent, "0", None)
|
||||
if isinstance(prev, WanDistZeroPad2d):
|
||||
split_padding = prev.split_padding
|
||||
deferred_padding = prev.padding
|
||||
elif isinstance(prev, nn.ZeroPad2d):
|
||||
pad = prev.padding
|
||||
split_padding = (int(pad[2]), int(pad[3])) if split_dim == "height" else (int(pad[0]), int(pad[1]))
|
||||
deferred_padding = tuple(int(value) for value in prev.padding)
|
||||
setattr(
|
||||
parent,
|
||||
name,
|
||||
WanDistConv2d(
|
||||
child,
|
||||
group,
|
||||
split_dim=split_dim,
|
||||
split_padding=split_padding,
|
||||
deferred_padding=deferred_padding,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -691,21 +743,20 @@ def _replace_child(
|
||||
def _patch_decoder_modules(
|
||||
module: nn.Module,
|
||||
group: dist.ProcessGroup,
|
||||
split_dim: str,
|
||||
inside_attention: bool = False,
|
||||
) -> None:
|
||||
if module.__class__.__name__ == "WanAttentionBlock":
|
||||
_patch_attention_block(module, group, split_dim)
|
||||
_patch_attention_block(module, group)
|
||||
inside_attention = True
|
||||
|
||||
for name, child in list(module.named_children()):
|
||||
if child.__class__.__name__ == "WanCausalConv3d":
|
||||
_replace_child(module, name, child, group, split_dim)
|
||||
_replace_child(module, name, child, group)
|
||||
continue
|
||||
if isinstance(child, nn.Conv2d) and not inside_attention:
|
||||
_replace_child(module, name, child, group, split_dim)
|
||||
_replace_child(module, name, child, group)
|
||||
continue
|
||||
_patch_decoder_modules(child, group, split_dim, inside_attention=inside_attention)
|
||||
_patch_decoder_modules(child, group, inside_attention=inside_attention)
|
||||
|
||||
|
||||
def _decoder_upsample_count(decoder: nn.Module) -> int:
|
||||
@@ -716,16 +767,25 @@ def _decoder_upsample_count(decoder: nn.Module) -> int:
|
||||
return count
|
||||
|
||||
|
||||
def _clear_wan_spatial_shard_runtime_buffers(vae: Any) -> None:
|
||||
decoder = getattr(vae, "decoder", None)
|
||||
if decoder is None:
|
||||
return
|
||||
for module in decoder.modules():
|
||||
if "_halo_recv_top_buf" in module._buffers:
|
||||
module._halo_recv_top_buf = None
|
||||
if "_halo_recv_bottom_buf" in module._buffers:
|
||||
module._halo_recv_bottom_buf = None
|
||||
|
||||
|
||||
def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_dim: str = "height") -> None:
|
||||
"""Patch ``vae.decoder`` once for spatially-sharded decode.
|
||||
|
||||
This mutates the already-loaded decoder in place by swapping its spatial
|
||||
convolutions/padding for halo-exchanging variants and wrapping
|
||||
``decoder.forward``. The patch is permanent for the lifetime of the VAE
|
||||
instance and is applied only once (subsequent calls are no-ops). A given
|
||||
instance is bound to a single ``split_dim``; switching between
|
||||
``"height"`` and ``"width"`` requires a fresh VAE instance and raises here
|
||||
otherwise.
|
||||
convolutions/padding for context-aware variants and attaching a dedicated
|
||||
spatial forward. Outside that forward the replacements execute their exact
|
||||
local PyTorch behavior, so a later auto-selected tile request remains safe.
|
||||
The same installed decoder can alternate height and width sharding.
|
||||
|
||||
Only group-relative rank 0 assembles the final decoded frame, mirroring the
|
||||
distributed tiled-decode ``broadcast_result=False`` contract; the other ranks
|
||||
@@ -733,18 +793,12 @@ def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_d
|
||||
"""
|
||||
_spatial_dim(split_dim)
|
||||
if getattr(vae, "_vllm_omni_wan_spatial_shard_installed", False):
|
||||
installed_split_dim = getattr(vae, "_vllm_omni_wan_spatial_shard_split_dim", "height")
|
||||
if installed_split_dim != split_dim:
|
||||
raise ValueError(
|
||||
"Wan spatial-shard VAE decoder was already patched for "
|
||||
f"{installed_split_dim!r} split; create a fresh VAE instance to use {split_dim!r} split."
|
||||
)
|
||||
return
|
||||
decoder = getattr(vae, "decoder", None)
|
||||
if decoder is None:
|
||||
raise ValueError("Wan spatial-shard VAE decode requires a decoder module.")
|
||||
|
||||
_patch_decoder_modules(decoder, group, split_dim)
|
||||
_patch_decoder_modules(decoder, group)
|
||||
upsample_count = _decoder_upsample_count(decoder)
|
||||
orig_forward = decoder.forward
|
||||
|
||||
@@ -754,6 +808,8 @@ def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_d
|
||||
feat_cache: list[torch.Tensor] | None = None,
|
||||
feat_idx: list[int] | None = None,
|
||||
first_chunk: bool = False,
|
||||
*,
|
||||
split_dim: str,
|
||||
) -> torch.Tensor:
|
||||
if feat_idx is None:
|
||||
feat_idx = [0]
|
||||
@@ -781,10 +837,9 @@ def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_d
|
||||
_SPATIAL_SHARD_CONTEXT.reset(token)
|
||||
return gather_and_trim_extent(out, expected_extent=expected_extent, split_dim=split_dim, group=group, dst=0)
|
||||
|
||||
decoder.forward = MethodType(_forward, decoder)
|
||||
decoder._vllm_omni_spatial_shard_forward = MethodType(_forward, decoder)
|
||||
vae._vllm_omni_wan_spatial_shard_installed = True
|
||||
vae._vllm_omni_wan_spatial_shard_split_dim = split_dim
|
||||
logger.info("Installed Wan VAE %s-sharded decode.", split_dim)
|
||||
logger.info("Installed Wan VAE dynamic-axis spatial-shard decode.")
|
||||
|
||||
|
||||
def spatial_shard_decode(
|
||||
@@ -813,11 +868,12 @@ def spatial_shard_decode(
|
||||
decoded_chunks = []
|
||||
for i in range(z.shape[2]):
|
||||
vae._conv_idx = [0]
|
||||
chunk = vae.decoder(
|
||||
chunk = vae.decoder._vllm_omni_spatial_shard_forward(
|
||||
x[:, :, i : i + 1, :, :],
|
||||
feat_cache=vae._feat_map,
|
||||
feat_idx=vae._conv_idx,
|
||||
first_chunk=(i == 0),
|
||||
split_dim=split_dim,
|
||||
)
|
||||
if produce_output:
|
||||
decoded_chunks.append(chunk)
|
||||
@@ -831,6 +887,10 @@ def spatial_shard_decode(
|
||||
out = z.new_zeros(0)
|
||||
finally:
|
||||
vae.clear_cache()
|
||||
# Halo buffers are request-shape scratch space allocated outside the
|
||||
# tagged weight pool. Drop live references so sleep/offload never has
|
||||
# to preserve or discard them as model state.
|
||||
_clear_wan_spatial_shard_runtime_buffers(vae)
|
||||
|
||||
if not return_dict:
|
||||
return (out,)
|
||||
|
||||
@@ -167,9 +167,6 @@ class OmniEngineArgs(EngineArgs):
|
||||
Required when single-stage mode is active.
|
||||
omni_master_port: TCP port for the OmniMasterServer registration
|
||||
socket. Required when single-stage mode is active.
|
||||
stage_configs_path: Optional path to a JSON/YAML file containing
|
||||
stage configurations for the multi-stage pipeline. If None,
|
||||
stage configs are resolved from the model's default configuration.
|
||||
output_modalities: Optional list of output modality names to enable
|
||||
(e.g. ["text", "audio"]). If None, all modalities supported by
|
||||
the model are used.
|
||||
@@ -231,7 +228,6 @@ class OmniEngineArgs(EngineArgs):
|
||||
omni_dp_size_local: int = 1
|
||||
omni_lb_policy: str = "random"
|
||||
omni_heartbeat_timeout: float = 30.0
|
||||
stage_configs_path: str | None = None
|
||||
output_modalities: list[str] | None = None
|
||||
log_stats: bool = False
|
||||
custom_pipeline_args: dict[str, Any] | None = None
|
||||
@@ -490,7 +486,6 @@ class OrchestratorArgs:
|
||||
ray_address: str | None = None
|
||||
|
||||
# === Config Files ===
|
||||
stage_configs_path: str | None = None
|
||||
deploy_config: str | None = None
|
||||
stage_overrides: str | None = None # raw JSON string; parsed downstream
|
||||
# Optional composable-parallel strategy.yaml; orchestrator reads it, overlays
|
||||
@@ -582,7 +577,6 @@ SHARED_FIELDS: frozenset[str] = frozenset(
|
||||
"model", # orch: detect model_type; engine: load weights
|
||||
"stage_id", # orch: route (headless); engine: identity
|
||||
"log_stats", # both want the flag
|
||||
"stage_configs_path", # orch: load legacy YAML; engine: may reference for validation
|
||||
"async_chunk", # orch: read from CLI, redistribute; engine: per-stage flag
|
||||
"tokenizer", # orch: detect model type; engine: tokenization
|
||||
}
|
||||
|
||||
@@ -9,7 +9,6 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
import dataclasses
|
||||
import json
|
||||
import queue
|
||||
import threading
|
||||
@@ -24,7 +23,6 @@ import janus
|
||||
import torch
|
||||
from omegaconf import OmegaConf
|
||||
from vllm import envs as vllm_envs
|
||||
from vllm.engine.arg_utils import EngineArgs
|
||||
from vllm.inputs import PromptType
|
||||
from vllm.logger import init_logger
|
||||
from vllm.v1.engine import EngineCoreRequest
|
||||
@@ -34,7 +32,6 @@ from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remot
|
||||
from vllm_omni.config.stage_config import (
|
||||
DuplexSessionRuntimeConfig,
|
||||
load_deploy_config,
|
||||
strip_parent_engine_args,
|
||||
)
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig, parse_attention_config
|
||||
from vllm_omni.diffusion.diffusion_engine import supports_audio_output
|
||||
@@ -69,7 +66,6 @@ from vllm_omni.engine.stage_runtime import (
|
||||
)
|
||||
from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
|
||||
from vllm_omni.entrypoints.utils import (
|
||||
is_new_format_deploy_config,
|
||||
load_and_resolve_stage_configs,
|
||||
parse_stage_overrides,
|
||||
)
|
||||
@@ -85,39 +81,6 @@ if TYPE_CHECKING:
|
||||
|
||||
_STARTUP_POLL_INTERVAL_S = 1.0
|
||||
_REQUEST_QUEUE_MAXSIZE = 256
|
||||
# ============================================================================
|
||||
# Parent-EngineArgs field-routing contracts (consumed by
|
||||
# AsyncOmniEngine._strip_parent_engine_args when ``stage_configs_path`` is set).
|
||||
# ============================================================================
|
||||
|
||||
# Fields that must survive the "equal to default → strip" filter because
|
||||
# diffusion stages need them even when equal to vllm's default value
|
||||
# (e.g. colocate worker setup relies on worker_extension_cls being forwarded).
|
||||
_PARENT_ARGS_KEEP: frozenset[str] = frozenset(
|
||||
{
|
||||
"worker_extension_cls",
|
||||
"allowed_local_media_path",
|
||||
"allowed_media_domains",
|
||||
# Legacy stage-config YAMLs may intentionally leave parallel or
|
||||
# distributed knobs unspecified at the stage level and rely on
|
||||
# top-level CLI values to fill them in during the per-stage merge.
|
||||
# Keep these fields so stages that omit them can inherit CLI values,
|
||||
# while stages with explicit YAML values still win because the legacy
|
||||
# stage-config loader prefers stage-local engine args.
|
||||
"tensor_parallel_size",
|
||||
}
|
||||
)
|
||||
|
||||
# Omni orchestrator-level fields consumed by ``_resolve_stage_configs`` that
|
||||
# must never leak into per-stage EngineArgs (``stage_configs_path`` would
|
||||
# trigger the ``create_model_config`` guard).
|
||||
_PARENT_ARGS_STRIP: frozenset[str] = frozenset({"stage_configs_path"})
|
||||
|
||||
|
||||
# Fields always populated by callers (via ``from_cli_args`` / ``asdict``) so
|
||||
# their presence as an override is never a surprise — suppress the
|
||||
# "override ignored" warning for these.
|
||||
_PARENT_ARGS_NO_WARN: frozenset[str] = frozenset({"model"})
|
||||
|
||||
|
||||
class AsyncOmniEngine:
|
||||
@@ -210,10 +173,8 @@ class AsyncOmniEngine:
|
||||
self._omni_master_port,
|
||||
)
|
||||
|
||||
# Stage resolution pops deploy_config, so get pipeline-wide settings
|
||||
# beforehand. The stage CLI exposes the same deploy YAML through
|
||||
# stage_configs_path.
|
||||
deploy_config_path = kwargs.get("deploy_config") or kwargs.get("stage_configs_path")
|
||||
# Stage resolution pops deploy_config, so get pipeline-wide settings beforehand.
|
||||
deploy_config_path = kwargs.get("deploy_config")
|
||||
# ``trust_remote_code`` is tri-state (bool | None): ``None`` means "not
|
||||
# specified" so stage-config resolution can defer to the deploy yaml's
|
||||
# per-stage value (see ``with_trust_remote_code_override``). The
|
||||
@@ -1116,40 +1077,6 @@ class AsyncOmniEngine:
|
||||
default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion"
|
||||
return default_stage_cfg
|
||||
|
||||
@staticmethod
|
||||
def _strip_single_engine_args(kwargs: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Remove parent ``EngineArgs`` fields from *kwargs*.
|
||||
|
||||
When ``stage_configs_path`` is set, per-stage engine args are defined
|
||||
in the YAML. Top-level single-engine fields (``compilation_config``,
|
||||
``tensor_parallel_size``, …) must not leak into per-stage configs via
|
||||
the ``base_engine_args`` merge in ``load_stage_configs_from_yaml`` —
|
||||
they can cause type errors (e.g. ``compilation_config`` as a JSON
|
||||
string rejected by ``VllmConfig``) or silently override YAML values.
|
||||
|
||||
Logs a warning for any parent field whose value differs from the
|
||||
dataclass default, so users know their explicit overrides are ignored.
|
||||
See the module-level ``_PARENT_ARGS_*`` constants for the routing
|
||||
contracts this method enforces.
|
||||
"""
|
||||
parent_fields: dict[str, dataclasses.Field] = {f.name: f for f in dataclasses.fields(EngineArgs)}
|
||||
result, overridden = strip_parent_engine_args(
|
||||
kwargs,
|
||||
parent_fields=parent_fields,
|
||||
keep_keys=_PARENT_ARGS_KEEP,
|
||||
strip_keys=_PARENT_ARGS_STRIP,
|
||||
no_warn_keys=_PARENT_ARGS_NO_WARN,
|
||||
)
|
||||
|
||||
if overridden:
|
||||
logger.warning(
|
||||
"stage_configs_path is set — the following top-level engine "
|
||||
"args are ignored (per-stage YAML takes precedence): %s",
|
||||
", ".join(sorted(overridden)),
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
def _apply_strategy_lb_policy(self, derived: str | None, kwargs: dict[str, Any]) -> None:
|
||||
"""Apply a strategy-derived ``omni_lb_policy`` to the engine.
|
||||
|
||||
@@ -1188,7 +1115,6 @@ class AsyncOmniEngine:
|
||||
) -> tuple[str, list[Any]]:
|
||||
"""Resolve stage configs and inject defaults shared by orchestrator/headless."""
|
||||
|
||||
stage_configs_path = kwargs.get("stage_configs_path", None)
|
||||
deploy_config_path = kwargs.pop("deploy_config", None)
|
||||
strategy_config_path = kwargs.pop("strategy_config", None)
|
||||
stage_overrides_json = kwargs.pop("stage_overrides", None)
|
||||
@@ -1196,26 +1122,15 @@ class AsyncOmniEngine:
|
||||
if explicit_stage_configs is not None:
|
||||
logger.warning(
|
||||
"`stage_configs` is not part of the public API. "
|
||||
"Ignoring it and resolving stages from stage_configs_path/model factory."
|
||||
"Ignoring it and resolving stages from deploy_config/model factory."
|
||||
)
|
||||
|
||||
# Legacy ``stage_args`` YAMLs own per-stage EngineArgs completely, so
|
||||
# strip parent CLI fields. New-format deploy YAMLs (also accepted via
|
||||
# ``--stage-configs-path``) merge CLI overrides the same way as
|
||||
# ``--deploy-config`` — do not strip, or flags like
|
||||
# ``interleave_mm_strings`` / ``media_io_kwargs`` are silently dropped.
|
||||
if stage_configs_path is not None and not is_new_format_deploy_config(stage_configs_path):
|
||||
base_kwargs = self._strip_single_engine_args(kwargs)
|
||||
else:
|
||||
base_kwargs = kwargs
|
||||
|
||||
# Parse --stage-overrides JSON string if provided
|
||||
stage_overrides = parse_stage_overrides(stage_overrides_json)
|
||||
|
||||
config_path, stage_configs, strategy_lb_policy = load_and_resolve_stage_configs(
|
||||
model,
|
||||
stage_configs_path,
|
||||
base_kwargs,
|
||||
kwargs,
|
||||
trust_remote_code=trust_remote_code,
|
||||
default_stage_cfg_factory=lambda: self._create_default_diffusion_stage_cfg(kwargs),
|
||||
deploy_config_path=deploy_config_path,
|
||||
|
||||
@@ -123,9 +123,8 @@ class AsyncOmni(EngineClient, OmniBase):
|
||||
Args:
|
||||
model: Model name or path to load.
|
||||
**kwargs: Additional keyword arguments.
|
||||
- stage_configs_path: Optional path to YAML file containing stage
|
||||
configurations. If None, configurations are resolved from model
|
||||
pipeline factory.
|
||||
- deploy_config: Optional path to a deploy YAML. If None,
|
||||
configurations are resolved from the model pipeline factory.
|
||||
- log_stats: Whether to enable statistics logging.
|
||||
- stage_init_timeout: Timeout for per-stage initialization.
|
||||
- init_timeout: Total timeout for orchestrator startup.
|
||||
|
||||
@@ -790,9 +790,10 @@ class OmniServeCommand(CLISubcommand):
|
||||
"--vae-parallel-mode",
|
||||
type=str,
|
||||
default="tile",
|
||||
choices=["tile", "spatial_shard_height", "spatial_shard_width"],
|
||||
choices=["tile", "auto", "spatial_shard_height", "spatial_shard_width"],
|
||||
help="VAE parallel decode strategy for diffusion models. "
|
||||
"'tile' (default) uses patch/tile parallel decode; "
|
||||
"'tile' (default) uses patch/tile decode; 'auto' selects the longer spatial-shard axis for "
|
||||
"validated Wan tiled decodes on a full DiT group, then falls back to patch/tile decode; "
|
||||
"'spatial_shard_height'/'spatial_shard_width' use spatially-sharded decode that splits "
|
||||
"decoder feature maps along height/width and exchanges halo regions. The "
|
||||
"'spatial_shard_*' modes require vae_patch_parallel_size to match the DiT group size. "
|
||||
@@ -945,8 +946,6 @@ def run_headless(args: TrackingNamespace) -> None:
|
||||
|
||||
config_path, stage_configs, _ = load_and_resolve_stage_configs(
|
||||
model,
|
||||
# The serve CLI no longer accepts legacy stage_args YAMLs.
|
||||
None,
|
||||
args_dict,
|
||||
# store_true cannot express an explicit False: absent maps to None
|
||||
# ("not specified") so the deploy yaml's per-stage value applies.
|
||||
|
||||
@@ -679,7 +679,7 @@ async def build_async_omni_from_stage_config(
|
||||
EngineClient instance (AsyncOmni) ready for use
|
||||
|
||||
Note:
|
||||
Stage configurations are loaded from args.stage_configs_path if provided,
|
||||
Stage configurations are loaded from ``args.deploy_config`` when provided,
|
||||
otherwise from the model's default configuration.
|
||||
"""
|
||||
|
||||
@@ -1156,7 +1156,7 @@ async def omni_init_app_state(
|
||||
state.openai_serving_duplex = None
|
||||
if state.openai_serving_chat is not None and should_enable_duplex_endpoint(
|
||||
state.stage_configs,
|
||||
config_path=getattr(args, "stage_configs_path", None) or getattr(args, "deploy_config", None),
|
||||
config_path=getattr(args, "deploy_config", None),
|
||||
):
|
||||
from vllm_omni.experimental.fullduplex.openai.serving import OmniDuplexSessionHandler
|
||||
|
||||
|
||||
@@ -6,7 +6,6 @@ from dataclasses import fields, is_dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, get_args, get_origin
|
||||
|
||||
import yaml
|
||||
from vllm.logger import init_logger
|
||||
from vllm.sampling_params import RequestOutputKind, SamplingParams
|
||||
from vllm.transformers_utils.config import get_config, get_hf_file_to_dict
|
||||
@@ -30,15 +29,6 @@ PROJECT_ROOT = Path(__file__).parent.parent.parent
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def is_new_format_deploy_config(path: str | None) -> bool:
|
||||
"""Return True when *path* is a deploy YAML (``stages:`` without legacy ``stage_args:``)."""
|
||||
if not path or not os.path.exists(path):
|
||||
return False
|
||||
with open(path, encoding="utf-8") as f:
|
||||
peek = yaml.safe_load(f) or {}
|
||||
return isinstance(peek, dict) and "stages" in peek and "stage_args" not in peek
|
||||
|
||||
|
||||
_DIFFUSERS_CLASS_TO_CONFIG: dict[str, str] = {
|
||||
"GlmImagePipeline": "glm_image",
|
||||
}
|
||||
@@ -606,7 +596,6 @@ def parse_stage_overrides(value: Any) -> dict[str, dict[str, Any]] | None:
|
||||
|
||||
def load_and_resolve_stage_configs(
|
||||
model: str,
|
||||
stage_configs_path: str | None,
|
||||
kwargs: dict | None,
|
||||
*,
|
||||
trust_remote_code: bool | None,
|
||||
@@ -615,17 +604,15 @@ def load_and_resolve_stage_configs(
|
||||
stage_overrides: dict[str, dict[str, Any]] | None = None,
|
||||
strategy_config_path: str | None = None,
|
||||
) -> tuple[str, list, str | None]:
|
||||
"""Load stage configurations from model or YAML file with fallback to defaults.
|
||||
"""Load stage configurations from a deploy YAML or model defaults.
|
||||
|
||||
Args:
|
||||
model: Model name or path
|
||||
stage_configs_path: Optional path to legacy YAML (stage_args format)
|
||||
kwargs: Engine arguments to merge with stage configs
|
||||
trust_remote_code: Whether to trust remote code while resolving the model config.
|
||||
default_stage_cfg_factory: Optional callable that takes no args and returns
|
||||
default stage config list when no configs are found
|
||||
deploy_config_path: Optional path to deploy YAML (new format).
|
||||
Mutually exclusive with ``stage_configs_path``.
|
||||
deploy_config_path: Optional path to a deploy YAML.
|
||||
stage_overrides: Per-stage overrides from ``--stage-overrides`` JSON.
|
||||
Keys are stage_id strings, values are dicts of overrides.
|
||||
strategy_config_path: Optional path to a composable-parallel
|
||||
@@ -636,33 +623,7 @@ def load_and_resolve_stage_configs(
|
||||
the strategy-derived pipeline-wide load-balancer policy (``None`` when no
|
||||
strategy set one), returned for the engine to apply.
|
||||
"""
|
||||
if stage_configs_path is not None and deploy_config_path is not None:
|
||||
raise ValueError(
|
||||
"--stage-configs-path and --deploy-config are mutually exclusive: "
|
||||
"they use different path resolution rules and loading paths. "
|
||||
"Use --deploy-config for new-format YAMLs (preferred); "
|
||||
"--stage-configs-path is kept only for the legacy `stage_args` format "
|
||||
"and will be removed in a future release."
|
||||
)
|
||||
if stage_configs_path is not None and deploy_config_path is None:
|
||||
if not os.path.exists(stage_configs_path):
|
||||
raise FileNotFoundError(
|
||||
f"--stage-configs-path {stage_configs_path!r} does not exist. "
|
||||
"Legacy `stage_configs/` yamls were replaced by `vllm_omni/deploy/<model>.yaml`; "
|
||||
"use --deploy-config. See docs/configuration/stage_configs.md."
|
||||
)
|
||||
if is_new_format_deploy_config(stage_configs_path):
|
||||
deploy_config_path = stage_configs_path
|
||||
stage_configs_path = None
|
||||
else:
|
||||
logger.warning(
|
||||
"--stage-configs-path is deprecated; migrate %r and use --deploy-config.",
|
||||
stage_configs_path,
|
||||
)
|
||||
|
||||
omni_lb_policy: str | None = None
|
||||
if deploy_config_path is not None:
|
||||
config_path = deploy_config_path
|
||||
config_path = deploy_config_path if deploy_config_path is not None else resolve_model_config_path(model)
|
||||
stage_configs, omni_lb_policy = load_stage_configs_from_model(
|
||||
model,
|
||||
trust_remote_code=trust_remote_code,
|
||||
@@ -677,24 +638,6 @@ def load_and_resolve_stage_configs(
|
||||
stage_configs = create_config(_convert_dataclasses_to_dict(default_stage_cfg))
|
||||
else:
|
||||
stage_configs = []
|
||||
elif stage_configs_path is None:
|
||||
config_path = resolve_model_config_path(model)
|
||||
stage_configs, omni_lb_policy = load_stage_configs_from_model(
|
||||
model,
|
||||
trust_remote_code=trust_remote_code,
|
||||
base_engine_args=kwargs,
|
||||
stage_overrides=stage_overrides,
|
||||
strategy_config_path=strategy_config_path,
|
||||
)
|
||||
if not stage_configs:
|
||||
if default_stage_cfg_factory is not None:
|
||||
default_stage_cfg = default_stage_cfg_factory()
|
||||
stage_configs = create_config(_convert_dataclasses_to_dict(default_stage_cfg))
|
||||
else:
|
||||
stage_configs = []
|
||||
else:
|
||||
config_path = stage_configs_path
|
||||
stage_configs = load_stage_configs_from_yaml(stage_configs_path, base_engine_args=kwargs)
|
||||
|
||||
stage_configs = filter_stages(config_path, stage_configs, kwargs)
|
||||
logger.debug(f"stage_configs: {stage_configs}")
|
||||
|
||||
@@ -730,7 +730,9 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin):
|
||||
has_encoder_input = self.model_config.is_encoder_decoder and num_encoder_reqs > 0
|
||||
|
||||
# Run forward pass
|
||||
clear_kv_metadata = self.speculative_config is None
|
||||
defer_kv_connector_finalize = self.speculative_config is not None and (
|
||||
get_pp_group().is_last_rank or self.broadcast_pp_output
|
||||
)
|
||||
with (
|
||||
record_function_or_nullcontext("forward"),
|
||||
set_ascend_forward_context(
|
||||
@@ -744,13 +746,12 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin):
|
||||
model_instance=self.model,
|
||||
skip_compiled=has_encoder_input,
|
||||
has_sinks=self._has_sinks,
|
||||
input_ids=input_ids,
|
||||
eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False,
|
||||
),
|
||||
self.maybe_get_kv_connector_output(
|
||||
scheduler_output,
|
||||
**(
|
||||
{"defer_finalize": not clear_kv_metadata}
|
||||
{"defer_finalize": defer_kv_connector_finalize}
|
||||
),
|
||||
) as kv_connector_output,
|
||||
):
|
||||
|
||||
@@ -381,7 +381,9 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin
|
||||
has_encoder_input = self.model_config.is_encoder_decoder and num_encoder_reqs > 0
|
||||
|
||||
# Run forward pass
|
||||
clear_kv_metadata = self.speculative_config is None
|
||||
defer_kv_connector_finalize = self.speculative_config is not None and (
|
||||
get_pp_group().is_last_rank or self.broadcast_pp_output
|
||||
)
|
||||
with (
|
||||
record_function_or_nullcontext("forward"),
|
||||
set_ascend_forward_context(
|
||||
@@ -395,12 +397,11 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin
|
||||
model_instance=self.model,
|
||||
skip_compiled=has_encoder_input,
|
||||
has_sinks=self._has_sinks,
|
||||
input_ids=input_ids,
|
||||
eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False,
|
||||
),
|
||||
self.maybe_get_kv_connector_output(
|
||||
scheduler_output,
|
||||
**({"defer_finalize": not clear_kv_metadata}),
|
||||
**({"defer_finalize": defer_kv_connector_finalize}),
|
||||
) as kv_connector_output,
|
||||
):
|
||||
if self.cache_config.mamba_cache_mode == "align":
|
||||
@@ -654,9 +655,10 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin
|
||||
num_active_loras: int = 0,
|
||||
profile_seq_lens: int | None = None,
|
||||
profile_cpp: bool = False,
|
||||
randomize_inputs: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
# only support eager mode and piecewise graph now
|
||||
assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes()
|
||||
assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.is_valid_runtime_mode()
|
||||
# If cudagraph_mode.decode_mode() == FULL and
|
||||
# cudagraph_mode.separate_routine(). This means that we are using
|
||||
# different graphs and/or modes for mixed prefill-decode batches vs.
|
||||
@@ -917,7 +919,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin
|
||||
batch_descriptor=batch_desc,
|
||||
model_instance=self.model,
|
||||
has_sinks=self._has_sinks,
|
||||
input_ids=input_ids,
|
||||
eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False,
|
||||
):
|
||||
outputs = self.model(
|
||||
|
||||
@@ -125,9 +125,10 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
|
||||
num_active_loras: int = 0,
|
||||
profile_seq_lens: int | None = None,
|
||||
profile_cpp: bool = False,
|
||||
randomize_inputs: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
# only support eager mode and piecewise graph now
|
||||
assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes()
|
||||
assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.is_valid_runtime_mode()
|
||||
# If cudagraph_mode.decode_mode() == FULL and
|
||||
# cudagraph_mode.separate_routine(). This means that we are using
|
||||
# different graphs and/or modes for mixed prefill-decode batches vs.
|
||||
@@ -367,7 +368,9 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
|
||||
if hasattr(self.drafter, "model") and hasattr(self.drafter.model, "compute_logits"):
|
||||
return self.drafter.model.compute_logits(hidden_states[dummy_indices])
|
||||
|
||||
with set_ascend_forward_context(
|
||||
with (
|
||||
self.maybe_randomize_inputs(input_ids, inputs_embeds, randomize_inputs=randomize_inputs),
|
||||
set_ascend_forward_context(
|
||||
attn_metadata,
|
||||
self.vllm_config,
|
||||
num_tokens=num_tokens_padded,
|
||||
@@ -378,8 +381,8 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
|
||||
batch_descriptor=batch_desc,
|
||||
model_instance=self.model,
|
||||
has_sinks=self._has_sinks,
|
||||
input_ids=input_ids,
|
||||
eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False,
|
||||
),
|
||||
):
|
||||
# ---------------------------------------Omni-new----------------------------------------------
|
||||
if getattr(self.model, "talker", None) is not None and self.has_talker_mtp:
|
||||
|
||||
Reference in new issue
Block a user