Merge origin/main into dev/rebase-v0170
Resolve conflict in omni_generation_scheduler.py: keep both num_cached_tokens tracking and additional_information caching. Made-with: Cursor Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
This commit is contained in:
tzhouam committed
commit
2398a7c1d7
40 files changed
+2209
-155
No files matched your search
@@ -1,5 +1,5 @@
|
||||
steps:
|
||||
- label: "Omni Model Test with H100"
|
||||
- label: ":full_moon: Omni Model Test with H100"
|
||||
timeout_in_minutes: 90
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
@@ -41,7 +41,7 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Omni Model Test"
|
||||
- label: ":full_moon: Omni Model Test"
|
||||
timeout_in_minutes: 60
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
@@ -56,18 +56,22 @@ steps:
|
||||
always-pull: true
|
||||
shm-size: "8gb"
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Omni Model Perf Test"
|
||||
timeout_in_minutes: 120
|
||||
- label: ":full_moon: Omni Model Perf Test with H100"
|
||||
key: nightly-performance
|
||||
timeout_in_minutes: 180
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export BENCHMARK_DIR=tests
|
||||
- pytest -s -v tests/perf/scripts/run_benchmark.py
|
||||
- buildkite-agent artifact upload "tests/*.json"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -96,3 +100,18 @@ steps:
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":email: Nightly Perf Collection & Email"
|
||||
key: nightly-perf-distribution
|
||||
depends_on: nightly-performance
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- pip install openpyxl
|
||||
- export DEFAULT_INPUT_DIR=tests
|
||||
- export DEFAULT_OUTPUT_DIR=tests
|
||||
- buildkite-agent artifact download "tests/*.json" . --step nightly-performance
|
||||
- python tools/nightly/generate_nightly_perf_excel.py
|
||||
- python tools/nightly/send_nightly_perf_email.py
|
||||
- buildkite-agent artifact upload "tests/*.xlsx"
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
This section lists the most common options for running vLLM-Omni.
|
||||
|
||||
For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.14.0/configuration/index.html)
|
||||
For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.16.0/configuration/index.html)
|
||||
|
||||
Currently, the main options are maintained by stage configs for each model.
|
||||
|
||||
@@ -19,3 +19,4 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
|
||||
- **[TeaCache Configuration](../user_guide/diffusion/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss
|
||||
- **[Cache-DiT Configuration](../user_guide/diffusion/cache_dit_acceleration.md)** - Enable Cache-DiT as cache acceleration backends for DiT models
|
||||
- **[Parallelism Configuration](../user_guide/diffusion/parallelism_acceleration.md)** - Enable parallelism (e.g., sequence parallelism) for for DiT models
|
||||
- **[CPU Offloading](../user_guide/diffusion/cpu_offload_diffusion.md)** - Enable CPU offloading (model-level and layerwise) for for DiT models
|
||||
@@ -0,0 +1,289 @@
|
||||
# VAE Patch Parallelism
|
||||
|
||||
This document describes how to add **VAE Patch Parallelism** support to a diffusion model.
|
||||
We use **Qwen-Image** as the reference implementation.
|
||||
|
||||
---
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Overview](#overview)
|
||||
- [Step-by-Step Implementation](#step-by-step-implementation)
|
||||
- [Testing](#testing)
|
||||
- [Reference Implementations](#reference-implementations)
|
||||
- [Summary](#summary)
|
||||
|
||||
---
|
||||
|
||||
## Overview
|
||||
|
||||
### What is Vae Patch parallel?
|
||||
|
||||
**VAE Patch Parallelism** is a decoding acceleration technique. Instead of decoding the entire latent tensor at once, the latent tensor is:
|
||||
|
||||
+ Split into multiple spatial tiles
|
||||
|
||||
+ Distributed across multiple ranks
|
||||
|
||||
+ Decoded in parallel
|
||||
|
||||
+ Merged to reconstruct the final output
|
||||
|
||||
This approach:
|
||||
|
||||
+ Distributes computation across multiple devices
|
||||
|
||||
+ Reduces peak memory usage per device
|
||||
|
||||
+ Accelerates decoding latency
|
||||
|
||||
### Architecture
|
||||
We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE decoding.
|
||||
|
||||
The executor is model-agnostic and accepts three function parameters:
|
||||
|
||||
+ split – Partition the latent into tiles
|
||||
|
||||
+ exec – Decode a single tile
|
||||
|
||||
+ merge – Combine decoded tiles into the final output
|
||||
|
||||
#### Execution Flow
|
||||
|
||||
+ Call split(z) to generate a list of TileTask and a GridSpec
|
||||
|
||||
+ Dispatch tasks across ranks using workload-based balancing
|
||||
|
||||
+ Each rank executes exec(task) on its assigned tiles
|
||||
|
||||
+ Gather decoded tile results to rank 0
|
||||
|
||||
+ Rank 0 performs merge(...)
|
||||
|
||||
+ (Optional) Broadcast final result to all ranks
|
||||
|
||||
This design separates:
|
||||
|
||||
+ Distributed execution logic
|
||||
|
||||
+ Model-specific tiling and merging logic
|
||||
|
||||
#### Why split / exec / merge is necessary?
|
||||
|
||||
The latent tensor cannot be arbitrarily partitioned.
|
||||
|
||||
During decoding:
|
||||
|
||||
+ Each output pixel may depend on neighboring pixels
|
||||
|
||||
+ The receptive field is model-dependent
|
||||
|
||||
Therefore:
|
||||
|
||||
+ Tiles must include overlap
|
||||
|
||||
+ Merge must perform blending to avoid seams
|
||||
|
||||
## Step-by-Step Implementation
|
||||
|
||||
### Step 1: Implement DistributedAutoencoderKLQwenImage
|
||||
`QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
|
||||
|
||||
|
||||
```
|
||||
class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMixin):
|
||||
@classmethod
|
||||
def from_pretrained(cls, *args: Any, **kwargs: Any):
|
||||
model = super().from_pretrained(*args, **kwargs)
|
||||
model.init_distributed()
|
||||
return model
|
||||
```
|
||||
**Key points**:
|
||||
+ Inherit both AutoencoderKLQwenImage and DistributedVaeMixin
|
||||
+ Call init_distributed() after loading weights
|
||||
|
||||
### Step 2: Implement split/exec/merge
|
||||
Reuse `AutoencoderKLQwenImage.tiled_decode` logic and divide it into three stages. And we need return tiles with `GridSpec` and `TileTask`:
|
||||
```
|
||||
class GridSpec:
|
||||
split_dims: tuple[int, ...] # Tensor dimensions being split (e.g., (2, 3) for (B, C, H, W))
|
||||
grid_shape: tuple[int, ...] # Tile grid layout (num_rows, num_cols)
|
||||
tile_spec: dict = field(default_factory=dict) # Metadata required for merging
|
||||
output_dtype: torch.dtype | None = None # Final output dtype
|
||||
```
|
||||
```
|
||||
class TileTask:
|
||||
tile_id: int # task id
|
||||
grid_coord: tuple[int, ...] # Tile position in grid
|
||||
tensor: torch.Tensor | list[torch.Tensor] # The tile tensor
|
||||
workload: int | float = 1 # Used for load balancing (e.g., tile area)
|
||||
```
|
||||
And tiled base split/exec/merge as follow:
|
||||
```
|
||||
def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
|
||||
# mostly copy from AutoencoderKL
|
||||
_, _, num_frames, height, width = z.shape
|
||||
sample_height = height * self.spatial_compression_ratio
|
||||
sample_width = width * self.spatial_compression_ratio
|
||||
|
||||
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
|
||||
tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
|
||||
tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
|
||||
tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
|
||||
|
||||
blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
|
||||
blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
|
||||
|
||||
# Split z into overlapping tiles and decode them separately.
|
||||
# The tiles have an overlap to avoid seams between tiles.
|
||||
tiletask_list = []
|
||||
for i in range(0, height, tile_latent_stride_height):
|
||||
for j in range(0, width, tile_latent_stride_width):
|
||||
time_list = []
|
||||
for k in range(num_frames):
|
||||
self._conv_idx = [0]
|
||||
tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
|
||||
time_list.append(tile)
|
||||
tiletask_list.append(
|
||||
TileTask(
|
||||
len(tiletask_list),
|
||||
(i // tile_latent_stride_height, j // tile_latent_stride_width),
|
||||
time_list,
|
||||
workload=time_list[0].shape[3] * time_list[0].shape[4],
|
||||
)
|
||||
)
|
||||
tile_spec = {
|
||||
"sample_height": sample_height,
|
||||
"sample_width": sample_width,
|
||||
"blend_height": blend_height,
|
||||
"blend_width": blend_width,
|
||||
}
|
||||
grid_spec = GridSpec(
|
||||
split_dims=(3, 4),
|
||||
grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
|
||||
tile_spec=tile_spec,
|
||||
output_dtype=self.dtype,
|
||||
)
|
||||
return tiletask_list, grid_spec
|
||||
|
||||
def tile_exec(self, task: TileTask) -> torch.Tensor:
|
||||
"""Decode a single latent tile into RGB space."""
|
||||
self.clear_cache()
|
||||
time = []
|
||||
for k in range(len(task.tensor)):
|
||||
self._conv_idx = [0]
|
||||
tile = self.post_quant_conv(task.tensor[k])
|
||||
decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx)
|
||||
time.append(decoded)
|
||||
result = torch.cat(time, dim=2)
|
||||
return result
|
||||
|
||||
def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
|
||||
"""Merge decoded tiles into a full image."""
|
||||
grid_h, grid_w = grid_spec.grid_shape
|
||||
result_rows = []
|
||||
self.clear_cache()
|
||||
|
||||
result_rows = []
|
||||
for i in range(grid_h):
|
||||
result_row = []
|
||||
for j in range(grid_w):
|
||||
tile = coord_tensor_map[(i, j)]
|
||||
if i > 0:
|
||||
tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
|
||||
if j > 0:
|
||||
tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
|
||||
result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width])
|
||||
result_rows.append(torch.cat(result_row, dim=-1))
|
||||
dec = torch.cat(result_rows, dim=3)[
|
||||
:, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"]
|
||||
]
|
||||
return dec
|
||||
```
|
||||
|
||||
### Step 3: Override tiled_decode
|
||||
We need to override tiled_decode, the main logic is:
|
||||
+ check distributed is enabled
|
||||
+ select split/exec/merge
|
||||
+ Invoke self.distributed_decoder.execute to decode
|
||||
```
|
||||
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)
|
||||
|
||||
logger.info("Decode run with distributed executor")
|
||||
result = self.distributed_decoder.execute(
|
||||
z,
|
||||
DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
|
||||
broadcast_result=True,
|
||||
)
|
||||
if not return_dict:
|
||||
return (result,)
|
||||
|
||||
return DecoderOutput(sample=result)
|
||||
```
|
||||
`broadcast_result` is set to True or False depending on the model; when enabled, the result will be used even on ranks other than 0.
|
||||
|
||||
### Step 4: Modify Pipeline
|
||||
Change vae model from AutoencoderKLQwenImage to DistributedAutoencoderKLQwenImage
|
||||
```
|
||||
class YourModelPipeline(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
od_config: OmniDiffusionConfig,
|
||||
prefix: str = "",
|
||||
):
|
||||
super().__init__()
|
||||
...
|
||||
- self.vae = AutoencoderKL.from_pretrained(
|
||||
- model, subfolder="vae", local_files_only=local_files_only).to(self.device)
|
||||
+ self.vae = DistributedAutoencoderKL.from_pretrained(
|
||||
+ model, subfolder="vae", local_files_only=local_files_only
|
||||
+ ).to(self.device)
|
||||
```
|
||||
|
||||
## Testing
|
||||
Verify numerical consistency between:
|
||||
+ vae_patch_parallel_size = 1
|
||||
|
||||
+ vae_patch_parallel_size = N
|
||||
|
||||
Example:
|
||||
torch.allclose(output_1, output_n, atol=1e-5)
|
||||
|
||||
Testing requirements:
|
||||
+ Fix random seed
|
||||
+ Use identical tiling strategy
|
||||
|
||||
```python
|
||||
m = Omni(
|
||||
model=model_name,
|
||||
vae_use_tiling=True,
|
||||
parallel_config=DiffusionParallelConfig(
|
||||
tensor_parallel_size=2,
|
||||
vae_patch_parallel_size=1, # or 2
|
||||
),
|
||||
)
|
||||
```
|
||||
When vae_patch_parallel_size is larger than the DiT world size, it will automatically fall back to using the DiT world size instead.
|
||||
|
||||
## Reference Implementations
|
||||
|
||||
Complete examples in the codebase:
|
||||
|
||||
| Model | Path | Notes |
|
||||
|-------|------|-------|
|
||||
| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | Distributed AutoencoderKL |
|
||||
| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | Distributed AutoencoderKLWan |
|
||||
| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | Distributed AutoencoderKLQwenImage |
|
||||
|
||||
---
|
||||
|
||||
## Summary
|
||||
|
||||
Adding Vae Patch Parallel support to diffusion model:
|
||||
|
||||
1. **Implement Distributed Vae** - mainly copy from `diffusers` tiled_decode, and refactor into split/exec/merge
|
||||
2. **Change vae model in pipeline to Distributed Vae**
|
||||
3. **Test** - Verify with `tensor_parallel_size=N` quality
|
||||
@@ -27,7 +27,7 @@ The following table shows which models are currently supported by parallelism me
|
||||
| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | ❌ |
|
||||
| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | ❌ |
|
||||
| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
@@ -51,9 +51,9 @@ The following table shows which models are currently supported by parallelism me
|
||||
|
||||
### VideoGen
|
||||
|
||||
| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP |
|
||||
|-------|------------------|:----------:|:--------------:|:---------------:|:----:|
|
||||
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ |
|
||||
| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP | VAE-Patch-Parallel |
|
||||
|-------|------------------|:----------:|:--------------:|:---------------:|:----:| :----:|
|
||||
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
|
||||
### Tensor Parallelism
|
||||
|
||||
@@ -85,7 +85,7 @@ outputs = omni.generate(
|
||||
VAE patch parallelism distributes the VAE decode workload across multiple ranks by splitting the latent spatially. It is configured via `DiffusionParallelConfig.vae_patch_parallel_size` and can be combined with other parallelism methods (e.g., TP).
|
||||
|
||||
!!! note "Enablement and feature gate"
|
||||
- VAE patch parallelism is currently **enabled only for validated pipelines** (currently: `Tongyi-MAI/Z-Image-Turbo`, `stabilityai/stable-diffusion-3.5`).
|
||||
- VAE patch parallelism is currently **enabled only for validated pipelines** (check [ImageGen](#imagegen) and [VideoGen](#videogen) for more information).
|
||||
- If `vae_patch_parallel_size > 1` is set for a validated pipeline, vLLM-Omni will automatically enable `vae_use_tiling` as a safety gate. (We use `vae_use_tiling` because it indicates the VAE supports diffusers tiling parameters like `tile_latent_min_size` and `tile_overlap_factor`.)
|
||||
|
||||
#### Offline Inference
|
||||
|
||||
@@ -58,8 +58,8 @@ The following table shows which models are currently supported by each accelerat
|
||||
| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||
| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||
| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
|
||||
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
@@ -67,13 +67,13 @@ The following table shows which models are currently supported by each accelerat
|
||||
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
|
||||
| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||
| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ |
|
||||
| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||
|
||||
### VideoGen
|
||||
|
||||
| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP |
|
||||
|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|
|
||||
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP | VAE-Patch-Parallel |
|
||||
|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|:----:|
|
||||
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
|
||||
### Quantization
|
||||
|
||||
|
||||
@@ -40,8 +40,7 @@ GLM-Image uses a 2-stage pipeline:
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model-path /path/to/glm-image \
|
||||
--config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--prompt "A beautiful sunset over the ocean with sailing boats" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
@@ -52,8 +51,7 @@ python end2end.py \
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model-path /path/to/glm-image \
|
||||
--config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--prompt "Transform this scene into a winter wonderland" \
|
||||
--image input.png \
|
||||
--output output_i2i.png
|
||||
@@ -64,7 +62,7 @@ python end2end.py \
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model-path /path/to/glm-image \
|
||||
--config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--prompt "A photorealistic cat sitting on a window sill" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
|
||||
@@ -178,9 +178,6 @@ def main(args: argparse.Namespace) -> None:
|
||||
print("=" * 60)
|
||||
|
||||
# Validate arguments
|
||||
if not args.model_path:
|
||||
raise ValueError("--model-path is required")
|
||||
|
||||
if not args.prompt:
|
||||
raise ValueError("--prompt is required")
|
||||
|
||||
@@ -365,7 +362,7 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument(
|
||||
"--model-path",
|
||||
type=str,
|
||||
required=True,
|
||||
default="zai-org/GLM-Image",
|
||||
help="Path to GLM-Image model directory or HuggingFace model ID",
|
||||
)
|
||||
parser.add_argument(
|
||||
|
||||
@@ -20,6 +20,7 @@ Usage:
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
@@ -105,6 +106,12 @@ def parse_args() -> argparse.Namespace:
|
||||
choices=[1, 2],
|
||||
help="Number of GPUs used for classifier free guidance parallel size.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vae-patch-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of GPUs used for VAE patch/tile parallelism (decode).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enforce-eager",
|
||||
action="store_true",
|
||||
@@ -173,6 +180,7 @@ def main():
|
||||
ring_degree=args.ring_degree,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
vae_patch_parallel_size=args.vae_patch_parallel_size,
|
||||
use_hsdp=args.use_hsdp,
|
||||
hsdp_shard_size=args.hsdp_shard_size,
|
||||
hsdp_replicate_size=args.hsdp_replicate_size,
|
||||
@@ -199,10 +207,14 @@ def main():
|
||||
print(f" Model: {args.model}")
|
||||
print(f" Inference steps: {args.num_inference_steps}")
|
||||
print(f" Frames: {args.num_frames}")
|
||||
print(f" Parallel configuration: cfg_parallel_size={args.cfg_parallel_size}")
|
||||
print(
|
||||
f" Parallel configuration: cfg_parallel_size={args.cfg_parallel_size},"
|
||||
f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size}"
|
||||
)
|
||||
print(f" Video size: {args.width}x{args.height}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
generation_start = time.perf_counter()
|
||||
# omni.generate() returns Generator[OmniRequestOutput, None, None]
|
||||
frames = omni.generate(
|
||||
{
|
||||
@@ -220,6 +232,11 @@ def main():
|
||||
num_frames=args.num_frames,
|
||||
),
|
||||
)
|
||||
generation_end = time.perf_counter()
|
||||
generation_time = generation_end - generation_start
|
||||
|
||||
# Print profiling results
|
||||
print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
|
||||
|
||||
# Extract video frames from OmniRequestOutput
|
||||
if isinstance(frames, list) and len(frames) > 0:
|
||||
|
||||
@@ -109,6 +109,12 @@ def parse_args() -> argparse.Namespace:
|
||||
default=1,
|
||||
help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vae-patch-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of GPUs used for VAE patch/tile parallelism (decode).",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -141,6 +147,7 @@ def main():
|
||||
ring_degree=args.ring_degree,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
vae_patch_parallel_size=args.vae_patch_parallel_size,
|
||||
)
|
||||
|
||||
# Check if profiling is requested via environment variable
|
||||
@@ -172,7 +179,9 @@ def main():
|
||||
print(f" Inference steps: {args.num_inference_steps}")
|
||||
print(f" Frames: {args.num_frames}")
|
||||
print(
|
||||
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}"
|
||||
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree},"
|
||||
f" cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size},"
|
||||
f" vae_patch_parallel_size={args.vae_patch_parallel_size}"
|
||||
)
|
||||
print(f" Video size: {args.width}x{args.height}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
@@ -160,10 +160,7 @@ esac
|
||||
echo "Running query type: $QUERY_TYPE"
|
||||
echo ""
|
||||
|
||||
|
||||
output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d @- <<EOF
|
||||
request_body=$(cat <<EOF
|
||||
{
|
||||
"model": "Qwen/Qwen2.5-Omni-7B",
|
||||
"sampling_params_list": $sampling_params_list,
|
||||
@@ -186,7 +183,12 @@ output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
|
||||
]
|
||||
}
|
||||
EOF
|
||||
)
|
||||
)
|
||||
|
||||
output=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \
|
||||
-X POST http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d "$request_body")
|
||||
|
||||
# Here it only shows the text content of the first choice. Audio content has many binaries, so it's not displayed here.
|
||||
echo "Output of request: $(echo "$output" | jq '.choices[0].message.content')"
|
||||
@@ -138,10 +138,7 @@ esac
|
||||
echo "Running query type: $QUERY_TYPE"
|
||||
echo ""
|
||||
|
||||
|
||||
output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d @- <<EOF
|
||||
request_body=$(cat <<EOF
|
||||
{
|
||||
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
||||
"sampling_params_list": $sampling_params_list,
|
||||
@@ -164,7 +161,12 @@ output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
|
||||
]
|
||||
}
|
||||
EOF
|
||||
)
|
||||
)
|
||||
|
||||
output=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \
|
||||
-X POST http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d "$request_body")
|
||||
|
||||
# Here it only shows the text content of the first choice. Audio content has many binaries, so it's not displayed here.
|
||||
echo "Output of request: $(echo "$output" | jq '.choices[0].message.content')"
|
||||
+2
-1
@@ -49,7 +49,8 @@ dev = [
|
||||
"imageio[ffmpeg]>=0.6.0",
|
||||
"opencv-python>=4.12.0.88",
|
||||
"mooncake-transfer-engine==0.3.8.post1",
|
||||
"av" # for ComfyUI tests
|
||||
"av", # for ComfyUI tests
|
||||
"openpyxl>=3.0.0", # for nightly CI
|
||||
]
|
||||
|
||||
docs = [
|
||||
|
||||
@@ -0,0 +1,165 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
|
||||
DistributedVaeExecutor,
|
||||
DistributedVaeMixin,
|
||||
GridSpec,
|
||||
TileTask,
|
||||
)
|
||||
|
||||
|
||||
class E2EOperator:
|
||||
"""tiles with (2, 3) -- (H,W)"""
|
||||
|
||||
def split(self, z):
|
||||
rows_num = 2
|
||||
cols_num = 3
|
||||
|
||||
h_size = z.shape[0]
|
||||
w_size = z.shape[1]
|
||||
|
||||
tasks = []
|
||||
for i in range(rows_num):
|
||||
for j in range(cols_num):
|
||||
tasks.append(
|
||||
TileTask(
|
||||
tile_id=len(tasks),
|
||||
grid_coord=(i, j),
|
||||
tensor=z[
|
||||
((i * h_size) // rows_num) : (((i + 1) * h_size) // rows_num),
|
||||
((j * w_size) // cols_num) : (((j + 1) * w_size) // cols_num),
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
grid_spec = GridSpec(
|
||||
split_dims=(0, 1),
|
||||
grid_shape=(2, 3),
|
||||
)
|
||||
return tasks, grid_spec
|
||||
|
||||
def exec(self, task: TileTask):
|
||||
return torch.full_like(task.tensor, fill_value=task.tile_id)
|
||||
|
||||
def merge(self, coord_tensor_map, grid_spec):
|
||||
tiles = []
|
||||
for r in range(grid_spec.grid_shape[0]):
|
||||
row_tiles = []
|
||||
for c in range(grid_spec.grid_shape[1]):
|
||||
coord = (r, c)
|
||||
row_tiles.append(coord_tensor_map[coord])
|
||||
tiles.append(torch.cat(row_tiles, dim=1))
|
||||
return torch.cat(tiles, dim=0)
|
||||
|
||||
|
||||
class DummyMixin(DistributedVaeMixin):
|
||||
def __init__(self):
|
||||
self.use_tiling = True
|
||||
self.distributed_decoder = MagicMock()
|
||||
self.distributed_decoder.parallel_size = 2
|
||||
self.distributed_decoder.group = None
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_dist():
|
||||
with (
|
||||
patch.object(dist, "get_world_size", return_value=2),
|
||||
patch.object(dist, "get_rank", return_value=0),
|
||||
patch.object(dist, "is_initialized", return_value=True),
|
||||
patch.object(dist, "gather", return_value=None),
|
||||
patch.object(dist, "broadcast", return_value=None),
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_dit_group():
|
||||
with patch(
|
||||
"vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor.get_dit_group",
|
||||
new=MagicMock(return_value=None),
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_dist_vae_executor():
|
||||
with (
|
||||
patch.object(DistributedVaeExecutor, "gather_tensors", side_effect=lambda x: [x]),
|
||||
patch.object(DistributedVaeExecutor, "broadcast_tensor", side_effect=lambda x: x),
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
# ============================
|
||||
# Unitest
|
||||
# ============================
|
||||
|
||||
|
||||
def test_balance_tasks():
|
||||
executor = DistributedVaeExecutor()
|
||||
workloads = [2, 5, 13, 8, 2, 4] # each is 17
|
||||
tasks = [TileTask(0, (i,), torch.tensor([i]), workload=load) for i, load in enumerate(workloads)]
|
||||
assigned = executor._balance_tasks(tasks, num_rank=2)
|
||||
assert len(assigned) == 2
|
||||
total_work = [sum(t.workload for t in group) for group in assigned]
|
||||
assert total_work[0] == total_work[1]
|
||||
|
||||
|
||||
def test_compute_global_padding_shape():
|
||||
executor = DistributedVaeExecutor()
|
||||
executor.rank = 0
|
||||
|
||||
local_results = [(0, torch.zeros((2, 3))), (1, torch.zeros((4, 2)))]
|
||||
shape = executor._compute_global_padding_shape(local_results, 2, "cpu")
|
||||
|
||||
assert shape == [2, 4, 3]
|
||||
|
||||
|
||||
def test_pack_and_unpack():
|
||||
executor = DistributedVaeExecutor()
|
||||
executor.world_size = 1
|
||||
|
||||
grid_spec = GridSpec(split_dims=(0, 1), grid_shape=(2, 2))
|
||||
|
||||
# ======================
|
||||
# pack
|
||||
# ======================
|
||||
local_results = [(0, torch.tensor([[1, 2], [3, 4]]))]
|
||||
|
||||
global_shape = [3, 3, 3] # (tiles, H, W)
|
||||
|
||||
tile_tensor, meta_tensor = executor._pack_local_tiles(
|
||||
local_results, global_shape, grid_spec, device="cpu", dtype=torch.int64
|
||||
)
|
||||
|
||||
# check pack
|
||||
assert tile_tensor.shape == torch.Size(global_shape)
|
||||
assert meta_tensor.shape == (global_shape[0], len(grid_spec.split_dims) + 1)
|
||||
assert meta_tensor[0, 0] == 0
|
||||
assert meta_tensor[0, 1] == 2
|
||||
assert meta_tensor[0, 2] == 2
|
||||
|
||||
# ======================
|
||||
# unpack
|
||||
# ======================
|
||||
meta_gather = [meta_tensor]
|
||||
tile_gather = [tile_tensor]
|
||||
|
||||
tid_coord_map = {0: (0, 0)}
|
||||
|
||||
coord_tensor_map = executor._unpack_tiles(meta_gather, tile_gather, grid_spec, tid_coord_map)
|
||||
|
||||
# check unpack
|
||||
assert torch.equal(coord_tensor_map[(0, 0)], torch.tensor([[1, 2], [3, 4]]))
|
||||
|
||||
|
||||
def test_is_distributed_enabled():
|
||||
mixin = DummyMixin()
|
||||
assert mixin.is_distributed_enabled() is True
|
||||
|
||||
mixin.use_tiling = False
|
||||
assert mixin.is_distributed_enabled() is False
|
||||
@@ -0,0 +1,154 @@
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
# ruff: noqa: E402
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
import time
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
# os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
all_model_configs = [
|
||||
{
|
||||
"model_name": "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
|
||||
"out_width": 1280,
|
||||
"out_height": 720,
|
||||
"out_frames": 16,
|
||||
},
|
||||
{
|
||||
"model_name": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
"out_width": 1280,
|
||||
"out_height": 720,
|
||||
"out_frames": 16,
|
||||
},
|
||||
{
|
||||
"model_name": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
|
||||
"out_width": 1280,
|
||||
"out_height": 720,
|
||||
"out_frames": 16,
|
||||
},
|
||||
{
|
||||
"model_name": "Qwen/Qwen-Image",
|
||||
"out_width": 3072,
|
||||
"out_height": 3072,
|
||||
"out_frames": 1,
|
||||
},
|
||||
# {
|
||||
# "model_name": "stepfun-ai/NextStep-1.1",
|
||||
# "out_width": 3072,
|
||||
# "out_height": 3072,
|
||||
# "out_frames": 1,
|
||||
# },
|
||||
]
|
||||
|
||||
|
||||
def is_nextstep_model(model_name: str) -> bool:
|
||||
"""Check if the model is a NextStep model by reading its config."""
|
||||
from vllm.transformers_utils.config import get_hf_file_to_dict
|
||||
|
||||
try:
|
||||
cfg = get_hf_file_to_dict("config.json", model_name)
|
||||
if cfg and cfg.get("model_type") == "nextstep":
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
return False
|
||||
|
||||
|
||||
def model_run(model_configs, tp, out_height, out_width, out_frames, using_tile, vae_patch_parallel_size=1):
|
||||
try:
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
tensor_parallel_size=tp,
|
||||
vae_patch_parallel_size=vae_patch_parallel_size,
|
||||
)
|
||||
|
||||
omni_kwargs = {
|
||||
"model": model_configs["model_name"],
|
||||
"vae_use_tiling": using_tile,
|
||||
"parallel_config": parallel_config,
|
||||
}
|
||||
use_nextstep = is_nextstep_model(model_configs["model_name"])
|
||||
if use_nextstep:
|
||||
# NextStep-1.1 requires explicit pipeline class
|
||||
omni_kwargs["model_class_name"] = "NextStep11Pipeline"
|
||||
m = Omni(**omni_kwargs)
|
||||
image = Image.new("RGB", (out_width, out_height), (0, 0, 0))
|
||||
start = time.perf_counter()
|
||||
outputs = m.generate(
|
||||
{
|
||||
"prompt": "A cat sitting on a table",
|
||||
"multi_modal_data": {"image": image},
|
||||
},
|
||||
sampling_params_list=OmniDiffusionSamplingParams(
|
||||
height=model_configs["out_height"],
|
||||
width=model_configs["out_width"],
|
||||
num_frames=model_configs["out_frames"],
|
||||
num_inference_steps=2,
|
||||
generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
|
||||
),
|
||||
)
|
||||
end = time.perf_counter()
|
||||
first_output = outputs[0]
|
||||
req_out = first_output.request_output[0]
|
||||
frames = req_out.images[0]
|
||||
if isinstance(frames, torch.Tensor):
|
||||
frames = frames.detach().cpu().numpy()
|
||||
elif isinstance(frames, Image.Image):
|
||||
frames = np.array(frames)
|
||||
# frames shape: (batch, num_frames, height, width, channels)
|
||||
cost = (end - start) * 1000
|
||||
return frames, cost
|
||||
finally:
|
||||
m.close()
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_configs", all_model_configs)
|
||||
def test_vae_parallel_model(model_configs: str):
|
||||
parallel_result, parallel_time = model_run(
|
||||
model_configs=model_configs,
|
||||
tp=2,
|
||||
out_width=1280,
|
||||
out_height=704,
|
||||
out_frames=5,
|
||||
using_tile=True,
|
||||
vae_patch_parallel_size=2,
|
||||
)
|
||||
non_parallel_result, non_parallel_time = model_run(
|
||||
model_configs=model_configs,
|
||||
tp=2,
|
||||
out_width=1280,
|
||||
out_height=704,
|
||||
out_frames=5,
|
||||
using_tile=True,
|
||||
vae_patch_parallel_size=1,
|
||||
)
|
||||
result_diff = np.abs(non_parallel_result - parallel_result)
|
||||
|
||||
mean_threshold = 3e-2
|
||||
max_threshold = 3e-2 # they should be totally same
|
||||
print(
|
||||
f"{model_configs['model_name']} TP = 2 (tile + parallel vs tile): "
|
||||
f"mean_abs_diff={result_diff.mean():.6e}, max_abs_diff={result_diff.max():.6e}; "
|
||||
f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; "
|
||||
f"parallel generate take time: {parallel_time:.2f} ms, non-parallel take time: {non_parallel_time:.2f} ms"
|
||||
)
|
||||
# assert non_parallel_time > parallel_time
|
||||
assert result_diff.mean() < mean_threshold
|
||||
assert result_diff.max() < max_threshold
|
||||
@@ -1,9 +1,5 @@
|
||||
import os
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
|
||||
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import threading
|
||||
from datetime import datetime
|
||||
@@ -14,6 +10,9 @@ import pytest
|
||||
|
||||
from tests.conftest import OmniServer, modify_stage_config
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
|
||||
|
||||
|
||||
def load_configs(config_path: str) -> list[dict[str, Any]]:
|
||||
try:
|
||||
@@ -123,6 +122,8 @@ def run_benchmark(args: list, test_name: str, flow, dataset_name: str, num_promp
|
||||
"--endpoint",
|
||||
"/v1/chat/completions",
|
||||
"--save-result",
|
||||
"--result-dir",
|
||||
os.environ.get("BENCHMARK_DIR", "tests"),
|
||||
"--result-filename",
|
||||
result_filename,
|
||||
]
|
||||
@@ -137,9 +138,9 @@ def run_benchmark(args: list, test_name: str, flow, dataset_name: str, num_promp
|
||||
for line in iter(process.stderr.readline, ""):
|
||||
print(line, end=" ")
|
||||
|
||||
if "--result-dir" in args:
|
||||
index = args.index("--result-dir")
|
||||
result_dir = args[index + 1]
|
||||
if "--result-dir" in command:
|
||||
index = command.index("--result-dir")
|
||||
result_dir = command[index + 1]
|
||||
else:
|
||||
result_dir = "./"
|
||||
|
||||
|
||||
@@ -0,0 +1,437 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Generate a nightly Excel performance report from JSON results.
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from collections.abc import Iterable, Sequence
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from openpyxl import Workbook
|
||||
from openpyxl.styles import PatternFill
|
||||
from openpyxl.utils import get_column_letter
|
||||
|
||||
LOGGER = logging.getLogger(__name__)
|
||||
|
||||
GREY_BLOCK_FILL = PatternFill(start_color="D3D3D3", fill_type="solid")
|
||||
|
||||
# Benchmark metric columns: grey the latest row's cell when value changed vs previous date.
|
||||
BENCHMARK_COLUMNS: tuple[str, ...] = (
|
||||
"num_prompts",
|
||||
"request_rate",
|
||||
"burstiness",
|
||||
"max_concurrency",
|
||||
"duration",
|
||||
"completed",
|
||||
"failed",
|
||||
"request_throughput",
|
||||
"output_throughput",
|
||||
"total_token_throughput",
|
||||
"mean_ttft_ms",
|
||||
"p99_ttft_ms",
|
||||
"mean_tpot_ms",
|
||||
"p99_tpot_ms",
|
||||
"mean_itl_ms",
|
||||
"p99_itl_ms",
|
||||
"mean_e2el_ms",
|
||||
"p99_e2el_ms",
|
||||
"mean_audio_rtf",
|
||||
"p99_audio_rtf",
|
||||
"mean_audio_duration_s",
|
||||
"p99_audio_duration_s",
|
||||
)
|
||||
# Columns that get float coercion and number format in Excel. Excludes request_rate ("inf" str)
|
||||
# and max_concurrency (null); leave those as-is. If they become float in the future, they are
|
||||
# still written correctly from JSON without coercion here.
|
||||
NUMERIC_FORMAT_COLUMNS: tuple[str, ...] = tuple(
|
||||
c for c in BENCHMARK_COLUMNS if c not in ("request_rate", "max_concurrency")
|
||||
)
|
||||
|
||||
_COLUMNS_FILENAME = "nightly_perf_summary_columns.txt"
|
||||
DEFAULT_INPUT_DIR = os.getenv("DEFAULT_INPUT_DIR") if os.getenv("DEFAULT_INPUT_DIR") else "tests"
|
||||
DEFAULT_OUTPUT_DIR = os.getenv("DEFAULT_OUTPUT_DIR") if os.getenv("DEFAULT_OUTPUT_DIR") else "tests"
|
||||
|
||||
|
||||
def _load_summary_columns(script_dir: str) -> list[str]:
|
||||
"""Load summary column names from a file next to this script; fallback to default if missing."""
|
||||
path = os.path.join(script_dir, _COLUMNS_FILENAME)
|
||||
default = [
|
||||
"date",
|
||||
"endpoint_type",
|
||||
"backend",
|
||||
"model_id",
|
||||
"tokenizer_id",
|
||||
"num_prompts",
|
||||
"request_rate",
|
||||
"burstiness",
|
||||
"max_concurrency",
|
||||
"duration",
|
||||
"completed",
|
||||
"failed",
|
||||
"request_throughput",
|
||||
"output_throughput",
|
||||
"total_token_throughput",
|
||||
"mean_ttft_ms",
|
||||
"p99_ttft_ms",
|
||||
"mean_tpot_ms",
|
||||
"p99_tpot_ms",
|
||||
"mean_itl_ms",
|
||||
"p99_itl_ms",
|
||||
"mean_e2el_ms",
|
||||
"p99_e2el_ms",
|
||||
"mean_audio_rtf",
|
||||
"p99_audio_rtf",
|
||||
"mean_audio_duration_s",
|
||||
"p99_audio_duration_s",
|
||||
"commit_sha",
|
||||
"build_id",
|
||||
"build_url",
|
||||
"source_file",
|
||||
]
|
||||
if not os.path.isfile(path):
|
||||
return default
|
||||
columns: list[str] = []
|
||||
with open(path, encoding="utf-8") as f:
|
||||
for line in f:
|
||||
s = line.strip()
|
||||
if s and not s.startswith("#"):
|
||||
columns.append(s)
|
||||
return columns if columns else default
|
||||
|
||||
|
||||
def _vllm_omni_root() -> str:
|
||||
"""Resolve vllm-omni repo root: directory that contains a 'tests' subdir (and usually 'tools')."""
|
||||
path = os.path.dirname(os.path.abspath(__file__))
|
||||
while path and path != os.path.dirname(path):
|
||||
if os.path.isdir(os.path.join(path, "tests")):
|
||||
return path
|
||||
path = os.path.dirname(path)
|
||||
return os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", ".."))
|
||||
|
||||
|
||||
def _default_input_dir() -> str:
|
||||
"""Default: vllm-omni root / DEFAULT_INPUT_DIR (where performance JSON files live)."""
|
||||
root = _vllm_omni_root()
|
||||
return os.path.join(root, DEFAULT_INPUT_DIR)
|
||||
|
||||
|
||||
def _default_output_file() -> str:
|
||||
"""Default: vllm-omni root / DEFAULT_OUTPUT_DIR / nightly_perf_<timestamp>.xlsx."""
|
||||
ts = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
|
||||
return os.path.join(_vllm_omni_root(), DEFAULT_OUTPUT_DIR, f"nightly_perf_{ts}.xlsx")
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
"""Parse command-line arguments."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Read performance JSON files from vllm-omni/tests/ and generate an Excel report."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input-dir",
|
||||
type=str,
|
||||
default=_default_input_dir(),
|
||||
help="Directory containing performance JSON files; default is <vllm-omni-root>/DEFAULT_INPUT_DIR.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-file",
|
||||
type=str,
|
||||
default=_default_output_file(),
|
||||
help="Output path of the Excel report; \
|
||||
default is <vllm-omni-root>/DEFAULT_OUTPUT_DIR/nightly_perf_<timestamp>.xlsx.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--commit-sha",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Optional commit SHA; defaults to environment variable BUILDKITE_COMMIT if unset.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--build-id",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Optional build ID; defaults to environment variable BUILDKITE_BUILD_ID if unset.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--build-url",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Optional build URL; defaults to environment variable BUILDKITE_BUILD_URL if unset.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _load_json_file(path: str) -> dict[str, Any] | None:
|
||||
"""Safely load a single JSON file; return None and log a warning on failure."""
|
||||
try:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
except (OSError, json.JSONDecodeError) as exc:
|
||||
LOGGER.warning("failed to load json '%s': %s", path, exc)
|
||||
return None
|
||||
|
||||
if not isinstance(data, dict):
|
||||
LOGGER.warning("json root in '%s' is not an object, skip", path)
|
||||
return None
|
||||
|
||||
return data
|
||||
|
||||
|
||||
def _iter_json_records(input_dir: str) -> Iterable[dict[str, Any]]:
|
||||
"""Iterate over JSON files in the input directory and yield normalized records.
|
||||
commit_sha/build_id/build_url are not set here; they are applied later only to
|
||||
rows with the latest date (see _apply_build_metadata_to_latest_only).
|
||||
"""
|
||||
if not os.path.isdir(input_dir):
|
||||
LOGGER.warning("input dir '%s' does not exist or is not a directory", input_dir)
|
||||
return
|
||||
|
||||
for entry in sorted(os.listdir(input_dir)):
|
||||
if not entry.endswith(".json"):
|
||||
continue
|
||||
full_path = os.path.join(input_dir, entry)
|
||||
if not os.path.isfile(full_path):
|
||||
continue
|
||||
|
||||
data = _load_json_file(full_path)
|
||||
if data is None:
|
||||
continue
|
||||
|
||||
record: dict[str, Any] = dict(data)
|
||||
record.setdefault("date", datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S"))
|
||||
record["source_file"] = os.path.basename(full_path)
|
||||
yield record
|
||||
|
||||
|
||||
def _collect_records(input_dir: str) -> list[dict[str, Any]]:
|
||||
"""Collect all JSON records into a list."""
|
||||
records: list[dict[str, Any]] = []
|
||||
for record in _iter_json_records(input_dir):
|
||||
records.append(record)
|
||||
return records
|
||||
|
||||
|
||||
def _apply_build_metadata_to_latest_only(
|
||||
records: Sequence[dict[str, Any]],
|
||||
commit_sha: str | None,
|
||||
build_id: str | None,
|
||||
build_url: str | None,
|
||||
) -> None:
|
||||
"""Set commit_sha, build_id, build_url only on rows with the latest date.
|
||||
Other rows get None so that build info is not duplicated for older benchmark data.
|
||||
"""
|
||||
if not records:
|
||||
return
|
||||
max_date = max((r.get("date") or "") for r in records)
|
||||
for r in records:
|
||||
if (r.get("date") or "") == max_date:
|
||||
r["commit_sha"] = commit_sha
|
||||
r["build_id"] = build_id
|
||||
r["build_url"] = build_url
|
||||
else:
|
||||
r["commit_sha"] = None
|
||||
r["build_id"] = None
|
||||
r["build_url"] = None
|
||||
|
||||
|
||||
def _sort_records_for_summary(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
"""Sort so that same model_id is grouped, newest date first within each group."""
|
||||
by_date_desc = sorted(records, key=lambda r: (r.get("date") or ""), reverse=True)
|
||||
return sorted(by_date_desc, key=lambda r: (r.get("model_id") or ""))
|
||||
|
||||
|
||||
def _values_differ(a: Any, b: Any) -> bool:
|
||||
"""Return True if two values are considered different; avoid direct float equality."""
|
||||
if a is None and b is None:
|
||||
return False
|
||||
if a is None or b is None:
|
||||
return True
|
||||
if isinstance(a, float) and isinstance(b, float):
|
||||
if a != a and b != b:
|
||||
return False
|
||||
if a != a or b != b:
|
||||
return True
|
||||
return abs(a - b) > 1e-9
|
||||
return a != b
|
||||
|
||||
|
||||
def _apply_benchmark_change_highlight(
|
||||
ws,
|
||||
summary_columns: Sequence[str],
|
||||
records: Sequence[dict[str, Any]],
|
||||
) -> None:
|
||||
"""Grey cells in the latest row of each model when a benchmark metric changed vs previous date."""
|
||||
if not records:
|
||||
return
|
||||
col_to_index = {c: i + 1 for i, c in enumerate(summary_columns)}
|
||||
# Walk by model_id blocks (records already sorted by model_id, date desc).
|
||||
i = 0
|
||||
while i < len(records):
|
||||
model_id = records[i].get("model_id")
|
||||
block_start = i
|
||||
while i < len(records) and records[i].get("model_id") == model_id:
|
||||
i += 1
|
||||
block_end = i
|
||||
newest_idx = block_start
|
||||
prev_idx = block_start + 1 if block_start + 1 < block_end else None
|
||||
if prev_idx is None:
|
||||
continue
|
||||
excel_row = newest_idx + 2
|
||||
for col in BENCHMARK_COLUMNS:
|
||||
if col not in col_to_index:
|
||||
continue
|
||||
cur_val = records[newest_idx].get(col)
|
||||
prev_val = records[prev_idx].get(col)
|
||||
if _values_differ(cur_val, prev_val):
|
||||
ws.cell(row=excel_row, column=col_to_index[col]).fill = GREY_BLOCK_FILL
|
||||
|
||||
|
||||
def _build_raw_columns(
|
||||
records: Sequence[dict[str, Any]],
|
||||
summary_columns: Sequence[str],
|
||||
) -> list[str]:
|
||||
"""Infer the column set for the raw sheet based on all records."""
|
||||
keys: set[str] = set()
|
||||
for record in records:
|
||||
keys.update(record.keys())
|
||||
# Ensure summary columns appear first; remaining columns sorted alphabetically.
|
||||
ordered_keys: list[str] = []
|
||||
for key in summary_columns:
|
||||
if key in keys:
|
||||
ordered_keys.append(key)
|
||||
keys.discard(key)
|
||||
ordered_keys.extend(sorted(keys))
|
||||
return ordered_keys
|
||||
|
||||
|
||||
def _to_float_if_numeric(value: Any) -> Any:
|
||||
"""Coerce to float when possible so Excel treats as number; avoid #### from narrow columns."""
|
||||
if value is None:
|
||||
return value
|
||||
if isinstance(value, (int, float)):
|
||||
return float(value) if isinstance(value, int) else value
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return float(value)
|
||||
except (ValueError, TypeError):
|
||||
return value
|
||||
return value
|
||||
|
||||
|
||||
def _write_sheet(
|
||||
ws,
|
||||
columns: Sequence[str],
|
||||
rows: Iterable[dict[str, Any]],
|
||||
numeric_columns: Sequence[str] = (),
|
||||
) -> None:
|
||||
"""Write column names and row data into the given worksheet."""
|
||||
numeric_set = set(numeric_columns)
|
||||
ws.append(list(columns))
|
||||
for record in rows:
|
||||
row_values = []
|
||||
for col in columns:
|
||||
v = record.get(col)
|
||||
if col in numeric_set:
|
||||
v = _to_float_if_numeric(v)
|
||||
row_values.append(v)
|
||||
ws.append(row_values)
|
||||
|
||||
|
||||
def _format_benchmark_columns(
|
||||
ws,
|
||||
columns: Sequence[str],
|
||||
num_data_rows: int,
|
||||
) -> None:
|
||||
"""Set number format and column width for numeric benchmark columns so values display (no ####)."""
|
||||
numeric_set = set(NUMERIC_FORMAT_COLUMNS)
|
||||
for c, col_name in enumerate(columns):
|
||||
if col_name not in numeric_set:
|
||||
continue
|
||||
col_letter = get_column_letter(c + 1)
|
||||
ws.column_dimensions[col_letter].width = 14
|
||||
for r in range(2, 2 + num_data_rows):
|
||||
cell = ws.cell(row=r, column=c + 1)
|
||||
if cell.value is not None and isinstance(cell.value, (int, float)):
|
||||
cell.number_format = "0.0000"
|
||||
elif isinstance(cell.value, str):
|
||||
try:
|
||||
cell.value = float(cell.value)
|
||||
cell.number_format = "0.0000"
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
|
||||
|
||||
def _ensure_parent_dir(path: str) -> None:
|
||||
"""Ensure that the parent directory of the output file exists."""
|
||||
parent = os.path.dirname(os.path.abspath(path))
|
||||
if not parent:
|
||||
return
|
||||
os.makedirs(parent, exist_ok=True)
|
||||
|
||||
|
||||
def generate_excel_report(
|
||||
input_dir: str,
|
||||
output_file: str,
|
||||
commit_sha: str | None,
|
||||
build_id: str | None,
|
||||
build_url: str | None,
|
||||
) -> None:
|
||||
"""Main logic: load JSON records and generate an Excel report."""
|
||||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
summary_columns = _load_summary_columns(script_dir)
|
||||
|
||||
records = _collect_records(input_dir)
|
||||
if not records:
|
||||
LOGGER.warning("no valid json records found under '%s'", input_dir)
|
||||
|
||||
sorted_records = _sort_records_for_summary(records)
|
||||
_apply_build_metadata_to_latest_only(sorted_records, commit_sha, build_id, build_url)
|
||||
|
||||
wb = Workbook()
|
||||
ws_summary = wb.active
|
||||
ws_summary.title = "summary"
|
||||
|
||||
_write_sheet(ws_summary, summary_columns, sorted_records, numeric_columns=NUMERIC_FORMAT_COLUMNS)
|
||||
_format_benchmark_columns(ws_summary, summary_columns, len(sorted_records))
|
||||
_apply_benchmark_change_highlight(ws_summary, summary_columns, sorted_records)
|
||||
|
||||
if sorted_records:
|
||||
raw_columns = _build_raw_columns(sorted_records, summary_columns)
|
||||
ws_raw = wb.create_sheet(title="raw")
|
||||
_write_sheet(ws_raw, raw_columns, sorted_records)
|
||||
|
||||
_ensure_parent_dir(output_file)
|
||||
wb.save(output_file)
|
||||
LOGGER.info("excel report saved to '%s'", output_file)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Command-line entrypoint."""
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
|
||||
)
|
||||
|
||||
args = parse_args()
|
||||
|
||||
commit_sha = args.commit_sha or os.getenv("BUILDKITE_COMMIT")
|
||||
build_id = args.build_id or os.getenv("BUILDKITE_BUILD_ID")
|
||||
build_url = args.build_url or os.getenv("BUILDKITE_BUILD_URL")
|
||||
|
||||
generate_excel_report(
|
||||
input_dir=args.input_dir,
|
||||
output_file=args.output_file,
|
||||
commit_sha=commit_sha,
|
||||
build_id=build_id,
|
||||
build_url=build_url,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,233 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Send the nightly performance Excel report by email to DAILY_EMAIL_LIST.
|
||||
|
||||
Reads SMTP and recipient config from environment variables. Use --dry-run to
|
||||
validate config and print subject/body without sending.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import os
|
||||
import smtplib
|
||||
import sys
|
||||
from email.mime.application import MIMEApplication
|
||||
from email.mime.multipart import MIMEMultipart
|
||||
from email.mime.text import MIMEText
|
||||
|
||||
LOGGER = logging.getLogger(__name__)
|
||||
|
||||
# Env keys for SMTP and recipients.
|
||||
ENV_SMTP_HOST = "SMTP_HOST"
|
||||
ENV_SMTP_PORT = "SMTP_PORT"
|
||||
ENV_SMTP_USERNAME = "SMTP_USERNAME"
|
||||
ENV_SMTP_PASSWORD = "SMTP_PASSWORD"
|
||||
ENV_DAILY_EMAIL_LIST = "DAILY_EMAIL_LIST"
|
||||
ENV_EMAIL_SENDER = "EMAIL_SENDER"
|
||||
ENV_EMAIL_SUBJECT_PREFIX = "EMAIL_SUBJECT_PREFIX"
|
||||
ENV_BUILD_URL = "BUILDKITE_BUILD_URL"
|
||||
ENV_COMMIT = "BUILDKITE_COMMIT"
|
||||
|
||||
DEFAULT_OUTPUT_DIR = os.getenv("DEFAULT_OUTPUT_DIR")
|
||||
|
||||
# Do not attach Excel if size >= this (bytes); body will suggest downloading from build URL.
|
||||
MAX_ATTACHMENT_BYTES = 20 * 1024 * 1024
|
||||
|
||||
SMTP_RETRIES = 3
|
||||
SMTP_RETRY_DELAY_SEC = 5
|
||||
|
||||
|
||||
def _get_required_env() -> dict[str, str]:
|
||||
"""Read required env vars; raise SystemExit with clear message if any missing."""
|
||||
required = {
|
||||
ENV_SMTP_HOST: os.environ.get(ENV_SMTP_HOST),
|
||||
ENV_SMTP_PORT: os.environ.get(ENV_SMTP_PORT),
|
||||
ENV_SMTP_USERNAME: os.environ.get(ENV_SMTP_USERNAME),
|
||||
ENV_SMTP_PASSWORD: os.environ.get(ENV_SMTP_PASSWORD),
|
||||
ENV_DAILY_EMAIL_LIST: os.environ.get(ENV_DAILY_EMAIL_LIST),
|
||||
}
|
||||
missing = [k for k, v in required.items() if not (v and str(v).strip())]
|
||||
if missing:
|
||||
raise SystemExit(f"Missing required env vars: {', '.join(missing)}. Set them (e.g. in Buildkite secrets).")
|
||||
return {k: str(v).strip() for k, v in required.items()}
|
||||
|
||||
|
||||
def _get_latest_file(folder_path: str) -> str:
|
||||
"""Get the latest modified file from the folder path."""
|
||||
files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.endswith(".xlsx")]
|
||||
if not files:
|
||||
raise SystemExit(f"No Excel files found in {folder_path}")
|
||||
return max(files, key=os.path.getmtime)
|
||||
|
||||
|
||||
def _recipients_list(comma_separated: str) -> list[str]:
|
||||
"""Parse DAILY_EMAIL_LIST into a list of addresses."""
|
||||
return [a.strip() for a in comma_separated.split(",") if a.strip()]
|
||||
|
||||
|
||||
def _build_body(
|
||||
date_str: str,
|
||||
commit_sha: str | None,
|
||||
build_url: str | None,
|
||||
attachment_skipped: bool,
|
||||
) -> str:
|
||||
"""Plain-text body with build metadata; note when attachment was skipped."""
|
||||
lines = [
|
||||
f"Nightly performance report date: {date_str}",
|
||||
"",
|
||||
f"Commit: {commit_sha or 'N/A'}",
|
||||
f"Build: {build_url or 'N/A'}",
|
||||
"",
|
||||
]
|
||||
if attachment_skipped:
|
||||
lines.append(
|
||||
"Report file was too large to attach. Please download the Excel from the build artifacts (Build URL above)."
|
||||
)
|
||||
lines.append("")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _build_subject(prefix: str | None, date_str: str) -> str:
|
||||
"""Subject line: optional prefix + date."""
|
||||
base = f"Nightly Perf {date_str}"
|
||||
if prefix and prefix.strip():
|
||||
return f"{prefix.strip()} {base}"
|
||||
return base
|
||||
|
||||
|
||||
def _send_mail(
|
||||
report_file: str,
|
||||
date_str: str,
|
||||
dry_run: bool,
|
||||
) -> None:
|
||||
"""Load config, build message, and send (or dry-run)."""
|
||||
cfg = _get_required_env()
|
||||
recipients = _recipients_list(cfg[ENV_DAILY_EMAIL_LIST])
|
||||
if not recipients:
|
||||
raise SystemExit("DAILY_EMAIL_LIST is empty after parsing.")
|
||||
|
||||
commit_sha = os.environ.get(ENV_COMMIT)
|
||||
build_url = os.environ.get(ENV_BUILD_URL)
|
||||
sender = os.environ.get(ENV_EMAIL_SENDER) or cfg[ENV_SMTP_USERNAME]
|
||||
prefix = os.environ.get(ENV_EMAIL_SUBJECT_PREFIX)
|
||||
|
||||
size = os.path.getsize(report_file) if os.path.isfile(report_file) else 0
|
||||
attach_excel = size < MAX_ATTACHMENT_BYTES and size > 0
|
||||
attachment_skipped = os.path.isfile(report_file) and not attach_excel
|
||||
|
||||
body = _build_body(
|
||||
commit_sha=commit_sha, build_url=build_url, date_str=date_str, attachment_skipped=attachment_skipped
|
||||
)
|
||||
subject = _build_subject(prefix=prefix, date_str=date_str)
|
||||
|
||||
msg = MIMEMultipart()
|
||||
msg["Subject"] = subject
|
||||
msg["From"] = sender
|
||||
msg["To"] = ", ".join(recipients)
|
||||
msg.attach(MIMEText(body, "plain", "utf-8"))
|
||||
|
||||
if attach_excel:
|
||||
with open(report_file, "rb") as f:
|
||||
part = MIMEApplication(f.read(), _subtype="vnd.openxmlformats-officedocument.spreadsheetml.sheet")
|
||||
part.add_header("Content-Disposition", "attachment", filename=os.path.basename(report_file))
|
||||
msg.attach(part)
|
||||
|
||||
if dry_run:
|
||||
LOGGER.info("dry-run: not sending mail")
|
||||
print("To:", recipients, file=sys.stderr)
|
||||
print("Subject:", subject, file=sys.stderr)
|
||||
print("Attachment:", "yes" if attach_excel else "no (size limit)", file=sys.stderr)
|
||||
print("Body preview:", body[:300] + ("..." if len(body) > 300 else ""), file=sys.stderr)
|
||||
return
|
||||
|
||||
port = int(cfg[ENV_SMTP_PORT], 10)
|
||||
last_err: Exception | None = None
|
||||
for attempt in range(SMTP_RETRIES):
|
||||
try:
|
||||
with smtplib.SMTP(cfg[ENV_SMTP_HOST], port=port, timeout=30) as smtp:
|
||||
smtp.starttls()
|
||||
smtp.login(cfg[ENV_SMTP_USERNAME], cfg[ENV_SMTP_PASSWORD])
|
||||
smtp.sendmail(sender, recipients, msg.as_string())
|
||||
LOGGER.info("sent nightly perf email to %d recipient(s)", len(recipients))
|
||||
return
|
||||
except Exception as e:
|
||||
last_err = e
|
||||
LOGGER.warning("SMTP attempt %d/%d failed: %s", attempt + 1, SMTP_RETRIES, e)
|
||||
if attempt < SMTP_RETRIES - 1:
|
||||
import time
|
||||
|
||||
time.sleep(SMTP_RETRY_DELAY_SEC)
|
||||
raise SystemExit(f"Failed to send email after {SMTP_RETRIES} attempts.") from last_err
|
||||
|
||||
|
||||
def _date_from_filename(path: str) -> str:
|
||||
"""Try to derive a date string from report filename (e.g. nightly_perf_20260211-020704.xlsx -> 2026-02-11)."""
|
||||
base = os.path.splitext(os.path.basename(path))[0]
|
||||
if base.startswith("nightly_perf_") and len(base) > 13:
|
||||
raw = base.replace("nightly_perf_", "")[:8]
|
||||
if len(raw) == 8 and raw.isdigit():
|
||||
return f"{raw[:4]}-{raw[4:6]}-{raw[6:8]}"
|
||||
return base or "unknown"
|
||||
|
||||
|
||||
def _vllm_omni_root() -> str:
|
||||
"""Resolve vllm-omni repo root: directory that contains a 'tests' subdir (and usually 'tools')."""
|
||||
path = os.path.dirname(os.path.abspath(__file__))
|
||||
while path and path != os.path.dirname(path):
|
||||
if os.path.isdir(os.path.join(path, "tests")):
|
||||
return path
|
||||
path = os.path.dirname(path)
|
||||
return os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", ".."))
|
||||
|
||||
|
||||
def _default_output_dir() -> str:
|
||||
"""Default: vllm-omni root / DEFAULT_OUTPUT_DIR (where performance .xlsx files live)."""
|
||||
root = _vllm_omni_root()
|
||||
return os.path.join(root, DEFAULT_OUTPUT_DIR)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
"""Parse command-line arguments."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Send nightly performance Excel report by email (config from env).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--report-file",
|
||||
type=str,
|
||||
default=_default_output_dir(),
|
||||
help="Folder/file path to the nightly_perf_*.xlsx file; default is DEFAULT_OUTPUT_DIR.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--date",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Date string for subject/body (default: inferred from report filename).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dry-run",
|
||||
action="store_true",
|
||||
help="Print recipient, subject, and body; do not send.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Entrypoint."""
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
|
||||
)
|
||||
args = parse_args()
|
||||
|
||||
report_file = _get_latest_file(args.report_file) if os.path.isdir(args.report_file) else args.report_file
|
||||
if not os.path.isfile(report_file):
|
||||
raise SystemExit(f"Report file not found: {report_file}")
|
||||
|
||||
date_str = args.date or _date_from_filename(report_file)
|
||||
_send_mail(report_file=report_file, date_str=date_str, dry_run=args.dry_run)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -54,6 +54,7 @@ class OmniGenerationScheduler(VLLMScheduler):
|
||||
scheduled_spec_decode_tokens: dict[str, list[int]] = {}
|
||||
scheduled_encoder_inputs: dict[str, list[int]] = {}
|
||||
cached_prompt_token_ids: dict[str, list[int]] = {}
|
||||
cached_additional_information: dict[str, dict | None] = {}
|
||||
|
||||
# Temporary queue: preserve waiting order, do not disturb non-diffusion requests
|
||||
skipped_waiting_requests = create_request_queue(self.policy)
|
||||
@@ -112,6 +113,7 @@ class OmniGenerationScheduler(VLLMScheduler):
|
||||
cached_prompt_token_ids[request.request_id] = request.prompt_token_ids
|
||||
if request.num_cached_tokens < 0:
|
||||
request.num_cached_tokens = num_computed_tokens
|
||||
cached_additional_information[request.request_id] = getattr(request, "additional_information", None)
|
||||
token_budget -= num_new_tokens
|
||||
scheduled_running_reqs.append(request)
|
||||
req_index += 1
|
||||
@@ -239,6 +241,7 @@ class OmniGenerationScheduler(VLLMScheduler):
|
||||
num_computed_tokens=cached_reqs_data.num_computed_tokens,
|
||||
num_output_tokens=cached_reqs_data.num_output_tokens,
|
||||
prompt_token_ids=cached_prompt_token_ids,
|
||||
additional_information=cached_additional_information,
|
||||
)
|
||||
|
||||
total_num_scheduled_tokens = sum(num_scheduled_tokens.values())
|
||||
|
||||
@@ -67,6 +67,7 @@ class OmniCachedRequestData(CachedRequestData):
|
||||
"""
|
||||
|
||||
prompt_token_ids: dict[str, list[int]]
|
||||
additional_information: dict[str, dict | None]
|
||||
|
||||
|
||||
@dataclass
|
||||
|
||||
Whitespace-only changes.
@@ -0,0 +1,210 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import math
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from diffusers.models.autoencoders import AutoencoderKL as Diffusers_AutoencoderKL
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
|
||||
DistributedOperator,
|
||||
DistributedVaeMixin,
|
||||
GridSpec,
|
||||
TileTask,
|
||||
)
|
||||
|
||||
# from vllm_omni.diffusion.models.nextstep_1_1.modeling_flux_vae import AutoencoderKL as Next_Step_AutoencoderKL
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
# We use base class because some model re-implement AutoencoderKL, but split is share.
|
||||
class DistributedAutoencoderKL_base(DistributedVaeMixin):
|
||||
@classmethod
|
||||
def from_pretrained(cls, *args: Any, **kwargs: Any):
|
||||
model = super().from_pretrained(*args, **kwargs)
|
||||
model.init_distributed()
|
||||
return model
|
||||
|
||||
def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
|
||||
# mostly copy from AutoencoderKL
|
||||
overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
|
||||
blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
|
||||
row_limit = self.tile_sample_min_size - blend_extent
|
||||
|
||||
# Split z into overlapping 64x64 tiles and decode them separately.
|
||||
# The tiles have an overlap to avoid seams between tiles.
|
||||
tiletask_list = []
|
||||
for i in range(0, z.shape[2], overlap_size):
|
||||
for j in range(0, z.shape[3], overlap_size):
|
||||
tile = z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
|
||||
tiletask_list.append(
|
||||
TileTask(
|
||||
len(tiletask_list),
|
||||
(i // overlap_size, j // overlap_size),
|
||||
tile,
|
||||
workload=tile.shape[2] * tile.shape[3],
|
||||
)
|
||||
)
|
||||
|
||||
tile_spec = {
|
||||
"blend_extent": blend_extent,
|
||||
"row_limit": row_limit,
|
||||
}
|
||||
grid_spec = GridSpec(
|
||||
split_dims=(2, 3),
|
||||
grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
|
||||
tile_spec=tile_spec,
|
||||
)
|
||||
return tiletask_list, grid_spec
|
||||
|
||||
def tile_exec(self, task: TileTask) -> torch.Tensor:
|
||||
"""Decode a single latent tile into RGB space."""
|
||||
tile = task.tensor
|
||||
if self.config.use_post_quant_conv:
|
||||
tile = self.post_quant_conv(tile)
|
||||
decoded = self.decoder(tile)
|
||||
return decoded
|
||||
|
||||
def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
|
||||
"""Merge decoded tiles into a full image."""
|
||||
|
||||
grid_h, grid_w = grid_spec.grid_shape
|
||||
result_rows = []
|
||||
for i in range(grid_h):
|
||||
result_row = []
|
||||
for j in range(grid_w):
|
||||
tile = coord_tensor_map[(i, j)]
|
||||
if i > 0:
|
||||
tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_extent"])
|
||||
if j > 0:
|
||||
tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_extent"])
|
||||
result_row.append(tile[:, :, : grid_spec.tile_spec["row_limit"], : grid_spec.tile_spec["row_limit"]])
|
||||
result_rows.append(torch.cat(result_row, dim=3))
|
||||
|
||||
dec = torch.cat(result_rows, dim=2)
|
||||
return dec
|
||||
|
||||
def patch_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
|
||||
overlap_latent = int(self.tile_latent_min_size * float(self.tile_overlap_factor))
|
||||
halo_base = max(0, overlap_latent // 2)
|
||||
|
||||
_, _, latent_h, latent_w = z.shape
|
||||
scale = int(2 ** (len(self.config.block_out_channels) - 1))
|
||||
max_parallel_size = self.distributed_decoder.parallel_size
|
||||
|
||||
root = int(math.sqrt(max_parallel_size))
|
||||
for rows in range(root, 0, -1):
|
||||
if max_parallel_size % rows == 0:
|
||||
grid_rows, grid_cols = rows, max_parallel_size // rows
|
||||
break
|
||||
tiletask_list = []
|
||||
halo_size = dict()
|
||||
for i in range(grid_rows):
|
||||
for j in range(grid_cols):
|
||||
h0 = (i * latent_h) // grid_rows
|
||||
h1 = ((i + 1) * latent_h) // grid_rows
|
||||
w0 = (j * latent_w) // grid_cols
|
||||
w1 = ((j + 1) * latent_w) // grid_cols
|
||||
|
||||
core_h = max(0, h1 - h0)
|
||||
core_w = max(0, w1 - w0)
|
||||
halo = max(halo_base, min(core_h, core_w) // 2)
|
||||
ph0 = max(0, h0 - halo)
|
||||
ph1 = min(latent_h, h1 + halo)
|
||||
pw0 = max(0, w0 - halo)
|
||||
pw1 = min(latent_w, w1 + halo)
|
||||
tile = z[:, :, ph0:ph1, pw0:pw1]
|
||||
tiletask_list.append(TileTask(len(tiletask_list), (i, j), tile, tile.shape[2] * tile.shape[3]))
|
||||
halo_size[(i, j)] = {
|
||||
"up": h0 - ph0,
|
||||
"down": ph1 - h1,
|
||||
"left": w0 - pw0,
|
||||
"right": pw1 - w1,
|
||||
}
|
||||
|
||||
tile_spec = {
|
||||
"halo_size": halo_size,
|
||||
"scale": scale,
|
||||
}
|
||||
grid_spec = GridSpec(
|
||||
split_dims=(2, 3),
|
||||
grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
|
||||
tile_spec=tile_spec,
|
||||
output_dtype=self.dtype,
|
||||
)
|
||||
return tiletask_list, grid_spec
|
||||
|
||||
def patch_exec(self, task: TileTask) -> torch.Tensor:
|
||||
return self.tile_exec(task)
|
||||
|
||||
def patch_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
|
||||
grid_h, grid_w = grid_spec.grid_shape
|
||||
result_rows = []
|
||||
for i in range(grid_h):
|
||||
result_row = []
|
||||
for j in range(grid_w):
|
||||
halo = grid_spec.tile_spec["halo_size"][(i, j)]
|
||||
scale = grid_spec.tile_spec["scale"]
|
||||
|
||||
tile = coord_tensor_map[(i, j)]
|
||||
halo_up = halo["up"] * scale
|
||||
halo_down = halo["down"] * scale
|
||||
halo_left = halo["left"] * scale
|
||||
halo_right = halo["right"] * scale
|
||||
|
||||
core_tile = tile[
|
||||
:,
|
||||
:,
|
||||
halo_up : (None if halo_down == 0 else -halo_down),
|
||||
halo_left : (None if halo_right == 0 else -halo_right),
|
||||
]
|
||||
result_row.append(core_tile)
|
||||
result_rows.append(torch.cat(result_row, dim=3))
|
||||
dec = torch.cat(result_rows, dim=2)
|
||||
return dec
|
||||
|
||||
def _strategy_select(self, z: torch.Tensor):
|
||||
tile_latent_min_size = getattr(self, "tile_latent_min_size", None)
|
||||
tile_overlap_factor = getattr(self, "tile_overlap_factor", None)
|
||||
if tile_latent_min_size is None or tile_overlap_factor is None:
|
||||
return None, None, None
|
||||
if z.shape[-1] > tile_latent_min_size or z.shape[-2] > tile_latent_min_size:
|
||||
return self.tile_split, self.tile_exec, self.tile_merge
|
||||
|
||||
return self.patch_split, self.patch_exec, self.patch_merge
|
||||
|
||||
# Normally, we should override tiled_decode. However, since we also need to
|
||||
# support the patch split strategy, we override decode instead.
|
||||
def decode(self, z: torch.Tensor, return_dict: bool = True, *args: Any, **kwargs: Any):
|
||||
if not self.is_distributed_enabled():
|
||||
return super().decode(z, return_dict=return_dict, *args, **kwargs)
|
||||
|
||||
split, exec, merge = self._strategy_select(z)
|
||||
|
||||
if split is not None:
|
||||
strategy = "tile" if split == self.tile_split else "patch"
|
||||
logger.info(f"Decode run with distributed executor, split strategy is {strategy}")
|
||||
result = self.distributed_decoder.execute(
|
||||
z, DistributedOperator(split=split, exec=exec, merge=merge), broadcast_result=False
|
||||
)
|
||||
if not return_dict:
|
||||
return (result,)
|
||||
|
||||
from diffusers.models.autoencoders.vae import DecoderOutput
|
||||
|
||||
return DecoderOutput(sample=result)
|
||||
else:
|
||||
return super().decode(z, return_dict=return_dict, *args, **kwargs)
|
||||
|
||||
|
||||
class DistributedAutoencoderKL(DistributedAutoencoderKL_base, Diffusers_AutoencoderKL):
|
||||
pass
|
||||
|
||||
|
||||
# Next_Step_AutoencoderKL not support tiling now, so we currently disable it.
|
||||
|
||||
# class DistributedAutoencoderKL_NextStep(DistributedAutoencoderKL_base, Next_Step_AutoencoderKL):
|
||||
# pass
|
||||
@@ -0,0 +1,120 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from diffusers.models.autoencoders import AutoencoderKLQwenImage
|
||||
from diffusers.models.autoencoders.vae import DecoderOutput
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
|
||||
DistributedOperator,
|
||||
DistributedVaeMixin,
|
||||
GridSpec,
|
||||
TileTask,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMixin):
|
||||
@classmethod
|
||||
def from_pretrained(cls, *args: Any, **kwargs: Any):
|
||||
model = super().from_pretrained(*args, **kwargs)
|
||||
model.init_distributed()
|
||||
return model
|
||||
|
||||
def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
|
||||
# mostly copy from AutoencoderKL
|
||||
_, _, num_frames, height, width = z.shape
|
||||
sample_height = height * self.spatial_compression_ratio
|
||||
sample_width = width * self.spatial_compression_ratio
|
||||
|
||||
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
|
||||
tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
|
||||
tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
|
||||
tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
|
||||
|
||||
blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
|
||||
blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
|
||||
|
||||
# Split z into overlapping tiles and decode them separately.
|
||||
# The tiles have an overlap to avoid seams between tiles.
|
||||
tiletask_list = []
|
||||
for i in range(0, height, tile_latent_stride_height):
|
||||
for j in range(0, width, tile_latent_stride_width):
|
||||
time_list = []
|
||||
for k in range(num_frames):
|
||||
self._conv_idx = [0]
|
||||
tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
|
||||
time_list.append(tile)
|
||||
tiletask_list.append(
|
||||
TileTask(
|
||||
len(tiletask_list),
|
||||
(i // tile_latent_stride_height, j // tile_latent_stride_width),
|
||||
time_list,
|
||||
workload=time_list[0].shape[3] * time_list[0].shape[4],
|
||||
)
|
||||
)
|
||||
tile_spec = {
|
||||
"sample_height": sample_height,
|
||||
"sample_width": sample_width,
|
||||
"blend_height": blend_height,
|
||||
"blend_width": blend_width,
|
||||
}
|
||||
grid_spec = GridSpec(
|
||||
split_dims=(3, 4),
|
||||
grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
|
||||
tile_spec=tile_spec,
|
||||
output_dtype=self.dtype,
|
||||
)
|
||||
return tiletask_list, grid_spec
|
||||
|
||||
def tile_exec(self, task: TileTask) -> torch.Tensor:
|
||||
"""Decode a single latent tile into RGB space."""
|
||||
self.clear_cache()
|
||||
time = []
|
||||
for k in range(len(task.tensor)):
|
||||
self._conv_idx = [0]
|
||||
tile = self.post_quant_conv(task.tensor[k])
|
||||
decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx)
|
||||
time.append(decoded)
|
||||
result = torch.cat(time, dim=2)
|
||||
return result
|
||||
|
||||
def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
|
||||
"""Merge decoded tiles into a full image."""
|
||||
grid_h, grid_w = grid_spec.grid_shape
|
||||
self.clear_cache()
|
||||
|
||||
result_rows = []
|
||||
for i in range(grid_h):
|
||||
result_row = []
|
||||
for j in range(grid_w):
|
||||
tile = coord_tensor_map[(i, j)]
|
||||
if i > 0:
|
||||
tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
|
||||
if j > 0:
|
||||
tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
|
||||
result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width])
|
||||
result_rows.append(torch.cat(result_row, dim=-1))
|
||||
dec = torch.cat(result_rows, dim=3)[
|
||||
:, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"]
|
||||
]
|
||||
return dec
|
||||
|
||||
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)
|
||||
|
||||
logger.info("Decode run with distributed executor")
|
||||
result = self.distributed_decoder.execute(
|
||||
z,
|
||||
DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
|
||||
broadcast_result=True,
|
||||
)
|
||||
if not return_dict:
|
||||
return (result,)
|
||||
|
||||
return DecoderOutput(sample=result)
|
||||
@@ -0,0 +1,142 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from diffusers.models.autoencoders import AutoencoderKLWan
|
||||
from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify
|
||||
from diffusers.models.autoencoders.vae import DecoderOutput
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
|
||||
DistributedOperator,
|
||||
DistributedVaeMixin,
|
||||
GridSpec,
|
||||
TileTask,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin):
|
||||
@classmethod
|
||||
def from_pretrained(cls, *args: Any, **kwargs: Any):
|
||||
model = super().from_pretrained(*args, **kwargs)
|
||||
model.init_distributed()
|
||||
return model
|
||||
|
||||
def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
|
||||
_, _, num_frames, height, width = z.shape
|
||||
sample_height = height * self.spatial_compression_ratio
|
||||
sample_width = width * self.spatial_compression_ratio
|
||||
|
||||
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
|
||||
tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
|
||||
tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
|
||||
tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
|
||||
tile_sample_stride_height = self.tile_sample_stride_height
|
||||
tile_sample_stride_width = self.tile_sample_stride_width
|
||||
if self.config.patch_size is not None:
|
||||
sample_height = sample_height // self.config.patch_size
|
||||
sample_width = sample_width // self.config.patch_size
|
||||
tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size
|
||||
tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size
|
||||
blend_height = self.tile_sample_min_height // self.config.patch_size - tile_sample_stride_height
|
||||
blend_width = self.tile_sample_min_width // self.config.patch_size - tile_sample_stride_width
|
||||
else:
|
||||
blend_height = self.tile_sample_min_height - tile_sample_stride_height
|
||||
blend_width = self.tile_sample_min_width - tile_sample_stride_width
|
||||
|
||||
tiletask_list = []
|
||||
for i in range(0, height, tile_latent_stride_height):
|
||||
for j in range(0, width, tile_latent_stride_width):
|
||||
time_list = []
|
||||
for k in range(num_frames):
|
||||
self._conv_idx = [0]
|
||||
tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
|
||||
time_list.append(tile)
|
||||
tiletask_list.append(
|
||||
TileTask(
|
||||
len(tiletask_list),
|
||||
(i // tile_latent_stride_height, j // tile_latent_stride_width),
|
||||
time_list,
|
||||
workload=time_list[0].shape[3] * time_list[0].shape[4],
|
||||
)
|
||||
)
|
||||
tile_spec = {
|
||||
"sample_height": sample_height,
|
||||
"sample_width": sample_width,
|
||||
"blend_height": blend_height,
|
||||
"blend_width": blend_width,
|
||||
"tile_sample_stride_height": tile_sample_stride_height,
|
||||
"tile_sample_stride_width": tile_sample_stride_width,
|
||||
}
|
||||
grid_spec = GridSpec(
|
||||
split_dims=(3, 4),
|
||||
grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
|
||||
tile_spec=tile_spec,
|
||||
output_dtype=self.dtype,
|
||||
)
|
||||
return tiletask_list, grid_spec
|
||||
|
||||
def tile_exec(self, task: TileTask) -> torch.Tensor:
|
||||
"""Decode a single latent tile into RGB space."""
|
||||
self.clear_cache()
|
||||
time = []
|
||||
for k in range(len(task.tensor)):
|
||||
self._conv_idx = [0]
|
||||
tile = self.post_quant_conv(task.tensor[k])
|
||||
decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0))
|
||||
time.append(decoded)
|
||||
result = torch.cat(time, dim=2)
|
||||
return result
|
||||
|
||||
def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
|
||||
"""Merge decoded tiles into a full image."""
|
||||
grid_h, grid_w = grid_spec.grid_shape
|
||||
result_rows = []
|
||||
self.clear_cache()
|
||||
for i in range(grid_h):
|
||||
result_row = []
|
||||
for j in range(grid_w):
|
||||
tile = coord_tensor_map[(i, j)]
|
||||
if i > 0:
|
||||
tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
|
||||
if j > 0:
|
||||
tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
|
||||
result_row.append(
|
||||
tile[
|
||||
:,
|
||||
:,
|
||||
:,
|
||||
: grid_spec.tile_spec["tile_sample_stride_height"],
|
||||
: grid_spec.tile_spec["tile_sample_stride_width"],
|
||||
]
|
||||
)
|
||||
result_rows.append(torch.cat(result_row, dim=-1))
|
||||
|
||||
dec = torch.cat(result_rows, dim=3)[
|
||||
:, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"]
|
||||
]
|
||||
|
||||
if self.config.patch_size is not None:
|
||||
dec = unpatchify(dec, patch_size=self.config.patch_size)
|
||||
|
||||
dec = torch.clamp(dec, min=-1.0, max=1.0)
|
||||
return dec
|
||||
|
||||
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)
|
||||
|
||||
logger.info("Decode run with distributed executor")
|
||||
result = self.distributed_decoder.execute(
|
||||
z,
|
||||
DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
|
||||
broadcast_result=False,
|
||||
)
|
||||
if not return_dict:
|
||||
return (result,)
|
||||
|
||||
return DecoderOutput(sample=result)
|
||||
@@ -0,0 +1,193 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.diffusion.distributed.parallel_state import get_dit_group
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GridSpec:
|
||||
"""The Grid shape split"""
|
||||
|
||||
split_dims: tuple[int, ...] # For example: (3,4) for (B, C, T, H, W), (2,3) for (B, C, H, W)
|
||||
grid_shape: tuple[int, ...] # For example: (nh, nw) for (B, C, T, H, W), (nh, nw) for (B, C, H, W)
|
||||
tile_spec: dict = field(default_factory=dict)
|
||||
output_dtype: torch.dtype | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class TileTask:
|
||||
tile_id: int
|
||||
grid_coord: tuple[int, ...] # The coordinate of the tile in GridSpec.grid_shape
|
||||
tensor: torch.Tensor | list[torch.Tensor] # The tile tensor
|
||||
workload: int | float = 1
|
||||
|
||||
|
||||
@dataclass
|
||||
class DistributedOperator:
|
||||
split: callable
|
||||
exec: callable
|
||||
merge: callable
|
||||
|
||||
|
||||
class DistributedVaeExecutor:
|
||||
"""
|
||||
Abstract util class for distributed patch/tile parallel VAE decoding.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.group = get_dit_group()
|
||||
self.world_size = dist.get_world_size(self.group)
|
||||
self.rank = dist.get_rank(self.group)
|
||||
self.parallel_size = 1
|
||||
|
||||
def set_parallel_size(self, parallel_size: int):
|
||||
self.parallel_size = parallel_size
|
||||
|
||||
def gather_tensors(self, tensor: torch.Tensor):
|
||||
gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)] if self.rank == 0 else None
|
||||
dist.gather(tensor, gather_list=gather_list, dst=0, group=self.group)
|
||||
return gather_list
|
||||
|
||||
def broadcast_tensor(self, tensor: torch.Tensor):
|
||||
dist.broadcast(tensor, src=0, group=self.group)
|
||||
return tensor
|
||||
|
||||
def _compute_global_padding_shape(self, local_results, output_ndim: int, device):
|
||||
local_tile_max_dims = [0] * output_ndim
|
||||
for _, tile_tensor in local_results:
|
||||
for dim_idx, dim_size in enumerate(tile_tensor.shape):
|
||||
local_tile_max_dims[dim_idx] = max(local_tile_max_dims[dim_idx], dim_size)
|
||||
local_shape_stat = torch.tensor([len(local_results), *local_tile_max_dims], device=device)
|
||||
dist.all_reduce(
|
||||
local_shape_stat,
|
||||
op=dist.ReduceOp.MAX,
|
||||
group=self.group,
|
||||
)
|
||||
return local_shape_stat.tolist()
|
||||
|
||||
def _pack_local_tiles(self, local_results, global_padding_shape, grid_spec, device, dtype):
|
||||
tile_tensor = torch.zeros(global_padding_shape, device=device, dtype=dtype)
|
||||
meta_tensor = torch.full(
|
||||
(global_padding_shape[0], len(grid_spec.split_dims) + 1), -1, device=device, dtype=torch.int64
|
||||
)
|
||||
for idx, (tid, t_tensor) in enumerate(local_results):
|
||||
meta_tensor[idx, 0] = tid
|
||||
for i, dim in enumerate(grid_spec.split_dims):
|
||||
meta_tensor[idx, i + 1] = t_tensor.shape[dim]
|
||||
slices = tuple(slice(0, s) for s in t_tensor.shape)
|
||||
tile_tensor[idx][slices] = t_tensor
|
||||
return tile_tensor, meta_tensor
|
||||
|
||||
def _unpack_tiles(self, meta_gather, tile_gather, grid_spec, tid_coord_map):
|
||||
coord_tensor_map = {}
|
||||
for r in range(self.world_size):
|
||||
meta_src = meta_gather[r]
|
||||
tiles_src = tile_gather[r]
|
||||
for idx in range(meta_src.shape[0]):
|
||||
tid = int(meta_src[idx, 0])
|
||||
if tid < 0:
|
||||
continue
|
||||
slices = [slice(None)] * tiles_src[idx].ndim
|
||||
for i, dim in enumerate(grid_spec.split_dims):
|
||||
slices[dim] = slice(0, int(meta_src[idx, i + 1]))
|
||||
slices = tuple(slices)
|
||||
coord_tensor_map[tid_coord_map[tid]] = tiles_src[idx][slices]
|
||||
|
||||
return coord_tensor_map
|
||||
|
||||
def _balance_tasks(self, task_list, num_rank):
|
||||
workloads = [0] * num_rank
|
||||
assigned = [[] for _ in range(num_rank)]
|
||||
|
||||
for task in sorted(task_list, key=lambda t: t.workload, reverse=True):
|
||||
r = workloads.index(min(workloads))
|
||||
assigned[r].append(task)
|
||||
workloads[r] += task.workload
|
||||
|
||||
return assigned
|
||||
|
||||
def execute(self, z: torch.Tensor, operator: DistributedOperator, broadcast_result: bool = True):
|
||||
pp_size = min(self.parallel_size, self.world_size)
|
||||
|
||||
# 1. Split into tiles
|
||||
tiletask_list, grid_spec = operator.split(z)
|
||||
tid_coord_map = {task.tile_id: task.grid_coord for task in tiletask_list}
|
||||
|
||||
# 2. local decode
|
||||
assigned = self._balance_tasks(tiletask_list, pp_size)
|
||||
local_tasks = assigned[self.rank] if pp_size <= self.world_size else []
|
||||
local_results = [(t.tile_id, operator.exec(t)) for t in local_tasks]
|
||||
|
||||
# 3. compute shape per rank
|
||||
global_padding_shape = self._compute_global_padding_shape(local_results, z.ndim, z.device)
|
||||
|
||||
# 5. prepare tile tensors
|
||||
output_dtype = grid_spec.output_dtype if grid_spec.output_dtype is not None else z.dtype
|
||||
local_tile_tensor, local_meta_tensor = self._pack_local_tiles(
|
||||
local_results, global_padding_shape, grid_spec, z.device, output_dtype
|
||||
)
|
||||
|
||||
# 6. gather tiles & meta
|
||||
meta_gather = self.gather_tensors(local_meta_tensor)
|
||||
tile_gather = self.gather_tensors(local_tile_tensor)
|
||||
|
||||
if self.rank != 0:
|
||||
result = torch.empty(0, device=z.device) # Dummy return for non-zero ranks
|
||||
else:
|
||||
# 7. reconstruct full tensor (rank 0)
|
||||
coord_tensor_map = self._unpack_tiles(meta_gather, tile_gather, grid_spec, tid_coord_map)
|
||||
result = operator.merge(coord_tensor_map, grid_spec)
|
||||
|
||||
if broadcast_result:
|
||||
result = self._sync_final_result(result, z.ndim, z.device, output_dtype)
|
||||
return result
|
||||
|
||||
def _sync_final_result(self, rank0_result, output_ndim, output_device, output_dtype):
|
||||
shape_tensor = torch.empty((output_ndim,), device=output_device, dtype=torch.int64)
|
||||
if self.rank == 0:
|
||||
shape_tensor.copy_(torch.tensor(tuple(rank0_result.shape), device=output_device, dtype=torch.int64))
|
||||
dist.broadcast(shape_tensor, src=0, group=self.group)
|
||||
|
||||
if self.rank != 0:
|
||||
sync_result = torch.empty(tuple(shape_tensor.tolist()), device=output_device, dtype=output_dtype)
|
||||
else:
|
||||
sync_result = rank0_result
|
||||
dist.broadcast(sync_result, src=0, group=self.group)
|
||||
return sync_result
|
||||
|
||||
|
||||
class DistributedVaeMixin:
|
||||
def init_distributed(self):
|
||||
self.distributed_decoder = DistributedVaeExecutor()
|
||||
|
||||
def set_parallel_size(self, parallel_size: int) -> bool:
|
||||
return self.distributed_decoder.set_parallel_size(parallel_size)
|
||||
|
||||
def is_distributed_enabled(self) -> bool:
|
||||
if (
|
||||
self.distributed_decoder.parallel_size <= 1
|
||||
or not dist.is_initialized()
|
||||
or not getattr(self, "use_tiling", False)
|
||||
):
|
||||
return False
|
||||
world_size = dist.get_world_size(group=self.distributed_decoder.group)
|
||||
pp_size = min(int(self.distributed_decoder.parallel_size), int(world_size))
|
||||
if pp_size <= 1:
|
||||
return False
|
||||
if self.distributed_decoder.parallel_size > pp_size:
|
||||
logger.warning(
|
||||
f"vae_patch_parallel_size={self.distributed_decoder.parallel_size} "
|
||||
f"is greater than dit_group={world_size};"
|
||||
f" using dit_group size={world_size}"
|
||||
)
|
||||
return True
|
||||
@@ -13,9 +13,6 @@ import numpy as np
|
||||
import torch
|
||||
import torch.distributed
|
||||
from diffusers.image_processor import VaeImageProcessor
|
||||
from diffusers.models.autoencoders.autoencoder_kl_qwenimage import (
|
||||
AutoencoderKLQwenImage,
|
||||
)
|
||||
from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
|
||||
FlowMatchEulerDiscreteScheduler,
|
||||
)
|
||||
@@ -25,6 +22,7 @@ from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer
|
||||
from vllm.model_executor.models.utils import AutoWeightsLoader
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
|
||||
@@ -269,9 +267,9 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin):
|
||||
self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
|
||||
model, subfolder="text_encoder", local_files_only=local_files_only
|
||||
)
|
||||
self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
|
||||
self.device
|
||||
)
|
||||
self.vae = DistributedAutoencoderKLQwenImage.from_pretrained(
|
||||
model, subfolder="vae", local_files_only=local_files_only
|
||||
).to(self.device)
|
||||
transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel)
|
||||
quant_config = get_vllm_quant_config_for_layers(od_config.quantization_config)
|
||||
self.transformer = QwenImageTransformer2DModel(
|
||||
|
||||
@@ -6,7 +6,6 @@ from collections.abc import Iterable
|
||||
|
||||
import torch
|
||||
from diffusers.image_processor import VaeImageProcessor
|
||||
from diffusers.models.autoencoders import AutoencoderKL
|
||||
from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
|
||||
FlowMatchEulerDiscreteScheduler,
|
||||
)
|
||||
@@ -16,6 +15,7 @@ from transformers import CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderMo
|
||||
from vllm.model_executor.models.utils import AutoWeightsLoader
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl import DistributedAutoencoderKL
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
@@ -174,9 +174,9 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin):
|
||||
)
|
||||
self.transformer = SD3Transformer2DModel(od_config=od_config)
|
||||
|
||||
self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
|
||||
self.device
|
||||
)
|
||||
self.vae = DistributedAutoencoderKL.from_pretrained(
|
||||
model, subfolder="vae", local_files_only=local_files_only
|
||||
).to(self.device)
|
||||
|
||||
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
|
||||
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
|
||||
|
||||
@@ -11,13 +11,13 @@ from typing import Any, cast
|
||||
|
||||
import PIL.Image
|
||||
import torch
|
||||
from diffusers import AutoencoderKLWan
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from torch import nn
|
||||
from transformers import AutoTokenizer, UMT5EncoderModel
|
||||
from vllm.model_executor.models.utils import AutoWeightsLoader
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
@@ -267,7 +267,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin):
|
||||
self.text_encoder = UMT5EncoderModel.from_pretrained(
|
||||
model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
|
||||
).to(self.device)
|
||||
self.vae = AutoencoderKLWan.from_pretrained(
|
||||
self.vae = DistributedAutoencoderKLWan.from_pretrained(
|
||||
model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
|
||||
).to(self.device)
|
||||
|
||||
|
||||
@@ -11,13 +11,13 @@ from typing import Any, cast
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
import torch
|
||||
from diffusers import AutoencoderKLWan
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from torch import nn
|
||||
from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, UMT5EncoderModel
|
||||
from vllm.model_executor.models.utils import AutoWeightsLoader
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
@@ -210,7 +210,7 @@ class Wan22I2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin):
|
||||
self.image_encoder = None
|
||||
|
||||
# VAE
|
||||
self.vae = AutoencoderKLWan.from_pretrained(
|
||||
self.vae = DistributedAutoencoderKLWan.from_pretrained(
|
||||
model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
|
||||
).to(self.device)
|
||||
|
||||
|
||||
@@ -24,7 +24,6 @@ from typing import Any
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from diffusers.image_processor import VaeImageProcessor
|
||||
from diffusers.models.autoencoders import AutoencoderKL
|
||||
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
|
||||
from diffusers.utils import logging
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
@@ -32,6 +31,7 @@ from transformers import AutoModel, AutoTokenizer
|
||||
from vllm.model_executor.models.utils import AutoWeightsLoader
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl import DistributedAutoencoderKL
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
from vllm_omni.diffusion.models.z_image.z_image_transformer import (
|
||||
@@ -171,9 +171,9 @@ class ZImagePipeline(nn.Module):
|
||||
self.text_encoder = AutoModel.from_pretrained(
|
||||
model, subfolder="text_encoder", local_files_only=local_files_only
|
||||
)
|
||||
self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
|
||||
self._execution_device
|
||||
)
|
||||
self.vae = DistributedAutoencoderKL.from_pretrained(
|
||||
model, subfolder="vae", local_files_only=local_files_only
|
||||
).to(self._execution_device)
|
||||
# Get vLLM quantization config for linear layers
|
||||
quant_config = get_vllm_quant_config_for_layers(od_config.quantization_config)
|
||||
self.transformer = ZImageTransformer2DModel(quant_config=quant_config)
|
||||
|
||||
@@ -8,6 +8,7 @@ from vllm.logger import init_logger
|
||||
from vllm.model_executor.models.registry import _LazyRegisteredModel, _ModelRegistry
|
||||
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin
|
||||
from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig, get_sp_plan_from_model
|
||||
from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel
|
||||
|
||||
@@ -123,13 +124,6 @@ DiffusionModelRegistry = _ModelRegistry(
|
||||
}
|
||||
)
|
||||
|
||||
_VAE_PATCH_PARALLEL_ALLOWLIST = {
|
||||
# Only enable for models we have validated end-to-end.
|
||||
"StableDiffusion3Pipeline",
|
||||
"ZImagePipeline",
|
||||
"NextStep11Pipeline",
|
||||
}
|
||||
|
||||
_NO_CACHE_ACCELERATION = {
|
||||
# Pipelines that do not support cache acceleration (cache_dit / tea_cache).
|
||||
"NextStep11Pipeline",
|
||||
@@ -161,17 +155,14 @@ def initialize_model(
|
||||
model = model_class(od_config=od_config)
|
||||
|
||||
vae_pp_size = od_config.parallel_config.vae_patch_parallel_size
|
||||
if vae_pp_size > 1 and od_config.model_class_name not in _VAE_PATCH_PARALLEL_ALLOWLIST:
|
||||
is_distributed_vae = hasattr(model, "vae") and isinstance(model.vae, DistributedVaeMixin)
|
||||
if vae_pp_size > 1 and not is_distributed_vae:
|
||||
logger.warning(
|
||||
"vae_patch_parallel_size=%d is set but VAE patch parallelism is only enabled for %s; ignoring.",
|
||||
"vae_patch_parallel_size=%d is set but VAE patch parallelism is NOT enabled for %s; ignoring.",
|
||||
vae_pp_size,
|
||||
sorted(_VAE_PATCH_PARALLEL_ALLOWLIST),
|
||||
od_config.model_class_name,
|
||||
)
|
||||
if (
|
||||
vae_pp_size > 1
|
||||
and od_config.model_class_name in _VAE_PATCH_PARALLEL_ALLOWLIST
|
||||
and not od_config.vae_use_tiling
|
||||
):
|
||||
if vae_pp_size > 1 and is_distributed_vae and not od_config.vae_use_tiling:
|
||||
logger.info(
|
||||
"vae_patch_parallel_size=%d requires vae_use_tiling; automatically enabling it.",
|
||||
vae_pp_size,
|
||||
@@ -184,20 +175,8 @@ def initialize_model(
|
||||
if hasattr(model.vae, "use_tiling"):
|
||||
model.vae.use_tiling = od_config.vae_use_tiling
|
||||
|
||||
if (
|
||||
vae_pp_size > 1
|
||||
and hasattr(model, "vae")
|
||||
and od_config.model_class_name in _VAE_PATCH_PARALLEL_ALLOWLIST
|
||||
and od_config.vae_use_tiling
|
||||
):
|
||||
from vllm_omni.diffusion.distributed.parallel_state import get_dit_group
|
||||
from vllm_omni.diffusion.distributed.vae_patch_parallel import maybe_wrap_vae_decode_with_patch_parallelism
|
||||
|
||||
maybe_wrap_vae_decode_with_patch_parallelism(
|
||||
model,
|
||||
vae_patch_parallel_size=vae_pp_size,
|
||||
group_getter=get_dit_group,
|
||||
)
|
||||
if is_distributed_vae:
|
||||
model.vae.set_parallel_size(vae_pp_size)
|
||||
|
||||
# Apply sequence parallelism if enabled
|
||||
# This follows diffusers' pattern where enable_parallelism() is called
|
||||
|
||||
@@ -155,6 +155,11 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
|
||||
|
||||
new_ids = payload_data.get("code_predictor_codes", [])
|
||||
request.prompt_token_ids = new_ids
|
||||
# Pass additional fields (like left_context_size) to the request
|
||||
# Only pass chunk context metadata in additional_information
|
||||
request.additional_information = {}
|
||||
if "left_context_size" in payload_data:
|
||||
request.additional_information["left_context_size"] = payload_data["left_context_size"]
|
||||
request.num_computed_tokens = 0
|
||||
|
||||
# Empty chunk with more data expected: keep polling.
|
||||
|
||||
@@ -322,6 +322,21 @@ def _to_dict(x: Any) -> dict[str, Any]:
|
||||
return {}
|
||||
|
||||
|
||||
def _resolve_model_to_local_path(model: str) -> str:
|
||||
"""Resolve an HF Hub model ID to its local cache snapshot path."""
|
||||
if os.path.isdir(model):
|
||||
return model
|
||||
|
||||
try:
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
# no network access is attempted, check local model path only
|
||||
return snapshot_download(model, local_files_only=True)
|
||||
except Exception:
|
||||
logger.warning(f"Could not resolve {model} to a local snapshot path; using as-is", exc_info=True)
|
||||
return model
|
||||
|
||||
|
||||
def _resolve_model_tokenizer_paths(
|
||||
model: str,
|
||||
engine_args: dict,
|
||||
@@ -331,30 +346,32 @@ def _resolve_model_tokenizer_paths(
|
||||
Some models (e.g., GLM-Image) have tokenizer in root and model in subdirectory.
|
||||
This function handles model_subdir and tokenizer_subdir engine_args.
|
||||
|
||||
When the base model path is an HF Hub ID rather than an absolute local path,
|
||||
the ID is first resolved to the local snapshot directory so that subdirectory
|
||||
joins produce valid filesystem paths.
|
||||
|
||||
Args:
|
||||
model: Base model path
|
||||
model: Base model path or HF Hub model ID
|
||||
engine_args: Engine arguments (modified in-place to remove subdir args
|
||||
and set tokenizer if needed)
|
||||
|
||||
Returns:
|
||||
Resolved model path (may be subdirectory of original)
|
||||
"""
|
||||
import os
|
||||
|
||||
model_subdir = engine_args.pop("model_subdir", None)
|
||||
tokenizer_subdir = engine_args.pop("tokenizer_subdir", None)
|
||||
base_model_path = model
|
||||
resolved_base = _resolve_model_to_local_path(model)
|
||||
|
||||
if model_subdir:
|
||||
model = os.path.join(model, model_subdir)
|
||||
model = os.path.join(resolved_base, model_subdir)
|
||||
logger.info(f"Using model subdirectory: {model}")
|
||||
|
||||
if tokenizer_subdir is not None:
|
||||
tokenizer_path = os.path.join(base_model_path, tokenizer_subdir) if tokenizer_subdir else base_model_path
|
||||
tokenizer_path = os.path.join(resolved_base, tokenizer_subdir) if tokenizer_subdir else resolved_base
|
||||
engine_args["tokenizer"] = tokenizer_path
|
||||
logger.info(f"Using tokenizer from: {tokenizer_path}")
|
||||
elif model_subdir and "tokenizer" not in engine_args:
|
||||
engine_args["tokenizer"] = base_model_path
|
||||
logger.info(f"Using tokenizer from base model path: {base_model_path}")
|
||||
engine_args["tokenizer"] = resolved_base
|
||||
logger.info(f"Using tokenizer from base model path: {resolved_base}")
|
||||
|
||||
return model
|
||||
@@ -5,6 +5,7 @@
|
||||
|
||||
from collections.abc import Iterable
|
||||
from functools import cached_property
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
@@ -160,6 +161,7 @@ class Qwen3OmniMoeForConditionalGeneration(
|
||||
self.requires_raw_input_tokens = True
|
||||
|
||||
elif self.model_stage == "code2wav":
|
||||
self.enable_update_additional_information = True
|
||||
self.thinker = None
|
||||
self.talker = None
|
||||
# Initialize code2wav (codec codes → audio waveform)
|
||||
@@ -254,7 +256,7 @@ class Qwen3OmniMoeForConditionalGeneration(
|
||||
codec: torch.Tensor | None = None,
|
||||
sampling_metadata: SamplingMetadata | None = None,
|
||||
logits_index: int | None = None,
|
||||
additional_information: dict[str, object] | None = None,
|
||||
runtime_additional_information: list[dict[str, Any]] | None = None,
|
||||
**kwargs: object,
|
||||
) -> torch.Tensor | IntermediateTensors | OmniOutput:
|
||||
"""
|
||||
@@ -359,7 +361,15 @@ class Qwen3OmniMoeForConditionalGeneration(
|
||||
codes = input_ids_flatten.reshape(1, 16, -1)
|
||||
|
||||
# Generate audio from codec codes
|
||||
audio_tensors = self.generate_audio(codes, voice_type, seq_token_counts)
|
||||
# Get every request's left_context_size from runtime_additional_information (passed via kwargs)
|
||||
left_context_size = []
|
||||
if runtime_additional_information is not None:
|
||||
for info in runtime_additional_information:
|
||||
if "left_context_size" in info:
|
||||
left_context_size.append(info["left_context_size"])
|
||||
else:
|
||||
logger.debug("No additional_information provided to code2wav stage.")
|
||||
audio_tensors = self.generate_audio(codes, voice_type, left_context_size, seq_token_counts)
|
||||
|
||||
return audio_tensors
|
||||
|
||||
@@ -440,6 +450,7 @@ class Qwen3OmniMoeForConditionalGeneration(
|
||||
self,
|
||||
code: torch.Tensor,
|
||||
voice_type: str,
|
||||
left_context_size: list[int] | None = None,
|
||||
seq_token_counts: list[int] | None = None,
|
||||
) -> list[torch.Tensor]:
|
||||
"""
|
||||
@@ -448,6 +459,7 @@ class Qwen3OmniMoeForConditionalGeneration(
|
||||
Args:
|
||||
code: [batch, num_quantizers, T] - RVQ codec codes
|
||||
voice_type: Voice type (not used in Qwen3, kept for compatibility)
|
||||
left_context_size: Left context size for streaming decode
|
||||
seq_token_counts: Token count for each request in batch
|
||||
|
||||
Returns:
|
||||
@@ -471,10 +483,10 @@ class Qwen3OmniMoeForConditionalGeneration(
|
||||
talker_codes = talker_codes.expand(1, 16, -1)
|
||||
|
||||
if self.vllm_config.model_config.async_chunk:
|
||||
# Only use left_context_size from additional information
|
||||
audio_tensors = self.code2wav.chunked_decode_streaming(
|
||||
talker_codes,
|
||||
chunk_size=25,
|
||||
left_context_size=25,
|
||||
left_context_size=left_context_size,
|
||||
seq_token_counts=seq_token_counts,
|
||||
)
|
||||
else:
|
||||
|
||||
@@ -213,8 +213,7 @@ class Qwen3OmniMoeCode2Wav(nn.Module):
|
||||
def chunked_decode_streaming(
|
||||
self,
|
||||
codes: torch.Tensor,
|
||||
chunk_size: int = 25,
|
||||
left_context_size: int = 25,
|
||||
left_context_size: list[int] | None = None,
|
||||
seq_token_counts: list[int] | None = None,
|
||||
) -> list[torch.Tensor]:
|
||||
"""
|
||||
@@ -222,9 +221,10 @@ class Qwen3OmniMoeCode2Wav(nn.Module):
|
||||
|
||||
Uses overlapping chunks with left context to avoid boundary artifacts.
|
||||
|
||||
No longer need chunk size here, which is different from chunked_decode
|
||||
|
||||
Args:
|
||||
codes: [batch, num_quantizers, seq_len] - num_quantizers-layer RVQ codes
|
||||
chunk_size: Number of codec frames per chunk
|
||||
left_context_size: Number of overlapping frames for context
|
||||
seq_token_counts: Token count for each request in batch
|
||||
|
||||
@@ -233,6 +233,13 @@ class Qwen3OmniMoeCode2Wav(nn.Module):
|
||||
codes. For ``batch_size == 1``, this is a list containing a
|
||||
single tensor with shape ``[1, waveform_len]``.
|
||||
"""
|
||||
if not (left_context_size and seq_token_counts and len(left_context_size) == len(seq_token_counts)):
|
||||
logger.warning(
|
||||
"chunked_decode_streaming: missing/invalid left_context_size or seq_token_counts; "
|
||||
"defaulting to left_context_size=zeros(len=codes.shape[0])."
|
||||
)
|
||||
left_context_size = [0] * codes.shape[0]
|
||||
# Decode chunk
|
||||
wavs = []
|
||||
batch_wav = self(codes)
|
||||
if seq_token_counts is not None:
|
||||
@@ -241,14 +248,10 @@ class Qwen3OmniMoeCode2Wav(nn.Module):
|
||||
# Fallback: assume all batch elements share the same sequence length.
|
||||
code_seq_lens = [codes.shape[-1]] * codes.shape[0]
|
||||
for idx, code_seq_len in enumerate(code_seq_lens):
|
||||
# TODO: need to optimize algorithms, current only support
|
||||
# chunk_size = left_context_size = 25
|
||||
if code_seq_len <= chunk_size:
|
||||
context_size = 0
|
||||
else:
|
||||
context_size = left_context_size
|
||||
# Remove context from output (context_size * total_upsample samples)
|
||||
wav_chunk = batch_wav[idx, :, context_size * self.total_upsample : code_seq_len * self.total_upsample]
|
||||
# Remove context from output (left_context_size * total_upsample samples)
|
||||
wav_chunk = batch_wav[
|
||||
idx, :, left_context_size[idx] * self.total_upsample : code_seq_len * self.total_upsample
|
||||
]
|
||||
wavs.append(wav_chunk)
|
||||
return wavs
|
||||
|
||||
|
||||
@@ -33,6 +33,7 @@ class Qwen3TTSCode2Wav(nn.Module):
|
||||
self.have_multimodal_outputs = True
|
||||
self.has_preprocess = False
|
||||
self.has_postprocess = False
|
||||
self.enable_update_additional_information = True
|
||||
# Generation-only stage (no logits / sampling).
|
||||
self.requires_raw_input_tokens = True
|
||||
|
||||
@@ -146,6 +147,7 @@ class Qwen3TTSCode2Wav(nn.Module):
|
||||
positions: torch.Tensor | None = None,
|
||||
intermediate_tensors: Any = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
runtime_additional_information: list[dict[str, Any]] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> OmniOutput:
|
||||
"""Decode codec codes into audio waveform.
|
||||
@@ -177,18 +179,28 @@ class Qwen3TTSCode2Wav(nn.Module):
|
||||
ids = input_ids.reshape(-1).to(dtype=torch.long)
|
||||
request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts"))
|
||||
|
||||
# Parse each request: extract ctx_frames, validate, reshape codes.
|
||||
# input_ids layout per request: [codec_context_frames, *flat_codes]
|
||||
# Parse each request: extract left_context_size, validate, reshape codes.
|
||||
# input_ids layout per request: [*flat_codes]
|
||||
# where flat_codes is codebook-major [q*F].
|
||||
parsed = [] # (ctx_frames, actual_frames)
|
||||
parsed = [] # (left_context_size, actual_frames)
|
||||
valid_codes = []
|
||||
valid_indices = []
|
||||
# Get left_context_size from runtime_additional_information (passed via kwargs)
|
||||
left_context_size = [0] * len(request_ids_list)
|
||||
if runtime_additional_information is not None:
|
||||
for i, info in enumerate(runtime_additional_information):
|
||||
if i >= len(left_context_size):
|
||||
break
|
||||
if "left_context_size" in info:
|
||||
left_context_size[i] = info["left_context_size"]
|
||||
else:
|
||||
logger.debug("No additional_information provided to code2wav stage.")
|
||||
for i, req_ids in enumerate(request_ids_list):
|
||||
if req_ids.numel() < 2:
|
||||
if req_ids.numel() < 1:
|
||||
parsed.append((0, 0))
|
||||
continue
|
||||
ctx_frames = int(req_ids[0].item())
|
||||
flat = req_ids[1:]
|
||||
ctx_frames = left_context_size[i]
|
||||
flat = req_ids
|
||||
n = flat.numel()
|
||||
# Warmup / dummy_run: not divisible by num_quantizers.
|
||||
if n == 0 or n % q != 0:
|
||||
|
||||
@@ -369,16 +369,17 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
|
||||
def _maybe_init_kv_cache(self, device: torch.device) -> None:
|
||||
if self._kv_cache is not None:
|
||||
return
|
||||
max_seq_len = int(getattr(self.config, "num_code_groups", 16) or 16)
|
||||
# Upper bound on batch size: vLLM scheduler max_num_seqs (fallback 8).
|
||||
max_batch = int(getattr(self._vllm_config.scheduler_config, "max_num_seqs", 8) or 8)
|
||||
max_batch = max(1, max_batch)
|
||||
self._kv_cache = _LocalPredictorKVCache(
|
||||
vllm_config=self._vllm_config,
|
||||
max_seq_len=max_seq_len,
|
||||
max_batch_size=max_batch,
|
||||
device=device,
|
||||
)
|
||||
with set_current_vllm_config(self._vllm_config):
|
||||
max_seq_len = int(getattr(self.config, "num_code_groups", 16) or 16)
|
||||
# Upper bound on batch size: vLLM scheduler max_num_seqs (fallback 8).
|
||||
max_batch = int(getattr(self._vllm_config.scheduler_config, "max_num_seqs", 8) or 8)
|
||||
max_batch = max(1, max_batch)
|
||||
self._kv_cache = _LocalPredictorKVCache(
|
||||
vllm_config=self._vllm_config,
|
||||
max_seq_len=max_seq_len,
|
||||
max_batch_size=max_batch,
|
||||
device=device,
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def reset_cache(self) -> None:
|
||||
|
||||
@@ -29,6 +29,9 @@ stage_args:
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
is_comprehension: true
|
||||
# Use named connector to apply runtime.connectors.extra.
|
||||
output_connectors:
|
||||
to_stage_1: connector_of_shared_memory
|
||||
default_sampling_params:
|
||||
temperature: 0.4
|
||||
top_p: 0.9
|
||||
@@ -60,6 +63,9 @@ stage_args:
|
||||
engine_input_source: [0]
|
||||
# final_output: true
|
||||
# final_output_type: text
|
||||
# Distributed connector configuration
|
||||
input_connectors:
|
||||
from_stage_0: connector_of_shared_memory
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
@@ -99,3 +105,13 @@ stage_args:
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.1
|
||||
|
||||
runtime:
|
||||
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
# Align with Omni: small chunks with sufficient context overlap.
|
||||
codec_chunk_frames: 25 # code2wav decode chunk size
|
||||
codec_left_context_frames: 25 # code2wav left context size
|
||||
@@ -220,6 +220,12 @@ def talker2code2wav_async_chunk(
|
||||
if "code_predictor_codes" not in pooling_output:
|
||||
return None
|
||||
|
||||
connector = getattr(transfer_manager, "connector", None)
|
||||
raw_cfg = getattr(connector, "config", {}) or {}
|
||||
cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
|
||||
chunk_size_config = int(cfg.get("codec_chunk_frames", 25))
|
||||
left_context_size_config = int(cfg.get("codec_left_context_frames", 25))
|
||||
|
||||
code_predictor_codes = pooling_output["code_predictor_codes"]
|
||||
|
||||
if code_predictor_codes is None:
|
||||
@@ -244,23 +250,27 @@ def talker2code2wav_async_chunk(
|
||||
return None
|
||||
|
||||
request_id = request.external_req_id
|
||||
chunk_size = left_context_size = 25
|
||||
transfer_manager.code_prompt_token_ids[request_id].append(codec_codes)
|
||||
length = len(transfer_manager.code_prompt_token_ids[request_id])
|
||||
chunk_length = length % chunk_size
|
||||
chunk_length = length % chunk_size_config
|
||||
if chunk_length != 0 and not is_finished:
|
||||
return None
|
||||
|
||||
context_length = chunk_length if chunk_length != 0 else chunk_size
|
||||
context_length = chunk_length if chunk_length != 0 else chunk_size_config
|
||||
# ensure left context does not exceed available length
|
||||
left_context_size = max(0, min(length - context_length, left_context_size_config))
|
||||
end_index = min(length, left_context_size + context_length)
|
||||
|
||||
codes = (
|
||||
torch.tensor(transfer_manager.code_prompt_token_ids[request_id][-end_index:])
|
||||
.transpose(0, 1)
|
||||
.reshape(-1)
|
||||
.tolist()
|
||||
)
|
||||
|
||||
info = {
|
||||
"code_predictor_codes": (
|
||||
torch.tensor(transfer_manager.code_prompt_token_ids[request_id][-end_index:])
|
||||
.transpose(0, 1)
|
||||
.reshape(-1)
|
||||
.tolist()
|
||||
),
|
||||
"code_predictor_codes": codes,
|
||||
"left_context_size": left_context_size,
|
||||
"finished": torch.tensor(is_finished, dtype=torch.bool),
|
||||
}
|
||||
return info
|
||||
|
||||
@@ -40,12 +40,12 @@ def talker2code2wav_async_chunk(
|
||||
connector = getattr(transfer_manager, "connector", None)
|
||||
raw_cfg = getattr(connector, "config", {}) or {}
|
||||
cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
|
||||
chunk_size = int(cfg.get("codec_chunk_frames", 25))
|
||||
left_context_size = int(cfg.get("codec_left_context_frames", 25))
|
||||
if chunk_size <= 0 or left_context_size < 0:
|
||||
chunk_size_config = int(cfg.get("codec_chunk_frames", 25))
|
||||
left_context_size_config = int(cfg.get("codec_left_context_frames", 25))
|
||||
if chunk_size_config <= 0 or left_context_size_config < 0:
|
||||
raise ValueError(
|
||||
f"Invalid codec chunk config: codec_chunk_frames={chunk_size}, "
|
||||
f"codec_left_context_frames={left_context_size}"
|
||||
f"Invalid codec chunk config: codec_chunk_frames={chunk_size_config}, "
|
||||
f"codec_left_context_frames={left_context_size_config}"
|
||||
)
|
||||
length = len(transfer_manager.code_prompt_token_ids[request_id])
|
||||
|
||||
@@ -59,22 +59,23 @@ def talker2code2wav_async_chunk(
|
||||
}
|
||||
return None
|
||||
|
||||
chunk_length = length % chunk_size
|
||||
chunk_length = length % chunk_size_config
|
||||
|
||||
if chunk_length != 0 and not finished:
|
||||
return None
|
||||
|
||||
context_length = chunk_length if chunk_length != 0 else chunk_size
|
||||
end_index = min(length, left_context_size + context_length)
|
||||
ctx_frames = max(0, int(end_index - context_length))
|
||||
context_length = chunk_length if chunk_length != 0 else chunk_size_config
|
||||
end_index = min(length, left_context_size_config + context_length)
|
||||
left_context_size = max(0, int(end_index - context_length))
|
||||
window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
|
||||
|
||||
# Pack context + chunk into codebook-major flat codes for adapter.
|
||||
code_predictor_codes = torch.tensor(window_frames).transpose(0, 1).reshape(-1).tolist()
|
||||
|
||||
# Build final prompt_token_ids with ctx_frames header for Qwen3-TTS Code2Wav.
|
||||
# The model expects input_ids layout: [ctx_frames, *flat_codes].
|
||||
# Build final prompt_token_ids and left_context_size header for Qwen3-TTS Code2Wav.
|
||||
# The model expects input_ids layout: [*flat_codes].
|
||||
return {
|
||||
"code_predictor_codes": [int(ctx_frames)] + code_predictor_codes,
|
||||
"code_predictor_codes": code_predictor_codes,
|
||||
"left_context_size": left_context_size,
|
||||
"finished": torch.tensor(finished, dtype=torch.bool),
|
||||
}
|
||||
@@ -1198,12 +1198,18 @@ class OmniGPUModelRunner(GPUModelRunner):
|
||||
# Prefill: overlay prompt_embeds and collect additional_information
|
||||
self._collect_additional_information_for_prefill(num_scheduled_tokens_np)
|
||||
|
||||
# Keep per-request additional_information in sync for both new and
|
||||
# cached requests. This is required for stages without preprocess
|
||||
# (e.g., code2wav) so runtime_additional_information can be refreshed
|
||||
# from scheduler cached infos on every step.
|
||||
if hasattr(self.model, "has_preprocess") or hasattr(self.model, "enable_update_additional_information"):
|
||||
if self.vllm_config.model_config.async_chunk:
|
||||
self._update_additional_information(scheduler_output)
|
||||
|
||||
if hasattr(self.model, "has_preprocess") and self.model.has_preprocess:
|
||||
# Overlay custom prompt_embeds per request for the prompt portion;
|
||||
# collect additional_information (tensor/list) for prefill portion only
|
||||
decode_req_ids = []
|
||||
if self.vllm_config.model_config.async_chunk:
|
||||
self._update_additional_information(scheduler_output)
|
||||
for req_index, req_id in enumerate(self.input_batch.req_ids):
|
||||
req_infos = self.model_intermediate_buffer.get(req_id, {})
|
||||
|
||||
|
||||
Reference in new issue
Block a user