From 821286794f1afaac7d44d7a75371e87527b30d22 Mon Sep 17 00:00:00 2001 From: lyj-jjj Date: Tue, 19 May 2026 00:35:30 +0800 Subject: [PATCH 01/38] [HY-Imgae3.0] support hunyuan image3 dit fa-fp8 on npu (#3540) Signed-off-by: lyj-jjj Co-authored-by: Cursor --- .../hunyuan_image3/end2end.py | 24 +++++++++++++++++++ .../hunyuan_image3_transformer.py | 16 ++++++++++++- 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/examples/offline_inference/hunyuan_image3/end2end.py b/examples/offline_inference/hunyuan_image3/end2end.py index 16f7d8f06..aed06bcb9 100644 --- a/examples/offline_inference/hunyuan_image3/end2end.py +++ b/examples/offline_inference/hunyuan_image3/end2end.py @@ -77,6 +77,24 @@ def parse_args(): parser.add_argument("--log-stats", action="store_true", default=False) parser.add_argument("--init-timeout", type=int, default=300, help="Initialization timeout in seconds.") parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.") + parser.add_argument( + "--diffusion-kv-cache-dtype", + type=str, + default=None, + help="Diffusion attention KV cache dtype, for example 'fp8'. Separate from vLLM --kv-cache-dtype.", + ) + parser.add_argument( + "--diffusion-kv-cache-skip-steps", + type=str, + default=None, + help="Denoising step selector to keep diffusion KV cache in native dtype, for example '0,1,4-6'.", + ) + parser.add_argument( + "--diffusion-kv-cache-skip-layers", + type=str, + default=None, + help="Transformer layer selector to keep diffusion KV cache in native dtype, for example '0-2,10'.", + ) parser.add_argument( "--additional-config", type=str, @@ -140,6 +158,9 @@ def main(): "init_timeout": args.init_timeout, "enforce_eager": args.enforce_eager, "mode": _MODALITY_MODE[args.modality], + "diffusion_kv_cache_dtype": args.diffusion_kv_cache_dtype, + "diffusion_kv_cache_skip_steps": args.diffusion_kv_cache_skip_steps, + "diffusion_kv_cache_skip_layers": args.diffusion_kv_cache_skip_layers, } if additional_config is not None: @@ -237,6 +258,9 @@ def main(): print(f" Inference steps: {args.steps}") print(f" Guidance scale: {args.guidance_scale}") print(f" Seed: {args.seed}") + print(f" diffusion_kv_cache_dtype: {args.diffusion_kv_cache_dtype}") + print(f" diffusion_kv_cache_skip_steps: {args.diffusion_kv_cache_skip_steps}") + print(f" diffusion_kv_cache_skip_layers: {args.diffusion_kv_cache_skip_layers}") if args.modality == "text2img": print(f" Output size: {args.width}x{args.height}") if args.image_path: diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py index 4edcfb6ca..c39965f67 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py @@ -72,6 +72,7 @@ from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelOutput, ) from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx from vllm_omni.diffusion.layers.rope import RotaryEmbedding from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import HunyuanFusedMoE from vllm_omni.model_executor.layers.timestep_embedding import timestep_embedding @@ -868,7 +869,15 @@ class ImageKVCacheManager: Manages specialized caching and updating of KV-Cache for image tokens in multimodal models. """ - def __init__(self, num_heads: int, num_kv_heads: int, head_dim: int, scaling: float, image_token_len: int = 4097): + def __init__( + self, + num_heads: int, + num_kv_heads: int, + head_dim: int, + scaling: float, + image_token_len: int = 4097, + prefix: str = "", + ): """ Args: image_token_len: Number of tokens per image (including special placeholders), @@ -892,6 +901,7 @@ class ImageKVCacheManager: causal=False, softmax_scale=self.scaling, num_kv_heads=self.num_kv_heads, + prefix=f"{prefix}.attn" if prefix else "", ) def _cache_prompt_kv( @@ -1650,6 +1660,7 @@ class HunYuanAttention(nn.Module): num_kv_heads=self.num_kv_heads, scaling=self.scaling, image_token_len=4097, + prefix=f"{prefix}.image_attn", ) self.image_rope2d_emb = HunYuanRotary2DEmbedder( num_heads=self.num_heads, @@ -3021,6 +3032,7 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline): with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): + set_forward_context_denoise_step_idx(i) if cfg_parallel_ready: # CFG parallel: each rank forwards its own branch (no batch doubling) latent_model_input = latents @@ -3109,6 +3121,8 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline): if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() + set_forward_context_denoise_step_idx(None) + if hasattr(self.vae.config, "scaling_factor") and self.vae.config.scaling_factor: latents = latents / self.vae.config.scaling_factor if hasattr(self.vae.config, "shift_factor") and self.vae.config.shift_factor: From ab59673f21d804729557ac53d4c839e6d7353afb Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Tue, 19 May 2026 02:59:23 +0800 Subject: [PATCH 02/38] [Bugfix][Qwen3-Omni] Handle short Code2Wav chunk outputs (#3687) Signed-off-by: Sy03 <1370724210@qq.com> Co-authored-by: amy-why-3459 --- .../qwen3_tts/test_cuda_graph_decoder.py | 37 +++++++++++++++++++ .../qwen3_tts/cuda_graph_decoder_wrapper.py | 21 ++++------- 2 files changed, 45 insertions(+), 13 deletions(-) diff --git a/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py b/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py index dbd217246..e1ee68991 100644 --- a/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py +++ b/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py @@ -74,6 +74,14 @@ class SyntheticDecoder(nn.Module): return self.out(x).clamp(min=-1, max=1) +class ShortOutputDecoder(SyntheticDecoder): + """Decoder variant that returns fewer samples than seq_len * total_upsample.""" + + def forward(self, codes): + out = super().forward(codes) + return out[..., :-5] if out.shape[-1] > 5 else out + + @pytest.fixture(scope="module") def decoder(): """Create a synthetic decoder on CUDA with fixed weights.""" @@ -228,6 +236,35 @@ def test_chunked_decode_output_survives_later_replay(wrapper): torch.testing.assert_close(graph_out, expected, atol=0, rtol=0) +def test_chunked_decode_preserves_short_chunk_concat_semantics(): + """Chunked decode must allow shorter-than-nominal chunk outputs. + + Qwen3-Omni non-async startup can hit a decoder path where a chunk returns + fewer waveform samples than ``code_len * total_upsample``. The wrapper must + preserve the original eager concat behavior instead of copying into exact + nominal slices. + """ + torch.manual_seed(123) + short_decoder = ShortOutputDecoder().to(DEVICE).eval() + short_wrapper = CUDAGraphDecoderWrapper( + decoder=short_decoder, + capture_sizes=[50], + capture_batch_sizes=[1], + num_quantizers=NUM_QUANTIZERS, + enabled=True, + ) + short_wrapper.warmup(DEVICE) + codes = _random_codes(100) + + with torch.no_grad(): + eager_out = _eager_chunked(short_decoder, codes, chunk_size=50, left_context_size=0) + graph_out = short_wrapper.chunked_decode_with_cudagraph(codes, chunk_size=50, left_context_size=0) + + assert graph_out.shape == eager_out.shape + assert graph_out.shape[-1] < 100 * TOTAL_UPSAMPLE + torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0) + + def test_batched_chunked_decode_variable_lengths_matches_per_request_eager(decoder, wrapper): """Variable-length chunk batching should match independent chunked decodes.""" long_codes = _random_codes(100) diff --git a/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py b/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py index a9cdc8948..1c25e0842 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py +++ b/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py @@ -554,10 +554,10 @@ class CUDAGraphDecoderWrapper: chunk_size: int = 300, left_context_size: int = 25, ) -> torch.Tensor: + wavs = [] start_index = 0 total_len = codes.shape[-1] total_upsample = self.decoder.total_upsample - wav_out = None while start_index < total_len: end_index = min(start_index + chunk_size, total_len) @@ -566,21 +566,16 @@ class CUDAGraphDecoderWrapper: codes_chunk = codes[..., start_index - context_size : end_index] wav_chunk = self._decode(codes_chunk, clone_graph_output=False) - if wav_out is None: - wav_out = torch.empty( - (*wav_chunk.shape[:-1], total_len * total_upsample), - dtype=wav_chunk.dtype, - device=wav_chunk.device, - ) - src_start = context_size * total_upsample - dst_start = start_index * total_upsample - dst_end = end_index * total_upsample - wav_out[..., dst_start:dst_end].copy_(wav_chunk[..., src_start:]) + # Keep origin/main's concat semantics: Qwen3-Omni can return a chunk + # that is shorter than the nominal code_len * total_upsample length. + # Clone each slice because graph outputs are static buffers that later + # replays may overwrite. + wavs.append(wav_chunk[..., context_size * total_upsample :].clone()) start_index = end_index - if wav_out is None: + if not wavs: return self.decoder(codes) - return wav_out + return torch.cat(wavs, dim=-1) def batched_chunked_decode_with_cudagraph( self, From 475a4002b0136235b4feb22d4a1e4b221ca5e112 Mon Sep 17 00:00:00 2001 From: "Chendi.Xue" Date: Mon, 18 May 2026 18:58:15 -0500 Subject: [PATCH 03/38] [XPU] set flash_attn as default diffusion attn backend and fix k_len for cross_attn (#3525) Signed-off-by: Chendi Xue --- vllm_omni/diffusion/attention/backends/flash_attn.py | 10 ++++++---- vllm_omni/platforms/xpu/platform.py | 8 ++++++++ 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index cf8d4224f..ab010f3ab 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -136,7 +136,9 @@ class FlashAttentionImpl(AttentionImpl): ) batch_size, q_len = query.size()[:2] - cu_seqlens = torch.arange(0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=query.device) + k_len = key.size(1) + cu_seqlens_q = torch.arange(0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=query.device) + cu_seqlens_k = torch.arange(0, (batch_size + 1) * k_len, step=k_len, dtype=torch.int32, device=query.device) # b s ... -> (b s) ... query = query.flatten(0, 1) key = key.flatten(0, 1) @@ -146,10 +148,10 @@ class FlashAttentionImpl(AttentionImpl): q=query, k=key, v=value, - cu_seqlens_q=cu_seqlens, - cu_seqlens_k=cu_seqlens, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, max_seqlen_q=q_len, - max_seqlen_k=q_len, + max_seqlen_k=k_len, causal=self.causal, softmax_scale=self.softmax_scale, ) diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py index 84707df08..5500e013f 100644 --- a/vllm_omni/platforms/xpu/platform.py +++ b/vllm_omni/platforms/xpu/platform.py @@ -36,12 +36,20 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform): selected_backend: str | None, head_size: int, ) -> str: + compute_capability = torch.xpu.get_device_capability() + # Intel Max 1100 and 1550 will not support flash_attn currently + flash_attn_supported = compute_capability["architecture"] not in [13136561920] + if selected_backend is not None: backend_upper = selected_backend.upper() backend = DiffusionAttentionBackendEnum[backend_upper] logger.debug("Using diffusion attention backend '%s'", backend_upper) return backend.get_path() + if flash_attn_supported: + logger.debug("Defaulting to diffusion attention backend FLASH_ATTN") + return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path() + logger.debug("Defaulting to diffusion attention backend SDPA") return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path() From 89f8819525589141fd825ce4f0d1e1be9cf3660b Mon Sep 17 00:00:00 2001 From: Rustam Khadipash <16683750+hadipash@users.noreply.github.com> Date: Tue, 19 May 2026 09:42:17 +0800 Subject: [PATCH 04/38] [Feature] Add support for Pipeline Parallel and integrate it into Wan 2.2 (#2322) Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com> --- docs/design/feature/pipeline_parallel.md | 356 ++++++++ .../diffusion/frame_interpolation.md | 1 - .../diffusion/parallelism/overview.md | 17 +- .../parallelism/pipeline_parallel.md | 205 +++++ docs/user_guide/diffusion_features.md | 95 ++- .../image_to_video/image_to_video.py | 76 +- .../text_to_video/text_to_video.py | 60 +- .../distributed/test_pipeline_parallel.py | 802 ++++++++++++++++++ .../models/wan2_2/test_wan22_ti2v_pipeline.py | 97 --- .../diffusion/cache/cache_dit_backend.py | 1 - .../diffusion/distributed/cfg_parallel.py | 8 +- .../distributed/group_coordinator.py | 72 ++ .../diffusion/distributed/parallel_state.py | 5 - .../distributed/pipeline_parallel.py | 291 +++++++ vllm_omni/diffusion/models/wan2_2/__init__.py | 8 - .../models/wan2_2/pipeline_wan2_2.py | 19 +- .../models/wan2_2/pipeline_wan2_2_i2v.py | 23 +- .../models/wan2_2/pipeline_wan2_2_ti2v.py | 731 ---------------- .../models/wan2_2/wan2_2_transformer.py | 122 ++- vllm_omni/engine/async_omni_engine.py | 1 + 20 files changed, 1981 insertions(+), 1009 deletions(-) create mode 100644 docs/design/feature/pipeline_parallel.md create mode 100644 docs/user_guide/diffusion/parallelism/pipeline_parallel.md create mode 100644 tests/diffusion/distributed/test_pipeline_parallel.py delete mode 100644 tests/diffusion/models/wan2_2/test_wan22_ti2v_pipeline.py create mode 100644 vllm_omni/diffusion/distributed/pipeline_parallel.py delete mode 100644 vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py diff --git a/docs/design/feature/pipeline_parallel.md b/docs/design/feature/pipeline_parallel.md new file mode 100644 index 000000000..3ed030288 --- /dev/null +++ b/docs/design/feature/pipeline_parallel.md @@ -0,0 +1,356 @@ +# Pipeline Parallel + +This section describes how to add Pipeline Parallelism (PP) to a diffusion pipeline. We use the Wan2.2 text-to-video and +image-to-video pipelines as the reference implementations. + +## Implementation Checklist + +Adding Pipeline Parallel support requires: + +1. ✅ **Inherit the mixin** - add `PipelineParallelMixin` before `CFGParallelMixin` in the pipeline class +2. ✅ **Make stage forward resumable** - support `intermediate_tensors` in the transformer forward path +3. ✅ **Return the right object type** - `IntermediateTensors` on non-last PP ranks, final model output on the last rank +4. ✅ **Use the standard denoising loop** - call `predict_noise_maybe_with_cfg()` and `scheduler_step_maybe_with_cfg()` +5. ✅ **Test parity** - compare PP results against the single-GPU baseline + +--- + +## Table of Contents + +- [Overview](#overview) +- [Step-by-Step Implementation](#step-by-step-implementation) +- [Testing](#testing) +- [Troubleshooting](#troubleshooting) +- [Reference Implementations](#reference-implementations) + +--- + +## Overview + +### What is Pipeline Parallelism? + +Pipeline Parallelism splits the denoising transformer into multiple sequential stages and places each stage on a +different rank. Instead of every rank holding the full DiT, each PP rank owns only a slice of the layers. + +For each denoising step: + +1. Rank 0 starts the forward pass using the current latents. +2. Each intermediate rank receives hidden states from the previous rank, runs its local layer slice, and forwards the + intermediate tensors downstream. +3. The last PP rank produces the final noise prediction. +4. The last PP rank applies the scheduler step and sends the updated latents back to rank 0 for the next timestep. + +This reduces per-rank model memory and enables larger diffusion transformers to run across multiple GPUs. It can also be +combined with CFG-Parallel, where each PP pipeline carries one CFG branch. + +### Architecture + +vLLM-Omni provides `PipelineParallelMixin` to encapsulate the PP communication pattern for diffusion pipelines. + +| Method | Purpose | Automatic Behavior | +|-----------------------------------|-------------------------------|---------------------------------------------------------------------------------| +| `diffuse()` | Denoising loop boundary | Wrapped by `PipelineParallelMixin` to flush pending async PP sends on exit | +| `predict_noise_maybe_with_cfg()` | Predict noise with PP support | Runs partial forwards on non-last PP ranks, combines with CFG logic when needed | +| `scheduler_step_maybe_with_cfg()` | Step scheduler with PP sync | Runs scheduler on the last PP rank, returns updated latents to rank 0 | +| `_sync_pp_send()` | Flush pending async sends | Waits for outstanding `isend` handles before later collectives or decode | + +`PipelineParallelMixin` is intentionally a pipeline-level abstraction. Your model-specific `predict_noise()` still +defines how a local stage executes. + +### How It Works + +`predict_noise_maybe_with_cfg()` automatically switches between these modes: + +- **PP disabled** (`pipeline_parallel_size == 1`): + - Falls back to `CFGParallelMixin.predict_noise_maybe_with_cfg()` +- **PP only** (`pipeline_parallel_size > 1`, `cfg_parallel_size == 1`): + - Rank 0 starts with the input latents + - Middle ranks receive `intermediate_tensors`, run their local layer range, and asynchronously send downstream + - The last rank returns the final noise prediction + - When CFG is enabled in this mode (sequential CFG), both the positive and negative branches run through the same + PP pipeline, doubling the communication volume per denoising step. Prefer PP + CFG-Parallel when + `cfg_parallel_size > 1` is available to avoid this overhead. +- **PP + CFG-Parallel** (`pipeline_parallel_size > 1`, `cfg_parallel_size > 1`): + - Each PP pipeline carries one CFG branch + - The last PP rank all-gathers across the CFG group + - CFG combination happens on every last PP rank in the CFG group, matching the non-PP CFG-parallel behavior + +`scheduler_step_maybe_with_cfg()` keeps the denoising loop consistent: + +- **PP disabled**: + - Falls back to `scheduler_step_maybe_with_cfg()` +- **PP enabled**: + - Only the last PP rank has `noise_pred` and runs the scheduler step + - The resulting latents are sent back to rank 0 + - Rank 0 receives an `AsyncLatents` wrapper that resolves only when the tensor is next consumed + +This asynchronous design avoids unnecessary blocking between denoising steps. + +When a pipeline class is defined, `PipelineParallelMixin` wraps its `diffuse()` method and calls `_sync_pp_send()` in a +`finally` block after `diffuse()` returns or raises. This keeps model code free of explicit PP cleanup while still +ensuring that the final non-blocking PP send completes before decode or later collectives. + +--- + +## Step-by-Step Implementation + +### Step 1: Inherit `PipelineParallelMixin` and `CFGParallelMixin` + +`PipelineParallelMixin` **requires** `CFGParallelMixin` and must appear before it in the class MRO. This is enforced at +class definition time via `__init_subclass__`: defining a pipeline that inherits `PipelineParallelMixin` without +`CFGParallelMixin`, or after `CFGParallelMixin`, raises a `TypeError` immediately on import. + +`PipelineParallelMixin` delegates noise prediction, CFG combination, and scheduler stepping to `CFGParallelMixin`, which +supplies `predict_noise()`, `predict_noise_maybe_with_cfg()`, `scheduler_step_maybe_with_cfg()`, and +`combine_cfg_noise()`. + +**Example:** + +```python +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin +from vllm_omni.diffusion.distributed.pipeline_parallel import PipelineParallelMixin +import torch.nn as nn + + +class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin): + ... +``` + +The order matters: `PipelineParallelMixin` must be listed before `CFGParallelMixin` so calls to +`predict_noise_maybe_with_cfg()` and `scheduler_step_maybe_with_cfg()` resolve to the PP-aware wrappers, while their +`super()` calls delegate to the CFG implementation when PP is disabled or after the last PP stage. + +### Step 2: Make model forward and `predict_noise()` PP-aware + +The PP mixin injects `intermediate_tensors` into the normal `predict_noise()` call. Your model forward path must support +two inputs: + +- Normal input from rank 0, typically passed as `hidden_states` or `x` +- `intermediate_tensors` from upstream PP ranks + +The standard model forward pattern is: + +1. If `intermediate_tensors` is present, read the local hidden state from it. +2. Run only this rank's layer slice. +3. Return `IntermediateTensors(...)` on non-last PP ranks. +4. Return the final model output on the last PP rank. `CFGParallelMixin.predict_noise()` already follows this contract + for common pipelines. + +**Minimal example:** + +```python +from vllm.sequence import IntermediateTensors +from vllm_omni.diffusion.distributed.parallel_state import get_pp_group + + +def forward(self, hidden_states=None, intermediate_tensors=None, **kwargs): + if intermediate_tensors is not None: + hidden_states = intermediate_tensors["hidden_states"] + + for i in range(self.start_layer, self.end_layer): + hidden_states = self.layers[i](hidden_states) + + pp_group = get_pp_group() + if not pp_group.is_last_rank: + return IntermediateTensors({"hidden_states": hidden_states}) + return (hidden_states,) +``` + +### Step 3: Partition the transformer layers + +The local module on each PP rank must expose only that rank's layer slice. In the reference tests and vLLM model +implementations, this is typically done with vLLM utilities such as `make_layers(...)`, which fill missing ranges with +`PPMissingLayer`. Besides partitioning the layers themselves, model authors should also wire +`make_empty_intermediate_tensors_factory(...)` for intermediate tensor allocation and `is_pp_missing_parameter(...)` for +PP-aware weight loading. + +To prepare a transformer for PP, implement the following pieces in order. + +#### 3.1 Split the transformer layers across PP ranks + +Each PP rank should own only its local layer range, typically exposed as `[start_layer, end_layer)`. In practice this is +usually done with `make_layers(...)`, which constructs local layers and fills missing ranges with `PPMissingLayer`. + +The goal is: + +- every PP rank knows its `[start_layer, end_layer)` range +- non-local layers do not execute on this rank +- the forward can resume from incoming `intermediate_tensors` + +By default, layers are distributed evenly across PP ranks using `get_pp_indices()`. When the number of layers is not +divisible by the PP size, remaining layers are assigned to middle partitions to balance compute and memory. You can +override this with the `VLLM_PP_LAYER_PARTITION` environment variable to specify an exact per-rank layer count: + +```bash +# Example: 40 layers across 4 PP ranks, assigning 8 / 12 / 12 / 8 layers +export VLLM_PP_LAYER_PARTITION=8,12,12,8 +``` + +The value must be a comma-separated list of integers whose length equals `pipeline_parallel_size` and whose sum equals +the total number of transformer layers. + +#### 3.2 Expose `make_empty_intermediate_tensors` + +Your transformer module should expose `self.make_empty_intermediate_tensors`, typically created with +`make_empty_intermediate_tensors_factory(...)`. + +This is important because PP ranks need a consistent way to allocate placeholder `IntermediateTensors` with the expected +keys and hidden dimension. + +**Example:** + +```python +from vllm.model_executor.models.utils import make_empty_intermediate_tensors_factory + +self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states"], + inner_dim, +) +``` + +For Wan2.2, the intermediate payload between PP stages is the token sequence stored under `"hidden_states"`, so the +factory is created with that key and the transformer hidden size. + +#### 3.3 Return `IntermediateTensors` on non-last PP ranks + +Your model `forward()` or custom `predict_noise()` implementation should consume `intermediate_tensors` on non-first +ranks and return `IntermediateTensors(...)` on non-last ranks. + +This lets each PP stage resume from the upstream hidden states and pass the local result to the next stage. + +#### 3.4 Skip non-local weights during `load_weights()` + +When a model is PP-partitioned, many parameters in the checkpoint belong to layers that do not exist on the current +rank. `load_weights()` must therefore skip parameters for missing PP stages using `is_pp_missing_parameter(...)`. + +If this is not done, weight loading will either fail or incorrectly try to load tensors into `PPMissingLayer` +placeholders. + +The Wan2.2 transformer is the best reference here: it uses `is_pp_missing_parameter(...)` before loading both remapped +and fused parameters. + +If your model has multiple transformer variants, PP still works as long as each selected transformer obeys the same +contract. + +### Step 4: Use the standard denoising contract + +vLLM-Omni diffusion pipelines already route denoising through `diffuse()`, `predict_noise_maybe_with_cfg()`, and +`scheduler_step_maybe_with_cfg()`. `PipelineParallelMixin` overrides those standard helpers when +`pipeline_parallel_size > 1`, so model integrations should not add separate PP-specific helper names or manual post-loop +synchronization. + +In PP mode: + +- Non-last PP ranks return `None` from `predict_noise_maybe_with_cfg()` after sending `IntermediateTensors` downstream. +- The last PP rank returns the final prediction, and with CFG-Parallel every last PP rank in the CFG group receives the + combined result. +- Rank 0 receives `AsyncLatents` from `scheduler_step_maybe_with_cfg()`, which resolves only when the latents are + consumed. +- Pending non-blocking sends are flushed automatically when `diffuse()` exits. + +--- + +## Testing + +Use an offline inference script with `pipeline_parallel_size > 1`. + +```bash +python examples/offline_inference/text_to_video/text_to_video.py \ +--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \ +--width=1280 \ +--height=704 \ +--guidance-scale=5.0 \ +--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \ +--output=t2v_5B_pp2.mp4 \ +--pipeline-parallel-size=2 +``` + +For PP + CFG-Parallel together: + +```bash +python examples/offline_inference/text_to_video/text_to_video.py \ +--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \ +--width=1280 \ +--height=704 \ +--guidance-scale=5.0 \ +--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \ +--output=t2v_5B_pp2_cfg2.mp4 \ +--pipeline-parallel-size=2 \ +--cfg-parallel-size=2 +``` + +**Verify:** + +1. The run completes without hangs at the PP stage boundary +2. Output quality matches the non-PP baseline within normal numerical variance +3. Peak memory per GPU drops compared with the single-rank model +4. There are no outstanding communication errors before decode + +--- + +## Troubleshooting + +### Issue: `TypeError` on import — `CFGParallelMixin` missing + +**Symptoms:** Importing a pipeline that inherits `PipelineParallelMixin` raises: + +``` +TypeError: YourPipeline inherits PipelineParallelMixin but not CFGParallelMixin. +``` + +or: + +``` +TypeError: YourPipeline must inherit PipelineParallelMixin before CFGParallelMixin ... +``` + +**Cause:** `PipelineParallelMixin` enforces via `__init_subclass__` that the subclass also inherits `CFGParallelMixin` +and lists `PipelineParallelMixin` first in the MRO. + +**Solution:** Add `CFGParallelMixin` to the pipeline's base classes after `PipelineParallelMixin`: + +```python +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin +from vllm_omni.diffusion.distributed.pipeline_parallel import PipelineParallelMixin + + +class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin): + ... +``` + +### Issue: Non-last PP ranks crash when calling `predict_noise` + +**Symptoms:** Shape or missing-input errors on ranks other than the first or last PP rank. + +**Cause:** The model forward path assumes direct input tensors and ignores `intermediate_tensors`. + +**Solution:** Update the transformer `forward()` or custom `predict_noise()` path to load hidden states from +`intermediate_tensors` when present. + +### Issue: PP outputs differ from single-GPU baseline + +**Symptoms:** The PP run finishes but produces numerically inconsistent outputs. + +**Causes & Solutions:** + +- **Local layer partitioning is wrong** + - Verify each rank runs only its `[start_layer, end_layer)` slice +- **Non-last ranks return a plain tensor instead of `IntermediateTensors`** + - Return `IntermediateTensors({...})` until the last PP stage +- **CFG branch wiring is incorrect** + - When CFG is enabled, confirm positive and negative kwargs are passed exactly as in the non-PP path + +--- + +## Reference Implementations + +Complete examples in the codebase: + +| Component | Path | Notes | +|-------------------------|------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------| +| `PipelineParallelMixin` | `vllm_omni/diffusion/distributed/pipeline_parallel.py` | Core PP communication and scheduler helpers | +| `CFGParallelMixin` | `vllm_omni/diffusion/distributed/cfg_parallel.py` | Default `predict_noise()` tuple normalization and CFG helper fallback | +| Wan2.2 transformer | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` | Reference for layer partitioning, `IntermediateTensors`, `make_empty_intermediate_tensors`, and PP-aware weight loading | +| Wan2.2 T2V pipeline | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py` | Reference PP + CFG integration for text-to-video | +| Wan2.2 I2V pipeline | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py` | Reference PP + CFG integration for image-to-video | +| PP tests | `tests/diffusion/distributed/test_pipeline_parallel.py` | Baseline parity and async communication tests | diff --git a/docs/user_guide/diffusion/frame_interpolation.md b/docs/user_guide/diffusion/frame_interpolation.md index 349af50c5..ab03363cf 100644 --- a/docs/user_guide/diffusion/frame_interpolation.md +++ b/docs/user_guide/diffusion/frame_interpolation.md @@ -28,7 +28,6 @@ Frame interpolation is currently supported for: - `WanPipeline` (Wan2.2 text-to-video) - `WanImageToVideoPipeline` -- `Wan22TI2VPipeline` ## Request Parameters diff --git a/docs/user_guide/diffusion/parallelism/overview.md b/docs/user_guide/diffusion/parallelism/overview.md index 90d0b9660..d482b5ea2 100644 --- a/docs/user_guide/diffusion/parallelism/overview.md +++ b/docs/user_guide/diffusion/parallelism/overview.md @@ -4,13 +4,14 @@ This guide covers the parallelism methods in vLLM-Omni for speeding up diffusion ## Supported Methods -| Method | Description | -|--------|-------------| -| **[Tensor Parallelism](tensor_parallel.md)** | Shards DiT weights across GPUs to reduce per-GPU memory | -| **[Sequence Parallelism](sequence_parallel.md)** | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos | -| **[CFG-Parallel](cfg_parallel.md)** | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation | -| **[VAE Patch Parallelism](vae_patch_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory | -| **[HSDP](hsdp.md)** | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs | -| **[Expert Parallelism](expert_parallel.md)** | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0) | +| Method | Description | +|----------------------------------------------------|---------------------------------------------------------------------------------------------------------------------| +| **[Tensor Parallelism](tensor_parallel.md)** | Shards DiT weights across GPUs to reduce per-GPU memory | +| **[Sequence Parallelism](sequence_parallel.md)** | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos | +| **[CFG-Parallel](cfg_parallel.md)** | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation | +| **[Pipeline Parallelism](pipeline_parallel.md)** | Splits the denoising transformer block-wise across sequential GPU stages to reduce per-GPU model memory | +| **[VAE Patch Parallelism](vae_patch_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory | +| **[HSDP](hsdp.md)** | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs | +| **[Expert Parallelism](expert_parallel.md)** | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0) | See [Supported Models](../../diffusion_features.md#supported-models) for per-model compatibility. diff --git a/docs/user_guide/diffusion/parallelism/pipeline_parallel.md b/docs/user_guide/diffusion/parallelism/pipeline_parallel.md new file mode 100644 index 000000000..0827e5efa --- /dev/null +++ b/docs/user_guide/diffusion/parallelism/pipeline_parallel.md @@ -0,0 +1,205 @@ +# Pipeline Parallelism Guide + +## Table of Content + +- [Overview](#overview) +- [Quick Start](#quick-start) +- [Example Script](#example-script) +- [Configuration Parameters](#configuration-parameters) +- [Best Practices](#best-practices) +- [Troubleshooting](#troubleshooting) +- [Summary](#summary) + +--- + +## Overview + +Pipeline Parallelism splits the denoising transformer block-wise into sequential stages across GPUs. Each rank owns only +part of the transformer, which reduces per-GPU model memory and enables larger diffusion models to run across multiple +devices. + +It can also be combined with other distributed methods such as CFG-Parallel, Tensor Parallelism, and Sequence +Parallelism. + +See supported models list in [Supported Models](../../diffusion_features.md#supported-models). + +--- + +## Quick Start + +### Basic Usage + +Simplest working example: + +```python +from vllm_omni import Omni +from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +omni = Omni( + model="Wan-AI/Wan2.2-TI2V-5B-Diffusers", + parallel_config=DiffusionParallelConfig( + pipeline_parallel_size=2, + ), +) + +outputs = omni.generate( + {"prompt": "A cinematic drone shot over snowy mountains"}, + OmniDiffusionSamplingParams( + num_inference_steps=40, + num_frames=81, + height=704, + width=1280, + ), +) +``` + +--- + +## Example Script + +### Offline Inference + +Use python scripts under: + +- `examples/offline_inference/text_to_video/text_to_video.py` +- `examples/offline_inference/image_to_video/image_to_video.py` + +Text-to-video example: + +```bash +python examples/offline_inference/text_to_video/text_to_video.py \ +--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \ +--width=1280 \ +--height=704 \ +--guidance-scale=5.0 \ +--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \ +--output=t2v_5B_pp2.mp4 \ +--pipeline-parallel-size=2 +``` + +Pipeline Parallelism can also be combined with CFG-Parallel: + +```bash +python examples/offline_inference/text_to_video/text_to_video.py \ +--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \ +--width=1280 \ +--height=704 \ +--guidance-scale=5.0 \ +--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \ +--output=t2v_5B_pp2_cfg2.mp4 \ +--pipeline-parallel-size=2 \ +--cfg-parallel-size=2 +``` + +### Online Serving + +Enable Pipeline Parallelism in online serving: + +```bash +# Default PP configuration +vllm serve Wan-AI/Wan2.2-TI2V-5B-Diffusers --omni --port 8091 --pipeline-parallel-size 2 + +# PP + CFG-Parallel +vllm serve Wan-AI/Wan2.2-TI2V-5B-Diffusers --omni --port 8091 \ + --pipeline-parallel-size 2 \ + --cfg-parallel-size 2 +``` + +--- + +## Configuration Parameters + +In `DiffusionParallelConfig` + +| Parameter | Type | Default | Description | +|--------------------------|------|---------|------------------------------------------------------------------------------------------------------------------| +| `pipeline_parallel_size` | int | 1 | Number of pipeline-parallel stages. Set to a value greater than 1 to split the denoising transformer across GPUs | + +> [!NOTE] +> Total GPU count is the product of all enabled distributed dimensions, for example +`pipeline_parallel_size * cfg_parallel_size * tensor_parallel_size * ulysses_degree * ring_degree`. + +### Manual Layer Partitioning + +By default, transformer layers are distributed evenly across PP ranks. You can override this with the +`VLLM_PP_LAYER_PARTITION` environment variable to assign a specific number of layers to each rank: + +```bash +# Example: 40 layers across 4 PP ranks, assigning 8 / 12 / 12 / 8 layers +export VLLM_PP_LAYER_PARTITION=8,12,12,8 +``` + +The value must be a comma-separated list of integers whose length equals `pipeline_parallel_size` and whose sum equals +the total number of transformer layers. This is useful when you want to balance memory or compute asymmetrically across +ranks. + + +--- + +## Best Practices + +### When to Use + +**Good for:** + +- Large diffusion transformers that do not fit comfortably on one GPU +- Multi-GPU setups where reducing per-GPU model memory is more important than minimizing communication +- Combining with CFG-Parallel or other distributed methods on supported models + +**Not for:** + +- Single GPU setups +- Models that do not support Pipeline Parallelism + (check [supported models](../../diffusion_features.md#supported-models)) +- Very small models where inter-stage communication overhead may outweigh the benefit + +### Expected Behavior + +Pipeline Parallelism primarily reduces per-GPU memory usage by splitting the transformer into stage-local blocks. +Depending on the model, topology, and resolution, it may also help execution fit into available hardware, but it is not +primarily a latency optimization. + +--- + +## Troubleshooting + +### Common Issue 1: No benefit from Pipeline Parallelism + +**Symptoms**: PP is enabled, but latency does not improve or becomes slightly worse. + +**Solutions**: + +1. **Check your goal:** + +```python +# PP is mainly for memory scaling, not guaranteed latency speedup +parallel_config = DiffusionParallelConfig(pipeline_parallel_size=2) +``` + +2. **Check model support:** + - Verify your model in [supported models](../../diffusion_features.md#supported-models) + - PP is currently validated only on selected pipelines + +3. **Combine with other methods when appropriate:** + - PP can be combined with CFG-Parallel, Tensor Parallelism, or Sequence Parallelism on supported models + +### Common Issue 2: PP pipeline fails at import + +**Symptoms**: Importing a custom pipeline raises a `TypeError` about `CFGParallelMixin` or mixin order. + +**Solutions**: + +1. Inherit both mixins. +2. Put `PipelineParallelMixin` before `CFGParallelMixin` in the class base list. +3. Use `class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin): ...` as the reference pattern. + +--- + +## Summary + +1. ✅ **Enable Pipeline Parallelism** - Set `pipeline_parallel_size > 1` in `DiffusionParallelConfig` +2. ✅ **Use Supported Models** - Verify your model supports PP in + [supported models](../../diffusion_features.md#supported-models) +3. ✅ **Combine When Needed** - PP can be combined with CFG-Parallel and other distributed methods on supported pipelines +4. ✅ **Scale for Memory** - Use PP primarily to reduce per-GPU model memory and fit larger transformers diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 94238f26f..26b4d9bf7 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -35,14 +35,15 @@ Cache methods trade minimal quality for significant speedup. Quality loss is typ Parallelism methods distribute computation across GPUs without quality loss (mathematically equivalent to single-GPU). -| Method | Description | Best For | -|--------|-------------|----------| -| **[Ulysses-SP](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via all-to-all communication | High-resolution images (>1536px) or long videos with 2-8 GPUs | -| **[Ring-Attention](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via ring-based communication | Videos, very long sequences, memory-constrained, with 2-8 GPUs | -| **[CFG-Parallel](diffusion/parallelism/cfg_parallel.md)** | Splits CFG positive/negative branches across devices | Image editing with CFG guidance (true_cfg_scale > 1) on 2 GPUs | -| **[Tensor Parallelism](diffusion/parallelism/tensor_parallel.md)** | Shards model weights across devices | Large models that don't fit in single GPU, with 2+ GPUs | -| **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP | -| **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) | +| Method | Description | Best For | +|------------------------------------------------------------------------|--------------------------------------------------------------------------|-------------------------------------------------------------------| +| **[Ulysses-SP](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via all-to-all communication | High-resolution images (>1536px) or long videos with 2-8 GPUs | +| **[Ring-Attention](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via ring-based communication | Videos, very long sequences, memory-constrained, with 2-8 GPUs | +| **[CFG-Parallel](diffusion/parallelism/cfg_parallel.md)** | Splits CFG positive/negative branches across devices | Image editing with CFG guidance (true_cfg_scale > 1) on 2 GPUs | +| **[Tensor Parallelism](diffusion/parallelism/tensor_parallel.md)** | Shards model weights across devices | Large models that don't fit in single GPU, with 2+ GPUs | +| **[Pipeline Parallelism](diffusion/parallelism/pipeline_parallel.md)** | Splits the denoising transformer block-wise across sequential GPU stages | Large diffusion transformers that need lower per-GPU model memory | +| **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP | +| **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) | #### Startup Optimization @@ -105,32 +106,32 @@ The following tables show which models support each feature: ### ImageGen -| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution | -|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:| -| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | -| **FLUX.1-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | -| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | -| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | -| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | -| **FLUX.2-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | -| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | -| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | -| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | -| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | -| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | -| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | -| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | -| **OmniGen2** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | -| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | -| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ | -| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ | -| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | -| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ | -| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | -| **SenseNova-U1** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | -| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ | -| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ (decode) | ✅ | ❌ | -| **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | +| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution | +|--------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:| +| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | +| **FLUX.1-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | +| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | +| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | +| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | +| **FLUX.2-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | +| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | +| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | +| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | +| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | +| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | +| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | +| **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | +| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ | +| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ | +| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | +| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ | +| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | +| **SenseNova-U1** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | +| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ | +| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ✅ | ❌ | ✅ (decode) | ✅ | ❌ | +| **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | > Notes: > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT. @@ -138,15 +139,15 @@ The following tables show which models support each feature: ### VideoGen -| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution | -|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:| -| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ | -| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | -| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | -| **LTX-2.3** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | -| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | -| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | -| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | +| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution | +|------------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-----------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:| +| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ | +| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | +| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | +| **LTX-2.3** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | +| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | +| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | **Frame Interpolation Support** @@ -155,9 +156,9 @@ The following tables show which models support each feature: ### AudioGen -| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution | -|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:| -| **Stable-Audio-Open** | ✅ | ❌ | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | +| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution | +|-----------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:| +| **Stable-Audio-Open** | ✅ | ❌ | ❓ | ❓ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ## Feature Compatibility @@ -257,7 +258,7 @@ Measured on NVIDIA H800: **Parallelism Methods:** -- **[Parallelism Overview](diffusion/parallelism/overview.md)** - Tensor Parallelism, Sequence Parallelism, CFG Parallelism, HSDP, and Expert Parallelism +- **[Parallelism Overview](diffusion/parallelism/overview.md)** - Tensor Parallelism, Sequence Parallelism, CFG Parallelism, Pipeline Parallelism, HSDP, and Expert Parallelism **Memory Optimization:** diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index 4964f10c4..9e2521233 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -142,6 +142,42 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable layerwise (blockwise) offloading on DiT modules.", ) + parser.add_argument( + "--enforce-eager", + action="store_true", + help="Disable torch.compile and force eager execution.", + ) + parser.add_argument( + "--audio-sample-rate", + type=int, + default=24000, + help="Sample rate for audio output when saved (default: 24000).", + ) + parser.add_argument( + "--cache-backend", + type=str, + default=None, + choices=["cache_dit", "tea_cache"], + help=( + "Cache backend to use for acceleration. " + "Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). " + "Default: None (no cache acceleration)." + ), + ) + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) + parser.add_argument( + "--quantization", + type=str, + default=None, + choices=["fp8", "mxfp8", "int8", "gguf"], + help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).", + ) + + # Distributed and parallel execution parser.add_argument( "--ulysses-degree", type=int, @@ -173,28 +209,6 @@ def parse_args() -> argparse.Namespace: default=1, help="Number of GPUs used for VAE patch/tile parallelism (decode).", ) - parser.add_argument( - "--enforce-eager", - action="store_true", - help="Disable torch.compile and force eager execution.", - ) - parser.add_argument( - "--audio-sample-rate", - type=int, - default=24000, - help="Sample rate for audio output when saved (default: 24000).", - ) - parser.add_argument( - "--cache-backend", - type=str, - default=None, - choices=["cache_dit", "tea_cache"], - help=( - "Cache backend to use for acceleration. " - "Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). " - "Default: None (no cache acceleration)." - ), - ) parser.add_argument( "--use-hsdp", action="store_true", @@ -219,16 +233,10 @@ def parse_args() -> argparse.Namespace: ), ) parser.add_argument( - "--quantization", - type=str, - default=None, - choices=["fp8", "mxfp8", "int8", "gguf"], - help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).", - ) - parser.add_argument( - "--enable-diffusion-pipeline-profiler", - action="store_true", - help="Enable diffusion pipeline profiler to display stage durations.", + "--pipeline-parallel-size", + type=int, + default=1, + help="Number of pipeline parallel stages.", ) parser.add_argument( "--profiler-config", @@ -326,6 +334,7 @@ def main(): use_hsdp=args.use_hsdp, hsdp_shard_size=args.hsdp_shard_size, hsdp_replicate_size=args.hsdp_replicate_size, + pipeline_parallel_size=args.pipeline_parallel_size, ) omni_kwargs = dict( model=args.model, @@ -366,7 +375,8 @@ def main(): print(f" diffusion_kv_cache_skip_layers(config): {args.diffusion_kv_cache_skip_layers}") 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}" + f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size}," + f" pipeline_parallel_size={args.pipeline_parallel_size}" ) print(f" Video size: {args.width}x{args.height}") print(f"{'=' * 60}\n") diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index b19f0095e..9e1775dd1 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -142,6 +142,32 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable layerwise (blockwise) offloading on DiT modules.", ) + parser.add_argument( + "--audio-sample-rate", + type=int, + default=24000, + help="Sample rate for audio output when saved (default: 24000).", + ) + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) + parser.add_argument( + "--profiler-config", + type=parse_profiler_config, + default=None, + help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', + ) + parser.add_argument( + "--quantization", + type=str, + default=None, + choices=["fp8", "mxfp8", "int8", "gguf"], + help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).", + ) + + # Distributed and parallel execution parser.add_argument( "--ulysses-mode", type=str, @@ -174,41 +200,23 @@ def parse_args() -> argparse.Namespace: default=1, help="Number of GPUs used for tensor parallelism (TP) inside the DiT.", ) - parser.add_argument( - "--audio-sample-rate", - type=int, - default=24000, - help="Sample rate for audio output when saved (default: 24000).", - ) 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( + "--pipeline-parallel-size", + type=int, + default=1, + help="Number of pipeline parallel stages.", + ) parser.add_argument( "--enable-expert-parallel", action="store_true", help="Enable expert parallelism for MoE layers.", ) - parser.add_argument( - "--enable-diffusion-pipeline-profiler", - action="store_true", - help="Enable diffusion pipeline profiler to display stage durations.", - ) - parser.add_argument( - "--profiler-config", - type=parse_profiler_config, - default=None, - help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', - ) - parser.add_argument( - "--quantization", - type=str, - default=None, - choices=["fp8", "mxfp8", "int8", "gguf"], - help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).", - ) parser.add_argument( "--use-hsdp", action="store_true", @@ -262,6 +270,7 @@ def main(): cfg_parallel_size=args.cfg_parallel_size, tensor_parallel_size=args.tensor_parallel_size, vae_patch_parallel_size=args.vae_patch_parallel_size, + pipeline_parallel_size=args.pipeline_parallel_size, enable_expert_parallel=args.enable_expert_parallel, ) @@ -307,7 +316,8 @@ def main(): print( 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}, enable_expert_parallel={args.enable_expert_parallel}" + f" vae_patch_parallel_size={args.vae_patch_parallel_size}, pipeline_parallel_size={args.pipeline_parallel_size}," + f" enable_expert_parallel={args.enable_expert_parallel}" ) print(f" Video size: {args.width}x{args.height}") print(f"{'=' * 60}\n") diff --git a/tests/diffusion/distributed/test_pipeline_parallel.py b/tests/diffusion/distributed/test_pipeline_parallel.py new file mode 100644 index 000000000..b76bf8c89 --- /dev/null +++ b/tests/diffusion/distributed/test_pipeline_parallel.py @@ -0,0 +1,802 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os +import socket + +import pytest +import torch +from vllm.model_executor.models.utils import PPMissingLayer, make_empty_intermediate_tensors_factory, make_layers +from vllm.sequence import IntermediateTensors +from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors + +import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin +from vllm_omni.diffusion.distributed.parallel_state import ( + destroy_distributed_env, + get_classifier_free_guidance_rank, + get_pp_group, + init_distributed_environment, + initialize_model_parallel, +) +from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.parallel] + + +def _find_free_port() -> str: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return str(s.getsockname()[1]) + + +def update_environment_variables(envs_dict: dict[str, str]) -> None: + for k, v in envs_dict.items(): + os.environ[k] = v + + +# --------------------------------------------------------------------------- +# Shared stubs used by both unit and distributed tests +# --------------------------------------------------------------------------- + + +class FakeWork: + """Drop-in for torch.distributed.Work that records whether wait() was called.""" + + def __init__(self): + self.waited = False + + def wait(self): + self.waited = True + + +class SimpleScheduler: + """Minimal diffusion-step scheduler: latents -= 0.1 * noise_pred.""" + + def step(self, noise_pred: torch.Tensor, t, latents: torch.Tensor, return_dict: bool = False): + return (latents - 0.1 * noise_pred,) + + +class FakeVAE: + def __init__(self, distributed_enabled: bool = False): + self.calls = 0 + self.distributed_enabled = distributed_enabled + + def decode(self, z: torch.Tensor): + """Original decode docstring.""" + self.calls += 1 + return (z + 1,) + + def is_distributed_enabled(self) -> bool: + return self.distributed_enabled + + +class MockPipelineParallel(PipelineParallelMixin, CFGParallelMixin): + """Minimal pipeline used to exercise PipelineParallelMixin. + + Uses vLLM's ``make_layers`` for layer partitioning — the same utility used + by real DiT models — so the PP layer-split logic is exercised faithfully. + + Each layer's weights are seeded by ``seed + layer_index`` so that layer ``i`` + is initialized identically on every rank regardless of which ranks are active, + allowing the distributed output to be compared against the single-GPU baseline. + + Args: + num_layers: Total number of Linear layers. + dim: Input / hidden dimension. + seed: Base RNG seed; layer ``i`` uses ``seed + i``. + device: Target device for layer weights (default: CPU). + dtype: Target dtype for layer weights (default: float32). + """ + + def __init__( + self, + num_layers: int = 4, + dim: int = 64, + seed: int = 42, + device: torch.device | None = None, + dtype: torch.dtype = torch.float32, + ): + self.start_layer, self.end_layer, self.layers = make_layers( + num_layers, + lambda prefix: torch.nn.Linear(dim, dim, bias=False), + prefix="layers", + ) + + for i, layer in enumerate(self.layers): + if not isinstance(layer, PPMissingLayer): + torch.manual_seed(seed + i) + torch.nn.init.normal_(layer.weight, mean=0.0, std=0.02) + + self.layers.to(device=device, dtype=dtype) + + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(["hidden_states"], dim) + self.scheduler = SimpleScheduler() + + def predict_noise(self, x=None, intermediate_tensors=None, **_kwargs) -> torch.Tensor | IntermediateTensors: + """Layer-split forward pass. + + * First PP rank: uses ``x`` from caller kwargs. + * Later PP ranks: overrides ``x`` with ``intermediate_tensors["hidden_states"]`` + (which transparently waits for the async receive). + * Non-last PP ranks return ``IntermediateTensors``; the last rank + returns the plain noise-prediction tensor. + """ + if intermediate_tensors is not None: + x = intermediate_tensors["hidden_states"] + + for i in range(self.start_layer, self.end_layer): + x = self.layers[i](x) + + pp_group = get_pp_group() + if not pp_group.is_last_rank: + return IntermediateTensors({"hidden_states": x}) + return x + + +# --------------------------------------------------------------------------- +# 1. AsyncLatents – unit tests (no distributed env required) +# --------------------------------------------------------------------------- + + +class TestAsyncLatents: + """Verifies the lazy-resolution behaviour of AsyncLatents without a real process group.""" + + pytestmark = [pytest.mark.cpu] + + def _make(self, tensor: torch.Tensor, handles: list | None = None, postproc: list | None = None) -> AsyncLatents: + return AsyncLatents({"latents": tensor}, handles or [], postproc or []) + + def test_resolve_returns_wrapped_tensor(self): + t = torch.randn(2, 4) + al = self._make(t) + assert al._resolve() is t + + def test_attribute_access_resolves(self): + t = torch.randn(2, 4) + al = self._make(t) + assert al.shape == t.shape + assert al.dtype == t.dtype + + def test_torch_function_protocol(self): + """torch ops that receive an AsyncLatents should see the underlying tensor.""" + t = torch.randn(2, 4) + al = self._make(t) + mask = torch.ones_like(t) + result = mask * al # triggers __torch_function__ + torch.testing.assert_close(result, mask * t) + + def test_torch_function_with_list_arg(self): + """__torch_function__ must unwrap AsyncLatents inside list/tuple args.""" + t = torch.randn(2, 4) + al = self._make(t) + result = torch.cat([al, al], dim=0) + torch.testing.assert_close(result, torch.cat([t, t], dim=0)) + + def test_torch_tensor_conversion(self): + """torch.as_tensor on an AsyncLatents must share storage with the underlying tensor (no copy).""" + t = torch.randn(2, 4) + al = self._make(t) + result = torch.as_tensor(al) + assert result.data_ptr() == t.data_ptr(), "torch.as_tensor copied the data instead of sharing storage" + + def test_handles_are_waited_before_resolve(self): + t = torch.randn(2, 4) + h1, h2 = FakeWork(), FakeWork() + al = self._make(t, handles=[h1, h2]) + _ = al.shape # trigger resolution + assert h1.waited and h2.waited, "Not all handles were waited on" + + def test_postproc_callbacks_invoked_on_resolve(self): + t = torch.randn(2, 4) + log: list[int] = [] + al = self._make(t, postproc=[lambda: log.append(1), lambda: log.append(2)]) + _ = al.shape + assert log == [1, 2], f"postproc not called in order: {log}" + + def test_idempotent_resolve(self): + """handle.wait() must not be called twice if _resolve() is called twice.""" + t = torch.randn(2, 4) + h = FakeWork() + al = self._make(t, handles=[h]) + _ = al.shape # first resolve + h.waited = False # reset sentinel + _ = al.dtype # second resolve + assert not h.waited, "handle.wait() was called a second time" + + +# --------------------------------------------------------------------------- +# 2. _sync_pp_send / diffuse wrapper – unit tests (no distributed env required) +# --------------------------------------------------------------------------- + + +class TestSyncPPSend: + """Verifies PipelineParallelMixin's internal PP-send flush.""" + + pytestmark = [pytest.mark.cpu] + + @staticmethod + def _make_pipeline() -> PipelineParallelMixin: + # Instantiate a bare mixin — no layers, no distributed env needed. + # _sync_pp_send only touches _pp_send_work, so this is sufficient. + class _BarePP(PipelineParallelMixin, CFGParallelMixin): + pass + + return _BarePP() + + def test_noop_when_work_list_empty(self): + pipeline = self._make_pipeline() + pipeline._sync_pp_send() + assert pipeline._pp_send_work == [] + + def test_waits_all_pending_handles(self): + pipeline = self._make_pipeline() + works = [FakeWork(), FakeWork(), FakeWork()] + pipeline._pp_send_work = works + pipeline._sync_pp_send() + assert all(w.waited for w in works), "Some handles were not waited on" + + def test_clears_work_list_after_sync(self): + pipeline = self._make_pipeline() + pipeline._pp_send_work = [FakeWork()] + pipeline._sync_pp_send() + assert pipeline._pp_send_work == [] + + +class TestDiffuseWrapper: + """Verifies that PipelineParallelMixin flushes pending sends when diffuse() exits.""" + + pytestmark = [pytest.mark.cpu] + + def test_diffuse_flushes_pending_sends_on_success(self): + work = FakeWork() + + class _DiffusePP(PipelineParallelMixin, CFGParallelMixin): + def diffuse(self): + self._pp_send_work = [work] + return "done" + + pipeline = _DiffusePP() + + assert pipeline.diffuse() == "done" + assert work.waited + assert pipeline._pp_send_work == [] + + def test_diffuse_flushes_pending_sends_on_exception(self): + work = FakeWork() + + class _DiffusePP(PipelineParallelMixin, CFGParallelMixin): + def diffuse(self): + self._pp_send_work = [work] + raise RuntimeError("boom") + + pipeline = _DiffusePP() + + with pytest.raises(RuntimeError, match="boom"): + pipeline.diffuse() + assert work.waited + assert pipeline._pp_send_work == [] + + def test_diffuse_wrapper_preserves_metadata(self): + class _DiffusePP(PipelineParallelMixin, CFGParallelMixin): + def diffuse(self): + """Original diffuse docstring.""" + return "done" + + assert _DiffusePP.diffuse.__name__ == "diffuse" + assert _DiffusePP.diffuse.__doc__ == "Original diffuse docstring." + + +class TestVaeDecodeGuard: + pytestmark = [pytest.mark.cpu] + + @staticmethod + def _make_pipeline(distributed_enabled: bool = False) -> PipelineParallelMixin: + class _DecodePP(PipelineParallelMixin, CFGParallelMixin): + def __init__(self): + self.vae = FakeVAE(distributed_enabled=distributed_enabled) + + return _DecodePP() + + @staticmethod + def _set_rank(monkeypatch, world_size: int, first_stage: bool) -> None: + monkeypatch.setattr(pp_module, "get_pipeline_parallel_world_size", lambda: world_size) + monkeypatch.setattr(pp_module, "is_pipeline_first_stage", lambda: first_stage) + + def test_calls_original_decode_when_pp_disabled(self, monkeypatch): + self._set_rank(monkeypatch, world_size=1, first_stage=True) + pipeline = self._make_pipeline() + z = torch.ones(2, 3) + + output = pipeline.vae.decode(z)[0] + + assert pipeline.vae.calls == 1 + torch.testing.assert_close(output, z + 1) + + def test_calls_original_decode_on_first_stage(self, monkeypatch): + self._set_rank(monkeypatch, world_size=2, first_stage=True) + pipeline = self._make_pipeline() + z = torch.ones(2, 3) + + output = pipeline.vae.decode(z)[0] + + assert pipeline.vae.calls == 1 + torch.testing.assert_close(output, z + 1) + + def test_skips_decode_on_non_first_stage(self, monkeypatch): + self._set_rank(monkeypatch, world_size=2, first_stage=False) + pipeline = self._make_pipeline() + z = torch.ones(2, 3) + + output = pipeline.vae.decode(z) + + assert pipeline.vae.calls == 0 + assert output == (None,) + + def test_calls_original_decode_when_distributed_vae_enabled(self, monkeypatch): + self._set_rank(monkeypatch, world_size=2, first_stage=False) + pipeline = self._make_pipeline(distributed_enabled=True) + z = torch.ones(2, 3) + + output = pipeline.vae.decode(z)[0] + + assert pipeline.vae.calls == 1 + torch.testing.assert_close(output, z + 1) + + def test_decode_wrapper_preserves_metadata(self): + pipeline = self._make_pipeline() + + assert pipeline.vae.decode.__name__ == "decode" + assert pipeline.vae.decode.__doc__ == "Original decode docstring." + + +@pytest.mark.cpu +def test_pipeline_parallel_requires_cfg_mixin(): + with pytest.raises(TypeError, match="inherits PipelineParallelMixin but not CFGParallelMixin"): + + class _MissingCFG(PipelineParallelMixin): + pass + + +@pytest.mark.cpu +def test_pipeline_parallel_requires_mro_before_cfg_mixin(): + with pytest.raises(TypeError, match="must inherit PipelineParallelMixin before CFGParallelMixin"): + + class _WrongOrder(CFGParallelMixin, PipelineParallelMixin): + pass + + +# --------------------------------------------------------------------------- +# Distributed test helpers +# --------------------------------------------------------------------------- + + +def init_dist(local_rank: int, world_size: int, master_port: str) -> torch.device: + """Initialise the distributed environment for a spawned worker.""" + device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") + current_omni_platform.set_device(device) + update_environment_variables( + { + "RANK": str(local_rank), + "LOCAL_RANK": str(local_rank), + "WORLD_SIZE": str(world_size), + "MASTER_ADDR": "localhost", + "MASTER_PORT": master_port, + } + ) + init_distributed_environment() + return device + + +def make_pipeline_and_inputs( + test_config: dict, dtype: torch.dtype, device: torch.device, do_true_cfg: bool = False +) -> tuple["MockPipelineParallel", dict, dict | None]: + """Create a MockPipelineParallel and seeded inputs from a test_config dict. + + Must be called after ``initialize_model_parallel`` so that ``make_layers`` + can read the PP group to determine this rank's layer slice. + + Returns ``(pipeline, positive_kwargs, negative_kwargs)``. + ``negative_kwargs`` is ``None`` when ``do_true_cfg=False``. + """ + pipeline = MockPipelineParallel( + num_layers=test_config["num_layers"], + dim=test_config["dim"], + seed=test_config["model_seed"], + device=device, + dtype=dtype, + ) + + torch.manual_seed(test_config["input_seed"]) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(test_config["input_seed"]) + pos_x = {"x": torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)} + + negative_kwargs = None + if do_true_cfg: + torch.manual_seed(test_config["input_seed"] + 1) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(test_config["input_seed"] + 1) + neg_x = torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device) + negative_kwargs = {"x": neg_x} + + return pipeline, pos_x, negative_kwargs + + +# --------------------------------------------------------------------------- +# 3. isend_tensor_dict / irecv_tensor_dict (2 GPUs) +# --------------------------------------------------------------------------- + + +def isend_irecv_worker(local_rank: int, world_size: int, master_port: str, result_queue): + device = init_dist(local_rank, world_size, master_port) + initialize_model_parallel(pipeline_parallel_size=world_size) + pp_group = get_pp_group() + + if pp_group.is_first_rank: + torch.manual_seed(77) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(77) + tensor = torch.randn(3, 5, dtype=torch.float32, device=device) + handles = pp_group.isend_tensor_dict({"t": tensor}) + for h in handles: + h.wait() + result_queue.put(("sent", tensor.cpu())) + else: + received = AsyncIntermediateTensors(*pp_group.irecv_tensor_dict()) + result_queue.put(("received", received["t"].cpu())) + + if torch.distributed.is_initialized(): + torch.distributed.barrier() + destroy_distributed_env() + + +@pytest.mark.gpu +@pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs") +@pytest.mark.parametrize("pp_size", [2]) +def test_isend_irecv_tensor_dict(pp_size: int): + """isend_tensor_dict / irecv_tensor_dict transfer a tensor dict without loss.""" + mp_context = torch.multiprocessing.get_context("spawn") + manager = mp_context.Manager() + q = manager.Queue() + + port = _find_free_port() + torch.multiprocessing.spawn(isend_irecv_worker, args=(pp_size, port, q), nprocs=pp_size) + + results = {label: tensor for label, tensor in [q.get(), q.get()]} + torch.testing.assert_close( + results["received"], results["sent"], rtol=0, atol=0, msg="isend/irecv transferred tensor incorrectly" + ) + + +# --------------------------------------------------------------------------- +# 4. predict_noise_maybe_with_cfg +# --------------------------------------------------------------------------- + +_baseline_cache: dict[tuple, torch.Tensor] = {} + + +def compute_single_gpu_baseline(test_config: dict, dtype: torch.dtype, do_true_cfg: bool) -> torch.Tensor: + """Compute expected single-GPU output using the same MockPipelineParallel. + + Initializes a trivial distributed env (world_size=1) so that ``make_layers`` and the PP/CFG mixins work normally. + Results are cached so identical configs are only computed once. + """ + key = ( + test_config["num_layers"], + test_config["dim"], + test_config["batch_size"], + test_config["model_seed"], + test_config["input_seed"], + test_config["cfg_scale"], + dtype, + do_true_cfg, + ) + if key in _baseline_cache: + return _baseline_cache[key] + + device = init_dist(0, 1, _find_free_port()) + initialize_model_parallel(pipeline_parallel_size=1) + + pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs( + test_config, dtype, device, do_true_cfg=do_true_cfg + ) + + with torch.inference_mode(): + noise_pred = pipeline.predict_noise_maybe_with_cfg( + do_true_cfg=do_true_cfg, + true_cfg_scale=test_config["cfg_scale"], + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + ) + + destroy_distributed_env() + + _baseline_cache[key] = noise_pred.cpu() + return _baseline_cache[key] + + +def predict_noise_worker( + local_rank: int, + world_size: int, + master_port: str, + pp_size: int, + cfg_size: int, + do_true_cfg: bool, + dtype: torch.dtype, + test_config: dict, + result_queue, +): + """Generic predict-noise worker parameterized by PP and CFG topology.""" + device = init_dist(local_rank, world_size, master_port) + initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size) + + pp_group = get_pp_group() + cfg_rank = get_classifier_free_guidance_rank() + + pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs( + test_config, dtype, device, do_true_cfg=do_true_cfg + ) + + with torch.inference_mode(): + noise_pred = pipeline.predict_noise_maybe_with_cfg( + do_true_cfg=do_true_cfg, + true_cfg_scale=test_config["cfg_scale"], + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + ) + # This worker exercises predict_noise_maybe_with_cfg directly, bypassing diffuse(). + # Flush the non-last PP rank's async send before barrier / process teardown. + pipeline._sync_pp_send() + + if pp_group.is_last_rank: + assert noise_pred is not None + if cfg_rank == 0: + result_queue.put(noise_pred.cpu()) + else: + assert noise_pred is None + + if torch.distributed.is_initialized(): + torch.distributed.barrier() + destroy_distributed_env() + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol", + [ + pytest.param( + 2, + 1, + False, + torch.float32, + 4, + 100, + 1e-5, + 1e-5, + marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"), + id="pp2-no_cfg-float32", + ), + pytest.param( + 2, + 1, + False, + torch.bfloat16, + 4, + 100, + 1e-2, + 1e-2, + marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"), + id="pp2-no_cfg-bfloat16", + ), + pytest.param( + 2, + 1, + True, + torch.bfloat16, + 4, + 100, + 1e-2, + 1e-2, + marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"), + id="pp2-seq_cfg-bfloat16", + ), + pytest.param( + 2, + 2, + True, + torch.bfloat16, + 4, + 100, + 1e-2, + 1e-2, + marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"), + id="pp2-cfg2-bfloat16", + ), + pytest.param( + 3, + 1, + False, + torch.bfloat16, + 6, + 100, + 1e-2, + 1e-2, + marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 3, reason="Need at least 3 GPUs"), + id="pp3-no_cfg-bfloat16", + ), + ], +) +def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol): + """predict_noise_maybe_with_cfg output matches the single-GPU baseline across PP / CFG topologies.""" + test_config = { + "num_layers": num_layers, + "dim": 64, + "batch_size": 2, + "cfg_scale": 7.5, + "model_seed": 42, + "input_seed": input_seed, + } + + baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg) + + mp_context = torch.multiprocessing.get_context("spawn") + manager = mp_context.Manager() + pp_q = manager.Queue() + + world_size = pp_size * cfg_size + port = _find_free_port() + torch.multiprocessing.spawn( + predict_noise_worker, + args=(world_size, port, pp_size, cfg_size, do_true_cfg, dtype, test_config, pp_q), + nprocs=world_size, + ) + + pp_out = pp_q.get() + + assert baseline_out.shape == pp_out.shape + torch.testing.assert_close( + pp_out, + baseline_out, + rtol=rtol, + atol=atol, + msg=f"PP={pp_size} cfg={cfg_size} {'with' if do_true_cfg else 'no'} CFG output differs from baseline ({dtype=})", + ) + + +# --------------------------------------------------------------------------- +# 5. scheduler_step_maybe_with_cfg +# --------------------------------------------------------------------------- + + +def compute_scheduler_step_baseline(test_config: dict, do_true_cfg: bool) -> torch.Tensor: + """Single-GPU reference: predict_noise + scheduler_step.""" + device = init_dist(0, 1, _find_free_port()) + initialize_model_parallel(pipeline_parallel_size=1) + + pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs( + test_config, torch.float32, device, do_true_cfg=do_true_cfg + ) + latents = positive_kwargs["x"] + t = torch.tensor(500, device=device) + + with torch.inference_mode(): + noise_pred = pipeline.predict_noise_maybe_with_cfg( + do_true_cfg=do_true_cfg, + true_cfg_scale=test_config["cfg_scale"], + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + ) + result = pipeline.scheduler_step_maybe_with_cfg( + noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg + ) + + destroy_distributed_env() + return result.cpu() + + +def scheduler_step_worker( + local_rank: int, + world_size: int, + master_port: str, + pp_size: int, + cfg_size: int, + do_true_cfg: bool, + test_config: dict, + result_queue, +): + device = init_dist(local_rank, world_size, master_port) + initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size) + + pp_group = get_pp_group() + cfg_rank = get_classifier_free_guidance_rank() + + pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs( + test_config, torch.float32, device, do_true_cfg=do_true_cfg + ) + latents = positive_kwargs["x"] + t = torch.tensor(500, device=device) + + with torch.inference_mode(): + noise_pred = pipeline.predict_noise_maybe_with_cfg( + do_true_cfg=do_true_cfg, + true_cfg_scale=test_config["cfg_scale"], + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + ) + latents = pipeline.scheduler_step_maybe_with_cfg( + noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg + ) + # This worker exercises scheduler_step_maybe_with_cfg directly, bypassing diffuse(). + # Flush the last PP rank's async latent send before barrier / process teardown. + pipeline._sync_pp_send() + + if pp_group.is_first_rank and cfg_rank == 0: + resolved = latents.contiguous() + result_queue.put(resolved.cpu()) + + if torch.distributed.is_initialized(): + torch.distributed.barrier() + destroy_distributed_env() + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "pp_size, cfg_size, do_true_cfg, input_seed", + [ + pytest.param( + 2, + 1, + False, + 300, + marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"), + id="pp2-no_cfg", + ), + pytest.param( + 2, + 2, + True, + 600, + marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"), + id="pp2-cfg2-true_cfg", + ), + ], +) +def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed): + """Rank 0 latents after scheduler_step match the single-GPU baseline across PP / CFG topologies.""" + test_config = { + "num_layers": 4, + "dim": 64, + "batch_size": 2, + "cfg_scale": 7.5, + "model_seed": 42, + "input_seed": input_seed, + } + + baseline = compute_scheduler_step_baseline(test_config, do_true_cfg) + + mp_context = torch.multiprocessing.get_context("spawn") + manager = mp_context.Manager() + q = manager.Queue() + + port = _find_free_port() + world_size = pp_size * cfg_size + torch.multiprocessing.spawn( + scheduler_step_worker, + args=(world_size, port, pp_size, cfg_size, do_true_cfg, test_config, q), + nprocs=world_size, + ) + + result = q.get() + torch.testing.assert_close( + result, + baseline, + rtol=0, + atol=0, + msg=f"PP={pp_size} CFG={cfg_size} scheduler step latents on rank 0 do not match single-GPU baseline", + ) diff --git a/tests/diffusion/models/wan2_2/test_wan22_ti2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_ti2v_pipeline.py deleted file mode 100644 index e611c37b6..000000000 --- a/tests/diffusion/models/wan2_2/test_wan22_ti2v_pipeline.py +++ /dev/null @@ -1,97 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from types import SimpleNamespace - -import pytest -import torch -from PIL import Image -from torch import nn - -from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar -from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_ti2v import ( - Wan22TI2VPipeline, - get_wan22_ti2v_pre_process_func, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] - - -def _make_ti2v_pipeline() -> Wan22TI2VPipeline: - pipeline = object.__new__(Wan22TI2VPipeline) - nn.Module.__init__(pipeline) - pipeline.device = torch.device("cpu") - pipeline.transformer = StubTransformer(in_channels=4, out_channels=4) - pipeline.vae = StubVAE(z_dim=4) - pipeline.vae_scale_factor_temporal = 4 - pipeline.vae_scale_factor_spatial = 8 - pipeline.progress_bar = noop_progress_bar - return pipeline - - -def test_ti2v_preprocess_uses_720p_area_for_image_condition() -> None: - preprocess = get_wan22_ti2v_pre_process_func(SimpleNamespace()) - request = SimpleNamespace( - prompts=[{"prompt": "p", "multi_modal_data": {"image": Image.new("RGB", (320, 160), "blue")}}], - sampling_params=SimpleNamespace(height=None, width=None), - ) - - result = preprocess(request) - - assert result.sampling_params.height == 672 - assert result.sampling_params.width == 1344 - assert result.prompts[0]["multi_modal_data"]["image"].size == (1344, 672) - - -def test_ti2v_diffuse_without_image_condition_expands_patch_timesteps() -> None: - pipeline = _make_ti2v_pipeline() - latents = torch.zeros(1, 4, 2, 4, 4) - calls = [] - - def fake_predict_noise_maybe_with_cfg(**kwargs): - calls.append(kwargs) - return torch.ones_like(latents) - - pipeline.predict_noise_maybe_with_cfg = fake_predict_noise_maybe_with_cfg # type: ignore[method-assign] - pipeline.scheduler_step_maybe_with_cfg = lambda noise, t, current, cfg: current + noise # type: ignore[method-assign] - - result = pipeline.diffuse( - latents=latents, - timesteps=torch.tensor([7]), - prompt_embeds=torch.zeros(1, 2, 3), - negative_prompt_embeds=torch.zeros(1, 2, 3), - guidance_scale=3.0, - dtype=torch.float32, - attention_kwargs={"a": "b"}, - num_latent_frames=2, - latent_height=4, - latent_width=4, - ) - - positive = calls[0]["positive_kwargs"] - assert calls[0]["do_true_cfg"] is True - assert positive["timestep"].shape == (1, 8) - torch.testing.assert_close(positive["timestep"], torch.full((1, 8), 7, dtype=positive["timestep"].dtype)) - torch.testing.assert_close(positive["hidden_states"], latents) - torch.testing.assert_close(result, torch.ones_like(latents)) - - -def test_ti2v_prepare_i2v_latents_encodes_condition_and_masks_first_frame() -> None: - pipeline = _make_ti2v_pipeline() - latents, latent_condition, first_frame_mask = pipeline.prepare_i2v_latents( - image=torch.zeros(1, 3, 16, 16), - batch_size=1, - num_channels_latents=4, - height=16, - width=16, - num_frames=5, - dtype=torch.float32, - device=torch.device("cpu"), - generator=None, - latents=torch.zeros(1, 4, 2, 2, 2), - ) - - torch.testing.assert_close(latents, torch.zeros(1, 4, 2, 2, 2)) - assert latent_condition.shape == (1, 4, 1, 2, 2) - assert first_frame_mask[:, :, 0].sum() == 0 - assert first_frame_mask[:, :, 1].sum() == 4 diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py index 5f9047fba..daf9b3939 100644 --- a/vllm_omni/diffusion/cache/cache_dit_backend.py +++ b/vllm_omni/diffusion/cache/cache_dit_backend.py @@ -1505,7 +1505,6 @@ CUSTOM_DIT_ENABLERS.update( { "Wan22Pipeline": enable_cache_for_wan22, "Wan22I2VPipeline": enable_cache_for_wan22, - "Wan22TI2VPipeline": enable_cache_for_wan22, "HunyuanImage3Pipeline": enable_cache_for_hunyuan_image3, "FluxPipeline": enable_cache_for_flux, "Flux2KleinPipeline": enable_cache_for_flux2_klein, diff --git a/vllm_omni/diffusion/distributed/cfg_parallel.py b/vllm_omni/diffusion/distributed/cfg_parallel.py index 98757006b..6810a447b 100644 --- a/vllm_omni/diffusion/distributed/cfg_parallel.py +++ b/vllm_omni/diffusion/distributed/cfg_parallel.py @@ -10,6 +10,7 @@ from typing import Any import torch from vllm.logger import init_logger +from vllm.sequence import IntermediateTensors from vllm_omni.diffusion.distributed.parallel_state import ( get_cfg_group, @@ -369,7 +370,7 @@ class CFGParallelMixin(metaclass=ABCMeta): results.append(comb) return _unwrap(tuple(results)) - def predict_noise(self, *args: Any, **kwargs: Any) -> torch.Tensor | tuple[torch.Tensor, ...]: + def predict_noise(self, *args: Any, **kwargs: Any) -> torch.Tensor | tuple[torch.Tensor, ...] | IntermediateTensors: """ Forward pass through transformer to predict noise. @@ -381,8 +382,11 @@ class CFGParallelMixin(metaclass=ABCMeta): multi-output models (e.g., video + audio). Multi-output models must also override combine_cfg_noise() and set self.scheduler to a composite scheduler that handles tuples. + Non-last Pipeline Parallel stages return `IntermediateTensors` + instead of final noise tensors wrapped in a tuple. """ - return self.transformer(*args, **kwargs)[0] + result = self.transformer(*args, **kwargs) + return result if isinstance(result, IntermediateTensors) else result[0] def diffuse( self, diff --git a/vllm_omni/diffusion/distributed/group_coordinator.py b/vllm_omni/diffusion/distributed/group_coordinator.py index 5294e6c9e..5c5b092ef 100644 --- a/vllm_omni/diffusion/distributed/group_coordinator.py +++ b/vllm_omni/diffusion/distributed/group_coordinator.py @@ -429,6 +429,78 @@ class GroupCoordinator: async_handle.wait() return tensor_dict + def isend_tensor_dict( + self, + tensor_dict: dict[str, torch.Tensor | Any], + dst: int | None = None, + ) -> list[torch.distributed.Work]: + """Non-blocking send of a tensor dictionary. + + Sends metadata via the Gloo CPU group (blocking) then starts a + non-blocking NCCL isend for each GPU tensor. Returns the list of + Work handles; the caller must call handle.wait() before the tensors + can be safely reused or freed. + + NOTE: `dst` is the group rank of the destination. + """ + if not torch.distributed.is_initialized() or self.world_size == 1: + return [] + + if dst is None: + dst = self.group_next_rank + assert dst < self.world_size, f"Invalid dst rank ({dst})" + + metadata_list, tensor_list = _split_tensor_dict(tensor_dict) + self.send_object(metadata_list, dst=dst) + + handles: list[torch.distributed.Work] = [] + for tensor in tensor_list: + if tensor.numel() == 0: + continue + group = self.cpu_group if tensor.is_cpu else self.device_group + handle = torch.distributed.isend(tensor, dst=self.ranks[dst], group=group) + if tensor.is_cuda: + # Keep allocator from reusing this CUDA buffer before the async send finishes. + tensor.record_stream(torch.cuda.current_stream(tensor.device)) + handles.append(handle) + return handles + + def irecv_tensor_dict( + self, + src: int | None = None, + ) -> tuple[dict[str, torch.Tensor | Any], list[torch.distributed.Work], list]: + """Non-blocking receive of a tensor dictionary. + + Receives metadata via the Gloo CPU group (blocking) then starts a + non-blocking NCCL irecv for each GPU tensor. Returns + ``(tensor_dict, comm_handles, comm_postprocess)`` matching the + interface expected by ``AsyncIntermediateTensors``. + + NOTE: `src` is the group rank of the source. + """ + if not torch.distributed.is_initialized() or self.world_size == 1: + return {}, [], [] + + if src is None: + src = self.group_prev_rank + assert src < self.world_size, f"Invalid src rank ({src})" + + recv_metadata_list = self.recv_object(src=src) + tensor_dict: dict[str, Any] = {} + handles: list[torch.distributed.Work] = [] + + for key, value in recv_metadata_list: + if isinstance(value, TensorMetadata): + tensor = torch.empty(value.size, dtype=value.dtype, device=value.device) + if tensor.numel() > 0: + group = self.cpu_group if tensor.is_cpu else self.device_group + handles.append(torch.distributed.irecv(tensor, src=self.ranks[src], group=group)) + _update_nested_dict(tensor_dict, key, tensor) + else: + _update_nested_dict(tensor_dict, key, value) + + return tensor_dict, handles, [] + def send_tensor_dict( self, tensor_dict: dict[str, torch.Tensor | Any], diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py index 82c9c1a45..163f9dd45 100644 --- a/vllm_omni/diffusion/distributed/parallel_state.py +++ b/vllm_omni/diffusion/distributed/parallel_state.py @@ -746,11 +746,6 @@ def initialize_model_parallel( f" but got {sequence_parallel_size} != {ring_degree} * {ulysses_degree}" ) - # FIXME: Since the async p2p communication operation of NPU is not same as cuda in torch, - # the pipefusion is not ready for npu yet - if current_omni_platform.is_npu(): - assert pipeline_parallel_size == 1, "Current pipefusion is not ready for NPU" - dit_parallel_size = ( data_parallel_size * cfg_parallel_size * sequence_parallel_size * pipeline_parallel_size * tensor_parallel_size ) diff --git a/vllm_omni/diffusion/distributed/pipeline_parallel.py b/vllm_omni/diffusion/distributed/pipeline_parallel.py new file mode 100644 index 000000000..b39d8052e --- /dev/null +++ b/vllm_omni/diffusion/distributed/pipeline_parallel.py @@ -0,0 +1,291 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from functools import wraps +from typing import Any + +import torch +from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors + +from vllm_omni.diffusion.distributed.parallel_state import ( + get_cfg_group, + get_classifier_free_guidance_rank, + get_classifier_free_guidance_world_size, + get_pipeline_parallel_world_size, + get_pp_group, + is_pipeline_first_stage, +) + + +class AsyncLatents: + """Transparent async wrapper returned by scheduler_step on rank 0. + + Wraps a pending ``irecv_tensor_dict`` and defers ``handle.wait()`` until the + underlying tensor is actually consumed — either via attribute access + (e.g. ``latents.to(dtype)``, ``latents.shape``) or via a torch operation + (e.g. ``mask * latents``). This keeps the first PP rank non-blocking after + posting the receive, matching the async philosophy used everywhere else in + the PP communication layer. + """ + + __slots__ = ("_tensor_dict", "_handles", "_postproc", "_tensor") + + def __init__( + self, + tensor_dict: dict[str, torch.Tensor], + handles: list[torch.distributed.Work], + postproc: list, + ): + self._tensor_dict = tensor_dict + self._handles = handles + self._postproc = postproc + self._tensor: torch.Tensor | None = None + + def _resolve(self) -> torch.Tensor: + if self._tensor is not None: + return self._tensor + for h in self._handles: + h.wait() + for fn in self._postproc: + fn() + self._tensor = self._tensor_dict["latents"] + return self._tensor + + # Attribute access (e.g. .shape, .to(), .dtype) delegates to the resolved tensor. + def __getattr__(self, name: str): + return getattr(self._resolve(), name) + + # Torch function protocol: any torch op involving an AsyncLatents resolves it first. + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + + def _unwrap(x): + if isinstance(x, AsyncLatents): + return x._resolve() + if isinstance(x, (list, tuple)): + return type(x)(_unwrap(item) for item in x) # type(x) return the class of x to preserve its type + return x + + args = tuple(_unwrap(a) for a in args) + kwargs = {k: _unwrap(v) for k, v in kwargs.items()} + return func(*args, **kwargs) + + +class PipelineParallelMixin: + """ + Mixin providing Pipeline Parallelism for diffusion pipelines. + + All PP ranks run the full denoising loop in `forward()`. + `predict_noise_maybe_with_cfg` and `scheduler_step_maybe_with_cfg` encapsulate all inter-rank communication. + + Communication pattern per denoising step: + Forward chain : rank 0 → 1 → … → N-1 via async isend/irecv (AsyncIntermediateTensors) + Next timestep : last rank → rank 0 via async isend/irecv (AsyncLatents) + + All communication is asynchronous using isend_tensor_dict/irecv_tensor_dict. + Only rank 0 needs updated latents for the next forward pass start. + + For sequential CFG (cfg_parallel_size=1) with PP, two full forward chains are + executed — one for the positive pass and one for the negative pass — so that each + PP stage operates on the correct encoder_hidden_states. + """ + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin + + if not issubclass(cls, CFGParallelMixin): + raise TypeError( + f"{cls.__name__} inherits PipelineParallelMixin but not CFGParallelMixin. " + "Pipeline Parallelism requires CFGParallelMixin for predict_noise(), " + "predict_noise_maybe_with_cfg(), scheduler_step_maybe_with_cfg(), and combine_cfg_noise(). " + "Add CFGParallelMixin to the base classes of your pipeline." + ) + mro = cls.mro() + if mro.index(PipelineParallelMixin) > mro.index(CFGParallelMixin): + raise TypeError( + f"{cls.__name__} must inherit PipelineParallelMixin before CFGParallelMixin so MRO selects " + f"PP-aware predict/scheduler wrappers and their `super()` calls delegate to CFGParallelMixin." + ) + + init = cls.__dict__.get("__init__") + if callable(init): + + @wraps(init) + def wrapped_init(self, *args: Any, **kwargs: Any) -> None: + init(self, *args, **kwargs) + vae = getattr(self, "vae", None) + if vae is not None and hasattr(vae, "decode"): + self._wrapped_vae_decode() + + cls.__init__ = wrapped_init + + diffuse = cls.__dict__.get("diffuse") + if callable(diffuse): + + @wraps(diffuse) + def wrapped_diffuse(self, *args: Any, **kwargs: Any) -> Any: + try: + latents = diffuse(self, *args, **kwargs) + if isinstance(latents, AsyncLatents): + latents = torch.as_tensor(latents) # avoid copying + return latents + finally: + self._sync_pp_send() + + cls.diffuse = wrapped_diffuse + + def _wrapped_vae_decode(self) -> None: + vae, orig_decode = self.vae, self.vae.decode + + @wraps(orig_decode) + def wrapped_decode(z: torch.Tensor, *args: Any, **kwargs: Any): + if hasattr(vae, "is_distributed_enabled") and vae.is_distributed_enabled(): + # Middle ranks (world size 3 or more) hold stale latents after the denoising loop. + # Broadcast from rank 0 so every rank splits identical tiles. + if get_pipeline_parallel_world_size() > 2: + z = get_pp_group().broadcast(z, src=0) + return orig_decode(z, *args, **kwargs) + elif is_pipeline_first_stage(): + return orig_decode(z, *args, **kwargs) + return (None,) # decoder returns a tuple + + self.vae.decode = wrapped_decode + + @property + def _pp_send_work(self) -> list[torch.distributed.Work]: + if not hasattr(self, "_pp_send_work_list"): + self._pp_send_work_list: list[torch.distributed.Work] = [] + return self._pp_send_work_list + + @_pp_send_work.setter + def _pp_send_work(self, work: list[torch.distributed.Work]) -> None: + self._pp_send_work_list = work + + def _sync_pp_send(self) -> None: + """ + Wait on all pending non-blocking PP sends. + + Must be called after the denoising loop so that the isend handles + from the last iteration are completed before any subsequent + collective (e.g. VAE decode broadcast) or tensor reuse. + """ + if self._pp_send_work: + for handle in self._pp_send_work: + handle.wait() + self._pp_send_work = [] + + def predict_noise_maybe_with_cfg( + self, + do_true_cfg: bool, + true_cfg_scale: float, + positive_kwargs: dict[str, Any], + negative_kwargs: dict[str, Any] | None, + cfg_normalize: bool = True, + output_slice: int | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, ...] | None: + """ + Drop-in replacement for predict_noise_maybe_with_cfg that also handles PP. + + Supports three modes: + - PP only, sequential CFG: both branches (cond and uncond) run through this PP pipeline. + This doubles communication volume per denoising step compared to PP + CFG-parallel. + - PP + CFG-parallel: each PP pipeline carries one branch. The last PP + rank all-gathers across the CFG group and combines, mirroring + CFGParallelMixin.predict_noise_maybe_with_cfg exactly. + - PP only, no CFG: cond branch only. + + Returns: + noise_pred on the last PP rank (all CFG ranks when CFG-parallel is active). + None on all other ranks. + """ + if get_pipeline_parallel_world_size() == 1: + return super().predict_noise_maybe_with_cfg( + do_true_cfg, true_cfg_scale, positive_kwargs, negative_kwargs, cfg_normalize, output_slice + ) + + self._sync_pp_send() + + pp_group = get_pp_group() + + cfg_parallel_ready = do_true_cfg and get_classifier_free_guidance_world_size() > 1 + if cfg_parallel_ready: + # Each PP pipeline carries exactly one CFG branch determined by cfg_rank. + all_kwargs = [positive_kwargs if get_classifier_free_guidance_rank() == 0 else negative_kwargs] + else: + # Sequential CFG (or no CFG): this PP pipeline handles all branches. + all_kwargs = [positive_kwargs] + ([negative_kwargs] if do_true_cfg else []) + + # Non-first ranks receive intermediate tensors asynchronously + n = len(all_kwargs) + its: list[AsyncIntermediateTensors | None] = [None] * n + if not pp_group.is_first_rank: + for i in range(n): + its[i] = AsyncIntermediateTensors(*pp_group.irecv_tensor_dict()) + + if not pp_group.is_last_rank: + # First / middle rank: run partial forwards and propagate ITs downstream. + for kwargs, it in zip(all_kwargs, its): + result = self.predict_noise(**kwargs, intermediate_tensors=it) + self._pp_send_work.extend(pp_group.isend_tensor_dict(result.tensors)) + return None + + # Last rank: run full forward + noise_preds = [self.predict_noise(**kwargs, intermediate_tensors=it) for kwargs, it in zip(all_kwargs, its)] + + if cfg_parallel_ready: + # All-gather the single-branch prediction across the CFG group and combine + # on all CFG ranks so every last PP rank has an identical noise_pred. + local_pred = noise_preds[0] + if output_slice is not None: + local_pred = local_pred[:, :output_slice] + gathered = get_cfg_group().all_gather(local_pred, separate_tensors=True) + return self.combine_cfg_noise(gathered[0], gathered[1], true_cfg_scale, cfg_normalize) + + # Sequential CFG or no-CFG path. + if do_true_cfg: + pos, neg = noise_preds[0], noise_preds[1] + if output_slice is not None: + pos = pos[:, :output_slice] + neg = neg[:, :output_slice] + return self.combine_cfg_noise(pos, neg, true_cfg_scale, cfg_normalize) + pred = noise_preds[0] + if output_slice is not None: + pred = pred[:, :output_slice] + return pred + + def scheduler_step_maybe_with_cfg( + self, + noise_pred: torch.Tensor | tuple[torch.Tensor, ...] | None, + t: torch.Tensor | tuple[torch.Tensor, ...], + latents: torch.Tensor | tuple[torch.Tensor, ...], + do_true_cfg: bool, + per_request_scheduler: Any | None = None, + generator: torch.Generator | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, ...] | AsyncLatents: + """ + Drop-in replacement for scheduler_step_maybe_with_cfg that also handles PP. + + Only the last rank runs the scheduler (it already has noise_pred); the result + is sent to rank 0 which needs it for the next forward pass. + + Returns a ``AsyncLatents`` on rank 0 that transparently defers + ``handle.wait()`` until the tensor is actually consumed (via attribute + access or a torch operation), keeping the rank non-blocking after the + ``irecv`` is posted. + """ + if get_pipeline_parallel_world_size() == 1: + return super().scheduler_step_maybe_with_cfg( + noise_pred, t, latents, do_true_cfg, per_request_scheduler, generator + ) + + pp_group = get_pp_group() + if pp_group.is_last_rank: + latents = super().scheduler_step_maybe_with_cfg( + noise_pred, t, latents, do_true_cfg, per_request_scheduler, generator + ) + self._pp_send_work = pp_group.isend_tensor_dict({"latents": latents}, dst=0) + elif pp_group.is_first_rank: + latents = AsyncLatents(*pp_group.irecv_tensor_dict(src=pp_group.world_size - 1)) + return latents diff --git a/vllm_omni/diffusion/models/wan2_2/__init__.py b/vllm_omni/diffusion/models/wan2_2/__init__.py index 2159a1184..4059c4fb5 100644 --- a/vllm_omni/diffusion/models/wan2_2/__init__.py +++ b/vllm_omni/diffusion/models/wan2_2/__init__.py @@ -22,11 +22,6 @@ from .pipeline_wan2_2_s2v import ( get_wan22_s2v_post_process_func, get_wan22_s2v_pre_process_func, ) -from .pipeline_wan2_2_ti2v import ( - Wan22TI2VPipeline, - get_wan22_ti2v_post_process_func, - get_wan22_ti2v_pre_process_func, -) from .pipeline_wan2_2_vace import ( Wan22VACEPipeline, get_wan22_vace_post_process_func, @@ -50,9 +45,6 @@ __all__ = [ "Wan22S2VPipeline", "get_wan22_s2v_post_process_func", "get_wan22_s2v_pre_process_func", - "Wan22TI2VPipeline", - "get_wan22_ti2v_post_process_func", - "get_wan22_ti2v_pre_process_func", "Wan22VACEPipeline", "get_wan22_vace_post_process_func", "get_wan22_vace_pre_process_func", diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index e3286cf4f..b4438c795 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -17,10 +17,12 @@ from torch import nn from transformers import AutoTokenizer, UMT5EncoderModel from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.sequence import IntermediateTensors 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.pipeline_parallel import AsyncLatents, PipelineParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader @@ -297,7 +299,9 @@ def get_wan22_pre_process_func( return pre_process_func -class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin): +class Wan22Pipeline( + nn.Module, PipelineParallelMixin, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin +): def __init__( self, *, @@ -460,7 +464,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe attention_kwargs: dict[str, Any], latent_condition: torch.Tensor | None = None, first_frame_mask: torch.Tensor | None = None, - ) -> torch.Tensor: + ) -> torch.Tensor | AsyncLatents: if attention_kwargs is None: attention_kwargs = {} with self.progress_bar(total=len(timesteps)) as pbar: @@ -858,7 +862,11 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None ) - def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor: + def predict_noise( + self, + current_model: nn.Module | None = None, + **kwargs: Any, + ) -> torch.Tensor | IntermediateTensors: """ Forward pass through transformer to predict noise. @@ -867,11 +875,12 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe **kwargs: Arguments to pass to the transformer Returns: - Predicted noise tensor + Predicted noise tensor or IntermediateTensors on non-last PP stages. """ if current_model is None: current_model = self.transformer - return current_model(**kwargs)[0] + result = current_model(**kwargs) + return result if isinstance(result, IntermediateTensors) else result[0] def encode_prompt( self, diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index 0651356cc..42c4eff6a 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -17,10 +17,12 @@ 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.sequence import IntermediateTensors 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.pipeline_parallel import AsyncLatents, PipelineParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader @@ -135,7 +137,12 @@ def get_wan22_i2v_pre_process_func( class Wan22I2VPipeline( - nn.Module, SupportImageInput, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin + nn.Module, + SupportImageInput, + PipelineParallelMixin, + CFGParallelMixin, + ProgressBarMixin, + DiffusionPipelineProfilerMixin, ): """ Wan2.2 Image-to-Video Pipeline. @@ -283,7 +290,7 @@ class Wan22I2VPipeline( attention_kwargs: dict[str, Any], condition: torch.Tensor, first_frame_mask: torch.Tensor, - ) -> torch.Tensor: + ) -> torch.Tensor | AsyncLatents: if attention_kwargs is None: attention_kwargs = {} with self.progress_bar(total=len(timesteps)) as pbar: @@ -348,7 +355,6 @@ class Wan22I2VPipeline( # Compute the previous noisy sample x_t -> x_t-1 with automatic CFG sync latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) - pbar.update() return latents @@ -667,7 +673,11 @@ class Wan22I2VPipeline( output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None ) - def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor: + def predict_noise( + self, + current_model: nn.Module | None = None, + **kwargs: Any, + ) -> torch.Tensor | IntermediateTensors: """ Forward pass through transformer to predict noise. @@ -676,11 +686,12 @@ class Wan22I2VPipeline( **kwargs: Arguments to pass to the transformer Returns: - Predicted noise tensor + Predicted noise tensor or IntermediateTensors on non-last PP stages. """ if current_model is None: current_model = self.transformer - return current_model(**kwargs)[0] + result = current_model(**kwargs) + return result if isinstance(result, IntermediateTensors) else result[0] def encode_prompt( self, diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py deleted file mode 100644 index 360eb83b9..000000000 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py +++ /dev/null @@ -1,731 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -""" -Wan2.2 TI2V (Text-Image-to-Video) Pipeline. - -This pipeline supports the unified TI2V-5B model that can generate videos from: -- Text only (T2V mode) -- Text + Image (I2V mode) - -The key difference from the MoE-based I2V pipeline is: -- Single transformer (not MoE with two transformers) -- Uses expand_timesteps mode for image conditioning -- No CLIP image encoder - only VAE encoding for image condition -""" - -from __future__ import annotations - -import logging -import os -from collections.abc import Iterable -from typing import Any, cast - -import numpy as np -import PIL.Image -import torch -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 OmniAutoencoderKLWan -from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin -from vllm_omni.diffusion.distributed.utils import get_local_device -from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx -from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader -from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders -from vllm_omni.diffusion.models.interface import SupportImageInput -from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin -from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( - build_wan_scheduler, - create_transformer_from_config, - load_transformer_config, - resolve_wan_flow_shift, - resolve_wan_sample_solver, - retrieve_latents, -) -from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel -from vllm_omni.diffusion.postprocess import interpolate_video_tensor -from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.inputs.data import OmniTextPrompt -from vllm_omni.platforms import current_omni_platform - -logger = logging.getLogger(__name__) - - -def get_wan22_ti2v_post_process_func( - od_config: OmniDiffusionConfig, -): - from diffusers.video_processor import VideoProcessor - - video_processor = VideoProcessor(vae_scale_factor=8) - - def post_process_func( - video: torch.Tensor, - output_type: str = "np", - sampling_params=None, - ): - if output_type == "latent": - return video - custom_output = {} - if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False): - video, multiplier = interpolate_video_tensor( - video, - exp=sampling_params.frame_interpolation_exp, - scale=sampling_params.frame_interpolation_scale, - model_path=sampling_params.frame_interpolation_model_path, - ) - custom_output["video_fps_multiplier"] = multiplier - return { - "video": video_processor.postprocess_video(video, output_type=output_type), - "custom_output": custom_output, - } - - return post_process_func - - -def get_wan22_ti2v_pre_process_func( - od_config: OmniDiffusionConfig, -): - """Pre-process function for TI2V: optionally load and resize input image.""" - - def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: - for i, prompt in enumerate(request.prompts): - multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None - raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None - if isinstance(prompt, str): - prompt = OmniTextPrompt(prompt=prompt) - if "additional_information" not in prompt: - prompt["additional_information"] = {} - - if raw_image is None: - raise ValueError( - """No image is provided. This model requires an image to run.""", - """Please correctly set `"multi_modal_data": {"image": , …}`""", - ) - if not isinstance(raw_image, (str, PIL.Image.Image)): - raise TypeError( - f"""Unsupported image format {raw_image.__class__}.""", - """Please correctly set `"multi_modal_data": {"image": , …}`""", - ) - image = PIL.Image.open(raw_image).convert("RGB") if isinstance(raw_image, str) else raw_image - - # Calculate dimensions based on aspect ratio if not provided - if request.sampling_params.height is None or request.sampling_params.width is None: - # Default max area for 720P (TI2V-5B default) - max_area = 720 * 1280 - aspect_ratio = image.height / image.width - - # Calculate dimensions maintaining aspect ratio - mod_value = 16 # Must be divisible by 16 - height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value - width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value - - if request.sampling_params.height is None: - request.sampling_params.height = height - if request.sampling_params.width is None: - request.sampling_params.width = width - - # Resize image to target dimensions - image = image.resize( - (request.sampling_params.width, request.sampling_params.height), # type: ignore # Above has ensured that width & height are not None - PIL.Image.Resampling.LANCZOS, - ) - prompt["multi_modal_data"]["image"] = image # type: ignore # key existence already checked above - - request.prompts[i] = prompt - return request - - return pre_process_func - - -class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, ProgressBarMixin): - """ - Wan2.2 Text-Image-to-Video (TI2V) Pipeline. - - This is a unified pipeline that supports both: - - Text-to-Video (T2V): when no image is provided - - Image-to-Video (I2V): when an image is provided - - Uses expand_timesteps mode for I2V conditioning where the first frame - is conditioned on the input image latent. - """ - - def __init__( - self, - *, - od_config: OmniDiffusionConfig, - prefix: str = "", - ): - super().__init__() - self.od_config = od_config - - self.device = get_local_device() - dtype = getattr(od_config, "dtype", torch.bfloat16) - - model = od_config.model - local_files_only = os.path.exists(model) - - # Set up weights sources for single transformer - self.weights_sources = [ - DiffusersPipelineLoader.ComponentSource( - model_or_path=od_config.model, - subfolder="transformer", - revision=None, - prefix="transformer.", - fall_back_to_pt=True, - ), - ] - - # See ``hub_prefetch.py`` for the transformers v5 subfolder race. - prefetch_subfolders( - model, - ["tokenizer", "text_encoder", "vae"], - local_files_only=local_files_only, - ) - - # Text encoder - self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) - self.text_encoder = UMT5EncoderModel.from_pretrained( - model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only - ).to(self.device) - - # VAE - self.vae = OmniAutoencoderKLWan.from_pretrained( - model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only - ).to(self.device) - - # Single transformer (TI2V uses dense 5B model, not MoE) - # Load config from model to get correct dimensions - transformer_config = load_transformer_config(model, "transformer", local_files_only) - self.transformer = self._create_transformer(transformer_config) - - self._sample_solver = "unipc" - self._flow_shift = od_config.flow_shift if od_config.flow_shift is not None else 5.0 - self.scheduler = build_wan_scheduler(self._sample_solver, self._flow_shift) - - # VAE scale factors - self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if hasattr(self.vae, "config") else 4 - self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if hasattr(self.vae, "config") else 8 - - # TI2V always uses expand_timesteps mode - self.expand_timesteps = True - - self._guidance_scale = None - self._num_timesteps = None - self._current_timestep = None - - def _create_transformer(self, config: dict) -> WanTransformer3DModel: - """Create a transformer from a config dict. Respects od_config.quantization_config.""" - quant_config = getattr(self.od_config, "quantization_config", None) - return create_transformer_from_config(config, quant_config=quant_config) - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def do_classifier_free_guidance(self): - return self._guidance_scale is not None and self._guidance_scale > 1.0 - - @property - def num_timesteps(self): - return self._num_timesteps - - @property - def current_timestep(self): - return self._current_timestep - - def diffuse( - self, - latents: torch.Tensor, - timesteps: torch.Tensor, - prompt_embeds: torch.Tensor, - negative_prompt_embeds: torch.Tensor | None, - guidance_scale: float, - dtype: torch.dtype, - attention_kwargs: dict[str, Any], - num_latent_frames: int, - latent_height: int, - latent_width: int, - latent_condition: torch.Tensor | None = None, - first_frame_mask: torch.Tensor | None = None, - ) -> torch.Tensor: - if attention_kwargs is None: - attention_kwargs = {} - with self.progress_bar(total=len(timesteps)) as pbar: - for step_idx, t in enumerate(timesteps): - self._current_timestep = t - set_forward_context_denoise_step_idx(step_idx) - - # Prepare latent input - if latent_condition is not None: - # I2V mode: blend condition with latents using mask - latent_model_input = (1 - first_frame_mask) * latent_condition + first_frame_mask * latents - latent_model_input = latent_model_input.to(dtype) - - # Expand timesteps for each patch (TI2V style) - temp_ts = (first_frame_mask[0][0][:, ::2, ::2] * t).flatten() - timestep = temp_ts.unsqueeze(0).expand(latents.shape[0], -1) - else: - # T2V mode: use latents directly - latent_model_input = latents.to(dtype) - - # Expand timesteps for TI2V model architecture - mask = torch.ones(1, 1, num_latent_frames, latent_height, latent_width, device=latents.device) - temp_ts = (mask[0][0][:, ::2, ::2] * t).flatten() - timestep = temp_ts.unsqueeze(0).expand(latents.shape[0], -1) - - do_true_cfg = guidance_scale > 1.0 and negative_prompt_embeds is not None - positive_kwargs = { - "hidden_states": latent_model_input, - "timestep": timestep, - "encoder_hidden_states": prompt_embeds, - "attention_kwargs": attention_kwargs, - "return_dict": False, - "current_model": self.transformer, - } - if do_true_cfg: - negative_kwargs = { - "hidden_states": latent_model_input, - "timestep": timestep, - "encoder_hidden_states": negative_prompt_embeds, - "attention_kwargs": attention_kwargs, - "return_dict": False, - "current_model": self.transformer, - } - else: - negative_kwargs = None - - noise_pred = self.predict_noise_maybe_with_cfg( - do_true_cfg=do_true_cfg, - true_cfg_scale=guidance_scale, - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - cfg_normalize=False, - ) - - latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) - pbar.update() - - return latents - - def forward( - self, - req: OmniDiffusionRequest, - prompt: str | None = None, - negative_prompt: str | None = None, - image: PIL.Image.Image | torch.Tensor | None = None, - height: int = 704, - width: int = 1280, - num_inference_steps: int = 40, - guidance_scale: float = 5.0, - frame_num: int = 81, - output_type: str | None = "np", - generator: torch.Generator | list[torch.Generator] | None = None, - prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - attention_kwargs: dict | None = None, - **kwargs, - ) -> DiffusionOutput: - # Get parameters from request or arguments - if len(req.prompts) > 1: - raise ValueError( - """This model only supports a single prompt, not a batched request.""", - """Please pass in a single prompt object or string, or a single-item list.""", - ) - if len(req.prompts) == 1: # If req.prompt is empty, default to prompt & neg_prompt in param list - prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt") - negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt") - if prompt is None and prompt_embeds is None: - raise ValueError("Prompt or prompt_embeds is required for Wan2.2 generation.") - - # Get image from request (optional for TI2V) - if image is None: - multi_modal_data = ( - req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None - ) - raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None - if isinstance(raw_image, list): - if len(raw_image) > 1: - logger.warning( - """Received a list of image. Only a single image is supported by this model.""" - """Taking only the first image for now.""" - ) - raw_image = raw_image[0] - if raw_image is None: - image = None - elif isinstance(raw_image, str): - image = PIL.Image.open(raw_image) - else: - image = cast(PIL.Image.Image | torch.Tensor, raw_image) - - # Default dimensions for TI2V-5B (720P) - height = req.sampling_params.height or height - width = req.sampling_params.width or width - num_frames = req.sampling_params.num_frames if req.sampling_params.num_frames else frame_num - num_steps = req.sampling_params.num_inference_steps or num_inference_steps - - # Respect per-request guidance_scale when explicitly provided. - if req.sampling_params.guidance_scale_provided: - guidance_scale = req.sampling_params.guidance_scale - - self._guidance_scale = guidance_scale - - # Validate inputs - self.check_inputs( - prompt=prompt, - negative_prompt=negative_prompt, - image=image, - height=height, - width=width, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - ) - - # Adjust num_frames to be compatible with VAE temporal scaling - if num_frames % self.vae_scale_factor_temporal != 1: - num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1 - num_frames = max(num_frames, 1) - - device = self.device - dtype = self.transformer.dtype - - # Generator setup - if generator is None: - generator = req.sampling_params.generator - if generator is None and req.sampling_params.seed is not None: - generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed) - - # Encode prompts - if prompt_embeds is None: - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, - do_classifier_free_guidance=guidance_scale > 1.0, - num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, - max_sequence_length=req.sampling_params.max_sequence_length or 512, - device=device, - dtype=dtype, - ) - else: - prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) - if negative_prompt_embeds is not None: - negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) - - batch_size = prompt_embeds.shape[0] - - sample_solver = resolve_wan_sample_solver(req, default=self._sample_solver) - flow_shift = resolve_wan_flow_shift(req, self.od_config) - if sample_solver != self._sample_solver or abs(flow_shift - self._flow_shift) > 1e-6: - self.scheduler = build_wan_scheduler(sample_solver, flow_shift) - self._sample_solver = sample_solver - self._flow_shift = flow_shift - - # Timesteps - self.scheduler.set_timesteps(num_steps, device=device) - timesteps = self.scheduler.timesteps - self._num_timesteps = len(timesteps) - - # Prepare latents - num_channels_latents = self.transformer.config.in_channels - num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 - latent_height = height // self.vae_scale_factor_spatial - latent_width = width // self.vae_scale_factor_spatial - - # Check if we have an image (I2V mode) or not (T2V mode) - if image is not None: - # I2V mode: prepare latents with image condition - from diffusers.video_processor import VideoProcessor - - video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - - if isinstance(image, PIL.Image.Image): - image_tensor = video_processor.preprocess(image, height=height, width=width) - else: - image_tensor = image - image_tensor = image_tensor.to(device=device, dtype=torch.float32) - - latents, latent_condition, first_frame_mask = self.prepare_i2v_latents( - image=image_tensor, - batch_size=batch_size, - num_channels_latents=num_channels_latents, - height=height, - width=width, - num_frames=num_frames, - dtype=torch.float32, - device=device, - generator=generator, - latents=req.sampling_params.latents, - ) - else: - # T2V mode: prepare random latents - latents = self.prepare_latents( - batch_size=batch_size, - num_channels_latents=num_channels_latents, - height=height, - width=width, - num_frames=num_frames, - dtype=torch.float32, - device=device, - generator=generator, - latents=req.sampling_params.latents, - ) - latent_condition = None - first_frame_mask = torch.ones( - 1, 1, num_latent_frames, latent_height, latent_width, dtype=torch.float32, device=device - ) - - if attention_kwargs is None: - attention_kwargs = {} - - latents = self.diffuse( - latents=latents, - timesteps=timesteps, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - guidance_scale=guidance_scale, - dtype=dtype, - attention_kwargs=attention_kwargs, - num_latent_frames=num_latent_frames, - latent_height=latent_height, - latent_width=latent_width, - latent_condition=latent_condition, - first_frame_mask=first_frame_mask, - ) - - # Wan2.2 is prone to out of memory errors when predicting large videos - # so we empty the cache here to avoid OOM before vae decoding. - if current_omni_platform.is_available(): - current_omni_platform.empty_cache() - self._current_timestep = None - - # For I2V mode, blend final latents with condition - if latent_condition is not None: - latents = (1 - first_frame_mask) * latent_condition + first_frame_mask * latents - - # Decode - if output_type == "latent": - output = latents - else: - latents = latents.to(self.vae.dtype) - latents_mean = ( - torch.tensor(self.vae.config.latents_mean) - .view(1, self.vae.config.z_dim, 1, 1, 1) - .to(latents.device, latents.dtype) - ) - latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( - latents.device, latents.dtype - ) - latents = latents / latents_std + latents_mean - output = self.vae.decode(latents, return_dict=False)[0] - - return DiffusionOutput(output=output) - - def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor: - """ - Forward pass through transformer to predict noise. - - Args: - current_model: The transformer model to use - **kwargs: Arguments to pass to the transformer - - Returns: - Predicted noise tensor - """ - if current_model is None: - current_model = self.transformer - return current_model(**kwargs)[0] - - def encode_prompt( - self, - prompt: str | list[str], - negative_prompt: str | list[str] | None = None, - do_classifier_free_guidance: bool = True, - num_videos_per_prompt: int = 1, - max_sequence_length: int = 512, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - ): - """Encode text prompts using T5 text encoder.""" - device = device or self.device - dtype = dtype or self.text_encoder.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - prompt_clean = [self._prompt_clean(p) for p in prompt] - batch_size = len(prompt_clean) - - text_inputs = self.tokenizer( - prompt_clean, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - add_special_tokens=True, - return_attention_mask=True, - return_tensors="pt", - ) - ids, mask = text_inputs.input_ids, text_inputs.attention_mask - seq_lens = mask.gt(0).sum(dim=1).long() - - prompt_embeds = self.text_encoder(ids.to(device), mask.to(device)).last_hidden_state - prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] - prompt_embeds = torch.stack( - [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0 - ) - - _, seq_len, _ = prompt_embeds.shape - prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) - prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) - - negative_prompt_embeds = None - if do_classifier_free_guidance: - negative_prompt = negative_prompt or "" - negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt - neg_text_inputs = self.tokenizer( - [self._prompt_clean(p) for p in negative_prompt], - padding="max_length", - max_length=max_sequence_length, - truncation=True, - add_special_tokens=True, - return_attention_mask=True, - return_tensors="pt", - ) - ids_neg, mask_neg = neg_text_inputs.input_ids, neg_text_inputs.attention_mask - seq_lens_neg = mask_neg.gt(0).sum(dim=1).long() - negative_prompt_embeds = self.text_encoder(ids_neg.to(device), mask_neg.to(device)).last_hidden_state - negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device) - negative_prompt_embeds = [u[:v] for u, v in zip(negative_prompt_embeds, seq_lens_neg)] - negative_prompt_embeds = torch.stack( - [ - torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) - for u in negative_prompt_embeds - ], - dim=0, - ) - negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_videos_per_prompt, 1) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) - - return prompt_embeds, negative_prompt_embeds - - @staticmethod - def _prompt_clean(text: str) -> str: - return " ".join(text.strip().split()) - - def prepare_latents( - self, - batch_size: int, - num_channels_latents: int, - height: int, - width: int, - num_frames: int, - dtype: torch.dtype | None, - device: torch.device | None, - generator: torch.Generator | list[torch.Generator] | None, - latents: torch.Tensor | None = None, - ) -> torch.Tensor: - """Prepare random latents for T2V mode.""" - if latents is not None: - return latents.to(device=device, dtype=dtype) - - num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 - shape = ( - batch_size, - num_channels_latents, - num_latent_frames, - int(height) // self.vae_scale_factor_spatial, - int(width) // self.vae_scale_factor_spatial, - ) - if isinstance(generator, list) and len(generator) != batch_size: - raise ValueError(f"Generator list length {len(generator)} does not match batch size {batch_size}.") - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - return latents - - def prepare_i2v_latents( - self, - image: torch.Tensor, - batch_size: int, - num_channels_latents: int, - height: int, - width: int, - num_frames: int, - dtype: torch.dtype | None, - device: torch.device | None, - generator: torch.Generator | list[torch.Generator] | None, - latents: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Prepare latents for I2V mode with image conditioning. - - Returns: - latents: Initial noise latents - latent_condition: Encoded first frame condition - first_frame_mask: Mask (0 for first frame, 1 for rest) - """ - num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 - latent_height = height // self.vae_scale_factor_spatial - latent_width = width // self.vae_scale_factor_spatial - - shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - latents = latents.to(device=device, dtype=dtype) - - # Prepare first frame condition - image = image.unsqueeze(2) # [batch, channels, 1, height, width] - image = image.to(device=device, dtype=self.vae.dtype) - - # Encode through VAE - latent_condition = retrieve_latents(self.vae.encode(image), sample_mode="argmax") - latent_condition = latent_condition.repeat(batch_size, 1, 1, 1, 1) - - # Normalize latents - latents_mean = ( - torch.tensor(self.vae.config.latents_mean) - .view(1, self.vae.config.z_dim, 1, 1, 1) - .to(latent_condition.device, latent_condition.dtype) - ) - latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( - latent_condition.device, latent_condition.dtype - ) - latent_condition = (latent_condition - latents_mean) * latents_std - latent_condition = latent_condition.to(dtype) - - # Create mask: 0 for first frame (condition), 1 for rest (to denoise) - first_frame_mask = torch.ones(1, 1, num_latent_frames, latent_height, latent_width, dtype=dtype, device=device) - first_frame_mask[:, :, 0] = 0 - - return latents, latent_condition, first_frame_mask - - def check_inputs( - self, - prompt, - negative_prompt, - image, - height, - width, - prompt_embeds=None, - negative_prompt_embeds=None, - ): - if height % 16 != 0 or width % 16 != 0: - raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.") - - if prompt is not None and prompt_embeds is not None: - raise ValueError("Cannot forward both `prompt` and `prompt_embeds`. Please provide only one.") - - if negative_prompt is not None and negative_prompt_embeds is not None: - raise ValueError( - "Cannot forward both `negative_prompt` and `negative_prompt_embeds`. Please provide only one." - ) - - if prompt is None and prompt_embeds is None: - raise ValueError("Provide either `prompt` or `prompt_embeds`.") - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Load weights using AutoWeightsLoader for vLLM integration.""" - loader = AutoWeightsLoader(self) - return loader.load_weights(weights) diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index b2dcec444..81889607a 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -21,9 +21,21 @@ from vllm.model_executor.layers.conv import Conv3dLayer from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.utils import ( + PPMissingLayer, + is_pp_missing_parameter, + make_empty_intermediate_tensors_factory, + make_layers, +) +from vllm.sequence import IntermediateTensors from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.distributed.parallel_state import ( + get_pipeline_parallel_world_size, + is_pipeline_first_stage, + is_pipeline_last_stage, +) from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, SequenceParallelOutput, @@ -871,14 +883,23 @@ class WanTransformer3DModel(nn.Module): inner_dim = num_attention_heads * attention_head_dim out_channels = out_channels or in_channels + if get_pipeline_parallel_world_size() == 0 and not is_pipeline_first_stage(): + raise RuntimeError( + "`initialize_model_parallel()` must be called before constructing `WanTransformer3DModel`" + ) + # 1. Patch & position embedding self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len) - self.patch_embedding = Conv3dLayer( - in_channels=in_channels, - out_channels=inner_dim, - kernel_size=patch_size, - stride=patch_size, - ) + # Patch embedding only on first PP stage; other stages receive hidden_states via P2P + if is_pipeline_first_stage(): + self.patch_embedding = Conv3dLayer( + in_channels=in_channels, + out_channels=inner_dim, + kernel_size=patch_size, + stride=patch_size, + ) + else: + self.patch_embedding = PPMissingLayer() # 2. Condition embeddings self.condition_embedder = WanTimeTextImageEmbedding( @@ -890,30 +911,39 @@ class WanTransformer3DModel(nn.Module): pos_embed_seq_len=pos_embed_seq_len, ) - # 3. Transformer blocks - self.blocks = nn.ModuleList( - [ - WanTransformerBlock( - inner_dim, - ffn_dim, - num_attention_heads, - eps, - added_kv_proj_dim, - cross_attn_norm, - quant_config=quant_config, - prefix=f"blocks.{layer_idx}", - ) - for layer_idx in range(num_layers) - ] + # 3. Transformer blocks — partitioned across PP stages via vLLM's `make_layers`. + # It computes the [start_layer, end_layer) slice for this rank and fills the remaining slots + # with PPMissingLayer so that weight names stay globally consistent. + self.start_layer, self.end_layer, self.blocks = make_layers( + num_layers, + lambda prefix: WanTransformerBlock( + inner_dim, + ffn_dim, + num_attention_heads, + eps, + added_kv_proj_dim, + cross_attn_norm, + quant_config=quant_config, + prefix=prefix, + ), + prefix="blocks", ) - # 4. Output norm & projection - self.norm_out = AdaLayerNorm(inner_dim, elementwise_affine=False, eps=eps) - self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size)) + # 4. Output norm & projection — only on the last PP stage + if is_pipeline_last_stage(): + self.norm_out = AdaLayerNorm(inner_dim, elementwise_affine=False, eps=eps) + self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size)) + else: + self.norm_out, self.proj_out = PPMissingLayer(), PPMissingLayer() # SP helper modules self.timestep_proj_prepare = TimestepProjPrepare() - self.output_scale_shift_prepare = OutputScaleShiftPrepare(inner_dim) + if is_pipeline_last_stage(): + self.output_scale_shift_prepare = OutputScaleShiftPrepare(inner_dim) + else: + self.output_scale_shift_prepare = PPMissingLayer() + + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(["hidden_states"], inner_dim) # ROPE helper self._cached_rope_emb = None @@ -930,9 +960,10 @@ class WanTransformer3DModel(nn.Module): timestep: torch.LongTensor, encoder_hidden_states: torch.Tensor, encoder_hidden_states_image: torch.Tensor | None = None, + intermediate_tensors: IntermediateTensors | None = None, return_dict: bool = True, attention_kwargs: dict[str, Any] | None = None, - ) -> torch.Tensor | Transformer2DModelOutput: + ) -> torch.Tensor | Transformer2DModelOutput | IntermediateTensors: batch_size, num_channels, num_frames, height, width = hidden_states.shape p_t, p_h, p_w = self.config.patch_size post_patch_num_frames = num_frames // p_t @@ -949,18 +980,25 @@ class WanTransformer3DModel(nn.Module): self._hidden_states_shape = hidden_states.shape self._cached_rope_emb = rotary_emb - # Patch embedding and flatten to sequence - # (hidden_states is sharded at blocks.0 input by _sp_plan) - hidden_states = self.patch_embedding(hidden_states) - hidden_states = hidden_states.flatten(2).transpose(1, 2) + if is_pipeline_first_stage(): + # Patch embedding and flatten to sequence + # (hidden_states is sharded at blocks.0 input by _sp_plan) + hidden_states = self.patch_embedding(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + else: + if intermediate_tensors is None: + raise RuntimeError("intermediate_tensors must be provided for non-first PP stages") + hidden_states = intermediate_tensors["hidden_states"] - # Handle timestep shape + # Handle timestep shape (2-D for TI2V, 1-D for T2V) if timestep.ndim == 2: ts_seq_len = timestep.shape[1] timestep = timestep.flatten() else: ts_seq_len = None + # Compute conditioning on all PP stages. + # Each stage needs temb/timestep_proj for scale-shift modulation in its local blocks. temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len ) @@ -996,14 +1034,13 @@ class WanTransformer3DModel(nn.Module): hidden_states_mask = None # Transformer blocks - for block in self.blocks: - hidden_states = block( - hidden_states, - encoder_hidden_states, - timestep_proj, - rotary_emb, - hidden_states_mask, - ) + for block in self.blocks[self.start_layer : self.end_layer]: + hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb, hidden_states_mask) + + if not is_pipeline_last_stage(): + # Non-last PP stage: hand the token sequence to the caller via IntermediateTensors. + # predict_noise will broadcast it to the next stage before calling that stage's forward. + return IntermediateTensors({"hidden_states": hidden_states}) # Output norm, projection & unpatchify shift, scale = self.output_scale_shift_prepare(temb) @@ -1074,7 +1111,8 @@ class WanTransformer3DModel(nn.Module): if weight_name not in original_name: continue lookup_name = original_name.replace(weight_name, param_name) - if lookup_name not in params_dict: + # Skip weights that belong to PP stages other than this one + if is_pp_missing_parameter(lookup_name, self) or lookup_name not in params_dict: break param = params_dict[lookup_name] weight_loader = param.weight_loader @@ -1098,6 +1136,10 @@ class WanTransformer3DModel(nn.Module): if modulation_alias in params_dict: lookup_name = modulation_alias + # Skip weights that belong to PP stages other than this one + if is_pp_missing_parameter(lookup_name, self): + continue + if lookup_name not in params_dict: logger.warning(f"Skipping weight {original_name} -> {lookup_name}") continue diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 687aa3bda..8a850f565 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1632,6 +1632,7 @@ class AsyncOmniEngine: data_parallel_size = normalized_kwargs.get("data_parallel_size") or 1 tensor_parallel_size = normalized_kwargs.get("tensor_parallel_size") or 1 cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1 + pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1 vae_patch_parallel_size = normalized_kwargs.get("vae_patch_parallel_size") or 1 enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False use_hsdp = normalized_kwargs.get("use_hsdp", False) From e2ed1c457455f8460182873111882b46829dc2df Mon Sep 17 00:00:00 2001 From: Daniel Huang Date: Mon, 18 May 2026 19:19:16 -0700 Subject: [PATCH 05/38] Disable sampler kernel for XPU test (#3718) Signed-off-by: Daniel Huang --- .buildkite/scripts/hardware_ci/run-xpu-test.sh | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.buildkite/scripts/hardware_ci/run-xpu-test.sh b/.buildkite/scripts/hardware_ci/run-xpu-test.sh index afed05852..ab8c2fd5d 100755 --- a/.buildkite/scripts/hardware_ci/run-xpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-xpu-test.sh @@ -52,5 +52,6 @@ time timeout -k 30 30m docker run \ pip install tblib==3.1.0 cd /workspace/vllm-omni pytest -v -s -m "core_model and xpu and B60" - pytest -v -s -m "advanced_model and xpu and B60" -k "not omni_expansion" + export VLLM_XPU_USE_SAMPLER_KERNEL=0 # NOTE: Remove this after vLLM v0.21.1 is merged. Fixes Qwen2-5 omni-expansion tests. + pytest -v -s -m "advanced_model and xpu and B60" ' From 2c6b1bb0c0b814aa562770737e8d0a6dd7c848f7 Mon Sep 17 00:00:00 2001 From: fan2956 Date: Tue, 19 May 2026 10:24:27 +0800 Subject: [PATCH 06/38] [Bugfix] Fix hunyuanimage3 dit quant storageshape mismatch error (#3694) Signed-off-by: fan2956 --- pyproject.toml | 1 + vllm_omni/platforms/npu/platform.py | 8 ++++++++ 2 files changed, 9 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 1cbf03dac..95e95fed8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -252,3 +252,4 @@ fullset = "fullset" Vai = "Vai" tockens = "tockens" CANN = "CANN" +ND = "ND" diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 8a30b4c7a..e371c73a6 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -31,6 +31,14 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): _omni_enum = OmniPlatformEnum.NPU dist_backend: str = "hccl" + @classmethod + def set_device(cls, device: torch.device) -> None: + super().set_device(device) + # Ascend quantized weights are converted from ND to FRACTAL_NZ + # after loading. Enable internal format so the NZ storage layout + # is preserved for fused NPU kernels. + torch.npu.config.allow_internal_format = True + @classmethod def get_omni_ar_worker_cls(cls) -> str: return "vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker" From 1fa734419ec6b578537aa5267c0d42f006499201 Mon Sep 17 00:00:00 2001 From: bjf-frz Date: Tue, 19 May 2026 11:30:33 +0800 Subject: [PATCH 07/38] [Refactor]Rename diffusion benchmark backend to endpoint (#3137) Signed-off-by: bjf-frz Signed-off-by: bjfwhite Co-authored-by: bjfwhite --- benchmarks/diffusion/README.md | 5 +- benchmarks/diffusion/backends.py | 34 +++- .../diffusion/diffusion_benchmark_serving.py | 81 ++++++---- .../qwen_image_serving_performance.md | 6 +- .../wan_2_2_serving_performance.md | 6 +- docs/.nav.yml | 3 +- ...on_chat_api.md => chat_completions_api.md} | 2 +- docs/serving/videos_api.md | 152 ++++++++++++++++++ .../examples/online_serving/glm_image.md | 2 +- .../examples/online_serving/image_to_image.md | 4 +- .../examples/online_serving/text_to_image.md | 4 +- recipes/Qwen/Qwen-Image.md | 2 +- .../test_diffusion_backends_metrics.py | 33 +++- .../scripts/diffusion_result_template.json | 5 +- .../perf/scripts/run_diffusion_benchmark.py | 54 ++++--- tests/test_generate_nightly_perf_excel.py | 5 + tests/test_generate_nightly_perf_html.py | 3 + tools/nightly/generate_nightly_perf_excel.py | 11 +- tools/nightly/generate_nightly_perf_html.py | 27 +++- 19 files changed, 352 insertions(+), 87 deletions(-) rename docs/serving/{diffusion_chat_api.md => chat_completions_api.md} (98%) create mode 100644 docs/serving/videos_api.md diff --git a/benchmarks/diffusion/README.md b/benchmarks/diffusion/README.md index 16a86c88f..54ed64261 100644 --- a/benchmarks/diffusion/README.md +++ b/benchmarks/diffusion/README.md @@ -30,7 +30,7 @@ python3 benchmarks/diffusion/diffusion_benchmark_serving.py \ **Notes** -- The benchmark talks to `http://:/v1/chat/completions`. +- By default, image tasks talk to `http://:/v1/chat/completions`; video tasks talk to `/v1/videos`. - If you run the server on another host or port, pass `--base-url` accordingly. ## 2. Supported Datasets @@ -95,8 +95,9 @@ You can point to your own trace using `--dataset-path`. ### Basic flags -- `--base-url`: Server address (the script calls `.../v1/chat/completions`). +- `--base-url`: Server address; `--endpoint` selects the path appended to this base URL. - `--model`: The OpenAI-compatible `model` field. +- `--endpoint`: API endpoint path. Leading `/` is optional, e.g. `/v1/videos` or `v1/videos`. - `--task`: Task type (e.g., `t2i`, `t2v`, `i2i`, `i2v`). - `--dataset`: Dataset mode (`vbench` / `trace` / `random`). - `--num-prompts`: Number of requests to send. diff --git a/benchmarks/diffusion/backends.py b/benchmarks/diffusion/backends.py index d33160f13..b7e533b2d 100644 --- a/benchmarks/diffusion/backends.py +++ b/benchmarks/diffusion/backends.py @@ -151,7 +151,7 @@ async def async_request_chat_completions( return output -async def async_request_openai_images( +async def async_request_openai_image_generations( input: RequestFuncInput, session: aiohttp.ClientSession, pbar: tqdm | None = None, @@ -348,12 +348,38 @@ async def async_request_v1_videos( return output +LEGACY_BACKEND_ENDPOINT_ALIASES = { + "vllm-omni": "/v1/chat/completions", + "openai": "/v1/images/generations", +} + + +def normalize_endpoint(value: str) -> str: + endpoint = str(value).strip() + if not endpoint: + raise ValueError("endpoint must not be empty.") + endpoint = LEGACY_BACKEND_ENDPOINT_ALIASES.get( + endpoint, + LEGACY_BACKEND_ENDPOINT_ALIASES.get(endpoint.lstrip("/"), endpoint), + ) + if not endpoint.startswith("/"): + endpoint = f"/{endpoint}" + return endpoint + + +def endpoint_filename_token(value: str) -> str: + token = normalize_endpoint(value).lstrip("/") + for bad in ("/", "\\", ":", "*", "?", '"', "<", ">", "|"): + token = token.replace(bad, "_") + return token or "endpoint" + + backends_function_mapping = { "2i": { - "vllm-omni": (async_request_chat_completions, "/v1/chat/completions"), - "openai": (async_request_openai_images, "/v1/images/generations"), + "/v1/chat/completions": (async_request_chat_completions, "/v1/chat/completions"), + "/v1/images/generations": (async_request_openai_image_generations, "/v1/images/generations"), }, "2v": { - "v1/videos": (async_request_v1_videos, "/v1/videos"), + "/v1/videos": (async_request_v1_videos, "/v1/videos"), }, } diff --git a/benchmarks/diffusion/diffusion_benchmark_serving.py b/benchmarks/diffusion/diffusion_benchmark_serving.py index f2b77ff8d..c59844f65 100644 --- a/benchmarks/diffusion/diffusion_benchmark_serving.py +++ b/benchmarks/diffusion/diffusion_benchmark_serving.py @@ -6,31 +6,31 @@ Benchmark online serving for diffusion models (Image/Video Generation). If you want to use i2v, i2i dataset, you should `uv pip install gdown` first -Supports multiple backends: - - vllm-omni: Uses /v1/chat/completions endpoint (default) - - openai: Uses /v1/images/generations endpoint - - v1/videos: Use /v1/videos endpoint +Supports multiple endpoints: + - /v1/chat/completions: OpenAI chat-compatible image requests + - /v1/images/generations: OpenAI image generation requests + - /v1/videos: Async video jobs Usage: - # Video (v1/videos backend) + # Video (/v1/videos endpoint) t2v: python3 benchmarks/diffusion/diffusion_benchmark_serving.py \ - --backend v1/videos --dataset vbench --task t2v --num-prompts 10 \ + --endpoint /v1/videos --dataset vbench --task t2v --num-prompts 10 \ --height 480 --width 640 --fps 16 --num-frames 80 i2v: python3 benchmarks/diffusion/diffusion_benchmark_serving.py \ - --backend v1/videos --dataset vbench --task i2v --num-prompts 10 + --endpoint /v1/videos --dataset vbench --task i2v --num-prompts 10 - # Image (vllm-omni backend) + # Image (/v1/chat/completions endpoint) t2i: python3 benchmarks/diffusion/diffusion_benchmark_serving.py \ - --backend vllm-omni --dataset vbench --task t2i --num-prompts 10 \ + --endpoint /v1/chat/completions --dataset vbench --task t2i --num-prompts 10 \ --height 1024 --width 1024 python3 benchmarks/diffusion/diffusion_benchmark_serving.py \ - --backend vllm-omni --dataset random --task t2i --num-prompts 1 \ + --endpoint /v1/chat/completions --dataset random --task t2i --num-prompts 1 \ --max-concurrency 1 --enable-negative-prompt \ --random-request-config '[ {"width":512,"height":512,"num_inference_steps":20,"weight":0.15}, @@ -41,18 +41,18 @@ Usage: i2i: python3 benchmarks/diffusion/diffusion_benchmark_serving.py \ - --backend vllm-omni --dataset vbench --task i2i --num-prompts 10 + --endpoint /v1/chat/completions --dataset vbench --task i2i --num-prompts 10 - # Image (openai backend) + # Image (/v1/images/generations endpoint) t2i: python3 benchmarks/diffusion/diffusion_benchmark_serving.py \ - --backend openai --dataset vbench --task t2i --num-prompts 10 \ + --endpoint /v1/images/generations --dataset vbench --task t2i --num-prompts 10 \ --height 1024 --width 1024 --port 3000 # Video (v1/videos) t2v: python3 benchmarks/diffusion/diffusion_benchmark_serving.py \ - --backend v1/videos --dataset random --task t2v --num-prompts 1 \ + --endpoint /v1/videos --dataset random --task t2v --num-prompts 1 \ --max-concurrency 1 --enable-negative-prompt \ --random-request-config '[ {"width":854,"height":480,"num_inference_steps":18,"num_frames":120,"fps":24,"weight":1} @@ -80,7 +80,12 @@ from typing import Any import aiohttp import numpy as np import requests -from backends import RequestFuncInput, RequestFuncOutput, backends_function_mapping +from backends import ( + RequestFuncInput, + RequestFuncOutput, + backends_function_mapping, + normalize_endpoint, +) from PIL import Image from tqdm.asyncio import tqdm @@ -883,6 +888,14 @@ def wait_for_service(base_url: str, timeout: int = 120) -> None: time.sleep(1) +def _default_endpoint_for_task(task: str) -> str: + if task in {"t2v", "i2v", "ti2v"}: + return "/v1/videos" + if task in {"t2i", "i2i", "ti2i"}: + return "/v1/chat/completions" + raise ValueError(f"Unsupported task for endpoint resolution: {task}") + + async def benchmark(args): # Construct base_url if not provided if args.base_url is None: @@ -902,18 +915,23 @@ async def benchmark(args): f"Valid image tasks: {sorted(IMAGE_TASKS)}" ) - valid_backends = sorted(backends_function_mapping[task_type].keys()) + raw_endpoint = args.endpoint if args.endpoint is not None else args.backend + if raw_endpoint is None: + raw_endpoint = _default_endpoint_for_task(args.task) + args.endpoint = normalize_endpoint(raw_endpoint) - if args.backend not in valid_backends: + valid_endpoints = sorted(backends_function_mapping[task_type].keys()) + + if args.endpoint not in valid_endpoints: logger.error( - f"Invalid backend '{args.backend}' for task '{args.task}' (task type: '{task_type}').\n" - f"Valid backends for this task type: {valid_backends}\n" - f"Example usage: --task {args.task} --backend {valid_backends[0]}" + f"Invalid endpoint '{args.endpoint}' for task '{args.task}' (task type: '{task_type}').\n" + f"Valid endpoints for this task type: {valid_endpoints}\n" + f"Example usage: --task {args.task} --endpoint {valid_endpoints[0]}" ) - raise ValueError("Backend validation failed. See log above for valid options.") + raise ValueError("Endpoint validation failed. See log above for valid options.") - # Setup API URL and request function based on backend - request_func, api_url = backends_function_mapping[task_type][args.backend] + # Setup API URL and request function based on endpoint. + request_func, api_url = backends_function_mapping[task_type][args.endpoint] api_url = f"{args.base_url}{api_url}" if args.dataset == "vbench": @@ -976,15 +994,15 @@ async def benchmark(args): metrics = calculate_metrics(outputs, total_duration, requests_list, args, args.slo) # Add configuration info to metrics for JSON output - metrics["backend"] = args.backend + metrics["endpoint"] = args.endpoint metrics["model"] = args.model metrics["dataset"] = args.dataset metrics["task"] = args.task - print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=60, c="=")) + print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="=")) # Section 1: Configuration - print("{:<40} {:<15}".format("Backend:", args.backend)) + print("{:<40} {:<15}".format("Endpoint:", args.endpoint)) print("{:<40} {:<15}".format("Model:", args.model)) print("{:<40} {:<15}".format("Dataset:", args.dataset)) print("{:<40} {:<15}".format("Task:", args.task)) @@ -1047,12 +1065,17 @@ if __name__ == "__main__": parser.add_argument("--host", type=str, default="localhost", help="Server host.") parser.add_argument("--port", type=int, default=8091, help="Server port.") parser.add_argument("--model", type=str, default="default", help="Model name.") + parser.add_argument( + "--endpoint", + type=str, + default=None, + help=("Endpoint path to target. Leading '/' is optional, e.g. /v1/videos or v1/videos."), + ) parser.add_argument( "--backend", type=str, - default="vllm-omni", - choices=["vllm-omni", "openai", "v1/videos"], - help="Backend to target the benchmark to.", + default=None, + help=argparse.SUPPRESS, ) parser.add_argument( "--dataset", diff --git a/benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md b/benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md index 172a8ed05..6d4a3c405 100644 --- a/benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md +++ b/benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md @@ -58,7 +58,7 @@ Record these parameters when reporting performance results. ```bash python benchmarks/diffusion/diffusion_benchmark_serving.py \ - --backend vllm-omni \ + --endpoint /v1/chat/completions \ --dataset \ --task t2i \ --num-prompts \ @@ -71,7 +71,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \ | Parameter | Description | | ---------------------- | --------------------------------- | -| `--backend` | Serving backend (use `vllm-omni`) | +| `--endpoint` | API endpoint (use `/v1/chat/completions`; leading `/` is optional) | | `--dataset` | Dataset name (`random` or custom) | | `--task` | Task type (e.g., `t2i`) | | `--num-prompts` | Total number of requests | @@ -129,7 +129,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \ ```bash python benchmarks/diffusion/diffusion_benchmark_serving.py \ - --backend vllm-omni \ + --endpoint /v1/chat/completions \ --dataset random \ --task t2i \ --num-prompts 1 \ diff --git a/benchmarks/diffusion/performance_dashboard/wan_2_2_serving_performance.md b/benchmarks/diffusion/performance_dashboard/wan_2_2_serving_performance.md index 9d6c40ece..d9b77fd99 100644 --- a/benchmarks/diffusion/performance_dashboard/wan_2_2_serving_performance.md +++ b/benchmarks/diffusion/performance_dashboard/wan_2_2_serving_performance.md @@ -53,7 +53,7 @@ Record these parameters when reporting performance results. ```bash python benchmarks/diffusion/diffusion_benchmark_serving.py \ - --backend v1/videos \ + --endpoint /v1/videos \ --dataset \ --task t2v \ --num-prompts \ @@ -66,7 +66,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \ | Parameter | Description | | ---------------------- | --------------------------------- | -| `--backend` | Serving backend (use `v1/videos`) | +| `--endpoint` | API endpoint (use `/v1/videos`; leading `/` is optional) | | `--dataset` | Dataset name (`random` or custom) | | `--task` | Task type (e.g., `t2v`) | | `--num-prompts` | Total number of requests | @@ -120,7 +120,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \ ```bash python benchmarks/diffusion/diffusion_benchmark_serving.py \ - --backend v1/videos \ + --endpoint /v1/videos \ --dataset random \ --task t2v \ --num-prompts 1 \ diff --git a/docs/.nav.yml b/docs/.nav.yml index 55283f0e8..a60a4368d 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -6,10 +6,11 @@ nav: - getting_started/installation/* - Serving: - OpenAI-Compatible API: - - Diffusion Chat API: serving/diffusion_chat_api.md + - Chat Completions API: serving/chat_completions_api.md - Image Generation: serving/image_generation_api.md - Image Edit: serving/image_edit_api.md - Text to Speech: serving/speech_api.md + - Videos: serving/videos_api.md - Audio Generation: serving/audio_generate_api.md - Streaming Video Input: serving/video_stream_api.md - Examples: diff --git a/docs/serving/diffusion_chat_api.md b/docs/serving/chat_completions_api.md similarity index 98% rename from docs/serving/diffusion_chat_api.md rename to docs/serving/chat_completions_api.md index d0e2990ad..d8f852f99 100644 --- a/docs/serving/diffusion_chat_api.md +++ b/docs/serving/chat_completions_api.md @@ -1,4 +1,4 @@ -# Diffusion Chat Completions API +# Chat Completions API vLLM-Omni supports generating and editing images via the `/v1/chat/completions` endpoint using diffusion models. This page explains how to pass generation diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md new file mode 100644 index 000000000..a8caee013 --- /dev/null +++ b/docs/serving/videos_api.md @@ -0,0 +1,152 @@ +# Videos API + +vLLM-Omni provides an OpenAI-compatible video generation API for diffusion +video models. The API supports asynchronous video jobs through `/v1/videos` and +a synchronous benchmark-oriented endpoint through `/v1/videos/sync`. + +Each server instance runs a single model specified at startup with +`vllm serve --omni`. + +## Quick Start + +### Start the Server + +```bash +vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 +``` + +### Create a Video Job + +```bash +create_response=$(curl -s http://localhost:8091/v1/videos \ + -F "prompt=A cinematic tracking shot of a mountain lake at sunrise" \ + -F "width=1280" \ + -F "height=720" \ + -F "num_frames=80" \ + -F "fps=16" \ + -F "num_inference_steps=40") + +video_id=$(echo "${create_response}" | jq -r '.id') +``` + +### Poll and Download + +```bash +curl -s "http://localhost:8091/v1/videos/${video_id}" | jq . +curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4 +``` + +## API Reference + +### Endpoints + +| Endpoint | Method | Description | +|----------|--------|-------------| +| `/v1/videos` | `POST` | Create an asynchronous video generation job | +| `/v1/videos/sync` | `POST` | Generate a video synchronously and return raw video bytes | +| `/v1/videos/{video_id}` | `GET` | Retrieve job status and metadata | +| `/v1/videos` | `GET` | List stored video jobs | +| `/v1/videos/{video_id}/content` | `GET` | Download generated video content | +| `/v1/videos/{video_id}` | `DELETE` | Delete a video job and stored output | + +### Request Parameters + +`POST /v1/videos` and `POST /v1/videos/sync` accept `multipart/form-data`. + +#### OpenAI-style fields + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `prompt` | string | **required** | Text prompt for video generation | +| `model` | string | server's model | Optional model name | +| `seconds` | string | null | Requested clip duration in seconds | +| `size` | string | null | Requested output size in `WIDTHxHEIGHT` format | +| `user` | string | null | Optional user identifier | + +#### vLLM-Omni extension fields + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `input_reference` | file | null | Uploaded reference image for image-to-video requests | +| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` | +| `width` | integer | model default | Output video width | +| `height` | integer | model default | Output video height | +| `num_frames` | integer | model default | Number of generated frames | +| `fps` | integer | model default | Output frames per second | +| `num_inference_steps` | integer | model default | Number of diffusion steps | +| `guidance_scale` | number | null | CFG guidance scale for the low-noise stage | +| `guidance_scale_2` | number | null | CFG guidance scale for the high-noise stage | +| `boundary_ratio` | number | null | Boundary split ratio for multi-stage denoising | +| `flow_shift` | number | null | Scheduler flow-shift value | +| `true_cfg_scale` | number | null | True CFG scale when supported by the model | +| `seed` | integer | null | Random seed for reproducibility | +| `negative_prompt` | string | null | Text describing what to avoid in the generated video | +| `enable_frame_interpolation` | boolean | null | Enable post-generation frame interpolation | +| `frame_interpolation_exp` | integer | null | Interpolation exponent; `1=2x`, `2=4x`, and so on | +| `frame_interpolation_scale` | number | null | RIFE inference scale | +| `frame_interpolation_model_path` | string | null | Local path or Hugging Face repo for the interpolation model | +| `lora` | string | null | JSON-encoded LoRA configuration object | +| `extra_params` | string | null | JSON-encoded object for additional model-specific parameters | + +### Create Response + +`POST /v1/videos` returns a job record: + +```json +{ + "id": "video-123", + "status": "queued", + "created_at": 1701234567 +} +``` + +The final content is available from `/v1/videos/{video_id}/content` after the +job status becomes `completed`. + +### Synchronous Response + +`POST /v1/videos/sync` blocks until generation finishes and returns raw video +bytes. It is useful for benchmarks and simple scripts that do not need job +storage or polling. + +## Examples + +### Image-to-Video + +```bash +curl -s http://localhost:8091/v1/videos \ + -F "prompt=animate this image with subtle camera movement" \ + -F "input_reference=@input.png" \ + -F "width=1280" \ + -F "height=720" \ + -F "num_frames=80" \ + -F "fps=16" +``` + +### Synchronous Generation + +```bash +curl -X POST http://localhost:8091/v1/videos/sync \ + -F "prompt=A small robot walking through a neon city" \ + -F "width=854" \ + -F "height=480" \ + -F "num_frames=80" \ + -F "fps=16" \ + -o output.mp4 +``` + +## Storage + +Set `VLLM_OMNI_STORAGE_PATH` to control where asynchronous video outputs are +stored: + +```bash +export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos +``` + +## Model-Specific Examples + +For complete text-to-video and image-to-video walkthroughs, see: + +- [Text-to-Video](../user_guide/examples/online_serving/text_to_video.md) +- [Image-to-Video](../user_guide/examples/online_serving/image_to_video.md) diff --git a/docs/user_guide/examples/online_serving/glm_image.md b/docs/user_guide/examples/online_serving/glm_image.md index 4cc49e846..410b1180a 100644 --- a/docs/user_guide/examples/online_serving/glm_image.md +++ b/docs/user_guide/examples/online_serving/glm_image.md @@ -101,7 +101,7 @@ guides. When using `/v1/chat/completions`, pass these inside `extra_body` in the curl JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the -[Diffusion Chat API guide](../../../serving/diffusion_chat_api.md)). +[Chat Completions API guide](../../../serving/chat_completions_api.md)). When using the dedicated [`/v1/images/generations`](../../../serving/image_generation_api.md) or [`/v1/images/edits`](../../../serving/image_edit_api.md) endpoints, pass the supported generation controls as top-level fields directly. For image diff --git a/docs/user_guide/examples/online_serving/image_to_image.md b/docs/user_guide/examples/online_serving/image_to_image.md index f40173c4f..89048ae71 100644 --- a/docs/user_guide/examples/online_serving/image_to_image.md +++ b/docs/user_guide/examples/online_serving/image_to_image.md @@ -208,7 +208,7 @@ Wrap generation parameters inside `extra_body` in the request JSON: ``` For details on how generation parameters are handled across different clients, see the - [Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md). + [Chat Completions API guide](../../../../serving/chat_completions_api.md). ### Layered Image Generation (Qwen-Image-Layered) @@ -361,7 +361,7 @@ Provide multiple images in `content` (order matters): When using `/v1/chat/completions`, pass these inside `extra_body` in the curl JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the -[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md)). +[Chat Completions API guide](../../../../serving/chat_completions_api.md)). When using the dedicated [`/v1/images/edits`](../../../../serving/image_edit_api.md) endpoint, pass the supported generation controls as top-level form fields directly. For image dimensions and count, use `size` and `n` rather than diff --git a/docs/user_guide/examples/online_serving/text_to_image.md b/docs/user_guide/examples/online_serving/text_to_image.md index 69c1480e3..43fb96594 100644 --- a/docs/user_guide/examples/online_serving/text_to_image.md +++ b/docs/user_guide/examples/online_serving/text_to_image.md @@ -216,7 +216,7 @@ Wrap generation parameters inside `extra_body` in the request JSON: ``` For details on how generation parameters are handled across different clients, see the - [Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md). + [Chat Completions API guide](../../../../serving/chat_completions_api.md). ### Multimodal Input (Text + Structured Content) @@ -237,7 +237,7 @@ Wrap generation parameters inside `extra_body` in the request JSON: When using `/v1/chat/completions`, pass these inside `extra_body` in the curl JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the -[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md)). +[Chat Completions API guide](../../../../serving/chat_completions_api.md)). When using the dedicated [`/v1/images/generations`](../../../../serving/image_generation_api.md) endpoint, pass the supported generation controls as top-level JSON fields directly. For image dimensions and count, use `size` and `n` rather than diff --git a/recipes/Qwen/Qwen-Image.md b/recipes/Qwen/Qwen-Image.md index d9d60bf0d..327160973 100644 --- a/recipes/Qwen/Qwen-Image.md +++ b/recipes/Qwen/Qwen-Image.md @@ -104,7 +104,7 @@ To replay the batching benefit with matched warmup: ```bash python benchmarks/diffusion/diffusion_benchmark_serving.py \ - --backend vllm-omni \ + --endpoint /v1/chat/completions \ --dataset vbench \ --task t2i \ --model Qwen/Qwen-Image \ diff --git a/tests/benchmarks/test_diffusion_backends_metrics.py b/tests/benchmarks/test_diffusion_backends_metrics.py index 2d51d0f1d..a89ab06be 100644 --- a/tests/benchmarks/test_diffusion_backends_metrics.py +++ b/tests/benchmarks/test_diffusion_backends_metrics.py @@ -1,6 +1,11 @@ import pytest -from benchmarks.diffusion.backends import RequestFuncInput, async_request_chat_completions +from benchmarks.diffusion.backends import ( + RequestFuncInput, + async_request_chat_completions, + endpoint_filename_token, + normalize_endpoint, +) class _MockResponse: @@ -29,6 +34,32 @@ class _MockSession: return _MockResponse(self._payload) +@pytest.mark.core_model +@pytest.mark.benchmark +@pytest.mark.cpu +def test_endpoint_normalization_accepts_optional_leading_slash(): + assert normalize_endpoint("v1/videos") == "/v1/videos" + assert normalize_endpoint("/v1/videos") == "/v1/videos" + assert normalize_endpoint("v1/chat/completions") == "/v1/chat/completions" + assert normalize_endpoint("v1/images/generations") == "/v1/images/generations" + + +@pytest.mark.core_model +@pytest.mark.benchmark +@pytest.mark.cpu +def test_endpoint_normalization_accepts_legacy_backend_aliases(): + assert normalize_endpoint("vllm-omni") == "/v1/chat/completions" + assert normalize_endpoint("openai") == "/v1/images/generations" + + +@pytest.mark.core_model +@pytest.mark.benchmark +@pytest.mark.cpu +def test_endpoint_filename_token_drops_leading_slash(): + assert endpoint_filename_token("/v1/videos") == "v1_videos" + assert endpoint_filename_token("v1/chat/completions") == "v1_chat_completions" + + @pytest.mark.core_model @pytest.mark.benchmark @pytest.mark.cpu diff --git a/tests/dfx/perf/scripts/diffusion_result_template.json b/tests/dfx/perf/scripts/diffusion_result_template.json index 86bdf1bc7..ba16a993d 100644 --- a/tests/dfx/perf/scripts/diffusion_result_template.json +++ b/tests/dfx/perf/scripts/diffusion_result_template.json @@ -1,7 +1,7 @@ [ { "test_name": null, - "backend": null, + "endpoint": null, "timestamp": null, "server_params": { "model": null, @@ -43,7 +43,7 @@ "stage_durations_mean": {}, "stage_durations_p50": {}, "stage_durations_p99": {}, - "backend": null, + "endpoint": null, "model": null, "dataset": null, "task": null @@ -51,6 +51,7 @@ "log_file": null, "Model": null, "Framework": null, + "API Endpoint": null, "Hardware": null, "Deployment": null, "Task": null, diff --git a/tests/dfx/perf/scripts/run_diffusion_benchmark.py b/tests/dfx/perf/scripts/run_diffusion_benchmark.py index 639f03623..9059904e4 100644 --- a/tests/dfx/perf/scripts/run_diffusion_benchmark.py +++ b/tests/dfx/perf/scripts/run_diffusion_benchmark.py @@ -5,16 +5,15 @@ This runner separates two concepts: 1. ``server_type``: how the serving process is started. Currently only ``vllm-omni`` is supported here. -2. ``benchmark_backend``: which serving API the benchmark client calls. - Examples: ``vllm-omni`` for ``/v1/chat/completions`` and ``v1/videos`` - for async video jobs. +2. ``benchmark_endpoint``: which serving API the benchmark client calls. + Examples: ``/v1/chat/completions`` and ``/v1/videos``. A config JSON file is REQUIRED via --test-config-file: pytest run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json All benchmark results for a session are consolidated into a single JSON file under BENCHMARK_RESULT_DIR (override via the DIFFUSION_BENCHMARK_DIR environment variable). -Each entry in the file contains the test metadata (test_name, backend, benchmark_params, +Each entry in the file contains the test metadata (test_name, endpoint, benchmark_params, timestamp) together with the raw metrics returned by the benchmark script. """ @@ -33,6 +32,8 @@ from typing import Any, cast import psutil import pytest +from benchmarks.diffusion.backends import endpoint_filename_token, normalize_endpoint + pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" @@ -400,7 +401,7 @@ def _unique_server_params(configs: list[dict[str, Any]]) -> list[dict[str, Any]] "model": cfg["server_params"]["model"], "serve_args_dict": serve_args_dict, "serve_args": _build_serve_args(serve_args_dict), - "benchmark_backend": cfg.get("benchmark_backend"), + "benchmark_endpoint": cfg.get("benchmark_endpoint", cfg.get("benchmark_backend")), "server_params": cfg["server_params"], } ) @@ -482,7 +483,7 @@ def run_benchmark( model: str, params: dict[str, Any], test_name: str, - backend: str = "vllm-omni", + endpoint: str = "/v1/chat/completions", server_cfg: dict[str, Any] | None = None, source_file: str = "", ) -> dict[str, Any]: @@ -490,18 +491,19 @@ def run_benchmark( The raw metrics are written to a temporary file by the subprocess. After the run completes the metrics are merged with full metadata (test_name, - backend, benchmark_params, timestamp, flat reporting fields) and appended + endpoint, benchmark_params, timestamp, flat reporting fields) and appended to the session-wide aggregated JSON file (AGGREGATED_RESULT_FILE). The temporary file is removed afterwards. Subprocess stdout/stderr are tee'd to a .log file under BENCHMARK_RESULT_DIR/logs/; its path is stored in the record. """ timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") + endpoint = normalize_endpoint(endpoint) log_dir = BENCHMARK_RESULT_DIR / "logs" log_dir.mkdir(parents=True, exist_ok=True) - backend_label = backend.replace("/", "_") - log_file = log_dir / f"{test_name}_{backend_label}_{timestamp}.log" + endpoint_label = endpoint_filename_token(endpoint) + log_file = log_dir / f"{test_name}_{endpoint_label}_{timestamp}.log" with tempfile.NamedTemporaryFile(mode="w", suffix=".json", prefix="diffusion_bench_tmp_", delete=False) as tmp: tmp_result_file = Path(tmp.name) @@ -517,8 +519,8 @@ def run_benchmark( str(port), "--model", model, - "--backend", - backend, + "--endpoint", + endpoint, "--dataset", params.get("dataset", "random"), "--task", @@ -545,7 +547,7 @@ def run_benchmark( # cause truncated or out-of-order log output when stdout is piped). cmd = [cmd[0], "-u"] + cmd[1:] - print(f"\nRunning benchmark (backend={backend}): {' '.join(cmd)}") + print(f"\nRunning benchmark (endpoint={endpoint}): {' '.join(cmd)}") print(f" Log file: {log_file}") # Redirect stdout + stderr directly to the log file at the OS level @@ -597,7 +599,7 @@ def run_benchmark( record: dict[str, Any] = { "test_name": test_name, - "backend": backend, + "endpoint": endpoint, "timestamp": timestamp, "server_params": server_cfg.get("server_params"), "benchmark_params": params, @@ -605,7 +607,7 @@ def run_benchmark( "log_file": str(log_file), "Model": model, "Framework": server_type, - "API Backend": backend, + "API Endpoint": endpoint, "Hardware": "", "Deployment": "", "Task": params.get("task", "t2i"), @@ -669,21 +671,21 @@ def assert_result(result: dict[str, Any], params: dict[str, Any]) -> None: assert current <= threshold, f"{metric}: {current:.4f} > baseline {threshold}" -def _default_benchmark_backend_for_task(task: str) -> str: - """Return the default client-side benchmark backend for a diffusion task.""" +def _default_benchmark_endpoint_for_task(task: str) -> str: + """Return the default client-side benchmark endpoint for a diffusion task.""" if task in {"t2v", "i2v", "ti2v"}: - return "v1/videos" + return "/v1/videos" if task in {"t2i", "i2i", "ti2i"}: - return "vllm-omni" - raise ValueError(f"Unsupported task for benchmark backend resolution: {task}") + return "/v1/chat/completions" + raise ValueError(f"Unsupported task for benchmark endpoint resolution: {task}") -def _resolve_benchmark_backend(server_cfg: dict[str, Any], params: dict[str, Any]) -> str: +def _resolve_benchmark_endpoint(server_cfg: dict[str, Any], params: dict[str, Any]) -> str: """Resolve which serving API the benchmark client should call.""" - configured = server_cfg.get("benchmark_backend") + configured = server_cfg.get("benchmark_endpoint") if configured: - return cast(str, configured) - return _default_benchmark_backend_for_task(cast(str, params.get("task", "t2i"))) + return normalize_endpoint(cast(str, configured)) + return _default_benchmark_endpoint_for_task(cast(str, params.get("task", "t2i"))) # --------------------------------------------------------------------------- @@ -711,7 +713,7 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params): test_name = benchmark_params["test_name"] params = benchmark_params["params"] server_cfg = getattr(diffusion_server, "server_cfg", {}) - backend = _resolve_benchmark_backend(server_cfg, params) + endpoint = _resolve_benchmark_endpoint(server_cfg, params) result = run_benchmark( host=diffusion_server.host, @@ -719,13 +721,13 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params): model=diffusion_server.model, params=params, test_name=test_name, - backend=backend, + endpoint=endpoint, server_cfg=server_cfg, source_file=cast(str, CONFIG_FILE_PATH), ) print(f"\n{'=' * 60}") - print(f"Results for {test_name} (server={diffusion_server.server_type}, backend={backend}):") + print(f"Results for {test_name} (server={diffusion_server.server_type}, endpoint={endpoint}):") for key in ( "throughput_qps", "latency_mean", diff --git a/tests/test_generate_nightly_perf_excel.py b/tests/test_generate_nightly_perf_excel.py index 9b05d6de0..6658b89b8 100644 --- a/tests/test_generate_nightly_perf_excel.py +++ b/tests/test_generate_nightly_perf_excel.py @@ -35,6 +35,8 @@ def test_generate_excel_report_with_perf_templates(tmp_path: Path): omni_record = json.loads(omni_template_path.read_text(encoding="utf-8")) diffusion_records = json.loads(diffusion_template_path.read_text(encoding="utf-8")) + diffusion_records[0]["endpoint"] = "/v1/videos" + diffusion_records[0]["result"]["endpoint"] = "/v1/videos" input_dir = tmp_path / "input" diffusion_input_dir = tmp_path / "diffusion_input" @@ -69,3 +71,6 @@ def test_generate_excel_report_with_perf_templates(tmp_path: Path): ws_omni_summary = wb["omni_summary"] assert _cell_value_by_header(ws_omni_summary, "commit_sha") == "test_commit_sha" assert _cell_value_by_header(ws_omni_summary, "build_id") == "test_build_id" + + ws_diff_summary = wb["diffusion_summary"] + assert _cell_value_by_header(ws_diff_summary, "endpoint") == "/v1/videos" diff --git a/tests/test_generate_nightly_perf_html.py b/tests/test_generate_nightly_perf_html.py index 4e77eb3ad..7a05165da 100644 --- a/tests/test_generate_nightly_perf_html.py +++ b/tests/test_generate_nightly_perf_html.py @@ -28,6 +28,8 @@ def test_generate_html_report_with_perf_templates(tmp_path: Path): omni_record = json.loads(omni_template_path.read_text(encoding="utf-8")) diffusion_records = json.loads(diffusion_template_path.read_text(encoding="utf-8")) + diffusion_records[0]["endpoint"] = "/v1/videos" + diffusion_records[0]["result"]["endpoint"] = "/v1/videos" input_dir = tmp_path / "input" diffusion_input_dir = tmp_path / "diffusion_input" @@ -52,3 +54,4 @@ def test_generate_html_report_with_perf_templates(tmp_path: Path): assert "Omni records 1" in html assert f"Diffusion records {len(diffusion_records)}" in html assert "const DIFF_DATA =" in html + assert '"endpoint": "/v1/videos"' in html diff --git a/tools/nightly/generate_nightly_perf_excel.py b/tools/nightly/generate_nightly_perf_excel.py index 6ba1d1eef..e76dc4346 100644 --- a/tools/nightly/generate_nightly_perf_excel.py +++ b/tools/nightly/generate_nightly_perf_excel.py @@ -18,6 +18,8 @@ from openpyxl import Workbook from openpyxl.styles import PatternFill from openpyxl.utils import get_column_letter +from benchmarks.diffusion.backends import normalize_endpoint + LOGGER = logging.getLogger(__name__) GREY_BLOCK_FILL = PatternFill(start_color="D3D3D3", fill_type="solid") @@ -60,7 +62,7 @@ DIFFUSION_SUMMARY_COLUMNS: tuple[str, ...] = ( "date", "test_name", "model", - "backend", + "endpoint", "dataset", "task", "completed_requests", @@ -528,6 +530,12 @@ def _process_diffusion_record(record: dict[str, Any]) -> dict[str, Any]: """Normalize a diffusion record by merging `result` and flattening stage metrics.""" flat = record.copy() flat.update(flat.pop("result", {})) + if flat.get("endpoint"): + flat["endpoint"] = normalize_endpoint(flat["endpoint"]) + elif flat.get("backend"): + flat["endpoint"] = normalize_endpoint(flat["backend"]) + flat.pop("backend", None) + flat.pop("API Backend", None) flat = _flatten_stage_durations(flat) flat.pop("benchmark_params", None) flat.pop("server_params", None) @@ -752,6 +760,7 @@ _OMNI_SUMMARY_WIDTHS = { _DIFFUSION_SUMMARY_WIDTHS = { "test_name": 30, "model": 15, + "endpoint": 24, } diff --git a/tools/nightly/generate_nightly_perf_html.py b/tools/nightly/generate_nightly_perf_html.py index a50a46255..73e245ab1 100644 --- a/tools/nightly/generate_nightly_perf_html.py +++ b/tools/nightly/generate_nightly_perf_html.py @@ -14,6 +14,8 @@ from collections.abc import Iterable, Sequence from datetime import datetime, timezone from typing import Any +from benchmarks.diffusion.backends import normalize_endpoint + LOGGER = logging.getLogger(__name__) _RESULT_JSON_PREFIX = "result_test_" @@ -222,6 +224,12 @@ def _iter_diffusion_json_records(input_dir: str) -> Iterable[dict[str, Any]]: result = flat.pop("result", None) if isinstance(result, dict): flat.update(result) + if flat.get("endpoint"): + flat["endpoint"] = normalize_endpoint(flat["endpoint"]) + elif flat.get("backend"): + flat["endpoint"] = normalize_endpoint(flat["backend"]) + flat.pop("backend", None) + flat.pop("API Backend", None) if "date" not in flat or not flat.get("date"): flat["date"] = filename_meta.get("date") or datetime.now( timezone.utc, @@ -1183,7 +1191,7 @@ function filterRows(rows, filters) {{ ) {{ return false; }} - if (filters.backend && String(row.backend || "") !== filters.backend) return false; + if (filters.backend && String(row[filters.backendKey] || "") !== filters.backend) return false; if ( filters.extraKey && filters.extraValue && @@ -1223,12 +1231,14 @@ function initSection(prefix, columns, data, numericCols, groups) {{ const modelKey = prefix === "diff" ? "model" : "model_id"; const datasetKey = prefix === "diff" ? "dataset" : "dataset_name"; + const backendKey = prefix === "diff" ? "endpoint" : "backend"; + const backendLabel = prefix === "diff" ? "Endpoint" : "Backend"; const extraKey = prefix === "diff" ? "" : "tokenizer_id"; const metaKeys = prefix === "diff" ? ["test_name"] : ["test_name", "max_concurrency", "num_prompts"]; const configFields = prefix === "diff" - ? ["test_name", "model", "backend", "dataset"] + ? ["test_name", "model", "endpoint", "dataset"] : [ "endpoint_type", "backend", @@ -1265,6 +1275,7 @@ function initSection(prefix, columns, data, numericCols, groups) {{ extraKey, modelKey, datasetKey, + backendKey, }}; }} @@ -1296,7 +1307,7 @@ function initSection(prefix, columns, data, numericCols, groups) {{ }} fillSelect( backendList, - uniqSorted(baseRows.map((row) => row.backend)), + uniqSorted(baseRows.map((row) => row[backendKey])), backendInput.value, ); if (extraList) {{ @@ -1378,7 +1389,7 @@ function initSection(prefix, columns, data, numericCols, groups) {{ }}); if (!backendInput.value.trim()) {{ - backendInput.value = firstValue(rows, "backend"); + backendInput.value = firstValue(rows, backendKey); }} rows = filterRows(data, {{ ...currentFilters(), @@ -1465,7 +1476,7 @@ function initSection(prefix, columns, data, numericCols, groups) {{ `
${{escapeHtml(filters.numPrompts || "All")}}
` ) }} -
Backend
${{escapeHtml( +
${{backendLabel}}
${{escapeHtml( filters.backend || "All", )}}
${{ @@ -1707,7 +1718,7 @@ window.addEventListener("load", () => {{ 'Diffusion
Image generation ' "benchmark history with throughput and latency trend views.
" '
' - "Grouping test_name + model + backend + dataset
" + "Grouping test_name + model + endpoint + dataset
" '
Snapshots single-point series ' "stay readable
" ), @@ -1736,7 +1747,7 @@ window.addEventListener("load", () => {{ " ", '
', ( - '
' + '
' '
' ), ( @@ -1803,7 +1814,7 @@ def generate_html_report(input_dir: str, diffusion_input_dir: str, output_file: "date", "test_name", "model", - "backend", + "endpoint", "dataset", "task", "duration", From 3ac739817f5afce9b5a291c2eddaccf5c1927cab Mon Sep 17 00:00:00 2001 From: JiaHong <2360655509@qq.com> Date: Tue, 19 May 2026 11:30:54 +0800 Subject: [PATCH 08/38] [Bugfix] Reject empty prompts in Flux2 Klein diffusion pipeline (#3711) Signed-off-by: MmMaiIIi <2360655509@qq.com> Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com> --- .../flux2/test_flux2_klein_validation.py | 80 +++++++++++++++++++ .../flux2_klein/pipeline_flux2_klein.py | 4 + 2 files changed, 84 insertions(+) create mode 100644 tests/diffusion/models/flux2/test_flux2_klein_validation.py diff --git a/tests/diffusion/models/flux2/test_flux2_klein_validation.py b/tests/diffusion/models/flux2/test_flux2_klein_validation.py new file mode 100644 index 000000000..de9ff6f39 --- /dev/null +++ b/tests/diffusion/models/flux2/test_flux2_klein_validation.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm_omni.diffusion.models.flux2_klein.pipeline_flux2_klein import ( + Flux2KleinPipeline, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _make_pipeline(): + pipeline = object.__new__(Flux2KleinPipeline) + pipeline.vae_scale_factor = 8 + pipeline.is_distilled = True + pipeline._guidance_scale = 0.0 + return pipeline + + +def _check(pipeline, prompt, prompt_embeds=None): + pipeline.check_inputs( + prompt=prompt, + height=512, + width=512, + prompt_embeds=prompt_embeds, + num_inference_steps=4, + guidance_scale=0.0, + ) + + +@pytest.mark.parametrize( + "prompt", + [ + "", + " ", + ], +) +def test_rejects_empty_or_whitespace_prompt(prompt): + pipe = _make_pipeline() + with pytest.raises(ValueError, match="empty or whitespace-only"): + _check(pipe, prompt) + + +@pytest.mark.parametrize( + "prompt", + [ + ["valid prompt", ""], + ["valid prompt", " "], + [" "], + ], +) +def test_rejects_list_with_empty_or_whitespace_element(prompt): + pipe = _make_pipeline() + with pytest.raises(ValueError, match="empty or whitespace-only"): + _check(pipe, prompt) + + +@pytest.mark.parametrize( + "prompt", + [ + "valid prompt", + ["valid prompt", "another valid prompt"], + ], +) +def test_accepts_valid_prompts(prompt): + pipe = _make_pipeline() + _check(pipe, prompt) + + +def test_rejects_none_without_embeds(): + pipe = _make_pipeline() + with pytest.raises(ValueError): + _check(pipe, None) + + +def test_accepts_none_with_prompt_embeds(): + pipe = _make_pipeline() + _check(pipe, None, prompt_embeds=torch.randn(1, 256, 4096)) diff --git a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py index 5543bc208..2ea267be1 100644 --- a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py +++ b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py @@ -701,6 +701,10 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi raise ValueError( "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." ) + elif isinstance(prompt, str) and not prompt.strip(): + raise ValueError("`prompt` cannot be empty or whitespace-only.") + elif isinstance(prompt, list) and any(isinstance(p, str) and not p.strip() for p in prompt): + raise ValueError("`prompt` cannot contain empty or whitespace-only strings.") elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") From ca9fd0b71ce04fa6283154c0ee7f32fcfc2eaf11 Mon Sep 17 00:00:00 2001 From: JiaHong <2360655509@qq.com> Date: Tue, 19 May 2026 11:52:41 +0800 Subject: [PATCH 09/38] Reject non-positive Flux2 Klein inference steps (#3717) Signed-off-by: MmMaiIIi <2360655509@qq.com> --- .../test_flux2_klein_num_inference_steps.py | 119 ++++++++++++++++++ .../flux2_klein/pipeline_flux2_klein.py | 6 +- 2 files changed, 124 insertions(+), 1 deletion(-) create mode 100644 tests/diffusion/models/flux2/test_flux2_klein_num_inference_steps.py diff --git a/tests/diffusion/models/flux2/test_flux2_klein_num_inference_steps.py b/tests/diffusion/models/flux2/test_flux2_klein_num_inference_steps.py new file mode 100644 index 000000000..3c6b5bada --- /dev/null +++ b/tests/diffusion/models/flux2/test_flux2_klein_num_inference_steps.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from unittest.mock import MagicMock + +import pytest + +from vllm_omni.diffusion.models.flux2_klein.pipeline_flux2_klein import ( + Flux2KleinPipeline, +) +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class _StopAfterCheckInputsError(Exception): + pass + + +def _make_pipeline(): + pipeline = object.__new__(Flux2KleinPipeline) + pipeline.vae_scale_factor = 8 + pipeline.is_distilled = True + pipeline._guidance_scale = 0.0 + return pipeline + + +def _make_minimal_request( + prompt="valid prompt", + *, + num_inference_steps=None, +): + """Build a minimal OmniDiffusionRequest-like object that forward() reads from.""" + params = OmniDiffusionSamplingParams( + height=512, + width=512, + num_inference_steps=num_inference_steps, + seed=42, + ) + req = MagicMock() + req.sampling_params = params + req.prompts = [prompt] + req.multi_modal_data = {} + return req + + +def _capture_num_inference_steps(pipe, steps): + """Run pipe.forward() and capture the num_inference_steps passed to check_inputs.""" + captured = {} + + def fake_check_inputs(**kwargs): + captured["num_inference_steps"] = kwargs.get("num_inference_steps") + raise _StopAfterCheckInputsError + + original = pipe.check_inputs + pipe.check_inputs = fake_check_inputs + try: + req = _make_minimal_request(num_inference_steps=steps) + pipe.forward(req) + except _StopAfterCheckInputsError: + pass + finally: + pipe.check_inputs = original + return captured.get("num_inference_steps") + + +# --- Forward resolution tests (the actual regression) --- + + +def test_forward_preserves_zero_not_default(): + """#3703: num_inference_steps=0 must reach validation as 0, not as default 50.""" + pipe = _make_pipeline() + captured = _capture_num_inference_steps(pipe, 0) + assert captured == 0, f"Expected 0, got {captured}" + + +def test_forward_preserves_negative(): + pipe = _make_pipeline() + captured = _capture_num_inference_steps(pipe, -1) + assert captured == -1 + + +def test_forward_none_uses_default(): + """None means unset — forward() should substitute the pipeline default.""" + pipe = _make_pipeline() + captured = _capture_num_inference_steps(pipe, None) + assert captured == 50 + + +def test_forward_preserves_positive(): + pipe = _make_pipeline() + captured = _capture_num_inference_steps(pipe, 9) + assert captured == 9 + + +# --- check_inputs validation tests --- + + +@pytest.mark.parametrize("steps", [0, -1]) +def test_check_inputs_rejects_non_positive(steps): + pipe = _make_pipeline() + with pytest.raises(ValueError): + pipe.check_inputs( + prompt="valid prompt", + height=512, + width=512, + num_inference_steps=steps, + ) + + +def test_check_inputs_accepts_none(): + pipe = _make_pipeline() + pipe.check_inputs(prompt="valid prompt", height=512, width=512, num_inference_steps=None) + + +@pytest.mark.parametrize("steps", [1, 2]) +def test_check_inputs_accepts_positive(steps): + pipe = _make_pipeline() + pipe.check_inputs(prompt="valid prompt", height=512, width=512, num_inference_steps=steps) diff --git a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py index 2ea267be1..9cb81fd37 100644 --- a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py +++ b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py @@ -865,7 +865,11 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi height = req.sampling_params.height or height width = req.sampling_params.width or width - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps + num_inference_steps = ( + req.sampling_params.num_inference_steps + if req.sampling_params.num_inference_steps is not None + else num_inference_steps + ) sigmas = req.sampling_params.sigmas or sigmas guidance_scale = ( req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale From e277feacaf859c1aa3f2f7354d6fc396cf06ba5d Mon Sep 17 00:00:00 2001 From: chickeyton Date: Tue, 19 May 2026 12:08:20 +0800 Subject: [PATCH 10/38] [large-scale-serving] Integrate OmniCoordinator into stage engine pipeline (#3569) Signed-off-by: chickeyton Signed-off-by: herotai214 Co-authored-by: herotai214 --- .../test_diffusion_engine_rpc_routing.py | 3 + .../omni_coordinator/test_load_balancer.py | 100 ++-- .../test_omni_coord_client_for_hub.py | 46 +- .../test_omni_coord_client_for_stage.py | 11 +- .../omni_coordinator/test_omni_coordinator.py | 68 +-- .../test_async_omni_engine_stage_init.py | 14 +- tests/engine/test_orchestrator.py | 2 +- tests/engine/test_single_stage_mode.py | 100 ++-- tests/entrypoints/test_serve.py | 446 +++++++++++------ vllm_omni/diffusion/stage_diffusion_client.py | 27 +- vllm_omni/diffusion/stage_diffusion_proc.py | 160 ++++++- .../omni_connectors/kv_transfer_manager.py | 9 +- .../omni_connectors/utils/initialization.py | 10 +- .../omni_connectors/utils/kv_utils.py | 36 +- .../distributed/omni_coordinator/__init__.py | 14 +- .../omni_coordinator/load_balancer.py | 103 ++-- .../distributed/omni_coordinator/messages.py | 46 +- .../omni_coord_client_for_hub.py | 68 +-- .../omni_coord_client_for_stage.py | 40 +- .../omni_coordinator/omni_coordinator.py | 152 +++--- .../distributed/omni_coordinator/runtime.py | 74 +++ vllm_omni/engine/arg_utils.py | 4 + vllm_omni/engine/async_omni_engine.py | 333 ++++++++++--- vllm_omni/engine/messages.py | 12 + .../engine/omni_core_engine_proc_manager.py | 162 +++++++ vllm_omni/engine/orchestrator.py | 260 +++++++++- vllm_omni/engine/stage_engine_core_client.py | 15 +- vllm_omni/engine/stage_engine_core_proc.py | 81 +++- vllm_omni/engine/stage_engine_startup.py | 449 +++++++++++++++--- vllm_omni/engine/stage_init_utils.py | 134 +++++- vllm_omni/engine/stage_pool.py | 387 +++++++++++++-- vllm_omni/entrypoints/cli/serve.py | 431 ++++++++++++++--- vllm_omni/entrypoints/stage_utils.py | 18 + 33 files changed, 3033 insertions(+), 782 deletions(-) create mode 100644 vllm_omni/distributed/omni_coordinator/runtime.py create mode 100644 vllm_omni/engine/omni_core_engine_proc_manager.py diff --git a/tests/diffusion/test_diffusion_engine_rpc_routing.py b/tests/diffusion/test_diffusion_engine_rpc_routing.py index f731f2472..b1b63a1a2 100644 --- a/tests/diffusion/test_diffusion_engine_rpc_routing.py +++ b/tests/diffusion/test_diffusion_engine_rpc_routing.py @@ -136,6 +136,7 @@ def _make_engine_with_loop( calls and a real ``RequestScheduler``. """ engine = DiffusionEngine.__new__(DiffusionEngine) + engine._closed = False engine.executor = _ConcurrencyTrackingExecutor(rpc_delay=rpc_delay) sched = RequestScheduler() @@ -371,6 +372,7 @@ def test_collective_rpc_before_loop_starts_calls_executor_directly(): on the caller's thread without enqueueing. """ engine = DiffusionEngine.__new__(DiffusionEngine) + engine._closed = False engine._loop_started = False engine._rpc_lock = threading.RLock() engine._cv = threading.Condition(engine._rpc_lock) @@ -391,6 +393,7 @@ def test_collective_rpc_before_loop_starts_serializes_concurrent_callers(): serialize them so they cannot race on the shared executor MQ pair. """ engine = DiffusionEngine.__new__(DiffusionEngine) + engine._closed = False engine._loop_started = False engine._rpc_lock = threading.RLock() engine._cv = threading.Condition(engine._rpc_lock) diff --git a/tests/distributed/omni_coordinator/test_load_balancer.py b/tests/distributed/omni_coordinator/test_load_balancer.py index 8350b33d3..1bf4575ea 100644 --- a/tests/distributed/omni_coordinator/test_load_balancer.py +++ b/tests/distributed/omni_coordinator/test_load_balancer.py @@ -6,18 +6,18 @@ from time import time import pytest from vllm_omni.distributed.omni_coordinator import ( - InstanceInfo, LeastQueueLengthBalancer, RandomBalancer, + ReplicaInfo, + ReplicaStatus, RoundRobinBalancer, - StageStatus, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def test_load_balancer_select_returns_valid_index(): - """Verify RandomBalancer.select() returns a valid index for instances.""" + """Verify RandomBalancer.select() returns a valid index for replicas.""" # Task structure mirrors async_omni; RandomBalancer ignores task contents. task: dict = { "request_id": "test", @@ -26,30 +26,30 @@ def test_load_balancer_select_returns_valid_index(): } now = time() - instances = [ - InstanceInfo( + replicas = [ + ReplicaInfo( input_addr="tcp://host:10001", output_addr="tcp://host:10001-out", stage_id=0, - status=StageStatus.UP, + status=ReplicaStatus.UP, queue_length=0, last_heartbeat=now, registered_at=now, ), - InstanceInfo( + ReplicaInfo( input_addr="tcp://host:10002", output_addr="tcp://host:10002-out", stage_id=0, - status=StageStatus.UP, + status=ReplicaStatus.UP, queue_length=1, last_heartbeat=now, registered_at=now, ), - InstanceInfo( + ReplicaInfo( input_addr="tcp://host:10003", output_addr="tcp://host:10003-out", stage_id=1, - status=StageStatus.UP, + status=ReplicaStatus.UP, queue_length=2, last_heartbeat=now, registered_at=now, @@ -58,38 +58,38 @@ def test_load_balancer_select_returns_valid_index(): balancer = RandomBalancer() - index = balancer.select(task, instances) + index = balancer.select(task, replicas) assert isinstance(index, int) - assert 0 <= index < len(instances) + assert 0 <= index < len(replicas) -def test_round_robin_balancer_cycles_instances(): +def test_round_robin_balancer_cycles_replicas(): now = time() - instances = [ - InstanceInfo( + replicas = [ + ReplicaInfo( input_addr="tcp://host:10001", output_addr="tcp://host:10001-out", stage_id=0, - status=StageStatus.UP, + status=ReplicaStatus.UP, queue_length=2, last_heartbeat=now, registered_at=now, ), - InstanceInfo( + ReplicaInfo( input_addr="tcp://host:10002", output_addr="tcp://host:10002-out", stage_id=0, - status=StageStatus.UP, + status=ReplicaStatus.UP, queue_length=1, last_heartbeat=now, registered_at=now, ), - InstanceInfo( + ReplicaInfo( input_addr="tcp://host:10003", output_addr="tcp://host:10003-out", stage_id=1, - status=StageStatus.UP, + status=ReplicaStatus.UP, queue_length=0, last_heartbeat=now, registered_at=now, @@ -97,35 +97,35 @@ def test_round_robin_balancer_cycles_instances(): ] balancer = RoundRobinBalancer() - results = [balancer.select({}, instances) for _ in range(5)] + results = [balancer.select({}, replicas) for _ in range(5)] # Default start_index=0 => 0,1,2,0,1 assert results == [0, 1, 2, 0, 1] -def test_round_robin_balancer_empty_instances_raises(): - with pytest.raises(ValueError, match="instances must not be empty"): +def test_round_robin_balancer_empty_replicas_raises(): + with pytest.raises(ValueError, match="replicas must not be empty"): RoundRobinBalancer().select({}, []) def test_round_robin_balancer_after_large_index_and_shorter_list(): - """Large start_index % len(instances) then counter wraps with shorter list.""" + """Large start_index % len(replicas) then counter wraps with shorter list.""" now = time() two = [ - InstanceInfo( + ReplicaInfo( input_addr="tcp://host:10001", output_addr="tcp://host:10001-out", stage_id=0, - status=StageStatus.UP, + status=ReplicaStatus.UP, queue_length=0, last_heartbeat=now, registered_at=now, ), - InstanceInfo( + ReplicaInfo( input_addr="tcp://host:10002", output_addr="tcp://host:10002-out", stage_id=0, - status=StageStatus.UP, + status=ReplicaStatus.UP, queue_length=0, last_heartbeat=now, registered_at=now, @@ -138,30 +138,30 @@ def test_round_robin_balancer_after_large_index_and_shorter_list(): def test_least_queue_length_balancer_picks_min_queue(): now = time() - instances = [ - InstanceInfo( + replicas = [ + ReplicaInfo( input_addr="tcp://host:10001", output_addr="tcp://host:10001-out", stage_id=0, - status=StageStatus.UP, + status=ReplicaStatus.UP, queue_length=2, last_heartbeat=now, registered_at=now, ), - InstanceInfo( + ReplicaInfo( input_addr="tcp://host:10002", output_addr="tcp://host:10002-out", stage_id=0, - status=StageStatus.UP, + status=ReplicaStatus.UP, queue_length=0, last_heartbeat=now, registered_at=now, ), - InstanceInfo( + ReplicaInfo( input_addr="tcp://host:10003", output_addr="tcp://host:10003-out", stage_id=1, - status=StageStatus.UP, + status=ReplicaStatus.UP, queue_length=5, last_heartbeat=now, registered_at=now, @@ -169,41 +169,41 @@ def test_least_queue_length_balancer_picks_min_queue(): ] balancer = LeastQueueLengthBalancer() - index = balancer.select({}, instances) + index = balancer.select({}, replicas) assert index == 1 -def test_least_queue_length_balancer_empty_instances_raises(): - with pytest.raises(ValueError, match="instances must not be empty"): +def test_least_queue_length_balancer_empty_replicas_raises(): + with pytest.raises(ValueError, match="replicas must not be empty"): LeastQueueLengthBalancer().select({}, []) def test_least_queue_length_balancer_equal_queues_uses_choice(mocker): now = time() - instances = [ - InstanceInfo( + replicas = [ + ReplicaInfo( input_addr="tcp://host:10001", output_addr="tcp://host:10001-out", stage_id=0, - status=StageStatus.UP, + status=ReplicaStatus.UP, queue_length=3, last_heartbeat=now, registered_at=now, ), - InstanceInfo( + ReplicaInfo( input_addr="tcp://host:10002", output_addr="tcp://host:10002-out", stage_id=0, - status=StageStatus.UP, + status=ReplicaStatus.UP, queue_length=3, last_heartbeat=now, registered_at=now, ), - InstanceInfo( + ReplicaInfo( input_addr="tcp://host:10003", output_addr="tcp://host:10003-out", stage_id=1, - status=StageStatus.UP, + status=ReplicaStatus.UP, queue_length=3, last_heartbeat=now, registered_at=now, @@ -214,21 +214,21 @@ def test_least_queue_length_balancer_equal_queues_uses_choice(mocker): "vllm_omni.distributed.omni_coordinator.load_balancer.random.choice", return_value=2, ) - assert balancer.select({}, instances) == 2 + assert balancer.select({}, replicas) == 2 def test_least_queue_length_balancer_negative_queue_raises(): now = time() - instances = [ - InstanceInfo( + replicas = [ + ReplicaInfo( input_addr="tcp://host:10001", output_addr="tcp://host:10001-out", stage_id=0, - status=StageStatus.UP, + status=ReplicaStatus.UP, queue_length=-1, last_heartbeat=now, registered_at=now, ), ] with pytest.raises(ValueError, match="queue_length must be non-negative"): - LeastQueueLengthBalancer().select({}, instances) + LeastQueueLengthBalancer().select({}, replicas) diff --git a/tests/distributed/omni_coordinator/test_omni_coord_client_for_hub.py b/tests/distributed/omni_coordinator/test_omni_coord_client_for_hub.py index 2fbd7c85b..542e3b06d 100644 --- a/tests/distributed/omni_coordinator/test_omni_coord_client_for_hub.py +++ b/tests/distributed/omni_coordinator/test_omni_coord_client_for_hub.py @@ -8,8 +8,8 @@ import pytest import zmq from vllm_omni.distributed.omni_coordinator import ( - InstanceList, OmniCoordClientForHub, + ReplicaList, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -32,8 +32,8 @@ def _wait_for_condition(cond, timeout: float = 2.0, interval: float = 0.01) -> b return False -def test_hub_client_caches_instance_list_from_pub(): - """Verify OmniCoordClientForHub receives instance list updates from OmniCoordinator and caches for get_instance_list().""" +def test_hub_client_caches_replica_list_from_pub(): + """Verify OmniCoordClientForHub receives replica list updates from OmniCoordinator and caches for get_replica_list().""" ctx, pub, endpoint = _bind_pub() client = OmniCoordClientForHub(endpoint) @@ -41,7 +41,7 @@ def test_hub_client_caches_instance_list_from_pub(): time.sleep(0.2) now = time.time() - instances_payload = [ + replicas_payload = [ { "input_addr": "tcp://stage:10001", "output_addr": "tcp://stage:10001-out", @@ -71,37 +71,37 @@ def test_hub_client_caches_instance_list_from_pub(): }, ] - payload = {"instances": instances_payload, "timestamp": now} + payload = {"replicas": replicas_payload, "timestamp": now} pub.send(json.dumps(payload).encode("utf-8")) - assert _wait_for_condition(lambda: len(client.get_instance_list().instances) == 3) + assert _wait_for_condition(lambda: len(client.get_replica_list().replicas) == 3) - inst_list = client.get_instance_list() - assert isinstance(inst_list, InstanceList) - assert len(inst_list.instances) == 3 + rep_list = client.get_replica_list() + assert isinstance(rep_list, ReplicaList) + assert len(rep_list.replicas) == 3 - for src, inst in zip(instances_payload, inst_list.instances, strict=True): - assert inst.input_addr == src["input_addr"] - assert inst.output_addr == src["output_addr"] - assert inst.stage_id == src["stage_id"] - assert inst.status.value == src["status"] + for src, rep in zip(replicas_payload, rep_list.replicas, strict=True): + assert rep.input_addr == src["input_addr"] + assert rep.output_addr == src["output_addr"] + assert rep.stage_id == src["stage_id"] + assert rep.status.value == src["status"] - stage0 = client.get_instances_for_stage(0) - stage1 = client.get_instances_for_stage(1) + stage0 = client.get_replicas_for_stage(0) + stage1 = client.get_replicas_for_stage(1) - assert all(inst.stage_id == 0 for inst in stage0.instances) - assert all(inst.stage_id == 1 for inst in stage1.instances) + assert all(rep.stage_id == 0 for rep in stage0.replicas) + assert all(rep.stage_id == 1 for rep in stage1.replicas) - # Send an updated list with fewer instances and verify cache refresh. + # Send an updated list with fewer replicas and verify cache refresh. updated_payload = { - "instances": instances_payload[:2], + "replicas": replicas_payload[:2], "timestamp": now + 1.0, } pub.send(json.dumps(updated_payload).encode("utf-8")) - assert _wait_for_condition(lambda: len(client.get_instance_list().instances) == 2) - updated_list = client.get_instance_list() - assert len(updated_list.instances) == 2 + assert _wait_for_condition(lambda: len(client.get_replica_list().replicas) == 2) + updated_list = client.get_replica_list() + assert len(updated_list.replicas) == 2 client.close() pub.close(0) diff --git a/tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py b/tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py index 0ba19c7ff..46eb2e0fb 100644 --- a/tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py +++ b/tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py @@ -9,7 +9,7 @@ import zmq from vllm_omni.distributed.omni_coordinator import ( OmniCoordClientForStage, - StageStatus, + ReplicaStatus, ) from vllm_omni.distributed.omni_coordinator import ( omni_coord_client_for_stage as stage_client_module, @@ -47,7 +47,7 @@ def test_stage_client_auto_register_on_init(): event = _recv_event(router) assert event["event_type"] == "update" - assert event["status"] == StageStatus.UP.value + assert event["status"] == ReplicaStatus.UP.value assert event["stage_id"] == stage_id assert event["input_addr"] == input_addr assert event["output_addr"] == output_addr @@ -70,13 +70,13 @@ def test_stage_client_update_info_sends_correct_event(): # Discard initial registration event. _recv_event(router) - client.update_info(status=StageStatus.ERROR) + client.update_info(status=ReplicaStatus.ERROR) client.update_info(queue_length=10) first = _recv_event(router) second = _recv_event(router) - assert first["status"] == StageStatus.ERROR.value + assert first["status"] == ReplicaStatus.ERROR.value assert first["stage_id"] == stage_id assert first["input_addr"] == input_addr assert first["output_addr"] == output_addr @@ -107,7 +107,7 @@ def test_stage_client_close_sends_down_status(): client.close() event = _recv_event(router) - assert event["status"] == StageStatus.DOWN.value + assert event["status"] == ReplicaStatus.DOWN.value assert event["stage_id"] == stage_id assert event["input_addr"] == input_addr assert event["output_addr"] == output_addr @@ -274,6 +274,7 @@ def test_heartbeat_loop_retries_after_transient_send_failure(): client._closed = False client._heartbeat_interval = 0.0 client._stop_event = _FakeStopEvent() + client._on_heartbeat = None calls = {"count": 0} diff --git a/tests/distributed/omni_coordinator/test_omni_coordinator.py b/tests/distributed/omni_coordinator/test_omni_coordinator.py index eff3d429e..81bd12454 100644 --- a/tests/distributed/omni_coordinator/test_omni_coordinator.py +++ b/tests/distributed/omni_coordinator/test_omni_coordinator.py @@ -11,14 +11,14 @@ from vllm.v1.utils import get_engine_client_zmq_addr from vllm_omni.distributed.omni_coordinator import ( OmniCoordClientForStage, OmniCoordinator, - StageStatus, + ReplicaStatus, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -def _recv_instance_list(sub: zmq.Socket, timeout_ms: int = 2000) -> dict | None: - """Receive InstanceList JSON from SUB socket. Returns None on timeout.""" +def _recv_replica_list(sub: zmq.Socket, timeout_ms: int = 2000) -> dict | None: + """Receive ReplicaList JSON from SUB socket. Returns None on timeout.""" sub.setsockopt(zmq.RCVTIMEO, timeout_ms) try: data = sub.recv() @@ -27,16 +27,16 @@ def _recv_instance_list(sub: zmq.Socket, timeout_ms: int = 2000) -> dict | None: return None -def _wait_for_instance_list( +def _wait_for_replica_list( sub: zmq.Socket, expected_count: int, timeout: float = 3.0, ) -> dict | None: - """Wait until received InstanceList with expected_count active instances.""" + """Wait until received ReplicaList with expected_count active replicas.""" start = time.time() while time.time() - start < timeout: - msg = _recv_instance_list(sub, timeout_ms=500) - if msg is not None and len(msg.get("instances", [])) == expected_count: + msg = _recv_replica_list(sub, timeout_ms=500) + if msg is not None and len(msg.get("replicas", [])) == expected_count: return msg return None @@ -45,7 +45,7 @@ def _drain_sub_messages(sub: zmq.Socket, max_seconds: float = 0.4) -> None: """Drain queued SUB messages for a short window.""" deadline = time.time() + max_seconds while time.time() < deadline: - _recv_instance_list(sub, timeout_ms=50) + _recv_replica_list(sub, timeout_ms=50) def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates(): @@ -81,7 +81,7 @@ def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates(): ) # Wait for initial registration broadcast and clear any queued messages. - msg = _wait_for_instance_list(sub, expected_count=1) + msg = _wait_for_replica_list(sub, expected_count=1) assert msg is not None _drain_sub_messages(sub) @@ -96,7 +96,7 @@ def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates(): deadline = time.time() + window_s recv_count = 0 while time.time() < deadline: - if _recv_instance_list(sub, timeout_ms=100) is not None: + if _recv_replica_list(sub, timeout_ms=100) is not None: recv_count += 1 assert recv_count < update_count // 2, ( @@ -110,8 +110,8 @@ def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates(): def test_omni_coordinator_registration_broadcast(): - """Verify that after multiple OmniCoordClientForStage instances register, - OmniCoordinator publishes an InstanceList containing all registered instances. + """Verify that after multiple OmniCoordClientForStage replicas register, + OmniCoordinator publishes a ReplicaList containing all registered replicas. """ router_addr = get_engine_client_zmq_addr( local_only=False, @@ -144,12 +144,12 @@ def test_omni_coordinator_registration_broadcast(): OmniCoordClientForStage(router_addr, "tcp://stage:10003", "tcp://stage:10003-out", 1), ] - msg = _wait_for_instance_list(sub, expected_count=3) - assert msg is not None, "Expected InstanceList with 3 instances" - assert len(msg["instances"]) == 3 + msg = _wait_for_replica_list(sub, expected_count=3) + assert msg is not None, "Expected ReplicaList with 3 replicas" + assert len(msg["replicas"]) == 3 assert isinstance(msg["timestamp"], (int, float)) - input_addrs = {inst["input_addr"] for inst in msg["instances"]} + input_addrs = {rep["input_addr"] for rep in msg["replicas"]} assert "tcp://stage:10001" in input_addrs assert "tcp://stage:10002" in input_addrs assert "tcp://stage:10003" in input_addrs @@ -162,7 +162,7 @@ def test_omni_coordinator_registration_broadcast(): def test_omni_coordinator_heartbeat_timeout_handling(): - """Verify that when a stage instance stops sending heartbeats, + """Verify that when a stage replica stops sending heartbeats, OmniCoordinator marks it as unhealthy and excludes it from the active list. """ router_addr = get_engine_client_zmq_addr( @@ -201,23 +201,23 @@ def test_omni_coordinator_heartbeat_timeout_handling(): "output_addr": "tcp://stage:c-out", "stage_id": 0, "event_type": "update", - "status": StageStatus.UP.value, + "status": ReplicaStatus.UP.value, "queue_length": 0, } dealer_c.send(json.dumps(reg_event).encode("utf-8")) - msg = _wait_for_instance_list(sub, expected_count=3) - assert msg is not None, "Expected initial 3 instances" - assert len(msg["instances"]) == 3 + msg = _wait_for_replica_list(sub, expected_count=3) + assert msg is not None, "Expected initial 3 replicas" + assert len(msg["replicas"]) == 3 # Wait for heartbeat timeout (timeout=5s, check interval ~2.5s). time.sleep(8.0) # Receive the update (C should be ERROR and excluded from active list). - msg_after_timeout = _wait_for_instance_list(sub, expected_count=2, timeout=5.0) - assert msg_after_timeout is not None, "Expected InstanceList with 2 instances after timeout" - instances = msg_after_timeout.get("instances", []) - input_addrs = {inst["input_addr"] for inst in instances} + msg_after_timeout = _wait_for_replica_list(sub, expected_count=2, timeout=5.0) + assert msg_after_timeout is not None, "Expected ReplicaList with 2 replicas after timeout" + replicas = msg_after_timeout.get("replicas", []) + input_addrs = {rep["input_addr"] for rep in replicas} assert "tcp://stage:a" in input_addrs assert "tcp://stage:b" in input_addrs @@ -232,8 +232,8 @@ def test_omni_coordinator_heartbeat_timeout_handling(): sub_ctx.term() -def test_omni_coordinator_instance_shutdown_handling(): - """Verify that when a stage instance sends status='down', +def test_omni_coordinator_replica_shutdown_handling(): + """Verify that when a stage replica sends status='down', OmniCoordinator removes it from the active list and broadcasts an updated list. """ router_addr = get_engine_client_zmq_addr( @@ -261,18 +261,18 @@ def test_omni_coordinator_instance_shutdown_handling(): client = OmniCoordClientForStage(router_addr, "tcp://stage:shutdown", "tcp://stage:shutdown-out", 0) - msg = _wait_for_instance_list(sub, expected_count=1) + msg = _wait_for_replica_list(sub, expected_count=1) assert msg is not None - assert len(msg["instances"]) == 1 - assert msg["instances"][0]["input_addr"] == "tcp://stage:shutdown" + assert len(msg["replicas"]) == 1 + assert msg["replicas"][0]["input_addr"] == "tcp://stage:shutdown" # Send down status (simulating graceful shutdown). - client.update_info(status=StageStatus.DOWN) + client.update_info(status=ReplicaStatus.DOWN) - # Receive updated list (should have 0 active instances). - msg = _wait_for_instance_list(sub, expected_count=0) + # Receive updated list (should have 0 active replicas). + msg = _wait_for_replica_list(sub, expected_count=0) assert msg is not None - assert len(msg["instances"]) == 0 + assert len(msg["replicas"]) == 0 client.close() coordinator.close() diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 51ed6b46c..b423aa407 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -171,18 +171,6 @@ def test_collect_initialized_clients_for_cleanup_deduplicates_clients(): assert cleanup_clients == [shared, extra] -def test_initialize_stages_rejects_non_diffusion_replicas_in_single_stage_mode(): - engine = object.__new__(AsyncOmniEngine) - engine.single_stage_mode = True - engine.stage_configs = [types.SimpleNamespace(stage_id=0, runtime={"num_replicas": 2})] - - with pytest.raises( - ValueError, - match="single_stage_mode only supports num_replicas > 1 for diffusion stages", - ): - engine._validate_single_stage_mode_replica_constraints() - - def test_initialize_diffusion_replica_restores_device_visibility_after_local_init(monkeypatch): import vllm_omni.engine.async_omni_engine as engine_mod from vllm_omni.platforms import current_omni_platform @@ -427,6 +415,7 @@ def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_rep engine.single_stage_mode = False engine._single_stage_id_filter = None engine._omni_master_server = None + engine._coordinator_runtime = None engine.stage_configs = [types.SimpleNamespace()] cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) @@ -471,6 +460,7 @@ def test_initialize_stages_cleans_up_late_successful_replicas_after_early_multi_ engine.single_stage_mode = False engine._single_stage_id_filter = None engine._omni_master_server = None + engine._coordinator_runtime = None engine.stage_configs = [types.SimpleNamespace()] cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index c84256f04..9e050a3f8 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -1052,7 +1052,7 @@ async def test_stage_pool_abort_requests_logs_when_binding_is_missing(caplog) -> target_logger.setLevel(prev_level) assert not stage0.abort_calls - assert "abort: no binding for req=missing-req in stage-0" in caplog.text + assert "abort: no live binding for req=missing-req in stage-0" in caplog.text @pytest.mark.asyncio diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index bcd910620..52187c341 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -203,6 +203,32 @@ class TestOmniMasterServerAllocation: server = OmniMasterServer(master_address="127.0.0.1", master_port=15004, stage_ids=[3]) assert server.get_allocation(3) is server._stage_routes[(3, 0)] + def test_next_free_replica_id_skips_head_local_slot_until_filled(self): + # Head pre-allocates slot (0, 0) for its own register_stage_with_omni_master + # call. A same-host headless that registers with auto-assign BEFORE the + # head's own registration must NOT be handed slot 0 — it should land on + # slot 1 instead. Without the head_local_replicas reservation, + # _next_free_replica_id would see (0, 0) absent from _stage_configs and + # return 0, colliding with the head's own bound sockets. + server = OmniMasterServer( + master_address="127.0.0.1", + master_port=15010, + stage_ids=[0], + head_local_replicas={0: [0]}, + ) + assert server._next_free_replica_id(0) == 1 + + def test_next_free_replica_id_uses_remote_slot_when_unowned(self): + # When the head pre-allocates a remote-only slot (the head's _initialize_* + # path waits on get_stage_config), auto-assign SHOULD fill it so the + # head's wait unblocks. This is the original behavior, preserved. + server = OmniMasterServer( + master_address="127.0.0.1", + master_port=15011, + stage_ids=[0], + ) + assert server._next_free_replica_id(0) == 0 + # --------------------------------------------------------------------------- # OmniMasterServer registration flow @@ -233,36 +259,6 @@ class TestOmniMasterServerRegistration: ctx.term() server.stop() - def test_server_handles_unknown_stage_id_gracefully(self): - import msgspec - import zmq - from vllm.utils.network_utils import get_open_port - - master_port = get_open_port() - server = OmniMasterServer(master_address="127.0.0.1", master_port=master_port, stage_ids=[0]) - server.start() - - ctx = zmq.Context() - bad_sock = None - good_sock = None - try: - bad_sock = ctx.socket(zmq.DEALER) - bad_sock.connect(f"tcp://127.0.0.1:{master_port}") - bad_sock.send(msgspec.msgpack.encode({"stage_id": 99})) - assert not bad_sock.poll(timeout=500) - - good_sock = ctx.socket(zmq.DEALER) - good_sock.connect(f"tcp://127.0.0.1:{master_port}") - good_sock.send(msgspec.msgpack.encode({"stage_id": 0})) - assert good_sock.poll(timeout=2_000) - good_sock.recv() - finally: - for sock in (bad_sock, good_sock): - if sock is not None: - sock.close(linger=0) - ctx.term() - server.stop() - def test_registration_stores_stage_config(self): import msgspec import zmq @@ -474,6 +470,8 @@ class TestSingleStageInitialization: engine._omni_master_address = "127.0.0.1" engine._omni_master_port = 26000 engine._omni_master_server = None + engine._omni_heartbeat_timeout = 30.0 + engine._coordinator_runtime = None engine.async_chunk = False engine.diffusion_batch_size = 2 return engine @@ -508,10 +506,16 @@ class TestSingleStageInitialization: def test_start_omni_master_server_uses_configured_stage_ids(self, mocker: MockerFixture): import vllm_omni.engine.async_omni_engine as engine_mod + from vllm_omni.distributed import omni_coordinator as omni_coord_mod engine = self._build_engine([], single_stage_mode=True, stage_id_filter=7) mock_oms = mocker.Mock(spec=OmniMasterServer) mocker.patch.object(engine_mod, "OmniMasterServer", return_value=mock_oms) + mocker.patch.object( + omni_coord_mod, + "OmniCoordinatorRuntime", + return_value=mocker.Mock(router_address="tcp://127.0.0.1:9999"), + ) stage_plans = [ _make_llm_plan(0, configured_stage_id=7, launch_mode="local"), @@ -520,12 +524,16 @@ class TestSingleStageInitialization: engine._start_omni_master_server(stage_plans) - engine_mod.OmniMasterServer.assert_called_once_with( - master_address="127.0.0.1", - master_port=26000, - stage_ids=[7, 11], - stage_replica_counts={7: 1, 11: 1}, - ) + call_kwargs = engine_mod.OmniMasterServer.call_args.kwargs + assert call_kwargs["master_address"] == "127.0.0.1" + assert call_kwargs["master_port"] == 26000 + assert call_kwargs["stage_ids"] == [7, 11] + assert call_kwargs["stage_replica_counts"] == {7: 1, 11: 1} + # head_local_replicas reserves slots that the head will register + # itself (launch_mode == "local"). Stage 11 is remote, so it must + # NOT appear in the head-owned set — that slot is for the headless + # to fill via auto-assign. + assert call_kwargs["head_local_replicas"] == {7: [0]} mock_oms.start.assert_called_once() def test_start_omni_master_server_duplicate_stage_ids_raise(self): @@ -545,7 +553,9 @@ class TestSingleStageInitialization: with pytest.raises(ValueError, match="requires both"): engine._start_omni_master_server([_make_llm_plan(0, configured_stage_id=7, launch_mode="local")]) - def test_build_logical_stage_init_plans_clears_runtime_cfg_in_single_stage_mode(self, mocker: MockerFixture): + def test_build_logical_stage_init_plans_preserves_runtime_cfg_for_local_llm_in_single_stage_mode( + self, mocker: MockerFixture + ): import vllm_omni.engine.async_omni_engine as engine_mod engine = self._build_engine([_make_stage_cfg(7)], single_stage_mode=True, stage_id_filter=7) @@ -574,7 +584,7 @@ class TestSingleStageInitialization: finally: monkeypatch.undo() - assert stage_plans[0].replicas[0].metadata.runtime_cfg is None + assert stage_plans[0].replicas[0].metadata.runtime_cfg == {"devices": "0"} def test_validate_single_stage_mode_allows_diffusion_replicas(self): stage_cfg = _make_stage_cfg(0, stage_type="diffusion") @@ -583,14 +593,6 @@ class TestSingleStageInitialization: engine._validate_single_stage_mode_replica_constraints() - def test_validate_single_stage_mode_rejects_llm_replicas(self): - stage_cfg = _make_stage_cfg(0, stage_type="llm") - stage_cfg.runtime.num_replicas = 2 - engine = self._build_engine([stage_cfg], single_stage_mode=True, stage_id_filter=0) - - with pytest.raises(ValueError, match="only supports num_replicas > 1 for diffusion"): - engine._validate_single_stage_mode_replica_constraints() - def test_build_logical_stage_init_plans_preserves_diffusion_runtime_cfg_in_single_stage_mode( self, mocker: MockerFixture ): @@ -809,6 +811,7 @@ class TestSingleStageReplicaInitialization: engine.model = "fake-model" engine.single_stage_mode = True engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) + engine._coordinator_runtime = None engine.stage_configs = [] fake_vllm_config = SimpleNamespace(parallel_config=SimpleNamespace()) @@ -896,6 +899,7 @@ class TestSingleStageReplicaInitialization: engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) engine._omni_master_server.address = "127.0.0.1" engine._omni_master_server.port = 25000 + engine._coordinator_runtime = None plan = _make_diffusion_plan(0, configured_stage_id=5, launch_mode="local").replicas[0] sentinel_client = SimpleNamespace() @@ -950,6 +954,9 @@ class TestSingleStageReplicaInitialization: handshake_address="tcp://hs", request_address="tcp://req", response_address="tcp://resp", + omni_coordinator_address=None, + omni_stage_id=5, + omni_replica_id=0, ) mock_handshake.assert_called_once_with(proc, "tcp://hs", 60) mock_from_addresses.assert_called_once_with( @@ -972,6 +979,7 @@ class TestSingleStageReplicaInitialization: engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) engine._omni_master_server.address = "127.0.0.1" engine._omni_master_server.port = 25000 + engine._coordinator_runtime = None plan = _make_diffusion_plan(0, configured_stage_id=5, launch_mode="local").replicas[0] proc = mocker.Mock() diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 1267a189a..86c311a60 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -3,6 +3,8 @@ from __future__ import annotations import argparse +from types import SimpleNamespace +from typing import Any import pytest from pytest_mock import MockerFixture @@ -31,143 +33,118 @@ def test_serve_parser_accepts_no_async_chunk() -> None: assert args.async_chunk is False -def _make_headless_args() -> argparse.Namespace: - return argparse.Namespace( +# --------------------------------------------------------------------------- +# run_headless validation +# --------------------------------------------------------------------------- + + +def _make_headless_args(**overrides: Any) -> argparse.Namespace: + """Build an argparse.Namespace shaped like the headless CLI passes in. + + Defaults pass every validation gate so individual tests can mutate just + the field they're exercising. + """ + defaults = dict( model="fake-model", - stage_id=3, + stage_id=0, replica_id=0, omni_master_address="127.0.0.1", omni_master_port=26000, + omni_replica_address=None, + omni_dp_size_local=1, api_server_count=0, worker_backend="multi_process", stage_configs_path=None, + deploy_config=None, log_stats=False, disable_log_stats=False, stage_init_timeout=600, + tokenizer=None, ) + defaults.update(overrides) + return argparse.Namespace(**defaults) -def test_run_headless_registers_stage_once_and_launches_all_local_engines(mocker: MockerFixture) -> None: - args = _make_headless_args() - stage_cfg = mocker.Mock(stage_id=3) - stage_cfgs = [stage_cfg] - parallel_config = mocker.Mock( - data_parallel_size_local=2, - data_parallel_rank=4, - data_parallel_rank_local=1, - node_rank_within_dp=0, - ) - vllm_config = mocker.Mock(parallel_config=parallel_config) - vllm_config.needs_dp_coordinator = False - executor_class = mocker.Mock() - engine_manager = mocker.Mock() +def test_run_headless_requires_stage_id() -> None: + args = _make_headless_args(stage_id=None) + with pytest.raises(ValueError, match="--stage-id is required"): + run_headless(args) + +def test_run_headless_requires_master_address() -> None: + args = _make_headless_args(omni_master_address=None) + with pytest.raises(ValueError, match="--omni-master-address and --omni-master-port"): + run_headless(args) + + +def test_run_headless_requires_master_port() -> None: + args = _make_headless_args(omni_master_port=None) + with pytest.raises(ValueError, match="--omni-master-address and --omni-master-port"): + run_headless(args) + + +def test_run_headless_rejects_multi_api_server_count() -> None: + args = _make_headless_args(api_server_count=2) + with pytest.raises(ValueError, match="api_server_count can't be set"): + run_headless(args) + + +def test_run_headless_rejects_non_multiprocess_worker_backend() -> None: + args = _make_headless_args(worker_backend="ray") + with pytest.raises(ValueError, match="worker_backend=multi_process"): + run_headless(args) + + +def test_run_headless_raises_when_stage_id_not_in_configs(mocker: MockerFixture) -> None: + """Headless looks up its assigned stage_id in the loaded deploy YAML and + fails fast when the launcher's --stage-id doesn't match any entry.""" + other_stage = SimpleNamespace(stage_id=99) mocker.patch( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", - return_value=("/fake/stages.yaml", stage_cfgs), + return_value=("/fake/stages.yaml", [other_stage]), ) - mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") - mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=mocker.Mock()) - mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={}) - mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={}) - mocker.patch( - "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage", - return_value=(None, None, None), - ) - mock_build_vllm_config = mocker.patch( - "vllm_omni.engine.stage_init_utils.build_vllm_config", - return_value=(vllm_config, executor_class), - ) - mock_register = mocker.patch( - "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", - return_value="tcp://127.0.0.1:26001", - ) - mock_manager_cls = mocker.patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager) - mocker.patch("signal.signal") - run_headless(args) - mock_build_vllm_config.assert_called_once_with( - stage_cfg, - "fake-model", - stage_connector_spec={}, - engine_args_dict={}, - headless=True, - ) - mock_register.assert_called_once_with( - omni_master_address="127.0.0.1", - omni_master_port=26000, - omni_stage_id=3, - omni_stage_config=stage_cfg, - coordinator=None, - replica_id=0, - ) - mock_manager_cls.assert_called_once() - manager_kwargs = mock_manager_cls.call_args.kwargs - assert manager_kwargs["local_engine_count"] == 2 - assert manager_kwargs["start_index"] == 4 - assert manager_kwargs["local_start_index"] == 0 - assert manager_kwargs["local_client"] is False - assert manager_kwargs["handshake_address"] == "tcp://127.0.0.1:26001" - assert manager_kwargs["log_stats"] is False - engine_manager.monitor_engine_liveness.assert_called_once_with() - engine_manager.shutdown.assert_called_once_with() + args = _make_headless_args(stage_id=0) + with pytest.raises(ValueError, match="No stage config found for stage_id=0"): + run_headless(args) -def test_run_headless_honors_explicit_log_stats_flag(mocker: MockerFixture) -> None: - args = _make_headless_args() - args.log_stats = True - stage_cfg = mocker.Mock(stage_id=3) - stage_cfgs = [stage_cfg] - parallel_config = mocker.Mock( - data_parallel_size_local=2, - data_parallel_rank=4, - data_parallel_rank_local=1, - node_rank_within_dp=0, - ) - vllm_config = mocker.Mock(parallel_config=parallel_config) - executor_class = mocker.Mock() - engine_manager = mocker.Mock() - - mocker.patch( - "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", - return_value=("/fake/stages.yaml", stage_cfgs), - ) - mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") - mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=mocker.Mock()) - mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={}) - mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={}) - mocker.patch( - "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage", - return_value=(None, None, None), - ) - mocker.patch( - "vllm_omni.engine.stage_init_utils.build_vllm_config", - return_value=(vllm_config, executor_class), - ) - mocker.patch( - "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", - return_value="tcp://127.0.0.1:26001", - ) - mock_manager_cls = mocker.patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager) - mocker.patch("signal.signal") - run_headless(args) - - manager_kwargs = mock_manager_cls.call_args.kwargs - assert manager_kwargs["log_stats"] is True +# --------------------------------------------------------------------------- +# run_headless happy paths +# --------------------------------------------------------------------------- -def test_run_headless_registers_llm_replica_id(mocker: MockerFixture) -> None: - args = _make_headless_args() - args.replica_id = 2 - stage_cfg = mocker.Mock(stage_id=3) - parallel_config = mocker.Mock( +def _make_stage_cfg(stage_id: int, stage_type: str) -> SimpleNamespace: + """Build a stage config that satisfies every attribute run_headless reads. + + Notably ``engine_args`` is a real dict (not a Mock) so + ``get_stage_devices_per_replica`` can call ``.get("tensor_parallel_size")`` + and feed the result through ``int()`` without TypeError. + """ + return SimpleNamespace( + stage_id=stage_id, + stage_type=stage_type, + # No "devices" key -> split_devices_for_replicas skipped, each replica + # inherits the launcher's CUDA_VISIBLE_DEVICES. + runtime=None, + engine_args={}, + ) + + +def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: MockerFixture) -> None: + """LLM headless: each loop iteration registers with auto-assigned + replica_id (master picks a free slot) and spawns one + ``OmniCoreEngineProcManager`` per local replica.""" + from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse + + stage_cfg = _make_stage_cfg(0, stage_type="llm") + parallel_config = SimpleNamespace( data_parallel_size_local=1, data_parallel_rank=0, data_parallel_rank_local=0, node_rank_within_dp=0, ) - vllm_config = mocker.Mock(parallel_config=parallel_config) - vllm_config.needs_dp_coordinator = False + vllm_config = SimpleNamespace(parallel_config=parallel_config, needs_dp_coordinator=False) engine_manager = mocker.Mock() mocker.patch( @@ -175,82 +152,247 @@ def test_run_headless_registers_llm_replica_id(mocker: MockerFixture) -> None: return_value=("/fake/stages.yaml", [stage_cfg]), ) mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") - mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=mocker.Mock()) - mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={}) - mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={}) + mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None) mocker.patch( "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage", return_value=(None, None, None), ) + mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={}) + mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={}) mocker.patch( "vllm_omni.engine.stage_init_utils.build_vllm_config", - return_value=(vllm_config, mocker.Mock()), + return_value=(vllm_config, object), ) mock_register = mocker.patch( "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", - return_value="tcp://127.0.0.1:26001", + return_value=StageRegistrationResponse( + handshake_address="tcp://127.0.0.1:26001", + input_address="tcp://127.0.0.1:26002", + output_address="tcp://127.0.0.1:26003", + replica_id=0, + coordinator_router_address="tcp://127.0.0.1:26100", + ), + ) + mock_manager_cls = mocker.patch( + "vllm_omni.engine.omni_core_engine_proc_manager.OmniCoreEngineProcManager", + return_value=engine_manager, ) - mocker.patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager) mocker.patch("signal.signal") - run_headless(args) + run_headless(_make_headless_args(stage_id=0)) - assert mock_register.call_args.kwargs["replica_id"] == 2 + # The launcher must request auto-assignment (replica_id=None) and the + # full response so it can wire the master-allocated coordinator into the + # spawned subprocess. ``replica_binds_sockets=False`` is required for LLM + # because the head binds all three sockets (handshake, input, output). + assert mock_register.call_count == 1 + kwargs = mock_register.call_args.kwargs + assert kwargs["omni_master_address"] == "127.0.0.1" + assert kwargs["omni_master_port"] == 26000 + assert kwargs["omni_stage_id"] == 0 + assert kwargs["omni_stage_config"] is stage_cfg + assert kwargs["replica_id"] is None + assert kwargs["return_full_response"] is True + assert kwargs["replica_binds_sockets"] is False + + assert mock_manager_cls.call_count == 1 + mgr_kwargs = mock_manager_cls.call_args.kwargs + assert mgr_kwargs["local_engine_count"] == 1 + assert mgr_kwargs["local_client"] is False + assert mgr_kwargs["handshake_address"] == "tcp://127.0.0.1:26001" + assert mgr_kwargs["omni_stage_id"] == 0 + assert mgr_kwargs["omni_coordinator_address"] == "tcp://127.0.0.1:26100" + assert mgr_kwargs["omni_replica_base_id"] == 0 + + engine_manager.monitor_engine_liveness.assert_called_once_with() + engine_manager.shutdown.assert_called_once_with() -def test_run_headless_launches_diffusion_stage_via_omni_master(mocker: MockerFixture) -> None: - args = _make_headless_args() - args.replica_id = 1 - stage_cfg = mocker.Mock(stage_id=3, stage_type="diffusion") - stage_cfg.engine_args = mocker.Mock() - stage_cfg.engine_input_source = [] - stage_cfgs = [stage_cfg] - metadata = mocker.Mock(stage_id=3, stage_type="diffusion") +def test_run_headless_llm_launches_one_manager_per_omni_dp_size_local(mocker: MockerFixture) -> None: + """``--omni-dp-size-local=N`` must spawn N managers, each with its own + master-assigned replica_id, and join all of them before returning.""" + from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse + + stage_cfg = _make_stage_cfg(0, stage_type="llm") + parallel_config = SimpleNamespace( + data_parallel_size_local=1, + data_parallel_rank=0, + data_parallel_rank_local=0, + node_rank_within_dp=0, + ) + vllm_config = SimpleNamespace(parallel_config=parallel_config, needs_dp_coordinator=False) + manager_a = mocker.Mock() + manager_b = mocker.Mock() + + mocker.patch( + "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", + return_value=("/fake/stages.yaml", [stage_cfg]), + ) + mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") + mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None) + mocker.patch( + "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage", + return_value=(None, None, None), + ) + mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={}) + mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={}) + mocker.patch( + "vllm_omni.engine.stage_init_utils.build_vllm_config", + return_value=(vllm_config, object), + ) + mocker.patch( + "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", + side_effect=[ + StageRegistrationResponse( + handshake_address=f"tcp://127.0.0.1:2700{idx}", + input_address=f"tcp://127.0.0.1:2710{idx}", + output_address=f"tcp://127.0.0.1:2720{idx}", + replica_id=idx, + coordinator_router_address=None, + ) + for idx in (0, 1) + ], + ) + mock_manager_cls = mocker.patch( + "vllm_omni.engine.omni_core_engine_proc_manager.OmniCoreEngineProcManager", + side_effect=[manager_a, manager_b], + ) + mocker.patch("signal.signal") + + run_headless(_make_headless_args(stage_id=0, omni_dp_size_local=2)) + + assert mock_manager_cls.call_count == 2 + assigned_ids = [call.kwargs["omni_replica_base_id"] for call in mock_manager_cls.call_args_list] + assert assigned_ids == [0, 1] + + # Multi-replica path joins the monitor threads instead of calling + # ``monitor_engine_liveness`` synchronously on the main thread, but every + # manager must still be shut down in the finally block. + manager_a.shutdown.assert_called_once_with() + manager_b.shutdown.assert_called_once_with() + + +def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture) -> None: + """Diffusion headless: registers as auto-assign, spawns a single + ``StageDiffusionProc`` per local replica, and waits for it via + ``multiprocessing.connection.wait``.""" + from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse + + stage_cfg = _make_stage_cfg(1, stage_type="diffusion") od_config = mocker.Mock() - proc = mocker.Mock() - proc.exitcode = 0 + proc = mocker.Mock(sentinel=object(), exitcode=0) proc.is_alive.return_value = False mocker.patch( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", - return_value=("/fake/stages.yaml", stage_cfgs), + return_value=("/fake/stages.yaml", [stage_cfg]), ) mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") - mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=mocker.Mock()) + mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None) mocker.patch( "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage", return_value=(None, None, None), ) - mocker.patch("vllm_omni.engine.stage_init_utils.extract_stage_metadata", return_value=metadata) - mock_inject_stage_info = mocker.patch("vllm_omni.engine.stage_init_utils.inject_kv_stage_info") + mocker.patch( + "vllm_omni.engine.stage_init_utils.extract_stage_metadata", + return_value=SimpleNamespace(stage_id=1, stage_type="diffusion"), + ) + mock_inject = mocker.patch("vllm_omni.engine.stage_init_utils.inject_kv_stage_info") mocker.patch("vllm_omni.engine.stage_init_utils.build_diffusion_config", return_value=od_config) mock_register = mocker.patch( "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", - return_value=("tcp://127.0.0.1:26001", "tcp://127.0.0.1:26002", "tcp://127.0.0.1:26003"), + return_value=StageRegistrationResponse( + handshake_address="tcp://127.0.0.1:26001", + input_address="tcp://127.0.0.1:26002", + output_address="tcp://127.0.0.1:26003", + replica_id=0, + coordinator_router_address="tcp://127.0.0.1:26100", + ), ) mock_spawn = mocker.patch( "vllm_omni.diffusion.stage_diffusion_proc.spawn_diffusion_proc", - return_value=(proc, "tcp://127.0.0.1:26001", "tcp://127.0.0.1:26002", "tcp://127.0.0.1:26003"), + return_value=(proc, None, None, None), ) mock_handshake = mocker.patch("vllm_omni.diffusion.stage_diffusion_proc.complete_diffusion_handshake") + # Replace the blocking wait with one that returns the only proc's sentinel + # immediately so the test does not hang. + mocker.patch( + "multiprocessing.connection.wait", + side_effect=lambda sentinels: [sentinels[0]], + ) + mocker.patch("vllm_omni.engine.stage_init_utils.terminate_alive_proc") mocker.patch("signal.signal") - run_headless(args) - mock_inject_stage_info.assert_called_once_with(stage_cfg, 3) - mock_register.assert_called_once_with( - omni_master_address="127.0.0.1", - omni_master_port=26000, - omni_stage_id=3, - omni_stage_config=stage_cfg, - return_addresses=True, - replica_id=1, - ) - mock_spawn.assert_called_once_with( - "fake-model", - od_config, - handshake_address="tcp://127.0.0.1:26001", - request_address="tcp://127.0.0.1:26002", - response_address="tcp://127.0.0.1:26003", - ) + run_headless(_make_headless_args(stage_id=1)) + + mock_inject.assert_called_once_with(stage_cfg, 1) + + reg_kwargs = mock_register.call_args.kwargs + assert reg_kwargs["omni_master_address"] == "127.0.0.1" + assert reg_kwargs["omni_master_port"] == 26000 + assert reg_kwargs["omni_stage_id"] == 1 + assert reg_kwargs["omni_stage_config"] is stage_cfg + assert reg_kwargs["replica_id"] is None + assert reg_kwargs["return_full_response"] is True + + spawn_kwargs = mock_spawn.call_args.kwargs + assert spawn_kwargs["handshake_address"] == "tcp://127.0.0.1:26001" + assert spawn_kwargs["request_address"] == "tcp://127.0.0.1:26002" + assert spawn_kwargs["response_address"] == "tcp://127.0.0.1:26003" + assert spawn_kwargs["omni_coordinator_address"] == "tcp://127.0.0.1:26100" + assert spawn_kwargs["omni_stage_id"] == 1 + assert spawn_kwargs["omni_replica_id"] == 0 + mock_handshake.assert_called_once_with(proc, "tcp://127.0.0.1:26001", 600) - proc.join.assert_called_once_with() + + +def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixture) -> None: + """A diffusion replica that exits with a non-zero code must surface as a + RuntimeError from ``run_headless`` (the head needs the signal to roll + back its own stage init).""" + from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse + + stage_cfg = _make_stage_cfg(1, stage_type="diffusion") + proc = mocker.Mock(sentinel=object(), exitcode=137, name="proc-stage1-rep0") + proc.is_alive.return_value = False + + mocker.patch( + "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", + return_value=("/fake/stages.yaml", [stage_cfg]), + ) + mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") + mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None) + mocker.patch( + "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage", + return_value=(None, None, None), + ) + mocker.patch( + "vllm_omni.engine.stage_init_utils.extract_stage_metadata", + return_value=SimpleNamespace(stage_id=1, stage_type="diffusion"), + ) + mocker.patch("vllm_omni.engine.stage_init_utils.inject_kv_stage_info") + mocker.patch("vllm_omni.engine.stage_init_utils.build_diffusion_config", return_value=mocker.Mock()) + mocker.patch( + "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", + return_value=StageRegistrationResponse( + handshake_address="tcp://127.0.0.1:26001", + input_address="tcp://127.0.0.1:26002", + output_address="tcp://127.0.0.1:26003", + replica_id=0, + coordinator_router_address=None, + ), + ) + mocker.patch( + "vllm_omni.diffusion.stage_diffusion_proc.spawn_diffusion_proc", + return_value=(proc, None, None, None), + ) + mocker.patch("vllm_omni.diffusion.stage_diffusion_proc.complete_diffusion_handshake") + mocker.patch( + "multiprocessing.connection.wait", + side_effect=lambda sentinels: [sentinels[0]], + ) + mocker.patch("vllm_omni.engine.stage_init_utils.terminate_alive_proc") + mocker.patch("signal.signal") + + with pytest.raises(RuntimeError, match=r"exited with code 137"): + run_headless(_make_headless_args(stage_id=1)) diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index 476b44b5b..6ab44b4b0 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/vllm_omni/diffusion/stage_diffusion_client.py @@ -17,6 +17,7 @@ from threading import Thread from typing import TYPE_CHECKING, Any import zmq +import zmq.asyncio from vllm.logger import init_logger from vllm.v1.engine.exceptions import EngineDeadError @@ -124,6 +125,11 @@ class StageDiffusionClient(StageClientBase): self.engine_input_source = getattr(metadata, "engine_input_source", []) self._proc = proc self._owns_process = proc is not None + # Expose the ZMQ addresses on the instance so callers (e.g. + # ``StagePool._client_input_addr``) can identify the diffusion + # replica by its bound address. + self.request_address = request_address + self.response_address = response_address self._zmq_ctx = zmq.Context() self._request_socket = self._zmq_ctx.socket(zmq.PUSH) @@ -131,6 +137,9 @@ class StageDiffusionClient(StageClientBase): self._response_socket = self._zmq_ctx.socket(zmq.PULL) self._response_socket.connect(response_address) + self._response_poller = zmq.asyncio.Poller() + self._response_poller.register(self._response_socket, zmq.POLLIN) + self._encoder = OmniMsgpackEncoder() self._decoder = OmniMsgpackDecoder() @@ -468,17 +477,27 @@ class StageDiffusionClient(StageClientBase): try: while True: self._drain_responses() - if rpc_id in self._rpc_results: - return self._rpc_results.pop(rpc_id) + result = self._rpc_results.pop(rpc_id, None) + if result is not None: + return result if self._engine_dead or (self._owns_process and self._proc is not None and not self._proc.is_alive()): self._engine_dead = True raise EngineDeadError( f"StageDiffusionProc died while waiting for " f"collective_rpc '{method}' (exit code {self._proc.exitcode})" ) - if deadline and time.monotonic() > deadline: + if deadline is not None and time.monotonic() > deadline: raise TimeoutError(f"collective_rpc_async '{method}' timed out after {timeout}s") - await asyncio.sleep(0.01) + # Block (async) until data arrives on the ZMQ response + # socket or until the timeout expires, then loop back to + # drain and check. + if deadline is not None: + poll_timeout_ms = max(int((deadline - time.monotonic()) * 1000), 0) + else: + poll_timeout_ms = 100 + # no exception raised on timeout (capped at 100ms so the + # engine-dead check still fires regularly). + await self._response_poller.poll(timeout=min(poll_timeout_ms, 100)) finally: self._pending_rpcs.discard(rpc_id) diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index 2f3ad5594..871a29729 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/vllm_omni/diffusion/stage_diffusion_proc.py @@ -7,6 +7,7 @@ communicating with StageDiffusionClient via ZMQ (PUSH/PULL). from __future__ import annotations import asyncio +import contextlib import signal import time from concurrent.futures import ThreadPoolExecutor @@ -30,6 +31,7 @@ from vllm_omni.distributed.omni_connectors.utils.serialization import ( OmniMsgpackDecoder, OmniMsgpackEncoder, ) +from vllm_omni.distributed.omni_coordinator import OmniCoordClientForStage from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -54,6 +56,56 @@ class StageDiffusionProc: self._engine: DiffusionEngine | None = None self._executor: ThreadPoolExecutor | None = None self._closed = False + # Set by ``run_loop`` to the live dispatch task dict so + # :attr:`queue_length` can report in-flight requests for the + # OmniCoordinator heartbeat hook. + self._active_tasks: dict[str, asyncio.Task] | None = None + # Set when a request-handler detects that the engine's multiproc + # executor has died (e.g. a worker process crashed and the executor's + # monitor thread closed it). Once set, ``run_loop`` breaks out so the + # outer ``except``/``finally`` can send ``DIFFUSION_PROC_DEAD`` and + # ``ReplicaStatus.DOWN``, then the subprocess exits non-zero. Without + # this, the run_loop would swallow per-request errors and keep + # serving 500s indefinitely while heartbeats still report UP. + self._fatal_event: asyncio.Event | None = None + + @property + def queue_length(self) -> int: + """Number of in-flight diffusion requests. + + Returns 0 before :meth:`run_loop` starts and after it exits. + """ + tasks = self._active_tasks + return 0 if tasks is None else len(tasks) + + def _is_executor_dead(self) -> bool: + """True iff the multiproc executor has been closed or marked failed. + + Detects the "workers died but the diffusion proc is still pulling + requests" case: ``MultiprocDiffusionExecutor`` sets ``_closed = True`` + and ``is_failed = True`` from its worker-monitor thread the moment any + worker process exits; every subsequent ``execute_request`` / + ``collective_rpc`` then raises ``RuntimeError("DiffusionExecutor is + closed.")`` inside the engine. Callers in ``run_loop`` use this to + decide whether a per-request failure is recoverable or fatal. + """ + if self._engine is None: + return False + executor = getattr(self._engine, "executor", None) + if executor is None: + return False + return bool(getattr(executor, "_closed", False) or getattr(executor, "is_failed", False)) + + def _signal_fatal_engine_failure(self, reason: str) -> None: + """Idempotently signal ``run_loop`` to tear down on a fatal engine error.""" + if self._fatal_event is None or self._fatal_event.is_set(): + return + logger.error( + "[StageDiffusionProc] fatal engine failure detected (%s); " + "signaling run_loop to send DIFFUSION_PROC_DEAD and exit.", + reason, + ) + self._fatal_event.set() # ------------------------------------------------------------------ # Initialization @@ -309,6 +361,14 @@ class StageDiffusionProc: decoder = OmniMsgpackDecoder() tasks: dict[str, asyncio.Task] = {} + # Expose the live task dict so :attr:`queue_length` (used by the + # OmniCoordinator heartbeat hook) can read the in-flight count. + self._active_tasks = tasks + # Wakes the main recv loop when a request-handler detects a fatal + # engine failure so we tear down promptly instead of swallowing + # "DiffusionExecutor is closed" on every subsequent request. + fatal_event = asyncio.Event() + self._fatal_event = fatal_event async def _dispatch_request( request_id: str, @@ -342,12 +402,40 @@ class StageDiffusionProc: } ) ) + # Per-request errors are usually recoverable, but a closed + # executor means every future request will get the same + # "DiffusionExecutor is closed" error. Signal the main loop + # to send DIFFUSION_PROC_DEAD and exit so the head's hub + # demotes this replica instead of waiting on the heartbeat + # timeout (~30 s by default). + if self._is_executor_dead(): + self._signal_fatal_engine_failure(f"add_request {request_id}: {e!s}") finally: tasks.pop(request_id, None) try: while True: - raw = await request_socket.recv() + # Await recv and fatal_event concurrently so the loop wakes + # up immediately when a per-request handler signals a fatal + # engine failure — even if no fresh ZMQ frame arrives. + recv_task: asyncio.Task = asyncio.ensure_future(request_socket.recv()) + fatal_task: asyncio.Task = asyncio.ensure_future(fatal_event.wait()) + try: + done, pending = await asyncio.wait( + [recv_task, fatal_task], + return_when=asyncio.FIRST_COMPLETED, + ) + finally: + for waiter in (recv_task, fatal_task): + if not waiter.done(): + waiter.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await waiter + if fatal_event.is_set(): + raise RuntimeError( + "StageDiffusionProc executor reported permanent failure; tearing down the diffusion subprocess." + ) + raw = recv_task.result() msg = decoder.decode(raw) msg_type = msg.get("type") @@ -397,6 +485,11 @@ class StageDiffusionProc: } ) ) + # Same rationale as the single-request path: a + # closed executor turns every subsequent batch + # into a 500, so escalate now. + if self._is_executor_dead(): + self._signal_fatal_engine_failure(f"add_batch_request {rid}: {e!s}") finally: tasks.pop(rid, None) @@ -446,6 +539,13 @@ class StageDiffusionProc: } ) ) + # Collective RPCs run through the same multiproc + # executor — a closed executor means every future + # RPC fails the same way, so tear down promptly. + if self._is_executor_dead(): + self._signal_fatal_engine_failure( + f"collective_rpc {msg['method']} (rpc_id={rpc_id}): {e!s}" + ) elif msg_type == "shutdown": break @@ -466,6 +566,8 @@ class StageDiffusionProc: if tasks: await asyncio.gather(*tasks.values(), return_exceptions=True) + self._active_tasks = None + self._fatal_event = None request_socket.close() response_socket.close() ctx.term() @@ -504,8 +606,22 @@ class StageDiffusionProc: handshake_address: str, request_address: str, response_address: str, + *, + omni_coordinator_address: str | None = None, + omni_stage_id: int | None = None, + omni_replica_id: int = 0, ) -> None: - """Entry point for the diffusion subprocess.""" + """Entry point for the diffusion subprocess. + + Omni-specific kwargs (mirroring :meth:`StageEngineCoreProc.run_stage_core`): + - ``omni_coordinator_address``: ROUTER address of the head-side + OmniCoordinator. When set, a :class:`OmniCoordClientForStage` + reports the diffusion replica's status + queue length. + - ``omni_stage_id``: logical stage id; required when + ``omni_coordinator_address`` is set. + - ``omni_replica_id``: cluster-unique replica id within the + stage (logging / metrics only). + """ shutdown_requested = False def signal_handler(signum: int, frame: Any) -> None: @@ -518,6 +634,7 @@ class StageDiffusionProc: signal.signal(signal.SIGINT, signal_handler) proc = cls(model, od_config) + coord_client: OmniCoordClientForStage | None = None try: proc.initialize() @@ -529,6 +646,32 @@ class StageDiffusionProc: handshake_socket.close() handshake_ctx.term() + # Wire OmniCoordClientForStage *after* READY so that the head + # has bound its head-side request/response sockets — the + # address pair we report is the same pair this proc binds to + # (request/response addresses passed in). + if omni_coordinator_address is not None: + if omni_stage_id is None: + raise ValueError("omni_stage_id must be provided when omni_coordinator_address is set") + coord_client = OmniCoordClientForStage( + coord_zmq_addr=omni_coordinator_address, + input_addr=request_address, + output_addr=response_address, + stage_id=int(omni_stage_id), + ) + + def _refresh_queue_length() -> None: + coord_client._queue_length = proc.queue_length # type: ignore[union-attr] + + coord_client._on_heartbeat = _refresh_queue_length + + logger.info( + "StageDiffusionProc registered with OmniCoordinator (stage_id=%d replica_id=%d coord=%s)", + omni_stage_id, + omni_replica_id, + omni_coordinator_address, + ) + # Run async event loop asyncio.run(proc.run_loop(request_address, response_address)) @@ -539,6 +682,9 @@ class StageDiffusionProc: logger.exception("StageDiffusionProc encountered a fatal error.") raise finally: + if coord_client is not None: + with contextlib.suppress(RuntimeError): + coord_client.close() proc.close() @@ -551,10 +697,17 @@ def spawn_diffusion_proc( handshake_address: str | None = None, request_address: str | None = None, response_address: str | None = None, + *, + omni_coordinator_address: str | None = None, + omni_stage_id: int | None = None, + omni_replica_id: int = 0, ) -> tuple[BaseProcess, str, str, str]: """Spawn a StageDiffusionProc subprocess. Returns ``(proc, handshake_address, request_address, response_address)``. + + Pass ``omni_coordinator_address`` / ``omni_stage_id`` / ``omni_replica_id`` + to have the subprocess publish heartbeats to an OmniCoordinator. """ handshake_address = handshake_address or get_open_zmq_ipc_path() request_address = request_address or get_open_zmq_ipc_path() @@ -570,6 +723,9 @@ def spawn_diffusion_proc( "handshake_address": handshake_address, "request_address": request_address, "response_address": response_address, + "omni_coordinator_address": omni_coordinator_address, + "omni_stage_id": omni_stage_id, + "omni_replica_id": omni_replica_id, }, ) proc.start() diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index 7fdd66cdb..241a6070d 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -21,6 +21,7 @@ from .utils.kv_utils import ( build_rank_aware_send_keys, get_kv_target_ranks, get_local_tp_rank, + get_omni_replica_id, get_tp_world_size, kv_zmq_port, merge_received_rank_shards, @@ -415,7 +416,13 @@ class OmniKVTransferManager: stage_int = int(self.config.from_stage) if self.config.from_stage is not None else 0 except (TypeError, ValueError): stage_int = 0 - zmq_port = kv_zmq_port(base_port, stage_int, self._tp_topo.local_rank) + replica_id = get_omni_replica_id() + zmq_port = kv_zmq_port( + base_port, + stage_int, + self._tp_topo.local_rank, + replica_id=replica_id, + ) if self.config.need_send_cache: c_extra["role"] = "sender" diff --git a/vllm_omni/distributed/omni_connectors/utils/initialization.py b/vllm_omni/distributed/omni_connectors/utils/initialization.py index e6d68f9bc..a8cafba61 100644 --- a/vllm_omni/distributed/omni_connectors/utils/initialization.py +++ b/vllm_omni/distributed/omni_connectors/utils/initialization.py @@ -26,9 +26,17 @@ KV_TRANSFER_PORT_OFFSET = 100 # Port stride between TP ranks so each worker binds a unique ZMQ port # when TP > 1. Must be larger than the maximum number of pipeline stages. -# Formula: zmq_port = base + KV_TRANSFER_PORT_OFFSET + rank * STRIDE + stage +# Formula: +# zmq_port = base + KV_TRANSFER_PORT_OFFSET +# + replica * KV_REPLICA_PORT_STRIDE +# + rank * KV_RANK_PORT_STRIDE +# + stage KV_RANK_PORT_STRIDE = 16 +# Port stride between Omni replicas of the same stage. This reserves a +# comfortably sized block per replica for TP-rank and stage offsets. +KV_REPLICA_PORT_STRIDE = 1024 + def initialize_connectors_from_config( config_path: str | Path | None = None, diff --git a/vllm_omni/distributed/omni_connectors/utils/kv_utils.py b/vllm_omni/distributed/omni_connectors/utils/kv_utils.py index 12b9b3d4f..fe0df21f2 100644 --- a/vllm_omni/distributed/omni_connectors/utils/kv_utils.py +++ b/vllm_omni/distributed/omni_connectors/utils/kv_utils.py @@ -16,7 +16,7 @@ from vllm.distributed.parallel_state import ( ) from vllm.logger import init_logger -from .initialization import KV_RANK_PORT_STRIDE, KV_TRANSFER_PORT_OFFSET +from .initialization import KV_RANK_PORT_STRIDE, KV_REPLICA_PORT_STRIDE, KV_TRANSFER_PORT_OFFSET logger = init_logger(__name__) @@ -94,20 +94,42 @@ def get_tp_world_size() -> int: return 1 +def get_omni_replica_id() -> int: + """Return the Omni replica id for this worker process.""" + try: + replica_id = int(os.environ.get("VLLM_OMNI_REPLICA_ID", "0")) + except (ValueError, TypeError): + return 0 + return max(replica_id, 0) + + # ------------------------------------------------------------------ # # ZMQ port computation # ------------------------------------------------------------------ # -def kv_zmq_port(base_port: int, from_stage: int, local_rank: int = 0) -> int: +def kv_zmq_port( + base_port: int, + from_stage: int, + local_rank: int = 0, + replica_id: int | None = None, +) -> int: """Compute the ZMQ port for a KV-transfer connector. - Each TP rank gets its own port so that TP > 1 deployments do not - cause ``EADDRINUSE`` when multiple sender workers bind on the same - host. The formula is backward-compatible: rank 0 produces the same - port as the previous ``base + OFFSET + stage`` formula. + Each Omni replica and TP rank gets its own port so multi-replica or + TP > 1 deployments do not cause ``EADDRINUSE`` when multiple sender + workers bind on the same host. The formula is backward-compatible: + replica 0 / rank 0 produces the previous ``base + OFFSET + stage`` port. + """ - return base_port + KV_TRANSFER_PORT_OFFSET + local_rank * KV_RANK_PORT_STRIDE + from_stage + replica = get_omni_replica_id() if replica_id is None else max(int(replica_id), 0) + return ( + base_port + + KV_TRANSFER_PORT_OFFSET + + replica * KV_REPLICA_PORT_STRIDE + + local_rank * KV_RANK_PORT_STRIDE + + from_stage + ) # ------------------------------------------------------------------ # diff --git a/vllm_omni/distributed/omni_coordinator/__init__.py b/vllm_omni/distributed/omni_coordinator/__init__.py index 6894e3113..326fb9b9a 100644 --- a/vllm_omni/distributed/omni_coordinator/__init__.py +++ b/vllm_omni/distributed/omni_coordinator/__init__.py @@ -8,18 +8,21 @@ from .load_balancer import ( RandomBalancer, RoundRobinBalancer, Task, + build_load_balancer_factory, ) -from .messages import InstanceEvent, InstanceInfo, InstanceList, StageStatus +from .messages import ReplicaEvent, ReplicaInfo, ReplicaList, ReplicaStatus from .omni_coord_client_for_hub import OmniCoordClientForHub from .omni_coord_client_for_stage import OmniCoordClientForStage from .omni_coordinator import OmniCoordinator +from .runtime import OmniCoordinatorRuntime __all__ = [ "OmniCoordinator", - "StageStatus", - "InstanceEvent", - "InstanceInfo", - "InstanceList", + "OmniCoordinatorRuntime", + "ReplicaStatus", + "ReplicaEvent", + "ReplicaInfo", + "ReplicaList", "OmniCoordClientForStage", "OmniCoordClientForHub", "Task", @@ -28,4 +31,5 @@ __all__ = [ "RandomBalancer", "RoundRobinBalancer", "LeastQueueLengthBalancer", + "build_load_balancer_factory", ] diff --git a/vllm_omni/distributed/omni_coordinator/load_balancer.py b/vllm_omni/distributed/omni_coordinator/load_balancer.py index 41b03be16..bcd9822ec 100644 --- a/vllm_omni/distributed/omni_coordinator/load_balancer.py +++ b/vllm_omni/distributed/omni_coordinator/load_balancer.py @@ -6,17 +6,18 @@ from __future__ import annotations import random import threading from abc import ABC, abstractmethod +from collections.abc import Callable from enum import Enum from typing import Any, TypedDict -from .messages import InstanceInfo +from .messages import ReplicaInfo class Task(TypedDict, total=False): - """Task structure passed from async_omni (stage.submit(task)). + """Task structure passed to ``StagePool.pick`` / ``LoadBalancer.select``. - Mirrors the dict built in AsyncOmni with request_id, engine_inputs, - sampling_params. Future load-balancing policies may use these fields. + Mirrors the dict built around a stage submission with request_id and any + payload-related fields a future load-balancing policy might inspect. """ request_id: str @@ -28,10 +29,7 @@ class LoadBalancingPolicy(str, Enum): """Enumeration for load balancing policies. These policies are used by :class:`LoadBalancer` implementations to route - tasks to a subset of available instances. - - TODO(NumberWan): Map enum values to balancer classes when OmniCoordinator - integration lands. Tracked in https://github.com/vllm-project/vllm-omni/pull/2448 + tasks to a subset of available replicas. """ RANDOM = "random" @@ -42,74 +40,60 @@ class LoadBalancingPolicy(str, Enum): class LoadBalancer(ABC): """Abstract base class for load balancers. - Subclasses implement :meth:`select` to choose an instance for a given task. + Subclasses implement :meth:`select` to choose a replica for a given task. """ @abstractmethod - def select(self, task: Task, instances: list[InstanceInfo]) -> int: - """Route a task to one of the available instances. + def select(self, task: Task, replicas: list[ReplicaInfo]) -> int: + """Route a task to one of the available replicas. Args: task: The task to route. Not used by the random policy but reserved for future strategies that may inspect task metadata. - instances: List of available instances to choose from. + replicas: List of available replicas to choose from. Returns: - Index of the selected instance in ``instances``. + Index of the selected replica in ``replicas``. Raises: - ValueError: If ``instances`` is empty. + ValueError: If ``replicas`` is empty. """ raise NotImplementedError class RandomBalancer(LoadBalancer): - """Load balancer that selects an instance uniformly at random. + """Load balancer that selects a replica uniformly at random.""" - It intentionally ignores the task payload and chooses a random index from - the provided instance list. More sophisticated policies (e.g. round-robin, - least-queue-length) can be implemented as additional subclasses of - :class:`LoadBalancer`. - """ + def select(self, task: Task, replicas: list[ReplicaInfo]) -> int: # noqa: ARG002 + if not replicas: + raise ValueError("replicas must not be empty") - def select(self, task: Task, instances: list[InstanceInfo]) -> int: # noqa: ARG002 - if not instances: - raise ValueError("instances must not be empty") - - return random.randrange(len(instances)) + return random.randrange(len(replicas)) class RoundRobinBalancer(LoadBalancer): - """Load balancer that selects instances in a round-robin fashion. + """Load balancer that selects replicas in a round-robin fashion. - This implementation keeps a running index modulo ``len(instances)``. It - therefore depends on the **order and stable meaning** of the ``instances`` + This implementation keeps a running index modulo ``len(replicas)``. It + therefore depends on the **order and stable meaning** of the ``replicas`` list between calls. If the list length or ordering changes, the sequence of picks may skip or repeat entries relative to a fixed set of backends. - When instance membership changes dynamically, callers should reset routing - state—for example by constructing a new ``RoundRobinBalancer`` or resetting - ``_next_index``—similar to rebuilding ``itertools.cycle`` after mutating - the instance list (as in vLLM's disaggregated proxy examples). - - Concurrency: ``select`` is synchronous and is expected to run on the - coordinator asyncio event loop thread without ``await`` inside this - method, so a single invocation is not interleaved with another on that - thread. A :class:`threading.Lock` still serializes updates to - ``_next_index`` for callers that might invoke ``select`` from multiple - threads or alongside threaded infrastructure (e.g. ZMQ receive threads). + Concurrency: a ``threading.Lock`` serializes updates to ``_next_index`` + for callers that invoke ``select`` from multiple threads or alongside + threaded infrastructure (e.g. ZMQ receive threads). """ def __init__(self, start_index: int = 0) -> None: self._next_index = start_index self._lock = threading.Lock() - def select(self, task: Task, instances: list[InstanceInfo]) -> int: # noqa: ARG002 - if not instances: - raise ValueError("instances must not be empty") + def select(self, task: Task, replicas: list[ReplicaInfo]) -> int: # noqa: ARG002 + if not replicas: + raise ValueError("replicas must not be empty") - n = len(instances) + n = len(replicas) with self._lock: idx = self._next_index % n self._next_index = (self._next_index + 1) % n @@ -117,27 +101,43 @@ class RoundRobinBalancer(LoadBalancer): class LeastQueueLengthBalancer(LoadBalancer): - """Select the instance with the smallest ``queue_length``. + """Select the replica with the smallest ``queue_length``. - If multiple instances share the same minimum queue length, one of them is + If multiple replicas share the same minimum queue length, one of them is chosen uniformly at random. Raises: - ValueError: If any instance has a negative ``queue_length``. + ValueError: If any replica has a negative ``queue_length``. """ - def select(self, task: Task, instances: list[InstanceInfo]) -> int: # noqa: ARG002 - if not instances: - raise ValueError("instances must not be empty") + def select(self, task: Task, replicas: list[ReplicaInfo]) -> int: # noqa: ARG002 + if not replicas: + raise ValueError("replicas must not be empty") - queue_lengths = [inst.queue_length for inst in instances] + queue_lengths = [rep.queue_length for rep in replicas] if any(q < 0 for q in queue_lengths): - raise ValueError("queue_length must be non-negative for all instances") + raise ValueError("queue_length must be non-negative for all replicas") min_q = min(queue_lengths) candidates = [i for i, q in enumerate(queue_lengths) if q == min_q] return random.choice(candidates) +def build_load_balancer_factory(policy: str) -> Callable[[], LoadBalancer]: + """Translate ``--omni-lb-policy`` (string) into a per-pool LB factory.""" + try: + normalized = LoadBalancingPolicy(policy) + except ValueError as exc: + valid = ", ".join(p.value for p in LoadBalancingPolicy) + raise ValueError(f"unknown --omni-lb-policy {policy!r} (valid: {valid})") from exc + if normalized is LoadBalancingPolicy.RANDOM: + return RandomBalancer + if normalized is LoadBalancingPolicy.ROUND_ROBIN: + return RoundRobinBalancer + if normalized is LoadBalancingPolicy.LEAST_QUEUE_LENGTH: + return LeastQueueLengthBalancer + raise ValueError(f"unhandled load balancing policy {normalized!r}") + + __all__ = [ "Task", "LoadBalancingPolicy", @@ -145,4 +145,5 @@ __all__ = [ "RandomBalancer", "RoundRobinBalancer", "LeastQueueLengthBalancer", + "build_load_balancer_factory", ] diff --git a/vllm_omni/distributed/omni_coordinator/messages.py b/vllm_omni/distributed/omni_coordinator/messages.py index 2bb590139..824475195 100644 --- a/vllm_omni/distributed/omni_coordinator/messages.py +++ b/vllm_omni/distributed/omni_coordinator/messages.py @@ -7,55 +7,55 @@ from dataclasses import dataclass from enum import Enum -class StageStatus(str, Enum): - """Enumeration for stage instance status.""" +class ReplicaStatus(str, Enum): + """Enumeration for stage replica status.""" - UP = "up" # Instance is ready and available - DOWN = "down" # Instance is shutdown gracefully - ERROR = "error" # Instance encountered an error or timeout + UP = "up" # Replica is ready and available + DOWN = "down" # Replica is shutdown gracefully + ERROR = "error" # Replica encountered an error or timeout @dataclass -class InstanceEvent: +class ReplicaEvent: """Wire payload from OmniCoordClientForStage to OmniCoordinator. Schema for Stage → Coordinator events over ZMQ: input_addr, output_addr, stage_id, status, queue_length, event_type. """ - input_addr: str # Stage instance input ZMQ address (e.g., "tcp://host:port") - output_addr: str # Stage instance output ZMQ address (e.g., "tcp://host:port") + input_addr: str # Stage replica input ZMQ address (e.g., "tcp://host:port") + output_addr: str # Stage replica output ZMQ address (e.g., "tcp://host:port") stage_id: int # Stage ID event_type: str # "update" | "heartbeat" - status: StageStatus # Current status + status: ReplicaStatus # Current status queue_length: int # Current queue length @dataclass -class InstanceInfo: - """Metadata for a single stage instance. +class ReplicaInfo: + """Metadata for a single stage replica. This type is stored in OmniCoordinator's internal registry and is also - published to hubs via :class:`InstanceList`. + published to hubs via :class:`ReplicaList`. """ - input_addr: str # Stage instance input ZMQ address (e.g., "tcp://host:port") - output_addr: str # Stage instance output ZMQ address (e.g., "tcp://host:port") - stage_id: int # Stage ID of this instance - status: StageStatus # Current status of the instance - queue_length: int # Current queue length of this instance + input_addr: str # Stage replica input ZMQ address (e.g., "tcp://host:port") + output_addr: str # Stage replica output ZMQ address (e.g., "tcp://host:port") + stage_id: int # Stage ID of this replica + status: ReplicaStatus # Current status of the replica + queue_length: int # Current queue length of this replica last_heartbeat: float # Timestamp of the last heartbeat received (seconds) - registered_at: float # Timestamp when the instance was registered (seconds) + registered_at: float # Timestamp when the replica was registered (seconds) @dataclass -class InstanceList: - """Container for instance list updates. +class ReplicaList: + """Container for replica list updates. - OmniCoordinator publishes an :class:`InstanceList` whenever its view of - active instances changes. OmniCoordClientForHub caches the latest value + OmniCoordinator publishes a :class:`ReplicaList` whenever its view of + active replicas changes. OmniCoordClientForHub caches the latest value and exposes it to AsyncOmni and the load balancer. """ - instances: list[InstanceInfo] + replicas: list[ReplicaInfo] timestamp: float # Time when the list was last updated (seconds) diff --git a/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_hub.py b/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_hub.py index 9081e4591..f4b6e5f80 100644 --- a/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_hub.py +++ b/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_hub.py @@ -9,16 +9,16 @@ from typing import Any import zmq -from .messages import InstanceInfo, InstanceList, StageStatus +from .messages import ReplicaInfo, ReplicaList, ReplicaStatus logger = logging.getLogger(__name__) class OmniCoordClientForHub: - """Client for AsyncOmni side to receive instance list updates. + """Client for AsyncOmni side to receive replica list updates. This client maintains a SUB socket connected to OmniCoordinator's PUB - endpoint and caches the latest :class:`InstanceList` in memory for use by + endpoint and caches the latest :class:`ReplicaList` in memory for use by the load balancer and routing logic. """ @@ -28,7 +28,7 @@ class OmniCoordClientForHub: self._ctx = zmq.Context() self._lock = threading.Lock() - self._instance_list: InstanceList | None = None + self._replica_list: ReplicaList | None = None self._closed = False self._stop_event = threading.Event() self._init_done = threading.Event() @@ -41,29 +41,29 @@ class OmniCoordClientForHub: if self._init_error: raise RuntimeError(f"Failed to connect to coordinator at {self._coord_zmq_addr}") from self._init_error[0] - def _decode_instance_list(self, payload: dict[str, Any]) -> InstanceList: - """Convert a JSON-decoded dict into an :class:`InstanceList`.""" - instances_payload = payload.get("instances", []) - instances: list[InstanceInfo] = [] - for inst in instances_payload: - instances.append( - InstanceInfo( - input_addr=inst["input_addr"], - output_addr=inst["output_addr"], - stage_id=int(inst["stage_id"]), - status=StageStatus(inst["status"]), - queue_length=int(inst["queue_length"]), - last_heartbeat=float(inst["last_heartbeat"]), - registered_at=float(inst["registered_at"]), + def _decode_replica_list(self, payload: dict[str, Any]) -> ReplicaList: + """Convert a JSON-decoded dict into a :class:`ReplicaList`.""" + replicas_payload = payload.get("replicas", []) + replicas: list[ReplicaInfo] = [] + for rep in replicas_payload: + replicas.append( + ReplicaInfo( + input_addr=rep["input_addr"], + output_addr=rep["output_addr"], + stage_id=int(rep["stage_id"]), + status=ReplicaStatus(rep["status"]), + queue_length=int(rep["queue_length"]), + last_heartbeat=float(rep["last_heartbeat"]), + registered_at=float(rep["registered_at"]), ) ) timestamp = float(payload.get("timestamp", time())) - return InstanceList(instances=instances, timestamp=timestamp) + return ReplicaList(replicas=replicas, timestamp=timestamp) def _recv_loop(self) -> None: - """Background loop that receives and caches instance lists.""" - sub = None + """Background loop that receives and caches replica lists.""" + sub: zmq.Socket | None = None try: sub = self._ctx.socket(zmq.SUB) sub.setsockopt(zmq.SUBSCRIBE, b"") @@ -115,9 +115,9 @@ class OmniCoordClientForHub: try: payload = json.loads(data.decode("utf-8")) - inst_list = self._decode_instance_list(payload) + rep_list = self._decode_replica_list(payload) with self._lock: - self._instance_list = inst_list + self._replica_list = rep_list except ( json.JSONDecodeError, KeyError, @@ -125,7 +125,7 @@ class OmniCoordClientForHub: TypeError, AttributeError, ) as e: - logger.warning("Invalid instance list message, skipping: %s", e) + logger.warning("Invalid replica list message, skipping: %s", e) finally: try: if sub is not None: @@ -137,22 +137,22 @@ class OmniCoordClientForHub: except zmq.ZMQError: pass - def get_instance_list(self) -> InstanceList: - """Return the latest cached :class:`InstanceList`. + def get_replica_list(self) -> ReplicaList: + """Return the latest cached :class:`ReplicaList`. If no update has been received yet, returns an empty list with ``timestamp=0.0``. """ with self._lock: - if self._instance_list is None: - return InstanceList(instances=[], timestamp=0.0) - return self._instance_list + if self._replica_list is None: + return ReplicaList(replicas=[], timestamp=0.0) + return self._replica_list - def get_instances_for_stage(self, stage_id: int) -> InstanceList: - """Return instances filtered by ``stage_id``.""" - base = self.get_instance_list() - filtered = [inst for inst in base.instances if inst.stage_id == stage_id] - return InstanceList(instances=filtered, timestamp=base.timestamp) + def get_replicas_for_stage(self, stage_id: int) -> ReplicaList: + """Return replicas filtered by ``stage_id``.""" + base = self.get_replica_list() + filtered = [rep for rep in base.replicas if rep.stage_id == stage_id] + return ReplicaList(replicas=filtered, timestamp=base.timestamp) def close(self) -> None: """Close the SUB socket and stop the background thread.""" diff --git a/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_stage.py b/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_stage.py index cd3c99ab8..aee25e584 100644 --- a/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_stage.py +++ b/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_stage.py @@ -1,24 +1,26 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import contextlib import json import logging import threading import time +from collections.abc import Callable from dataclasses import asdict import zmq -from .messages import InstanceEvent, StageStatus +from .messages import ReplicaEvent, ReplicaStatus logger = logging.getLogger(__name__) class OmniCoordClientForStage: - """Client used by stage instances to send events to OmniCoordinator. + """Client used by stage replicas to send events to OmniCoordinator. This client maintains a DEALER socket connected to OmniCoordinator's - ROUTER endpoint and sends JSON-encoded events describing instance status. + ROUTER endpoint and sends JSON-encoded events describing replica status. """ def __init__( @@ -42,13 +44,18 @@ class OmniCoordClientForStage: self._socket.close() raise RuntimeError(f"Failed to connect to coordinator at {self._coord_zmq_addr}: {e}") from e - self._status = StageStatus.UP + self._status = ReplicaStatus.UP self._queue_length = 0 self._closed = False self._closing = False self._heartbeat_interval = 5.0 self._stop_event = threading.Event() self._send_lock = threading.RLock() + # Optional hook invoked from the heartbeat thread before each + # heartbeat send. Stages set this to refresh ``queue_length`` (or any + # other field) just-in-time. Exceptions raised by the hook are + # suppressed and logged. + self._on_heartbeat: Callable[[], None] | None = None self._send_event("update") @@ -100,11 +107,13 @@ class OmniCoordClientForStage: return False def _send_event(self, event_type: str) -> None: - """Send an InstanceEvent to OmniCoordinator. + """Send a ReplicaEvent to OmniCoordinator. Wire format: input_addr, output_addr, stage_id, status, queue_length, event_type. - For "update": includes status and queue_length from instance state. - For "heartbeat": status and queue_length are null. + For "update": includes status and queue_length from replica state. + For "heartbeat": includes the latest queue_length (refreshed by the + optional ``_on_heartbeat`` hook) so the coordinator can propagate + live load to load balancers between explicit ``update`` events. On send failure (ZMQError / RuntimeError), attempts to reconnect up to 3 times (5s sleep each) and retries the send once after a @@ -114,7 +123,7 @@ class OmniCoordClientForStage: if self._closed: raise RuntimeError("Client already closed") - event = InstanceEvent( + event = ReplicaEvent( input_addr=self._input_addr, output_addr=self._output_addr, stage_id=self._stage_id, @@ -147,10 +156,10 @@ class OmniCoordClientForStage: def update_info( self, - status: StageStatus | None = None, + status: ReplicaStatus | None = None, queue_length: int | None = None, ) -> None: - """Update instance information and notify OmniCoordinator. + """Update replica information and notify OmniCoordinator. At least one of ``status`` or ``queue_length`` must be provided. """ @@ -174,6 +183,15 @@ class OmniCoordClientForStage: if self._closed: break + # Invoke the optional pre-heartbeat hook so callers (e.g. the + # engine subprocess) can refresh ``queue_length`` from live state + # before the heartbeat is sent. Exceptions are swallowed so a + # buggy hook never breaks the heartbeat loop. + hook = self._on_heartbeat + if hook is not None: + with contextlib.suppress(Exception): + hook() + try: self._send_event("heartbeat") except (RuntimeError, zmq.ZMQError) as e: @@ -199,7 +217,7 @@ class OmniCoordClientForStage: self._closing = True # Mark status as DOWN and send one last update. - self._status = StageStatus.DOWN + self._status = ReplicaStatus.DOWN try: self._send_event("update") except (RuntimeError, zmq.ZMQError): diff --git a/vllm_omni/distributed/omni_coordinator/omni_coordinator.py b/vllm_omni/distributed/omni_coordinator/omni_coordinator.py index 2c7c8fbb9..8f75f841d 100644 --- a/vllm_omni/distributed/omni_coordinator/omni_coordinator.py +++ b/vllm_omni/distributed/omni_coordinator/omni_coordinator.py @@ -11,22 +11,22 @@ from typing import Any import zmq -from .messages import InstanceEvent, InstanceInfo, InstanceList, StageStatus +from .messages import ReplicaEvent, ReplicaInfo, ReplicaList, ReplicaStatus logger = logging.getLogger(__name__) class OmniCoordinator: - """Coordinator for stage instances and hub clients. + """Coordinator for stage replicas and hub clients. - This service receives instance events from :class:`OmniCoordClientForStage` - via a ZMQ ROUTER socket and publishes active instance lists to + This service receives replica events from :class:`OmniCoordClientForStage` + via a ZMQ ROUTER socket and publishes active replica lists to :class:`OmniCoordClientForHub` via a PUB socket. - The coordinator maintains an in-memory registry of all known instances, + The coordinator maintains an in-memory registry of all known replicas, including their status, queue length, and heartbeat timestamps. A background thread periodically checks for heartbeat timeouts and marks - unhealthy instances as ``StageStatus.ERROR``. + unhealthy replicas as ``ReplicaStatus.ERROR``. """ def __init__( @@ -40,7 +40,7 @@ class OmniCoordinator: Args: router_zmq_addr: ZMQ address to bind the ROUTER socket. pub_zmq_addr: ZMQ address to bind the PUB socket. - heartbeat_timeout: Seconds before an instance is considered + heartbeat_timeout: Seconds before a replica is considered unhealthy if no heartbeat / update is received. """ self._router_zmq_addr = router_zmq_addr @@ -55,7 +55,7 @@ class OmniCoordinator: self._pub = self._ctx.socket(zmq.PUB) self._pub.bind(self._pub_zmq_addr) - self._instances: dict[str, InstanceInfo] = {} + self._replicas: dict[str, ReplicaInfo] = {} self._lock = threading.Lock() self._pub_lock = threading.Lock() @@ -75,43 +75,43 @@ class OmniCoordinator: self._periodic_thread = threading.Thread(target=self._periodic_loop, daemon=True) self._periodic_thread.start() - def get_active_instances(self) -> InstanceList: - """Return an :class:`InstanceList` of active (UP) instances only.""" + def get_active_replicas(self) -> ReplicaList: + """Return a :class:`ReplicaList` of active (UP) replicas only.""" with self._lock: - active = [inst for inst in self._instances.values() if inst.status == StageStatus.UP] - return InstanceList(instances=active, timestamp=time()) + active = [rep for rep in self._replicas.values() if rep.status == ReplicaStatus.UP] + return ReplicaList(replicas=active, timestamp=time()) - def add_new_instance(self, event: InstanceEvent) -> None: - """Add a new instance based on an incoming event.""" + def add_new_replica(self, event: ReplicaEvent) -> None: + """Add a new replica based on an incoming event.""" with self._lock: - self._add_new_instance_locked(event) + self._add_new_replica_locked(event) self._schedule_broadcast() - def update_instance_info(self, event: InstanceEvent) -> None: - """Update an existing instance based on an incoming event.""" + def update_replica_info(self, event: ReplicaEvent) -> None: + """Update an existing replica based on an incoming event.""" with self._lock: - self._update_instance_info_locked(event) + self._update_replica_info_locked(event) self._schedule_broadcast() - def remove_instance(self, event: InstanceEvent) -> None: - """Mark an instance as removed / down based on an incoming event. + def remove_replica(self, event: ReplicaEvent) -> None: + """Mark a replica as removed / down based on an incoming event. - This marks the instance's status as DOWN or ERROR (depending on the + This marks the replica's status as DOWN or ERROR (depending on the event) but keeps it in the internal registry. It is removed from the - *active* instance list published to hubs. + *active* replica list published to hubs. """ with self._lock: - self._remove_instance_locked(event) + self._remove_replica_locked(event) self._schedule_broadcast() - def publish_instance_list_update(self) -> bool: - """Publish the current active instance list to all subscribers. + def publish_replica_list_update(self) -> bool: + """Publish the current active replica list to all subscribers. Returns: True if the PUB send succeeded, False if it was dropped (e.g. socket not ready when using ``zmq.NOBLOCK``). """ - active_list = self.get_active_instances() + active_list = self.get_active_replicas() payload = asdict(active_list) data = json.dumps(payload).encode("utf-8") @@ -133,12 +133,12 @@ class OmniCoordinator: with self._pending_lock: self._pending_broadcast = True - def _mark_instance_error_locked(self, info: InstanceInfo) -> None: - """Mark instance as ERROR (e.g. after heartbeat timeout).""" - info.status = StageStatus.ERROR + def _mark_replica_error_locked(self, info: ReplicaInfo) -> None: + """Mark replica as ERROR (e.g. after heartbeat timeout).""" + info.status = ReplicaStatus.ERROR def _check_heartbeat_timeouts(self) -> None: - """Mark instances as ERROR if their heartbeat has timed out.""" + """Mark replicas as ERROR if their heartbeat has timed out.""" now = time() timed_out = False gc_ttl = 600.0 # 10 minutes @@ -146,17 +146,17 @@ class OmniCoordinator: with self._lock: to_delete: list[str] = [] - for input_addr, info in self._instances.items(): - if info.status == StageStatus.UP and now - info.last_heartbeat > self._heartbeat_timeout: - self._mark_instance_error_locked(info) + for input_addr, info in self._replicas.items(): + if info.status == ReplicaStatus.UP and now - info.last_heartbeat > self._heartbeat_timeout: + self._mark_replica_error_locked(info) timed_out = True - elif info.status in (StageStatus.DOWN, StageStatus.ERROR) and now - info.last_heartbeat > gc_ttl: + elif info.status in (ReplicaStatus.DOWN, ReplicaStatus.ERROR) and now - info.last_heartbeat > gc_ttl: to_delete.append(input_addr) for input_addr in to_delete: - del self._instances[input_addr] + del self._replicas[input_addr] if timed_out: - # Instance liveness changed; request broadcast. + # Replica liveness changed; request broadcast. self._schedule_broadcast() def close(self) -> None: @@ -187,22 +187,22 @@ class OmniCoordinator: except zmq.ZMQError: pass - def _parse_instance_event(self, data: dict[str, Any]) -> InstanceEvent | None: - """Parse wire payload dict into InstanceEvent. Returns None if invalid.""" + def _parse_replica_event(self, data: dict[str, Any]) -> ReplicaEvent | None: + """Parse wire payload dict into ReplicaEvent. Returns None if invalid.""" try: - return InstanceEvent( + return ReplicaEvent( input_addr=str(data["input_addr"]), output_addr=str(data["output_addr"]), stage_id=int(data["stage_id"]), event_type=str(data["event_type"]), - status=StageStatus(data.get("status")), + status=ReplicaStatus(data.get("status")), queue_length=data.get("queue_length"), ) except (KeyError, ValueError, TypeError): return None def _recv_loop(self) -> None: - """Background loop that receives and processes instance events.""" + """Background loop that receives and processes replica events.""" while self._running: try: frames = self._router.recv_multipart() @@ -219,12 +219,12 @@ class OmniCoordinator: payload = frames[-1] try: data = json.loads(payload.decode("utf-8")) - event = self._parse_instance_event(data) + event = self._parse_replica_event(data) except json.JSONDecodeError as e: - logger.warning("Invalid JSON in instance event, dropping: %s", e) + logger.warning("Invalid JSON in replica event, dropping: %s", e) continue if event is None: - logger.warning("Malformed instance event, dropping") + logger.warning("Malformed replica event, dropping") continue self._handle_event(event) @@ -234,7 +234,10 @@ class OmniCoordinator: Heartbeat timeouts are checked on their original cadence, while all broadcast requests are coalesced and flushed at most once per - ``_publish_min_interval``. + ``_publish_min_interval``. The heartbeat-check tick also schedules a + keepalive broadcast so late-joining hubs (which miss any PUB sends + that happened before their SUB connected) catch up within at most + ``heartbeat_interval`` seconds. """ heartbeat_interval = max(1.0, min(self._heartbeat_timeout / 2.0, 5.0)) loop_interval = self._publish_min_interval @@ -245,6 +248,13 @@ class OmniCoordinator: if now - last_heartbeat_check >= heartbeat_interval: self._check_heartbeat_timeouts() + # Keepalive broadcast: ZMQ PUB doesn't queue for late + # subscribers, so an OmniCoordClientForHub that connects + # after the initial UP events miss them entirely and would + # never see the current replica list otherwise. Scheduling a + # broadcast on every heartbeat tick caps that staleness at + # ``heartbeat_interval`` without flooding the wire. + self._schedule_broadcast() last_heartbeat_check = now with self._pending_lock: @@ -256,43 +266,51 @@ class OmniCoordinator: continue # Publish outside lock. Clear pending only on success. - if self.publish_instance_list_update(): + if self.publish_replica_list_update(): with self._pending_lock: self._pending_broadcast = False if self._stop_event.wait(timeout=loop_interval): break - def _handle_event(self, event: InstanceEvent) -> None: + def _handle_event(self, event: ReplicaEvent) -> None: """Dispatch an incoming event to the appropriate handler.""" try: input_addr = event.input_addr - # Heartbeat: only update last_heartbeat; if previously ERROR, - # promote back to UP and broadcast once. + # Heartbeat: refresh last_heartbeat and queue_length. The stage + # client refreshes queue_length just-in-time via its + # ``_on_heartbeat`` hook, so heartbeats are the only periodic + # source of live load for LeastQueueLengthBalancer; failing to + # propagate it here would let the policy route on stale data. + # If previously ERROR, promote back to UP and broadcast once. if event.event_type == "heartbeat": promote = False + queue_changed = False with self._lock: - info = self._instances.get(input_addr) + info = self._replicas.get(input_addr) if info is not None: info.last_heartbeat = time() - if info.status == StageStatus.ERROR: - info.status = StageStatus.UP + if event.queue_length is not None and info.queue_length != event.queue_length: + info.queue_length = event.queue_length + queue_changed = True + if info.status == ReplicaStatus.ERROR: + info.status = ReplicaStatus.UP promote = True - if promote: + if promote or queue_changed: self._schedule_broadcast() return # Check-and-act under single lock to avoid TOCTOU race (duplicate - # registration when concurrent events arrive for the same instance). + # registration when concurrent events arrive for the same replica). with self._lock: - if input_addr not in self._instances: - self._add_new_instance_locked(event) + if input_addr not in self._replicas: + self._add_new_replica_locked(event) else: - if event.status == StageStatus.DOWN: - self._remove_instance_locked(event) + if event.status == ReplicaStatus.DOWN: + self._remove_replica_locked(event) else: - self._update_instance_info_locked(event) + self._update_replica_info_locked(event) # Any non-heartbeat state change that affects the active list # is coalesced and flushed via the periodic loop. @@ -300,7 +318,7 @@ class OmniCoordinator: except (KeyError, ValueError, TypeError) as e: logger.warning("Dropping malformed event: %s", e) - def _add_new_instance_locked(self, event: InstanceEvent) -> None: + def _add_new_replica_locked(self, event: ReplicaEvent) -> None: input_addr = event.input_addr if not input_addr: raise KeyError("input_addr required") @@ -309,7 +327,7 @@ class OmniCoordinator: raise KeyError("stage_id required and must be non-negative") now = time() - info = InstanceInfo( + info = ReplicaInfo( input_addr=input_addr, output_addr=event.output_addr, stage_id=stage_id, @@ -318,11 +336,11 @@ class OmniCoordinator: last_heartbeat=now, registered_at=now, ) - self._instances[input_addr] = info + self._replicas[input_addr] = info - def _update_instance_info_locked(self, event: InstanceEvent) -> None: + def _update_replica_info_locked(self, event: ReplicaEvent) -> None: input_addr = event.input_addr - info = self._instances[input_addr] + info = self._replicas[input_addr] if event.status is not None: info.status = event.status @@ -330,10 +348,10 @@ class OmniCoordinator: if event.queue_length is not None: info.queue_length = event.queue_length - def _remove_instance_locked(self, event: InstanceEvent) -> None: + def _remove_replica_locked(self, event: ReplicaEvent) -> None: input_addr = event.input_addr - info = self._instances.get(input_addr) + info = self._replicas.get(input_addr) if info is None: return - info.status = StageStatus.DOWN + info.status = ReplicaStatus.DOWN diff --git a/vllm_omni/distributed/omni_coordinator/runtime.py b/vllm_omni/distributed/omni_coordinator/runtime.py new file mode 100644 index 000000000..71df7c052 --- /dev/null +++ b/vllm_omni/distributed/omni_coordinator/runtime.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Lifecycle wrapper around :class:`OmniCoordinator`. + +``OmniCoordinatorRuntime`` is the single-purpose owner of the head-side +coordinator process artifacts: it picks two free TCP ports, constructs an +:class:`OmniCoordinator` bound to them, exposes the resulting addresses, and +provides a single ``close()`` method to tear everything down. + +The ROUTER address is later handed to :class:`OmniMasterServer` so it can be +published to registering replicas; the PUB address is handed to the +``Orchestrator``, which constructs its :class:`OmniCoordClientForHub` against +it. +""" + +from __future__ import annotations + +import logging + +from vllm.utils.network_utils import get_open_ports_list + +from .omni_coordinator import OmniCoordinator + +logger = logging.getLogger(__name__) + + +class OmniCoordinatorRuntime: + """Own one :class:`OmniCoordinator` and the two ports it binds. + + Constructor binds; :meth:`close` tears down. The class deliberately does + not expose the coordinator instance — callers should consume the + coordinator only via its wire protocol through + :class:`OmniCoordClientForStage` and :class:`OmniCoordClientForHub`. + """ + + def __init__( + self, + *, + host: str, + heartbeat_timeout: float, + ) -> None: + if not host: + raise ValueError("host must be a non-empty string") + if heartbeat_timeout <= 0: + raise ValueError("heartbeat_timeout must be positive") + + router_port, pub_port = get_open_ports_list(count=2) + self.router_address: str = f"tcp://{host}:{router_port}" + self.pub_address: str = f"tcp://{host}:{pub_port}" + + self._closed = False + self._coordinator = OmniCoordinator( + router_zmq_addr=self.router_address, + pub_zmq_addr=self.pub_address, + heartbeat_timeout=heartbeat_timeout, + ) + + logger.info( + "[OmniCoordinatorRuntime] Started (router=%s pub=%s heartbeat_timeout=%.1fs)", + self.router_address, + self.pub_address, + heartbeat_timeout, + ) + + def close(self) -> None: + """Tear down the underlying coordinator. Idempotent.""" + if self._closed: + return + self._closed = True + try: + self._coordinator.close() + except Exception: + logger.exception("[OmniCoordinatorRuntime] coordinator close failed") diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 4a2f46b01..14d0d43e8 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -155,6 +155,10 @@ class OmniEngineArgs(EngineArgs): omni_master_address: str | None = None omni_master_port: int | None = None + # OmniCoordinator integration knobs (process-local). + omni_dp_size_local: int = 1 + omni_lb_policy: str = "random" + omni_heartbeat_timeout: float = 30.0 stage_configs_path: str | None = None output_modalities: list[str] | None = None log_stats: bool = False diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 8a850f565..420f808dc 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -18,7 +18,7 @@ import threading import time import uuid import weakref -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from contextlib import ExitStack from dataclasses import asdict, dataclass from typing import TYPE_CHECKING, Any, Literal, cast @@ -45,6 +45,10 @@ from vllm_omni.diffusion.stage_diffusion_proc import ( from vllm_omni.distributed.omni_connectors.utils.initialization import ( resolve_omni_kv_config_for_stage, ) +from vllm_omni.distributed.omni_coordinator import ( + LoadBalancer, + build_load_balancer_factory, +) from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.messages import ( AbortRequestMessage, @@ -53,6 +57,7 @@ from vllm_omni.engine.messages import ( CollectiveRPCResultMessage, EngineQueueMessage, ErrorMessage, + RegisterRemoteReplicaMessage, ShutdownRequestMessage, StageSubmissionMessage, ) @@ -77,6 +82,7 @@ from vllm_omni.engine.stage_engine_startup import ( from vllm_omni.engine.stage_init_utils import ( LogicalStageInitPlan, ReplicaInitPlan, + StageRemoteFactoryContext, _inject_inferred_kv_tp_topology, acquire_device_locks, acquire_diffusion_device_locks, @@ -85,6 +91,7 @@ from vllm_omni.engine.stage_init_utils import ( build_llm_stage_output_processor, build_stage0_input_processor, build_vllm_config, + capture_stage_factory_contexts, compute_replica_layout, extract_stage_metadata, get_stage_connector_spec, @@ -148,6 +155,7 @@ _PARENT_ARGS_KEEP: frozenset[str] = frozenset( # trigger the ``create_model_config`` guard). _PARENT_ARGS_STRIP: frozenset[str] = frozenset({"stage_configs_path"}) + # Fields always populated by callers (via ``from_cli_args`` / ``asdict``) so # their presence as an override is never a surprise — suppress the # "override ignored" warning for these. @@ -229,10 +237,9 @@ def _weak_shutdown_async_omni_engine( pass for q in (request_queue, output_queue, rpc_output_queue): - if q is None: - continue try: - q.close() + if q is not None: + q.close() except Exception: pass @@ -301,6 +308,23 @@ class AsyncOmniEngine: self._omni_master_port: int | None = kwargs.get("omni_master_port") self._omni_master_server: OmniMasterServer | None = None + # New omni-coordinator flags. Consumed only in single_stage_mode. + # ``omni_dp_size_local`` is process-local: each invocation (head and + # every headless) launches that many replicas for its own stage. + self._omni_dp_size_local: int = int(kwargs.get("omni_dp_size_local") or 1) + if self._omni_dp_size_local < 1: + raise ValueError(f"--omni-dp-size-local must be >= 1, got {self._omni_dp_size_local}") + self._omni_lb_policy: str = str(kwargs.get("omni_lb_policy") or "random") + self._omni_heartbeat_timeout: float = float(kwargs.get("omni_heartbeat_timeout") or 30.0) + if self._omni_heartbeat_timeout <= 0: + raise ValueError(f"--omni-heartbeat-timeout must be > 0, got {self._omni_heartbeat_timeout}") + # Coordinator runtime (head-distributed only). + self._coordinator_runtime: Any | None = None + # Per-stage construction context, captured after _initialize_stages + # and used by ``_build_remote_replica`` (the RemoteReplicaFactory + # passed to Orchestrator) when a headless replica registers. + self._stage_remote_factory_contexts: dict[int, StageRemoteFactoryContext] = {} + if single_stage_mode: logger.info( "[AsyncOmniEngine] Single-stage mode enabled (stage_id_filter=%s, master=%s:%s)", @@ -321,9 +345,14 @@ class AsyncOmniEngine: self.supported_tasks: tuple[str, ...] = ("generate",) self.default_sampling_params_list: list[OmniSamplingParams] = [] self.stage_metadata: list[StageRuntimeInfo] = [] - self.request_queue: janus.Queue[EngineQueueMessage] | None = None - self.output_queue: janus.Queue[EngineQueueMessage] | None = None - self.rpc_output_queue: janus.Queue[EngineQueueMessage] | None = None + # Janus queues are constructed eagerly here (not deferred to the + # orchestrator thread) so the master server's ROUTER thread always + # sees a non-None ``self.request_queue`` when on_register fires. + # ``async_q`` lazily binds to whatever event loop first awaits on + # it (the orchestrator loop), so cross-thread use stays correct. + self.request_queue: janus.Queue[EngineQueueMessage] = janus.Queue() + self.output_queue: janus.Queue[EngineQueueMessage] = janus.Queue() + self.rpc_output_queue: janus.Queue[EngineQueueMessage] = janus.Queue() self._shutdown_called = False self._weak_finalizer: weakref.finalize | None = None self._rpc_lock = threading.Lock() @@ -443,30 +472,37 @@ class AsyncOmniEngine: ) def _validate_single_stage_mode_replica_constraints(self) -> None: - """Reject unsupported replica fan-out in single-stage mode.""" + """Apply --omni-dp-size-local to the local stage's runtime.num_replicas. + + In the previous revision this method rejected LLM stages with + ``num_replicas > 1``. The whole point of ``--omni-dp-size-local`` is + to lift that restriction for the *local* stage, so the rejection is + gone. We now use this hook to write ``--omni-dp-size-local`` onto + the self-stage's runtime config so downstream code + (``compute_replica_layout`` → ``_build_logical_stage_init_plans``) + sees a consistent view. + """ if not self.single_stage_mode: return + target_stage_id = self._single_stage_id_filter + if target_stage_id is None: + return - unsupported: list[tuple[int, int]] = [] for idx, stage_cfg in enumerate(self.stage_configs): - runtime_cfg = getattr(stage_cfg, "runtime", {}) - num_replicas = int( - runtime_cfg.get("num_replicas", 1) - if hasattr(runtime_cfg, "get") - else getattr(runtime_cfg, "num_replicas", 1) - ) - if num_replicas <= 1: - continue - if getattr(stage_cfg, "stage_type", "llm") == "diffusion": - continue stage_id = int(getattr(stage_cfg, "stage_id", idx)) - unsupported.append((stage_id, num_replicas)) - - if unsupported: - raise ValueError( - "single_stage_mode only supports num_replicas > 1 for diffusion stages; " - f"found non-diffusion stages {unsupported}" - ) + runtime_cfg = getattr(stage_cfg, "runtime", None) + if runtime_cfg is None: + continue + if stage_id == target_stage_id: + # Self stage: take --omni-dp-size-local from this process. + try: + runtime_cfg.num_replicas = self._omni_dp_size_local + except Exception: + if hasattr(runtime_cfg, "__setitem__"): + runtime_cfg["num_replicas"] = self._omni_dp_size_local + # Other stages keep their config-declared num_replicas; in + # head-distributed mode they will be launched as ``launch_mode + # == "remote"`` with the configured count. def _build_logical_stage_init_plans( self, @@ -540,9 +576,18 @@ class AsyncOmniEngine: replica_metadata = extract_stage_metadata(replica_cfg) replica_metadata.replica_id = replica_id - if self.single_stage_mode: - if replica_metadata.stage_type != "diffusion": - replica_metadata.runtime_cfg = None + # In single_stage_mode the head only owns its self stage's + # replicas; the remote-stage metadata exists only so the + # orchestrator can route requests through StagePool. Wiping + # ``runtime_cfg`` there makes sense (we don't manage those + # devices). For the *self stage* we MUST keep the + # per-replica runtime so ``setup_stage_devices`` can apply + # the device split from ``compute_replica_layout`` — + # otherwise every replica inherits the parent's full + # CUDA_VISIBLE_DEVICES and stacks on cuda:0 (OOM with any + # model whose footprint exceeds ~1/(2N) of the card). + if launch_mode == "remote" and replica_metadata.stage_type != "diffusion": + replica_metadata.runtime_cfg = None replicas.append( ReplicaInitPlan( @@ -578,6 +623,11 @@ class AsyncOmniEngine: all_stage_ids: list[int] = [] stage_replica_counts: dict[int, int] = {} + # Slots that the head itself will register (launch_mode == "local") + # — auto-assign on the master must not hand these out to remote + # headless registrations even within the race window before the + # head's own register_stage_with_omni_master call completes. + head_local_replicas: dict[int, list[int]] = {} seen_stage_ids: set[int] = set() for plan in stage_plans: stage_id = plan.configured_stage_id @@ -588,12 +638,27 @@ class AsyncOmniEngine: seen_stage_ids.add(stage_id) all_stage_ids.append(stage_id) stage_replica_counts[stage_id] = len(plan.replicas) + local_rids = [rep.replica_id for rep in plan.replicas if rep.launch_mode == "local"] + if local_rids: + head_local_replicas[stage_id] = local_rids + + # Start the OmniCoordinator runtime first so its router address is + # available to publish in every registration reply. + from vllm_omni.distributed.omni_coordinator import OmniCoordinatorRuntime + + self._coordinator_runtime = OmniCoordinatorRuntime( + host=self._omni_master_address, + heartbeat_timeout=self._omni_heartbeat_timeout, + ) self._omni_master_server = OmniMasterServer( master_address=self._omni_master_address, master_port=self._omni_master_port, stage_ids=all_stage_ids, stage_replica_counts=stage_replica_counts, + coordinator_router_address=self._coordinator_runtime.router_address, + on_register=self._dispatch_master_register, + head_local_replicas=head_local_replicas, ) self._omni_master_server.start() logger.info( @@ -601,6 +666,147 @@ class AsyncOmniEngine: all_stage_ids, ) + # ------------------------------------------------------------------ + # Remote replica factory (head-side client construction) + # ------------------------------------------------------------------ + + async def _build_remote_replica(self, stage_id: int, replica_id: int) -> Any: + """Construct a head-side stage client for a newly-registered remote replica. + + Used by :class:`Orchestrator` as its ``remote_replica_factory``. + The orchestrator awaits this from its own asyncio loop, so the + client is created in the same loop that owns ZMQ sockets — no + cross-thread setup is required. + + Raises if the stage is not known or if its construction context + was not captured (e.g. the stage was empty at bring-up time). + """ + ctx = self._stage_remote_factory_contexts.get(stage_id) + if ctx is None: + raise RuntimeError( + f"no factory context captured for stage {stage_id}; " + f"known stages: {sorted(self._stage_remote_factory_contexts.keys())}" + ) + if self._omni_master_server is None: + raise RuntimeError("OmniMasterServer is not running; cannot build remote replica") + + alloc = self._omni_master_server.get_allocation(stage_id, replica_id=replica_id) + + # Build a per-replica copy of the base metadata so ``replica_id`` + # is correct (StageMetadata is a plain dataclass-like). + metadata = copy.copy(ctx.base_metadata) + try: + metadata.replica_id = replica_id + except Exception: + # Best-effort: if the metadata object is frozen / unusual, the + # downstream client will fall back to ``replica_id = 0``. + pass + + if ctx.stage_type == "diffusion": + client = StageDiffusionClient.from_addresses( + metadata, + request_address=alloc.input_bind_address, + response_address=alloc.output_bind_address, + batch_size=ctx.diffusion_batch_size, + ) + logger.info( + "[AsyncOmniEngine] Built remote diffusion client for stage=%d replica=%d (req=%s resp=%s)", + stage_id, + replica_id, + alloc.input_bind_address, + alloc.output_bind_address, + ) + return client + + # LLM path + if ctx.vllm_config is None or ctx.executor_class is None: + raise RuntimeError(f"stage {stage_id} factory context is missing vllm_config / executor_class") + + # The headless's StageEngineCoreProc subprocess calls + # vllm.v1.engine.core.startup_handshake at boot and blocks until the + # head's handshake ROUTER answers — without this step it hits the + # built-in 5-minute timeout and exits. The bootstrap (pre-allocated) + # path runs `connect_remote_engine_cores` to perform that + # handshake; dynamic attach must do the same, otherwise every + # replica that comes in via `on_register` (auto-assigned or + # explicit-id-beyond-pre-alloc) deadlocks. Run the blocking + # handshake in a thread so the orchestrator loop stays responsive, + # then build the async client on this loop where `make_async_mp_client` + # expects to be invoked. + master_server = self._omni_master_server + ctx_vllm_config = ctx.vllm_config + ctx_executor_class = ctx.executor_class + + def _run_handshake() -> Any: + with connect_remote_engine_cores( + vllm_config=ctx_vllm_config, + omni_master_server=master_server, + stage_id=stage_id, + replica_id=replica_id, + ) as remote_resources: + _engine_manager, _coordinator, addresses, _ = remote_resources + return _engine_manager, _coordinator, addresses + + engine_manager, coordinator, addresses = await asyncio.to_thread(_run_handshake) + + client_addresses: dict[str, str] = { + "input_address": addresses.inputs[0], + "output_address": addresses.outputs[0], + } + if addresses.frontend_stats_publish_address is not None: + client_addresses["stats_update_address"] = addresses.frontend_stats_publish_address + + client = StageEngineCoreClientBase.make_async_mp_client( + vllm_config=ctx_vllm_config, + executor_class=ctx_executor_class, + metadata=metadata, + client_addresses=client_addresses, + proc=None, + engine_manager=engine_manager, + coordinator=coordinator, + ) + logger.info( + "[AsyncOmniEngine] Built remote LLM client for stage=%d replica=%d (input=%s)", + stage_id, + replica_id, + client_addresses["input_address"], + ) + return client + + # ------------------------------------------------------------------ + # OmniCoordinator on_register proxy + # ------------------------------------------------------------------ + + def _dispatch_master_register(self, stage_id: int, replica_id: int, alloc: Any) -> None: + """Forward a master-server registration to the orchestrator queue. + + Called on :class:`OmniMasterServer`'s ROUTER thread (not the + orchestrator loop). Must return promptly. ``self.request_queue`` + is initialized in :meth:`_initialize_janus_queues` (running inside + the orchestrator thread) before the master server starts, so it + is non-None here; ``janus.Queue.sync_q`` is thread-safe by + construction. + """ + if self.request_queue is None: + logger.warning( + "[AsyncOmniEngine] request_queue not initialized; dropping register_remote_replica stage=%d replica=%d", + stage_id, + replica_id, + ) + return + msg = RegisterRemoteReplicaMessage( + stage_id=int(stage_id), + replica_id=int(replica_id), + ) + try: + self.request_queue.sync_q.put_nowait(msg) + except Exception: + logger.exception( + "[AsyncOmniEngine] Failed to enqueue register_remote_replica for stage=%d replica=%d", + stage_id, + replica_id, + ) + def _initialize_llm_replica( self, plan: ReplicaInitPlan, @@ -686,6 +892,11 @@ class AsyncOmniEngine: stage_init_timeout, ) if self.single_stage_mode and self._omni_master_server is not None: + coord_router_addr: str | None = ( + self._coordinator_runtime.router_address + if self._coordinator_runtime is not None + else None + ) engine_manager, coordinator, addresses = launch_stack.enter_context( launch_omni_core_engines( vllm_config=vllm_config, @@ -695,6 +906,7 @@ class AsyncOmniEngine: stage_id=plan.metadata.stage_id, stage_config=stage_cfg, replica_id=plan.replica_id, + omni_coordinator_address=coord_router_addr, ) ) else: @@ -830,12 +1042,20 @@ class AsyncOmniEngine: "[AsyncOmniEngine] Stage %s diffusion registration completed", plan.metadata.stage_id, ) + coord_router_addr: str | None = ( + self._coordinator_runtime.router_address + if self._coordinator_runtime is not None + else None + ) proc, _, _, _ = spawn_diffusion_proc( self.model, od_config, handshake_address=handshake_address, request_address=request_address, response_address=response_address, + omni_coordinator_address=coord_router_addr, + omni_stage_id=plan.metadata.stage_id, + omni_replica_id=plan.replica_id, ) complete_diffusion_handshake(proc, handshake_address, stage_init_timeout) logger.info( @@ -1049,7 +1269,13 @@ class AsyncOmniEngine: replicas_per_stage, replica_devices_map, ) + # Capture per-stage context now (before _start_omni_master_server) + # so the on_register proxy can build head-side clients for + # registrations that arrive immediately after the server starts. if self.single_stage_mode: + self._stage_remote_factory_contexts = capture_stage_factory_contexts( + stage_plans, diffusion_batch_size=self.diffusion_batch_size + ) self._start_omni_master_server(stage_plans) stage_pools: list[StagePool] = [] @@ -1085,6 +1311,14 @@ class AsyncOmniEngine: self._omni_master_server.stop() except Exception: logger.exception("[AsyncOmniEngine] Failed to stop OmniMasterServer during stage-init cleanup") + if self._coordinator_runtime is not None: + try: + self._coordinator_runtime.close() + except Exception: + logger.exception( + "[AsyncOmniEngine] Failed to close OmniCoordinatorRuntime during stage-init cleanup" + ) + self._coordinator_runtime = None raise self.stage_pools = stage_pools @@ -1104,13 +1338,6 @@ class AsyncOmniEngine: supported_tasks.add("speech") self.supported_tasks = tuple(supported_tasks) if supported_tasks else ("generate",) - def _initialize_janus_queues(self) -> None: - """Initialize janus queues inside orchestrator thread loop context.""" - self.request_queue = janus.Queue() - self.output_queue = janus.Queue() - self.rpc_output_queue = janus.Queue() - logger.debug("[AsyncOmniEngine] janus queues initialized in orchestrator thread loop") - def _bootstrap_orchestrator( self, stage_init_timeout: int, @@ -1122,10 +1349,15 @@ class AsyncOmniEngine: asyncio.set_event_loop(loop) async def _run_orchestrator() -> None: - self._initialize_janus_queues() - self._initialize_stages(stage_init_timeout) pd_config = self._detect_pd_config() + coordinator_pub_address: str | None = None + load_balancer_factory: Callable[[], LoadBalancer] | None = None + remote_replica_factory: Callable[[int, int], Awaitable[Any]] | None = None + if self._coordinator_runtime is not None: + coordinator_pub_address = self._coordinator_runtime.pub_address + load_balancer_factory = build_load_balancer_factory(self._omni_lb_policy) + remote_replica_factory = self._build_remote_replica orchestrator = Orchestrator( request_async_queue=self.request_queue.async_q, output_async_queue=self.output_queue.async_q, @@ -1133,6 +1365,9 @@ class AsyncOmniEngine: stage_pools=self.stage_pools, async_chunk=self.async_chunk, pd_config=pd_config, + coordinator_pub_address=coordinator_pub_address, + load_balancer_factory=load_balancer_factory, + remote_replica_factory=remote_replica_factory, ) if not startup_future.done(): startup_future.set_result(asyncio.get_running_loop()) @@ -1948,8 +2183,6 @@ class AsyncOmniEngine: reasoning_ended=reasoning_ended, resumable=resumable, ) - if self.request_queue is None: - raise RuntimeError("request_queue is not initialized") self.request_queue.sync_q.put_nowait(msg) # CFG companion expansion: create and enqueue companion requests @@ -2017,8 +2250,6 @@ class AsyncOmniEngine: resumable=resumable, message_type="streaming_update", ) - if self.request_queue is None: - raise RuntimeError("request_queue is not initialized") self.request_queue.sync_q.put_nowait(msg) async def add_streaming_update_async( @@ -2045,8 +2276,6 @@ class AsyncOmniEngine: def try_get_output(self, timeout: float = 0.001) -> EngineQueueMessage | None: """Read one output message from the Orchestrator output queue.""" - if self.output_queue is None: - return None try: return self.output_queue.sync_q.get(timeout=timeout) except queue.Empty: @@ -2056,8 +2285,6 @@ class AsyncOmniEngine: async def try_get_output_async(self) -> EngineQueueMessage | None: """Async read from the Orchestrator output queue.""" - if self.output_queue is None: - return None try: return self.output_queue.sync_q.get_nowait() except queue.Empty: @@ -2092,11 +2319,6 @@ class AsyncOmniEngine: This uses a dedicated RPC output queue so control-plane messages do not race with the normal request output polling loop. """ - if self.request_queue is None: - raise RuntimeError("request_queue is not initialized") - if self.rpc_output_queue is None: - raise RuntimeError("rpc_output_queue is not initialized") - rpc_id = uuid.uuid4().hex msg = CollectiveRPCRequestMessage( rpc_id=rpc_id, @@ -2184,8 +2406,6 @@ class AsyncOmniEngine: logger.warning("[AsyncOmniEngine] Orchestrator thread did not exit in time") for q in (self.request_queue, self.output_queue, self.rpc_output_queue): - if q is None: - continue try: q.close() except Exception: @@ -2198,6 +2418,13 @@ class AsyncOmniEngine: logger.exception("[AsyncOmniEngine] Failed to stop OmniMasterServer during shutdown") self._omni_master_server = None + if self._coordinator_runtime is not None: + try: + self._coordinator_runtime.close() + except Exception: + logger.exception("[AsyncOmniEngine] Failed to close OmniCoordinatorRuntime during shutdown") + self._coordinator_runtime = None + def _try_shutdown(self, *args, **kwargs) -> None: try: self.shutdown() diff --git a/vllm_omni/engine/messages.py b/vllm_omni/engine/messages.py index 28a5c721c..0e55105f1 100644 --- a/vllm_omni/engine/messages.py +++ b/vllm_omni/engine/messages.py @@ -56,6 +56,18 @@ class ShutdownRequestMessage(EngineQueueMessage, kw_only=True): type: Literal["shutdown"] = "shutdown" +class RegisterRemoteReplicaMessage(EngineQueueMessage, kw_only=True): + type: Literal["register_remote_replica"] = "register_remote_replica" + stage_id: int + replica_id: int + + +class UnregisterRemoteReplicaMessage(EngineQueueMessage, kw_only=True): + type: Literal["unregister_remote_replica"] = "unregister_remote_replica" + stage_id: int + input_addr: str + + class ErrorMessage(EngineQueueMessage, kw_only=True): type: Literal["error"] = "error" error: str diff --git a/vllm_omni/engine/omni_core_engine_proc_manager.py b/vllm_omni/engine/omni_core_engine_proc_manager.py new file mode 100644 index 000000000..8df5fd053 --- /dev/null +++ b/vllm_omni/engine/omni_core_engine_proc_manager.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Process manager for omni stage engine subprocesses. + +This is a drop-in replacement for vLLM's :class:`CoreEngineProcManager` that +spawns :meth:`StageEngineCoreProc.run_stage_core` instead of the upstream +``EngineCoreProc.run_engine_core``, and forwards omni-specific kwargs +(coordinator address, stage id, per-rank replica id). + +Each spawned subprocess corresponds to exactly one omni *replica*: it has its +own ZMQ allocation from :class:`OmniMasterServer` and (when an +``omni_coordinator_address`` is provided) its own +:class:`OmniCoordClientForStage` reporting heartbeat / status. + +Liveness monitoring and shutdown are inherited from +:class:`CoreEngineProcManager` unchanged. +""" + +from __future__ import annotations + +import contextlib +import threading +import weakref +from multiprocessing.process import BaseProcess +from multiprocessing.queues import Queue + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.utils import numa_utils +from vllm.utils.system_utils import get_mp_context +from vllm.v1.engine.utils import CoreEngineProcManager +from vllm.v1.executor import Executor +from vllm.v1.utils import shutdown + +from vllm_omni.engine.stage_engine_core_proc import StageEngineCoreProc + +logger = init_logger(__name__) + +try: + # ``set_device_control_env_var`` lives next to CoreEngineProcManager and + # is only required for non-CUDA DP, so we tolerate its absence on + # older / future vLLM revisions. + from vllm.v1.engine.utils import set_device_control_env_var # type: ignore +except ImportError: # pragma: no cover - depends on vLLM build + set_device_control_env_var = None # type: ignore[assignment] + + +class OmniCoreEngineProcManager(CoreEngineProcManager): + """Spawn :class:`StageEngineCoreProc` subprocesses with omni kwargs. + + The body mirrors :class:`CoreEngineProcManager.__init__` because the + upstream class hardcodes ``target=EngineCoreProc.run_engine_core`` and + does not expose an extensibility hook. The differences from upstream are: + + * ``target`` is :meth:`StageEngineCoreProc.run_stage_core`. + * Per-rank ``omni_replica_id`` is computed as + ``base_replica_id + rank_idx`` and added to each subprocess's kwargs. + * ``omni_coordinator_address`` (if provided) and ``omni_stage_id`` are + added to every subprocess's kwargs. + """ + + def __init__( + self, + local_engine_count: int, + start_index: int, + local_start_index: int, + vllm_config: VllmConfig, + local_client: bool, + handshake_address: str, + executor_class: type[Executor], + log_stats: bool, + *, + omni_stage_id: int, + omni_coordinator_address: str | None = None, + omni_replica_base_id: int = 0, + client_handshake_address: str | None = None, + tensor_queue: Queue | None = None, + ) -> None: + # NOTE: we intentionally do not call ``super().__init__`` — the + # parent's body hardcodes the wrong target. We re-implement it here + # while reusing the parent's instance methods (shutdown, monitor). + if local_engine_count <= 0: + raise ValueError(f"local_engine_count must be > 0, got {local_engine_count}") + + context = get_mp_context() + common_kwargs: dict[str, object] = { + "vllm_config": vllm_config, + "local_client": local_client, + "handshake_address": handshake_address, + "executor_class": executor_class, + "log_stats": log_stats, + "tensor_queue": tensor_queue, + "omni_stage_id": int(omni_stage_id), + "omni_coordinator_address": omni_coordinator_address, + } + + if client_handshake_address: + common_kwargs["client_handshake_address"] = client_handshake_address + + # Intra-replica vLLM DP mesh (i.e. ``data_parallel_size`` ranks sharing + # one engine, one DPCoordinator, one set of weights). Distinct from + # the omni-level notion of multiple independent replicas of a stage — + # those each spawn their own OmniCoreEngineProcManager and never join + # a vLLM DP group across replicas. + has_intra_replica_dp = vllm_config.parallel_config.data_parallel_size > 1 + + self.processes: list[BaseProcess] = [] + local_dp_ranks: list[int] = [] + for index in range(local_engine_count): + local_index = local_start_index + index + global_index = start_index + index + # Each spawned subprocess is one omni replica. The replica id + # is contiguous within this manager; the master server may have + # pre-allocated a contiguous block starting at ``omni_replica_base_id``. + omni_replica_id = omni_replica_base_id + index + + local_dp_ranks.append(local_index) + self.processes.append( + context.Process( + target=StageEngineCoreProc.run_stage_core, + name=( + f"StageEngineCoreProc_stage{omni_stage_id}" + f"_replica{omni_replica_id}" + (f"_DP{global_index}" if has_intra_replica_dp else "") + ), + kwargs=common_kwargs + | { + "dp_rank": global_index, + "local_dp_rank": local_index, + "omni_replica_id": omni_replica_id, + }, + ) + ) + + self._finalizer = weakref.finalize(self, shutdown, self.processes) + self.manager_stopped = threading.Event() + self.failed_proc_name: str | None = None + + try: + for proc, local_dp_rank in zip(self.processes, local_dp_ranks): + device_control_context: contextlib.AbstractContextManager[None] = contextlib.nullcontext() + if ( + has_intra_replica_dp + and set_device_control_env_var is not None + and (not current_platform.is_cuda_alike() or vllm_config.parallel_config.use_ray) + ): + device_control_context = set_device_control_env_var(vllm_config, local_dp_rank) + + with ( + device_control_context, + numa_utils.configure_subprocess( + vllm_config, + local_rank=0, + dp_local_rank=local_dp_rank, + process_kind="EngineCore", + ), + ): + proc.start() + finally: + if self.finished_procs(): + self.shutdown() diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index ea615d579..27bb44ad7 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -4,12 +4,20 @@ Orchestrator for vLLM-Omni multi-stage runtime. Runs inside a background thread with its own asyncio event loop. Owns logical request progression across stage pools and handles stage-to-stage transfer logic. + +In distributed mode (``coordinator_pub_address`` provided), it also +owns the single :class:`OmniCoordClientForHub`, runs a +:meth:`_watch_replica_list` task that converts replica disappearances +into ``unregister_remote_replica`` control messages, and handles the +``register_remote_replica`` / ``unregister_remote_replica`` flow that +attaches / detaches head-side stage clients for headless replicas. """ from __future__ import annotations import asyncio import time as _time +from collections.abc import Awaitable, Callable from dataclasses import dataclass, field from typing import Any @@ -22,6 +30,12 @@ from vllm.sampling_params import SamplingParams from vllm.v1.engine import EngineCoreOutputs from vllm.v1.engine.exceptions import EngineDeadError +from vllm_omni.distributed.omni_coordinator import ( + LoadBalancer, + OmniCoordClientForHub, + RandomBalancer, + ReplicaStatus, +) from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker from vllm_omni.engine.messages import ( @@ -32,13 +46,26 @@ from vllm_omni.engine.messages import ( EngineQueueMessage, ErrorMessage, OutputMessage, + RegisterRemoteReplicaMessage, + ShutdownRequestMessage, StageMetricsMessage, StageSubmissionMessage, + UnregisterRemoteReplicaMessage, ) from vllm_omni.engine.serialization import serialize_additional_information from vllm_omni.engine.stage_pool import StagePool from vllm_omni.outputs import OmniRequestOutput +# Factory signature for building a head-side stage client for a +# *dynamically attached* (auto-assigned) remote replica. +# +# Receives ``(stage_id, replica_id)`` and returns an awaitable yielding the +# constructed client (any type — it must satisfy the shape expected by the +# matching :class:`StagePool`, i.e. expose ``client_addresses["input_address"]`` +# or ``request_address``, plus the usual ``add_request_async`` / +# ``get_output_async`` / ``shutdown`` surface). +RemoteReplicaFactory = Callable[[int, int], Awaitable[Any]] + logger = init_logger(__name__) @@ -124,6 +151,10 @@ class StreamingInputState: class Orchestrator: """Runs inside a background thread's asyncio event loop.""" + # Cadence at which the replica-list watcher polls for disappearances. + _WATCH_REPLICA_INTERVAL_S: float = 0.5 + _WATCH_REPLICA_IDLE_INTERVAL_S: float = 1.0 + def __init__( self, request_async_queue: janus.AsyncQueue[EngineQueueMessage], @@ -133,6 +164,9 @@ class Orchestrator: *, async_chunk: bool = False, pd_config: dict[str, Any] | None = None, + coordinator_pub_address: str | None = None, + load_balancer_factory: Callable[[], LoadBalancer] | None = None, + remote_replica_factory: RemoteReplicaFactory | None = None, ) -> None: self.request_async_queue = request_async_queue self.output_async_queue = output_async_queue @@ -159,6 +193,28 @@ class Orchestrator: self._fatal_error: str | None = None self._fatal_error_stage_id: int | None = None + # Background tasks for fire-and-forget message handlers (currently + # only ``register_remote_replica`` and ``unregister_remote_replica``). + # Held as a set so each task's reference survives the loop and the + # task can self-deregister on completion. + self._membership_tasks: set[asyncio.Task[None]] = set() + + # Distributed-mode wiring. The hub is constructed on the + # orchestrator's asyncio loop because it spawns a SUB background + # thread; building it from another thread would race the + # ``_init_done`` event. + self._hub: OmniCoordClientForHub | None = ( + OmniCoordClientForHub(coordinator_pub_address) if coordinator_pub_address is not None else None + ) + self._remote_replica_factory = remote_replica_factory + # Inject hub + per-pool LB into each StagePool so they can run + # distributed dispatch via ``StagePool.pick``. + if self._hub is not None: + factory = load_balancer_factory or RandomBalancer + for pool in self.stage_pools: + pool.attach_hub(self._hub) + pool.attach_load_balancer(factory()) + async def run(self) -> None: """Main entry point for the Orchestrator event loop.""" logger.info("[Orchestrator] Starting event loop") @@ -168,9 +224,16 @@ class Orchestrator: self._orchestration_output_handler(), name="orchestrator-stage-output-handler", ) + # The replica watcher only runs in distributed mode. It's still + # created in both cases so ``run()`` has a uniform task graph; + # ``_watch_replica_list`` is a no-op poll when ``self._hub`` is None. + watch_task = asyncio.create_task( + self._watch_replica_list(), + name="orchestrator-replica-watcher", + ) try: - await asyncio.gather(request_task, output_task) + await asyncio.gather(request_task, output_task, watch_task) except asyncio.CancelledError: raise except EngineDeadError as e: @@ -185,11 +248,11 @@ class Orchestrator: raise finally: self._shutdown_event.set() - for task in (request_task, output_task): + for task in (request_task, output_task, watch_task): if not task.done(): task.cancel() try: - await asyncio.gather(request_task, output_task, return_exceptions=True) + await asyncio.gather(request_task, output_task, watch_task, return_exceptions=True) except Exception: pass @@ -199,8 +262,32 @@ class Orchestrator: if self._fatal_error is not None: await self._drain_pending_requests_on_fatal() + # Wait briefly for any in-flight membership handlers (register / + # unregister remote replica) to finish so they don't leave the + # head-side pool in a half-attached state. Cancel anything that + # hasn't completed in time; the generic pending-task sweep below + # will collect the cancellations. + if self._membership_tasks: + try: + await asyncio.wait_for( + asyncio.gather(*self._membership_tasks, return_exceptions=True), + timeout=10.0, + ) + except (asyncio.TimeoutError, Exception): + for t in self._membership_tasks: + if not t.done(): + t.cancel() + self._shutdown_stages() + # Close the hub last so any in-flight dispatch still has access. + if self._hub is not None: + try: + self._hub.close() + except RuntimeError: + pass + self._hub = None + loop = asyncio.get_running_loop() pending = [t for t in asyncio.all_tasks(loop) if t is not asyncio.current_task() and not t.done()] for task in pending: @@ -208,6 +295,27 @@ class Orchestrator: if pending: await asyncio.gather(*pending, return_exceptions=True) + # ---- Background task helpers ---- + + def _spawn_membership_task(self, coro: Awaitable[None], *, label: str) -> None: + """Run a fire-and-forget membership-change coroutine. + + Holds a strong reference until completion (asyncio would otherwise + garbage-collect a bare task), and logs any uncaught exception. + """ + task = asyncio.create_task(coro, name=f"orchestrator-{label}") + self._membership_tasks.add(task) + + def _on_done(t: asyncio.Task[None]) -> None: + self._membership_tasks.discard(t) + if t.cancelled(): + return + exc = t.exception() + if exc is not None: + logger.error("[Orchestrator] %s task crashed", label, exc_info=exc) + + task.add_done_callback(_on_done) + # ---- Request handling ---- async def _request_handler(self) -> None: @@ -226,7 +334,20 @@ class Orchestrator: await self._handle_abort(msg) elif msg_type == "collective_rpc": await self._handle_collective_rpc(msg) - elif msg_type == "shutdown": + elif isinstance(msg, RegisterRemoteReplicaMessage): + # Dynamic-attach involves a ~5s blocking handshake (run in a + # thread by ``_build_remote_replica``); ``await`` here would + # block the queue and stall the next ``add_request`` until + # the attach finishes. Dispatch as a background task so the + # main message loop keeps draining. + self._spawn_membership_task(self._handle_register_remote_replica(msg), label="register_remote_replica") + elif isinstance(msg, UnregisterRemoteReplicaMessage): + # Symmetric with register: keep the main queue flowing. + self._spawn_membership_task( + self._handle_unregister_remote_replica(msg), + label="unregister_remote_replica", + ) + elif isinstance(msg, ShutdownRequestMessage): logger.info("[Orchestrator] Received shutdown signal") self._shutdown_event.set() # Pre-mark stage clients as shutting down to prevent @@ -415,7 +536,7 @@ class Orchestrator: results: list[Any] = [] stage_ids: list[int] = [] for pool in target_pools: - for replica_id in range(pool.num_replicas): + for replica_id in pool.live_replica_ids(): stage_result = await pool.collective_rpc( replica_id=replica_id, method=method, @@ -451,7 +572,7 @@ class Orchestrator: idle = True for stage_id in range(self.num_stages): pool = self.stage_pools[stage_id] - for replica_id in range(pool.num_replicas): + for replica_id in pool.live_replica_ids(): if self._shutdown_event.is_set(): return @@ -1105,14 +1226,137 @@ class Orchestrator: ) self.request_states.pop(req_id, None) + # ---- Distributed-mode replica attach / detach ---- + + async def _watch_replica_list(self) -> None: + """Convert hub replica disappearances into unregister control messages.""" + last_up: set[tuple[int, str]] = set() + while not self._shutdown_event.is_set(): + if self._hub is None: + # No coordinator wired up; sleep coarsely and re-check shutdown. + try: + await asyncio.sleep(self._WATCH_REPLICA_IDLE_INTERVAL_S) + except asyncio.CancelledError: + raise + continue + + try: + snap = self._hub.get_replica_list() + current = {(rep.stage_id, rep.input_addr) for rep in snap.replicas if rep.status == ReplicaStatus.UP} + for stage_id, addr in last_up - current: + await self.request_async_queue.put( + UnregisterRemoteReplicaMessage( + stage_id=stage_id, + input_addr=addr, + ) + ) + last_up = current + except asyncio.CancelledError: + raise + except Exception: + logger.exception("[Orchestrator] _watch_replica_list iteration failed") + + try: + await asyncio.sleep(self._WATCH_REPLICA_INTERVAL_S) + except asyncio.CancelledError: + raise + + async def _handle_register_remote_replica(self, msg: RegisterRemoteReplicaMessage) -> None: + """Bind a head-side client for a newly registered remote replica.""" + stage_id = int(msg.stage_id) + replica_id = int(msg.replica_id) + if not (0 <= stage_id < self.num_stages): + logger.warning( + "[Orchestrator] register_remote_replica: stage_id %d out of range (num_stages=%d)", + stage_id, + self.num_stages, + ) + return + if self._remote_replica_factory is None: + logger.warning( + "[Orchestrator] register_remote_replica received for stage=%d replica=%d but no factory installed", + stage_id, + replica_id, + ) + return + + try: + await self._attach_remote_replica(stage_id, replica_id) + except Exception: + logger.exception( + "[Orchestrator] failed to attach remote replica stage=%d replica=%d", + stage_id, + replica_id, + ) + + async def _handle_unregister_remote_replica(self, msg: UnregisterRemoteReplicaMessage) -> None: + """Tear down the head-side client for a vanished remote replica.""" + stage_id = int(msg.stage_id) + input_addr = str(msg.input_addr) + if not (0 <= stage_id < self.num_stages): + return + pool = self.stage_pools[stage_id] + affected = pool.invalidate_addr(input_addr) + self._detach_remote_replica(stage_id, input_addr) + if affected: + await self._cleanup_request_ids(affected, abort=True) + for req_id in affected: + await self.output_async_queue.put( + ErrorMessage( + error="stage replica disappeared", + request_id=req_id, + stage_id=stage_id, + ) + ) + + async def _attach_remote_replica(self, stage_id: int, replica_id: int) -> None: + """Build a head-side stage client via the injected factory and register it.""" + factory = self._remote_replica_factory + if factory is None: + return + pool = self.stage_pools[stage_id] + client = await factory(stage_id, replica_id) + input_addr = StagePool._client_input_addr(client) + if input_addr is None: + raise RuntimeError( + f"remote replica factory for stage {stage_id} produced a client without a discoverable input address" + ) + pool.add_client(input_addr, client) + logger.info( + "[Orchestrator] attached remote replica stage=%d replica=%d addr=%s", + stage_id, + replica_id, + input_addr, + ) + + def _detach_remote_replica(self, stage_id: int, input_addr: str) -> None: + """Shut down + remove the head-side client at ``input_addr``.""" + pool = self.stage_pools[stage_id] + client = pool.remove_client(input_addr) + if client is None: + return + try: + client.shutdown() + except Exception: + logger.exception( + "[Orchestrator] failed to shutdown client for stage=%d addr=%s", + stage_id, + input_addr, + ) + logger.info( + "[Orchestrator] detached remote replica stage=%d addr=%s", + stage_id, + input_addr, + ) + def _shutdown_stages(self) -> None: """Shutdown all stage pools.""" if self._stages_shutdown: return self._stages_shutdown = True - total = sum(pool.num_replicas for pool in self.stage_pools) + total = sum(pool.live_num_replicas for pool in self.stage_pools) logger.info("[Orchestrator] Shutting down all %d client(s)", total) for pool in self.stage_pools: - for replica_id in range(pool.num_replicas): + for replica_id in pool.live_replica_ids(): pool.shutdown_replica(replica_id) diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index 25e4294e5..e9e2944a7 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -23,6 +23,7 @@ from vllm.v1.engine.exceptions import EngineDeadError from vllm_omni.distributed.omni_connectors.utils.initialization import ( KV_TRANSFER_PORT_OFFSET, ) +from vllm_omni.distributed.omni_connectors.utils.kv_utils import kv_zmq_port from vllm_omni.engine.stage_client import StageClientBase from vllm_omni.engine.stage_init_utils import StageMetadata @@ -344,7 +345,12 @@ class StageEngineCoreClientBase(StageClientBase): try: # Orchestrator always reports rank-0's port; receiver # workers add their own local_rank * KV_RANK_PORT_STRIDE. - sender_port = int(base_port) + KV_TRANSFER_PORT_OFFSET + int(from_stage) + sender_port = kv_zmq_port( + int(base_port), + int(from_stage), + local_rank=0, + replica_id=self.replica_id, + ) except (TypeError, ValueError): logger.warning( "[StageEngineCoreClient] stage-%s [rep-%s] could not resolve sender_zmq_port " @@ -386,7 +392,12 @@ class StageEngineCoreClientBase(StageClientBase): # rank-0 base port; receiver workers adjust per KV_RANK_PORT_STRIDE. return { "host": self._kv_sender_host, - "zmq_port": base_port + kv_transfer_port_offset + int(self.stage_id), + "zmq_port": kv_zmq_port( + base_port - KV_TRANSFER_PORT_OFFSET + kv_transfer_port_offset, + int(self.stage_id), + local_rank=0, + replica_id=self.replica_id, + ), } def set_engine_outputs(self, engine_outputs: EngineCoreOutput) -> None: diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index 2ab8b37dd..0b3dadcbd 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -7,6 +7,8 @@ busy loop in a subprocess, communicating with StageEngineCoreClient via ZMQ. from __future__ import annotations +import contextlib +import os import signal from multiprocessing.process import BaseProcess from typing import TYPE_CHECKING, Any @@ -33,6 +35,8 @@ from vllm.v1.engine.utils import ( ) from vllm.v1.utils import shutdown +from vllm_omni.distributed.omni_coordinator import OmniCoordClientForStage + if TYPE_CHECKING: from vllm.config import VllmConfig from vllm.v1.executor import Executor @@ -53,28 +57,41 @@ class StageEngineCoreProc(EngineCoreProc): *args: Any, dp_rank: int = 0, local_dp_rank: int = 0, + omni_coordinator_address: str | None = None, + omni_stage_id: int | None = None, + omni_replica_id: int = 0, **kwargs: Any, ) -> None: - """Launch StageEngineCoreProc busy loop in background process.""" + """Launch StageEngineCoreProc busy loop in background process. + + Omni-specific kwargs: + - ``omni_coordinator_address``: ROUTER address of the head-side + :class:`OmniCoordinator`. When provided, this subprocess + instantiates an :class:`OmniCoordClientForStage` after the + HELLO/INIT/READY handshake completes and reports its status + + queue length via heartbeats. The hook is wired so each + heartbeat refreshes ``queue_length`` from the live scheduler. + - ``omni_stage_id``: logical stage id this replica belongs to. + Required when ``omni_coordinator_address`` is provided. + - ``omni_replica_id``: cluster-unique replica id within the + stage (assigned by :class:`OmniMasterServer`). Used for + logging / metrics only. + """ signal_callback: SignalCallback | None = None maybe_register_config_serialize_by_value() engine_core: StageEngineCoreProc | None = None + coord_client: OmniCoordClientForStage | None = None try: - vllm_config: VllmConfig = kwargs["vllm_config"] - parallel_config = vllm_config.parallel_config + # NOTE: previous revisions hardcoded data_parallel_size=1 here + # (TODO referencing issue #984). The hardcoding has been removed + # so the DP fields propagate through from the caller exactly + # like upstream vLLM. - set_process_title(f"StageEngineCoreProc_DP{dp_rank}") + stage_label = f"stage{omni_stage_id}" if omni_stage_id is not None else "noid" + set_process_title(f"StageEngineCoreProc_{stage_label}_replica{omni_replica_id}_DP{dp_rank}") decorate_logs() - - # the current vllm-omni does not support data parallelism, - # so we set the data parallel size to 1. - # [TODO] support data parallelism in the future. - # https://github.com/vllm-project/vllm-omni/issues/984 - parallel_config.data_parallel_size = 1 - parallel_config.data_parallel_size_local = 1 - parallel_config.data_parallel_rank = 0 - parallel_config.data_parallel_index = dp_rank + os.environ["VLLM_OMNI_REPLICA_ID"] = str(max(int(omni_replica_id), 0)) engine_core = StageEngineCoreProc( *args, @@ -82,6 +99,41 @@ class StageEngineCoreProc(EngineCoreProc): **kwargs, ) + # Each subprocess corresponds to exactly one omni replica with + # its own OmniMasterServer allocation, so the heartbeat client + # runs unconditionally — there is no dp_rank-based gating. + if omni_coordinator_address is not None: + if omni_stage_id is None: + raise ValueError("omni_stage_id must be provided when omni_coordinator_address is set") + addresses: EngineZmqAddresses = engine_core.addresses + if not addresses.inputs or not addresses.outputs: + raise RuntimeError( + "EngineCore handshake did not populate input/output addresses; " + "cannot start OmniCoordClientForStage" + ) + coord_client = OmniCoordClientForStage( + coord_zmq_addr=omni_coordinator_address, + input_addr=addresses.inputs[0], + output_addr=addresses.outputs[0], + stage_id=int(omni_stage_id), + ) + + def _refresh_queue_length() -> None: + """Pre-heartbeat hook: refresh queue_length from scheduler.""" + scheduler = getattr(engine_core, "scheduler", None) + if scheduler is None: + return + try: + coord_client._queue_length = int( # type: ignore[union-attr] + scheduler.get_num_unfinished_requests() + ) + except Exception: + # Live scheduler stats are best-effort — heartbeats + # must not fail because of a stats lookup error. + pass + + coord_client._on_heartbeat = _refresh_queue_length + def wakeup_engine() -> None: engine_core.input_queue.put_nowait((EngineCoreRequestType.WAKEUP, None)) @@ -111,6 +163,9 @@ class StageEngineCoreProc(EngineCoreProc): signal.signal(signal.SIGINT, signal.SIG_DFL) if signal_callback is not None: signal_callback.stop() + if coord_client is not None: + with contextlib.suppress(RuntimeError): + coord_client.close() if engine_core is not None: engine_core.shutdown() diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py index 05bcdf7d1..a2744eea4 100644 --- a/vllm_omni/engine/stage_engine_startup.py +++ b/vllm_omni/engine/stage_engine_startup.py @@ -4,8 +4,9 @@ from __future__ import annotations import contextlib import dataclasses +import socket import threading -from collections.abc import Iterator +from collections.abc import Callable, Iterator from dataclasses import dataclass from typing import Any @@ -31,6 +32,15 @@ logger = init_logger(__name__) StageRoute = tuple[int, int] +# Sentinel that signals "auto-assign me a replica_id" on the wire. Negative +# values are not valid replica ids, so any sub-zero value works equivalently. +AUTO_ASSIGN_REPLICA_ID = -1 + +# Callback signature for OmniMasterServer.on_register. Fires only for +# auto-assigned replicas (new, headless-launched). The arguments are +# (stage_id, replica_id, allocation). +OnRegisterCallback = Callable[[int, int, "StageAllocation"], None] + # Poll period (ms) used by the registration/handshake loop. _POLL_PERIOD_MS = 5_000 # Default timeout (s) for a stage to send READY. @@ -110,6 +120,10 @@ class OmniMasterServer: master_port: int, stage_ids: list[int], stage_replica_counts: dict[int, int] | None = None, + *, + coordinator_router_address: str | None = None, + on_register: OnRegisterCallback | None = None, + head_local_replicas: dict[int, list[int]] | None = None, ) -> None: self._address = master_address self._port = master_port @@ -117,25 +131,41 @@ class OmniMasterServer: self._stage_configs: dict[StageRoute, Any] = {} self._stage_coordinator_addresses: dict[StageRoute, StageCoordinatorAddresses] = {} self._stage_config_events: dict[StageRoute, threading.Event] = {} + # Coordinator ROUTER address echoed back in every registration reply + # so OmniCoordClientForStage knows where to connect from inside the + # engine subprocess. + self._coordinator_router_address = coordinator_router_address + # Fires only for *newly assigned* (auto-assigned) replicas, not for + # head-side pre-allocated slots that already have head-side clients. + self._on_register = on_register + # Per-stage allocation lock + auto-assign cursor, so concurrent + # registrations from multiple headless processes for the same stage + # don't race on the routing table. + self._alloc_lock = threading.Lock() + self._stage_ids_known: set[int] = set(int(sid) for sid in stage_ids) self._thread: threading.Thread | None = None self._stop_event = threading.Event() stage_replica_counts = dict(stage_replica_counts or {}) + # Slots the *head* itself will fill via ``launch_omni_core_engines`` + # / its own ``register_stage_with_omni_master`` call. Auto-assigning + # headless registrations must skip these even when they appear + # ``_stage_configs``-unfilled — otherwise a fast headless on the same + # host can race the head's own registration and steal slot 0. + self._head_local_slots: set[StageRoute] = set() + for sid, rids in (head_local_replicas or {}).items(): + for rid in rids: + self._head_local_slots.add((int(sid), int(rid))) + for sid in stage_ids: - replica_count = max(1, int(stage_replica_counts.get(sid, 1))) + replica_count = int(stage_replica_counts.get(sid, 1)) + # Allow 0 explicitly so non-self stages (head distributed mode) + # can declare "no local replicas; remote ones will register + # dynamically". + if replica_count < 0: + raise ValueError(f"stage_replica_counts[{sid}] must be >= 0, got {replica_count}") for replica_id in range(replica_count): - route = (sid, replica_id) - self._stage_config_events[route] = threading.Event() - self._stage_coordinator_addresses[route] = StageCoordinatorAddresses() - hs_port, inp_port, out_port = get_open_ports_list(count=3) - self._stage_routes[route] = StageAllocation( - handshake_bind_address=f"tcp://{master_address}:{hs_port}", - handshake_connect_address=f"tcp://{master_address}:{hs_port}", - input_bind_address=f"tcp://{master_address}:{inp_port}", - input_connect_address=f"tcp://{master_address}:{inp_port}", - output_bind_address=f"tcp://{master_address}:{out_port}", - output_connect_address=f"tcp://{master_address}:{out_port}", - ) + self._allocate_route_locked(sid, replica_id) logger.info( "[OmniMasterServer] Pre-allocated addresses for stages %s (master=%s:%d)", @@ -157,10 +187,84 @@ class OmniMasterServer: """Return the registration port exposed to stage launchers.""" return self._port + @property + def coordinator_router_address(self) -> str | None: + """Return the OmniCoordinator ROUTER address echoed to replicas.""" + return self._coordinator_router_address + def get_allocation(self, stage_id: int, replica_id: int = 0) -> StageAllocation: """Return the full address allocation for *stage_id*.""" return self._stage_routes[(stage_id, replica_id)] + # ------------------------------------------------------------------ + # Allocation + # ------------------------------------------------------------------ + + def _allocate_route_locked(self, stage_id: int, replica_id: int) -> StageAllocation: + """Allocate handshake/input/output ports for ``(stage_id, replica_id)``. + + Idempotent: if the route already exists, returns the existing + allocation unchanged. Caller is responsible for holding + ``self._alloc_lock`` when needed. + """ + route = (stage_id, replica_id) + existing = self._stage_routes.get(route) + if existing is not None: + return existing + + self._stage_config_events[route] = threading.Event() + self._stage_coordinator_addresses[route] = StageCoordinatorAddresses() + hs_port, inp_port, out_port = get_open_ports_list(count=3) + alloc = StageAllocation( + handshake_bind_address=f"tcp://{self._address}:{hs_port}", + handshake_connect_address=f"tcp://{self._address}:{hs_port}", + input_bind_address=f"tcp://{self._address}:{inp_port}", + input_connect_address=f"tcp://{self._address}:{inp_port}", + output_bind_address=f"tcp://{self._address}:{out_port}", + output_connect_address=f"tcp://{self._address}:{out_port}", + ) + self._stage_routes[route] = alloc + return alloc + + def _next_free_replica_id(self, stage_id: int) -> int: + """Return the next replica id to assign for an auto-assign registration. + + Strategy: prefer filling a pre-allocated-but-unfilled slot (one that + ``__init__`` reserved in ``_stage_routes`` but no registration has + completed yet) so the head's bootstrap path — which waits on + ``_stage_config_events[(stage_id, replica_id)]`` for specific + pre-allocated ids — unblocks. Only when every pre-allocated slot for + this stage has been filled do we allocate a fresh id. + + Slots in ``_head_local_slots`` are reserved for the head's own + ``launch_omni_core_engines`` registration. Auto-assign must skip + them even when ``_stage_configs`` shows them unfilled — otherwise a + same-host headless that registers before the head's own + ``register_stage_with_omni_master`` call would steal slot 0. + + Without this, a headless contributor using ``--omni-dp-size-local > 1`` + (auto-assign mode) would skip past pre-allocated slot 0 and pick ids + beyond ``num_replicas``, deadlocking the head's + ``connect_remote_engine_cores`` wait. + """ + # Pre-allocated slots that haven't received a registration yet are + # tracked by absence from ``_stage_configs``. Head-owned slots are + # not auto-assignable. + for sid, rid in sorted(self._stage_routes): + if sid != stage_id: + continue + if (sid, rid) in self._head_local_slots: + continue + if (sid, rid) not in self._stage_configs: + return rid + # Every pre-allocated slot is filled (or head-owned); allocate a + # fresh id past the existing routes. + used = {rid for (sid, rid) in self._stage_routes if sid == stage_id} + rid = 0 + while rid in used: + rid += 1 + return rid + def register_stage_config( self, stage_id: int, @@ -280,12 +384,21 @@ class OmniMasterServer: poller = zmq.Poller() poller.register(reg_socket, zmq.POLLIN) + # The server runs until ``stop()`` is called so that headless replicas + # spawned after the head finished its initial bring-up can still + # register dynamically. ``pending`` is kept around purely for + # debug-level logging of which pre-allocated slots have not yet + # registered; once empty it does not terminate the loop. pending: set[StageRoute] = set(self._stage_routes.keys()) - while pending and not self._stop_event.is_set(): + while not self._stop_event.is_set(): events: list[tuple[zmq.Socket, int]] = poller.poll(_POLL_PERIOD_MS) # type: ignore[assignment] if not events: - logger.debug("[OmniMasterServer] Still waiting for registration from stages: %s", pending) + if pending: + logger.debug( + "[OmniMasterServer] Still waiting for registration from pre-allocated slots: %s", + pending, + ) continue for sock, _ in events: @@ -296,12 +409,12 @@ class OmniMasterServer: # Cleanup reg_socket.close(linger=0) - logger.info("[OmniMasterServer] All stages registered; server thread exiting.") + logger.info("[OmniMasterServer] Server thread exiting.") def _handle_registration(self, reg_socket: zmq.Socket) -> StageRoute | None: # type: ignore[type-arg] """Receive a stage registration and reply with the handshake address. - Returns the registered stage_id on success, or None on failure. + Returns ``(stage_id, replica_id)`` on success or ``None`` on failure. """ frames = reg_socket.recv_multipart() if len(frames) < 2: @@ -318,45 +431,186 @@ class OmniMasterServer: logger.warning("[OmniMasterServer] Failed to decode registration message: %s", exc) return None - stage_id: int | None = msg.get("stage_id") - replica_id = int(msg.get("replica_id", 0) or 0) - key = (stage_id, replica_id) - if key not in self._stage_routes: + stage_id_raw = msg.get("stage_id") + if not isinstance(stage_id_raw, int) or stage_id_raw < 0: logger.warning( - "[OmniMasterServer] Received registration for unknown stage_id=%s replica_id=%s", - stage_id, - replica_id, + "[OmniMasterServer] Registration missing or invalid stage_id: %r", + stage_id_raw, ) return None + stage_id: int = stage_id_raw - self.register_stage_config( - stage_id, - msg.get("stage_config"), - coordinator_addresses=StageCoordinatorAddresses( - coordinator_input=msg.get("coordinator_input"), - coordinator_output=msg.get("coordinator_output"), - frontend_stats_publish_address=msg.get("frontend_stats_publish_address"), - ), - replica_id=replica_id, - ) + incoming_replica_id = int(msg.get("replica_id", 0) or 0) + was_auto_assigned = incoming_replica_id < 0 + + # Distinguish two registration shapes: + # - Pre-allocated slots (concrete replica_id >= 0): the head built + # this slot during _initialize_stages. Just confirm it; do NOT + # fire ``on_register`` (the head already has a head-side client). + # - Auto-assigned slots (replica_id == AUTO_ASSIGN_REPLICA_ID): + # a *new* replica from a headless launcher. Allocate, then + # fire ``on_register`` so the orchestrator attaches. + with self._alloc_lock: + if was_auto_assigned: + replica_id = self._next_free_replica_id(stage_id) + # When auto-assign picks a slot the head pre-allocated (and + # is therefore waiting on in ``connect_remote_engine_cores``), + # the head's bootstrap path builds the head-side client. We + # must NOT also fire ``on_register`` for it; otherwise the + # orchestrator would build a duplicate client and overwrite + # the bootstrap-built one in the pool, leaking it. + preexisting_slot = (stage_id, replica_id) in self._stage_routes + alloc = self._allocate_route_locked(stage_id, replica_id) + if preexisting_slot: + was_auto_assigned = False + else: + replica_id = incoming_replica_id + if (stage_id, replica_id) not in self._stage_routes: + # Tolerate explicit replica_ids that haven't been + # pre-allocated (e.g. headless that wants a specific id). + alloc = self._allocate_route_locked(stage_id, replica_id) + was_auto_assigned = True + else: + alloc = self._stage_routes[(stage_id, replica_id)] + + # Cross-host override: when the registering replica advertised + # its own bind address + ports, rewrite the StageAllocation so + # each socket is rooted on the host that actually binds it + # (the master's pre-allocated ports are unreachable from a + # remote replica's host). + # + # Diffusion and LLM stages have different binder ownership: + # + # Diffusion remote replica (StageDiffusionProc): + # - handshake: replica binds -> rewrite to replica IP + # - input : replica binds -> rewrite to replica IP + # - output : replica binds -> rewrite to replica IP + # + # LLM remote replica (CoreClient on head): + # - handshake: head binds (``connect_remote_engine_cores``) + # -> keep on master IP, worker TCP-connects + # - input : head binds (``CoreClient`` ROUTER) + # -> keep on master IP, worker TCP-connects + # - output : head binds (``CoreClient`` PULL — default + # bind=True for PULL in ``make_zmq_socket``) + # -> keep on master IP, worker TCP-connects + # + # The registrant indicates which case via the boolean + # ``replica_binds_sockets`` payload flag. It defaults to + # True (the diffusion / single-host case) so older callers + # still get the previous full-rewrite semantics. For LLM + # remote replicas, the master keeps every address on its + # own host and the remote worker establishes 3 outbound + # TCP connections to the master. + new_bind_address = msg.get("replica_bind_address") + if new_bind_address: + replica_binds_sockets = bool(msg.get("replica_binds_sockets", True)) + if replica_binds_sockets: + hs_port = int(msg["replica_handshake_port"]) + inp_port = int(msg["replica_input_port"]) + out_port = int(msg["replica_output_port"]) + hs_bind_addr = f"tcp://{new_bind_address}:{hs_port}" + inp_bind_addr = f"tcp://{new_bind_address}:{inp_port}" + out_bind_addr = f"tcp://{new_bind_address}:{out_port}" + alloc = StageAllocation( + handshake_bind_address=hs_bind_addr, + handshake_connect_address=hs_bind_addr, + input_bind_address=inp_bind_addr, + input_connect_address=inp_bind_addr, + output_bind_address=out_bind_addr, + output_connect_address=out_bind_addr, + ) + self._stage_routes[(stage_id, replica_id)] = alloc + logger.info( + "[OmniMasterServer] Stage %d replica %d cross-host bind (sockets bound on %s; replica_ip=%s)", + stage_id, + replica_id, + "replica" if replica_binds_sockets else "master", + new_bind_address, + ) + + # Mark the slot as filled *inside* the lock. Without this, + # concurrent auto-assign registrations from a second headless + # could call ``_next_free_replica_id`` between the lock + # releasing above and the ``register_stage_config`` call + # below, observe the slot as unfilled, and hand the same + # pre-allocated handshake/input/output addresses to two + # different replicas — which then collide on + # ``zmq_socket_ctx(handshake_address, ROUTER, bind=True)``. + self.register_stage_config( + stage_id, + msg.get("stage_config"), + coordinator_addresses=StageCoordinatorAddresses( + coordinator_input=msg.get("coordinator_input"), + coordinator_output=msg.get("coordinator_output"), + frontend_stats_publish_address=msg.get("frontend_stats_publish_address"), + ), + replica_id=replica_id, + ) + + # Fire on_register only for genuinely new (auto-assigned or newly + # allocated) replicas, on the ROUTER thread. Callback is expected to + # be cheap and non-blocking (e.g. enqueue onto an asyncio queue). + if was_auto_assigned and self._on_register is not None: + try: + self._on_register(stage_id, replica_id, alloc) + except Exception: + logger.exception( + "[OmniMasterServer] on_register callback failed for stage=%d replica=%d", + stage_id, + replica_id, + ) - alloc = self._stage_routes[key] response = msgspec.msgpack.encode( { "handshake_address": alloc.handshake_connect_address, "input_address": alloc.input_bind_address, "output_address": alloc.output_bind_address, + "replica_id": replica_id, + "coordinator_router_address": self._coordinator_router_address, } ) # ROUTER-DEALER: reply is [identity, payload] (no empty delimiter). reg_socket.send_multipart([identity, response]) logger.info( - "[OmniMasterServer] Stage %d replica %d registered; assigned handshake=%s", + "[OmniMasterServer] Stage %d replica %d registered (auto=%s); handshake=%s", stage_id, replica_id, + was_auto_assigned, alloc.handshake_connect_address, ) - return key + return (stage_id, replica_id) + + +@dataclass(frozen=True) +class StageRegistrationResponse: + """Reply payload returned by :class:`OmniMasterServer` after a successful registration.""" + + handshake_address: str + input_address: str + output_address: str + replica_id: int + coordinator_router_address: str | None + + +def _detect_local_bind_address(master_address: str, master_port: int) -> str: + """Return the local IP the kernel would use to reach the master. + + Uses a connected UDP socket as a routing-table probe: ``connect()`` on + SOCK_DGRAM sends no packets but forces a route lookup, after which + ``getsockname()[0]`` exposes the source IP that an outbound packet to + ``(master_address, master_port)`` would carry. For a co-located master + this returns the loopback or eth0 IP (same effect as the legacy + ``self._address`` behaviour); for a remote master it returns the + NIC IP that's actually reachable from the master — which is exactly + the address the headless's per-stage ZMQ sockets must bind on. + """ + s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + try: + s.connect((master_address, master_port)) + return s.getsockname()[0] + finally: + s.close() def register_stage_with_omni_master( @@ -367,23 +621,37 @@ def register_stage_with_omni_master( omni_stage_config: Any = None, coordinator: DPCoordinator | None = None, return_addresses: bool = False, - replica_id: int = 0, -) -> str | tuple[str, str, str]: + replica_id: int | None = 0, + return_full_response: bool = False, + replica_bind_address: str | None = None, + replica_binds_sockets: bool = True, +) -> str | tuple[str, str, str] | StageRegistrationResponse: """Register a stage with the omni master server. Returns the per-stage handshake address by default. When ``return_addresses`` is true, also returns the stage input/output - addresses allocated by the master. + addresses allocated by the master. When ``return_full_response`` is + true, returns the full :class:`StageRegistrationResponse` including the + assigned ``replica_id`` and the OmniCoordinator ROUTER address (if + published by the master). + + Pass ``replica_id=None`` to request auto-assignment of a free replica + id by the master (used by headless launchers). """ + if replica_id is None: + wire_replica_id = AUTO_ASSIGN_REPLICA_ID + else: + wire_replica_id = int(replica_id) + reg_ctx = zmq.Context() try: reg_sock: zmq.Socket = reg_ctx.socket(zmq.DEALER) # type: ignore[attr-defined] try: reg_sock.connect(f"tcp://{omni_master_address}:{omni_master_port}") - payload = { + payload: dict[str, Any] = { "stage_id": omni_stage_id, - "replica_id": replica_id, + "replica_id": wire_replica_id, "stage_config": _serialize_stage_config(omni_stage_config), } if coordinator is not None: @@ -392,6 +660,31 @@ def register_stage_with_omni_master( payload["coordinator_output"] = coordinator_output payload["frontend_stats_publish_address"] = coordinator.get_stats_publish_address() + # Always advertise THIS host's local bind address + 3 locally + # free ports so the master can root the per-stage socket + # allocation on the replica's own interface. For a co-located + # replica the detected IP matches the master's address and + # the override is a no-op semantically; for a cross-host + # replica it's what makes the headless's ROUTER bind succeed + # (otherwise the master would hand back ``tcp://:port`` + # and ``zmq.bind`` would EADDRNOTAVAIL on the remote host). + if replica_bind_address is None: + replica_bind_address = _detect_local_bind_address(omni_master_address, omni_master_port) + hs_port, inp_port, out_port = get_open_ports_list(count=3) + payload["replica_bind_address"] = replica_bind_address + payload["replica_handshake_port"] = hs_port + payload["replica_input_port"] = inp_port + payload["replica_output_port"] = out_port + # ``False`` only for LLM headless replicas: the head's + # ``connect_remote_engine_cores`` is the binder for the + # handshake ROUTER, and ``CoreClient.__init__`` binds the + # input ROUTER and the output PULL (``make_zmq_socket`` + # defaults bind=True for PULL). The master must keep all + # three addresses on the master's host so the head can + # ``bind`` them; the remote LLM worker TCP-connects across + # hosts on all three. + payload["replica_binds_sockets"] = bool(replica_binds_sockets) + reg_sock.send(msgspec.msgpack.encode(payload)) timeout_ms = _DEFAULT_STARTUP_TIMEOUT_S * 1_000 if not reg_sock.poll(timeout=timeout_ms): @@ -402,13 +695,16 @@ def register_stage_with_omni_master( f"for stage {omni_stage_id}." ) response_bytes = reg_sock.recv() - response = msgspec.msgpack.decode(response_bytes) - handshake_address: str = response["handshake_address"] - input_address: str = response["input_address"] - output_address: str = response["output_address"] + response_msg = msgspec.msgpack.decode(response_bytes) + handshake_address: str = response_msg["handshake_address"] + input_address: str = response_msg["input_address"] + output_address: str = response_msg["output_address"] + assigned_replica_id: int = int(response_msg.get("replica_id", wire_replica_id)) + coord_router_addr: str | None = response_msg.get("coordinator_router_address") logger.info( - "Stage %d registered; handshake_address=%s", + "Stage %d replica %d registered; handshake_address=%s", omni_stage_id, + assigned_replica_id, handshake_address, ) finally: @@ -416,6 +712,14 @@ def register_stage_with_omni_master( finally: reg_ctx.term() + if return_full_response: + return StageRegistrationResponse( + handshake_address=handshake_address, + input_address=input_address, + output_address=output_address, + replica_id=assigned_replica_id, + coordinator_router_address=coord_router_addr, + ) if return_addresses: return handshake_address, input_address, output_address return handshake_address @@ -543,8 +847,16 @@ def launch_omni_core_engines( stage_id: int, stage_config: Any = None, replica_id: int = 0, + *, + omni_coordinator_address: str | None = None, ) -> Iterator[tuple[CoreEngineProcManager, DPCoordinator | None, EngineZmqAddresses]]: - """Launch local engine cores using the omni registration flow.""" + """Launch local engine cores using the omni registration flow. + + When ``omni_coordinator_address`` is provided, the spawned engine + subprocesses use :class:`OmniCoreEngineProcManager` and each + instantiates an :class:`OmniCoordClientForStage` after the handshake + completes so the head's :class:`OmniCoordinator` knows about them. + """ addresses = omni_master_server.get_zmq_addresses(stage_id, replica_id=replica_id) parallel_config = vllm_config.parallel_config # Determine the number of local engines and their ranks. @@ -608,16 +920,35 @@ def launch_omni_core_engines( handshake_bind_address = omni_master_server.get_allocation(stage_id, replica_id=replica_id).handshake_bind_address with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket: - local_engine_manager = CoreEngineProcManager( - local_engine_count=local_engine_count, - start_index=start_index, - local_start_index=local_start_index, - vllm_config=vllm_config, - local_client=True, - handshake_address=handshake_address, - executor_class=executor_class, - log_stats=log_stats, - ) + if omni_coordinator_address is not None: + # Use the omni subclass so each spawned subprocess instantiates + # an OmniCoordClientForStage and heartbeats to the coordinator. + from vllm_omni.engine.omni_core_engine_proc_manager import OmniCoreEngineProcManager + + local_engine_manager: CoreEngineProcManager = OmniCoreEngineProcManager( + local_engine_count=local_engine_count, + start_index=start_index, + local_start_index=local_start_index, + vllm_config=vllm_config, + local_client=True, + handshake_address=handshake_address, + executor_class=executor_class, + log_stats=log_stats, + omni_stage_id=stage_id, + omni_coordinator_address=omni_coordinator_address, + omni_replica_base_id=replica_id, + ) + else: + local_engine_manager = CoreEngineProcManager( + local_engine_count=local_engine_count, + start_index=start_index, + local_start_index=local_start_index, + vllm_config=vllm_config, + local_client=True, + handshake_address=handshake_address, + executor_class=executor_class, + log_stats=log_stats, + ) yield local_engine_manager, coordinator, addresses diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 7ab89400b..1a9f98fe5 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -61,6 +61,63 @@ class LogicalStageInitPlan: replicas: list[ReplicaInitPlan] +@dataclass +class StageRemoteFactoryContext: + """Per-stage context cached by AsyncOmniEngine for dynamic replica attach. + + Populated once during ``_bootstrap_orchestrator`` from the per-stage + init plans. ``_build_remote_replica`` consumes it to construct the + right head-side stage client when a headless replica registers. + """ + + stage_id: int + stage_type: str + stage_cfg: Any + base_metadata: Any + # LLM-only fields: + vllm_config: Any | None = None + executor_class: type | None = None + # Diffusion-only fields: + diffusion_batch_size: int = 1 + + +def capture_stage_factory_contexts( + stage_plans: Sequence[LogicalStageInitPlan], + diffusion_batch_size: int, +) -> dict[int, StageRemoteFactoryContext]: + """Snapshot per-stage construction context for dynamic replica attach. + + Called once after ``_initialize_stages`` finishes. The captured + context holds everything ``_build_remote_replica`` needs to build a + fresh head-side client when a new headless replica registers + (vllm_config / executor_class for LLM, batch_size for diffusion, + plus the base stage metadata). + + Per-replica fields like ``replica_id`` are filled in at build time, + not at capture time. + """ + contexts: dict[int, StageRemoteFactoryContext] = {} + for plan in stage_plans: + if not plan.replicas: + # Stage was declared but has zero replicas locally; we still + # want to be able to attach incoming headless ones, so use + # the stage_cfg-derived context if any replica plan exists. + continue + template = plan.replicas[0] + stage_id = int(plan.configured_stage_id) + stage_type = template.metadata.stage_type or "llm" + contexts[stage_id] = StageRemoteFactoryContext( + stage_id=stage_id, + stage_type=stage_type, + stage_cfg=template.stage_cfg, + base_metadata=template.metadata, + vllm_config=template.stage_vllm_config, + executor_class=template.executor_class, + diffusion_batch_size=diffusion_batch_size, + ) + return contexts + + def _resolve_model_to_local_path(model: str) -> str: """Resolve an HF Hub model ID to a local cache path.""" if os.path.isdir(model): @@ -420,32 +477,59 @@ def split_devices_for_replicas( """Split a devices string into per-replica subsets. When ``num_replicas`` is 1, returns ``[devices_str]`` unchanged. - Otherwise, the total number of device IDs must equal - ``num_replicas * tp_size``; each replica gets ``tp_size`` consecutive - device IDs. + Otherwise, two YAML shapes are accepted: - Example:: + 1. **Legacy / pool mode** — ``len(devices) == num_replicas * tp_size``: + the string enumerates the full per-stage pool. Each replica gets + ``tp_size`` consecutive entries. The values are logical indices + into the launcher's ``CUDA_VISIBLE_DEVICES``. - split_devices_for_replicas("1,2,3,4", num_replicas=2, tp_size=2, stage_id=1) - # → ["1,2", "3,4"] + ``split_devices_for_replicas("1,2,3,4", 2, 2, 1) → ["1,2", "3,4"]`` + + 2. **Template mode** — ``len(devices) == tp_size``: the YAML declares + a single per-replica template (the same shape one replica would + use), and is **dp-independent**. Each replica r gets the offsets + ``[r*tp_size + a for a in template]`` of the launcher's + ``CUDA_VISIBLE_DEVICES``. The template's entries must lie in + ``[0, tp_size)``. + + ``split_devices_for_replicas("0,1", 2, 2, 1) → ["0,1", "2,3"]`` + ``split_devices_for_replicas("0,1", 4, 2, 1) → ["0,1", "2,3", "4,5", "6,7"]`` + + This lets the same ``devices: "0,1"`` YAML work for any + ``--omni-dp-size-local``: the launcher's CVD scales, the YAML + does not. + + Any other length raises ``ValueError`` (the two modes are + length-disjoint for ``num_replicas > 1``). """ if num_replicas <= 1 or devices_str is None: return [devices_str] if devices_str is not None else [devices_str] device_list = [d.strip() for d in devices_str.split(",") if d.strip()] - required = num_replicas * tp_size - if len(device_list) != required: - raise ValueError( - f"Stage {stage_id}: num_replicas={num_replicas}, " - f"tensor_parallel_size={tp_size} requires " - f"{required} devices, got {len(device_list)}: {devices_str}" - ) - result: list[str] = [] - for r in range(num_replicas): - chunk = device_list[r * tp_size : (r + 1) * tp_size] - result.append(",".join(chunk)) - return result + if len(device_list) == num_replicas * tp_size: + return [",".join(device_list[r * tp_size : (r + 1) * tp_size]) for r in range(num_replicas)] + + if len(device_list) == tp_size: + try: + offsets = [int(a) for a in device_list] + except ValueError as e: + raise ValueError(f"Stage {stage_id}: template-mode devices must be ints, got {devices_str!r}") from e + bad = [a for a in offsets if not (0 <= a < tp_size)] + if bad: + raise ValueError( + f"Stage {stage_id}: template-mode device offset(s) {bad} " + f"out of range [0, {tp_size}); devices={devices_str!r}" + ) + return [",".join(str(r * tp_size + a) for a in offsets) for r in range(num_replicas)] + + raise ValueError( + f"Stage {stage_id}: devices={devices_str!r} has {len(device_list)} id(s); " + f"need either {tp_size} (template, dp-independent) or " + f"{num_replicas * tp_size} (pool / legacy). " + f"num_replicas={num_replicas}, tensor_parallel_size={tp_size}." + ) def get_stage_tp_size(stage_cfg: Any) -> int: @@ -479,9 +563,19 @@ def get_stage_devices_per_replica(stage_cfg: Any) -> int: def compute_replica_layout( stage_configs: Sequence[Any], + *, + allow_zero: bool = False, ) -> tuple[list[int], dict[int, list[str]]]: """Compute per-stage replica counts and device assignments. + Args: + stage_configs: per-stage config objects with a ``runtime`` sub-config + exposing ``num_replicas`` and ``devices``. + allow_zero: when True, ``num_replicas == 0`` is honored (used by + single-stage / head-distributed mode for non-self stages that + will be filled dynamically by remote registrations); when False + (default), the count is clamped to at least 1. + Returns: replicas_per_stage: num_replicas per logical stage. replica_devices_map: stage_idx -> per-replica device strings @@ -495,7 +589,9 @@ def compute_replica_layout( if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "num_replicas", 1) ) - replicas_per_stage.append(max(1, num_replicas)) + if num_replicas < 0: + raise ValueError(f"num_replicas must be >= 0, got {num_replicas}") + replicas_per_stage.append(num_replicas if allow_zero else max(1, num_replicas)) replica_devices_map: dict[int, list[str]] = {} for stage_id, stage_cfg in enumerate(stage_configs): diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 427368a02..68de78788 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -10,6 +10,13 @@ from typing import TYPE_CHECKING, Any, cast from vllm.logger import init_logger from vllm.v1.engine import EngineCoreOutputs +from vllm_omni.distributed.omni_coordinator import ( + LoadBalancer, + OmniCoordClientForHub, + ReplicaInfo, + ReplicaStatus, +) +from vllm_omni.distributed.omni_coordinator.load_balancer import Task from vllm_omni.engine.stage_client import ( StagePoolClient, StagePoolDiffusionClient, @@ -35,7 +42,33 @@ class _ReplicaMetrics: class StagePool: - """Replicas of one logical stage with RR + affinity selection.""" + """Replicas of one logical stage + per-stage routing (LB + affinity). + + The pool owns the head-side stage clients for one logical stage. It also + absorbs the per-stage dispatch responsibility (load balancing, affinity + tracking, bounded-wait pick) that used to live in a separate + ``StageDispatcher`` class — see the design doc for the rationale. + + In distributed mode (when an :class:`OmniCoordClientForHub` and a + :class:`LoadBalancer` are injected via :meth:`attach_hub` / + :meth:`attach_load_balancer`), :meth:`pick` consults the hub's cached + replica list and routes via the load balancer, sticking subsequent calls + for the same ``request_id`` to the same replica. + + In non-distributed mode (no hub attached), :meth:`pick` falls back to the + legacy ``select_replica_id`` round-robin path so the multi-stage + in-process invocation is unchanged. + + Dynamic replica membership: when a remote replica is added or removed + (driven by :class:`Orchestrator` via :meth:`add_client` / + :meth:`remove_client`), the pool keeps stable integer ``replica_id``s by + storing clients in a list whose entries can be ``None`` after a removal. + Iteration callers should use :meth:`live_replica_ids` rather than + ``range(pool.num_replicas)`` to skip the gaps. + """ + + DISPATCH_WAIT_TIMEOUT_S: float = 10.0 + DISPATCH_RETRY_INTERVAL_S: float = 0.1 def __init__( self, @@ -50,33 +83,69 @@ class StagePool: else: normalized_clients = [clients] - if not normalized_clients: - raise ValueError(f"StagePool for stage {stage_id} has no replicas") + # Allow empty pools when running in distributed head mode for a + # non-self stage; clients will arrive via add_client(...). self.stage_id = stage_id - self.clients: list[StagePoolClient] = normalized_clients + # Slots can become None after a dynamic remove_client (distributed mode); + # iterate via live_replica_ids() to skip holes. + self.clients: list[StagePoolClient | None] = list(normalized_clients) self._output_processor = output_processor self._stage_vllm_config = stage_vllm_config self._next_replica_id = 0 self._request_bindings: dict[str, int] = {} self._replica_metrics: list[_ReplicaMetrics] = [_ReplicaMetrics() for _ in self.clients] + # Distributed-mode state. Populated by add_client / remove_client. + self._addr_to_replica_id: dict[str, int] = {} + for replica_id, client in enumerate(self.clients): + if client is not None: + addr = self._client_input_addr(client) + if addr is not None: + self._addr_to_replica_id[addr] = replica_id + + # Distributed-mode dispatch hooks (injected by Orchestrator on bring-up). + self._hub: OmniCoordClientForHub | None = None + self._lb: LoadBalancer | None = None + # ``request_id`` → ``input_addr`` affinity (distributed mode only). + # Kept separate from the legacy ``_request_bindings`` so the two + # binding shapes do not collide. + self._affinity: dict[str, str] = {} + # ---- Stage-level properties ---- @property def num_replicas(self) -> int: + """Total slot count, including ``None`` holes from removed replicas. + + Use :meth:`live_replica_ids` to iterate only live entries. + """ return len(self.clients) + @property + def live_num_replicas(self) -> int: + """Number of currently live (non-None) replicas in this pool.""" + return sum(1 for c in self.clients if c is not None) + + def live_replica_ids(self) -> list[int]: + """Return the indices of currently live replicas in this pool.""" + return [i for i, c in enumerate(self.clients) if c is not None] + @property def stage_type(self) -> str | None: - return self.stage_client.stage_type + client = self.stage_client + return None if client is None else client.stage_type @property def final_output(self) -> bool: - return self.clients[0].final_output + client = self.stage_client + return False if client is None else bool(client.final_output) @property - def stage_client(self) -> StagePoolClient: - return self.clients[0] + def stage_client(self) -> StagePoolClient | None: + for client in self.clients: + if client is not None: + return client + return None @property def llm_stage_client(self) -> StagePoolLLMClient: @@ -90,11 +159,209 @@ class StagePool: def output_processor(self) -> Any: return self._output_processor - # ---- Route binding lifecycle ---- + @property + def is_distributed(self) -> bool: + """True iff a hub has been attached (i.e. running in head-distributed mode).""" + return self._hub is not None + + # ---- Distributed-mode dispatch hooks ---- + + def attach_hub(self, hub: OmniCoordClientForHub | None) -> None: + """Inject the shared :class:`OmniCoordClientForHub`. + + Called once by :class:`Orchestrator` after the hub is constructed. + ``hub=None`` keeps the pool in legacy mode (no behavior change). + """ + self._hub = hub + + def attach_load_balancer(self, lb: LoadBalancer | None) -> None: + """Inject the per-pool :class:`LoadBalancer` for distributed-mode pick.""" + self._lb = lb + + # ---- Dynamic membership (distributed mode) ---- + + @staticmethod + def _client_input_addr(client: Any) -> str | None: + """Return the input ZMQ address advertised by ``client`` if any. + + LLM clients expose ``client_addresses["input_address"]``; diffusion + clients expose ``request_address``. Both are stable strings used by + :class:`OmniCoordinator` to key replicas. + """ + request_address = getattr(client, "request_address", None) + if isinstance(request_address, str) and request_address: + return request_address + addrs = getattr(client, "client_addresses", None) + if isinstance(addrs, dict): + addr = addrs.get("input_address") + if isinstance(addr, str) and addr: + return addr + return None + + def add_client(self, input_addr: str, client: Any) -> int: + """Register a head-side client for ``input_addr``. + + Returns the assigned ``replica_id`` (index into :attr:`clients`). + If the address is already known, replaces the existing client and + returns its existing id (this should not happen in practice — the + master server assigns unique slots — but the contract is idempotent + to keep the dispatch layer robust). + """ + if not input_addr: + raise ValueError("input_addr must be a non-empty string") + + existing = self._addr_to_replica_id.get(input_addr) + if existing is not None: + self.clients[existing] = client + return existing + + replica_id = len(self.clients) + self.clients.append(client) + self._addr_to_replica_id[input_addr] = replica_id + self._replica_metrics.append(_ReplicaMetrics()) + return replica_id + + def remove_client(self, input_addr: str) -> Any | None: + """Remove the client at ``input_addr``. Returns the removed client or ``None``. + + Slot is marked ``None`` to preserve indices for outstanding bindings. + """ + replica_id = self._addr_to_replica_id.pop(input_addr, None) + if replica_id is None: + return None + client = self.clients[replica_id] + self.clients[replica_id] = None + return client + + def get_client_by_addr(self, input_addr: str) -> Any | None: + """Return the live client for ``input_addr`` if present.""" + replica_id = self._addr_to_replica_id.get(input_addr) + if replica_id is None: + return None + return self.clients[replica_id] + + def get_replica_id_by_addr(self, input_addr: str) -> int | None: + """Return the stable replica_id for ``input_addr`` if registered.""" + return self._addr_to_replica_id.get(input_addr) + + # ---- Per-request distributed dispatch ---- + + async def pick( + self, + request_id: str, + task: Task | None = None, + *, + affinity_request_id: str | None = None, + ) -> int: + """Return a replica id for ``request_id``. + + In distributed mode: consults the hub for UP replicas, runs the load + balancer, and records affinity so future picks for the same + ``request_id`` return the same replica. Bounded wait up to + ``DISPATCH_WAIT_TIMEOUT_S`` when no UP replica is currently usable. + + In non-distributed (legacy) mode: delegates to + :meth:`select_replica_id`. + """ + if self._hub is None or self._lb is None: + return self.select_replica_id(request_id, affinity_request_id=affinity_request_id) + + # 1. Sticky: previously bound and still serviceable? + bound_addr = self._affinity.get(request_id) + if bound_addr is not None: + replica_id = self._serviceable_replica_id_for_addr(bound_addr) + if replica_id is not None: + return replica_id + # Bound replica is gone or DOWN — fall through to re-select. + self._affinity.pop(request_id, None) + + # 2. Inherited affinity (CFG companion sharing a parent request_id). + if affinity_request_id is not None: + parent_addr = self._affinity.get(affinity_request_id) + if parent_addr is not None: + replica_id = self._serviceable_replica_id_for_addr(parent_addr) + if replica_id is not None: + self._affinity[request_id] = parent_addr + return replica_id + + # 3. Fresh pick: poll hub + LB with bounded wait. + task = task or Task(request_id=request_id) + deadline = _time.monotonic() + self.DISPATCH_WAIT_TIMEOUT_S + while True: + candidates = self._collect_serviceable_replicas() + if candidates: + # LB chose an index *into our candidates list*. + lb_idx = self._lb.select(task, [rep for rep, _ in candidates]) + replica_info, replica_id = candidates[lb_idx] + self._affinity[request_id] = replica_info.input_addr + return replica_id + + now = _time.monotonic() + if now >= deadline: + raise RuntimeError(f"no UP replica for stage {self.stage_id} after {self.DISPATCH_WAIT_TIMEOUT_S:.1f}s") + await asyncio.sleep(min(self.DISPATCH_RETRY_INTERVAL_S, deadline - now)) + + def _collect_serviceable_replicas(self) -> list[tuple[ReplicaInfo, int]]: + """Return list of ``(ReplicaInfo, replica_id)`` for UP, attached replicas.""" + if self._hub is None: + return [] + snap = self._hub.get_replicas_for_stage(self.stage_id) + out: list[tuple[ReplicaInfo, int]] = [] + for rep in snap.replicas: + if rep.status != ReplicaStatus.UP: + continue + replica_id = self._addr_to_replica_id.get(rep.input_addr) + if replica_id is None: + continue # Hub knows about it but head-side client not attached yet. + if self.clients[replica_id] is None: + continue + out.append((rep, replica_id)) + return out + + def _serviceable_replica_id_for_addr(self, input_addr: str) -> int | None: + """Return ``replica_id`` for ``input_addr`` iff currently UP + attached.""" + if self._hub is None: + return None + replica_id = self._addr_to_replica_id.get(input_addr) + if replica_id is None or self.clients[replica_id] is None: + return None + snap = self._hub.get_replicas_for_stage(self.stage_id) + for rep in snap.replicas: + if rep.input_addr == input_addr and rep.status == ReplicaStatus.UP: + return replica_id + return None + + def bind(self, request_id: str, input_addr: str) -> None: + """Explicitly record affinity (distributed mode).""" + self._affinity[request_id] = input_addr + + def release(self, request_id: str) -> None: + """Drop affinity (distributed mode) and legacy binding for ``request_id``.""" + self._affinity.pop(request_id, None) + self.release_binding(request_id) + + def invalidate_addr(self, input_addr: str) -> list[str]: + """Drop affinity rows pointing at ``input_addr``; return affected request ids.""" + affected: list[str] = [rid for rid, addr in self._affinity.items() if addr == input_addr] + for rid in affected: + self._affinity.pop(rid, None) + return affected + + # ---- Legacy (non-distributed) route binding ---- def get_bound_replica_id(self, request_id: str) -> int | None: - """Return the currently bound replica id for *request_id* if present.""" - return self._request_bindings.get(request_id) + """Return the currently bound replica id for *request_id* if present. + + In distributed mode the binding may have been recorded via + :meth:`pick`; we honor it transparently here. + """ + legacy = self._request_bindings.get(request_id) + if legacy is not None: + return legacy + addr = self._affinity.get(request_id) + if addr is None: + return None + return self._addr_to_replica_id.get(addr) def get_bound_client(self, request_id: str) -> StagePoolClient | None: """Return the currently bound client for *request_id* if present.""" @@ -113,6 +380,7 @@ class StagePool: def release_binding(self, request_id: str) -> None: """Drop the route binding for *request_id* in this stage.""" self._request_bindings.pop(request_id, None) + self._affinity.pop(request_id, None) def release_bindings(self, request_ids: list[str]) -> None: """Drop route bindings for the given request ids in this stage.""" @@ -125,27 +393,43 @@ class StagePool: *, affinity_request_id: str | None = None, ) -> int: - """Pick a replica id for *request_id* and cache the choice.""" + """Pick a replica id for *request_id* and cache the choice (legacy path).""" cached = self.get_bound_replica_id(request_id) - if cached is not None: + if cached is not None and self.clients[cached] is not None: return cached - chosen = self.get_bound_replica_id(affinity_request_id) if affinity_request_id is not None else None + chosen: int | None = None + if affinity_request_id is not None: + parent = self.get_bound_replica_id(affinity_request_id) + if parent is not None and self.clients[parent] is not None: + chosen = parent + if chosen is None: - if self.num_replicas == 1: - chosen = 0 + live = self.live_replica_ids() + if not live: + raise RuntimeError(f"stage {self.stage_id} has no live replicas") + if len(live) == 1: + chosen = live[0] else: - chosen = self._next_replica_id - self._next_replica_id = (self._next_replica_id + 1) % self.num_replicas + # Round-robin over live replicas only. + start = self._next_replica_id % len(live) + chosen = live[start] + self._next_replica_id = (self._next_replica_id + 1) % len(live) self._request_bindings[request_id] = chosen return chosen def _llm_client(self, replica_id: int) -> StagePoolLLMClient: - return cast(StagePoolLLMClient, self.clients[replica_id]) + client = self.clients[replica_id] + if client is None: + raise RuntimeError(f"stage {self.stage_id} replica {replica_id} is not attached") + return cast(StagePoolLLMClient, client) def _diffusion_client(self, replica_id: int) -> StagePoolDiffusionClient: - return cast(StagePoolDiffusionClient, self.clients[replica_id]) + client = self.clients[replica_id] + if client is None: + raise RuntimeError(f"stage {self.stage_id} replica {replica_id} is not attached") + return cast(StagePoolDiffusionClient, client) # ---- Metrics ---- @@ -206,7 +490,7 @@ class StagePool: params = params_override if params_override is not None else req_state.sampling_params_list[self.stage_id] submit_kwargs = dict(submit_kwargs or {}) if self.stage_type == "diffusion": - replica_id = self.select_replica_id( + replica_id = await self._pick_or_select( request_id, affinity_request_id=affinity_request_id, ) @@ -217,10 +501,13 @@ class StagePool: await client.add_request_async(request_id, request, params, **submit_kwargs) return replica_id - replica_id = self.select_replica_id( + replica_id = await self._pick_or_select( request_id, affinity_request_id=affinity_request_id, ) + client = self.clients[replica_id] + if client is None: + raise RuntimeError(f"stage {self.stage_id} replica {replica_id} is not attached") try: self.output_processor.add_request( request=request, @@ -262,8 +549,12 @@ class StagePool: """Submit a streaming update to an already admitted request.""" params = req_state.sampling_params_list[self.stage_id] replica_id = self.get_bound_replica_id(request_id) - if replica_id is None: - replica_id = self.select_replica_id(request_id) + if replica_id is None or self.clients[replica_id] is None: + replica_id = await self._pick_or_select(request_id) + + client = self.clients[replica_id] + if client is None: + raise RuntimeError(f"stage {self.stage_id} replica {replica_id} is not attached") if self.stage_type == "diffusion": await self._diffusion_client(replica_id).add_request_async(request_id, request, params) @@ -281,6 +572,17 @@ class StagePool: await self._llm_client(replica_id).add_request_async(request) return replica_id + async def _pick_or_select( + self, + request_id: str, + *, + affinity_request_id: str | None = None, + ) -> int: + """Bridge to ``pick`` in distributed mode or ``select_replica_id`` legacy.""" + if self.is_distributed: + return await self.pick(request_id, affinity_request_id=affinity_request_id) + return self.select_replica_id(request_id, affinity_request_id=affinity_request_id) + # ---- Stage-local polling ---- async def _poll_stage_raw(self, client: StagePoolLLMClient) -> EngineCoreOutputs | None: @@ -296,7 +598,10 @@ class StagePool: raw_outputs: EngineCoreOutputs, ) -> list[Any]: """Run the shared LLM output processor on one raw poll result.""" - client = self._llm_client(replica_id) + raw_client = self.clients[replica_id] + if raw_client is None: + return [] + client = cast(StagePoolLLMClient, raw_client) processor = self.output_processor processed = processor.process_outputs( raw_outputs.outputs, @@ -319,7 +624,10 @@ class StagePool: timeout_s: float = 0.001, ) -> EngineCoreOutputs | None: """Poll raw EngineCore outputs from one LLM replica once.""" - client = self._llm_client(replica_id) + raw_client = self.clients[replica_id] + if raw_client is None: + return None + client = cast(StagePoolLLMClient, raw_client) try: return await asyncio.wait_for( self._poll_stage_raw(client), @@ -339,7 +647,10 @@ class StagePool: def poll_diffusion_output(self, replica_id: int) -> Any | None: """Drain one ready diffusion output from the given replica if present.""" - return self._diffusion_client(replica_id).get_diffusion_output_nowait() + raw_client = self.clients[replica_id] + if raw_client is None: + return None + return cast(StagePoolDiffusionClient, raw_client).get_diffusion_output_nowait() # ---- Stage-local control plane ---- @@ -347,7 +658,7 @@ class StagePool: """Abort the given requests in this stage pool. Request-bound abort routing stays inside the pool because route affinity - (`request_id -> replica_id`) is pool-owned. + (``request_id -> replica_id``) is pool-owned. """ if not request_ids: return @@ -355,13 +666,16 @@ class StagePool: request_ids_by_replica: dict[int, list[str]] = {} for request_id in request_ids: replica_id = self.get_bound_replica_id(request_id) - if replica_id is None: - logger.debug("[StagePool] abort: no binding for req=%s in stage-%s", request_id, self.stage_id) + if replica_id is None or self.clients[replica_id] is None: + logger.debug("[StagePool] abort: no live binding for req=%s in stage-%s", request_id, self.stage_id) continue request_ids_by_replica.setdefault(replica_id, []).append(request_id) for replica_id, replica_request_ids in request_ids_by_replica.items(): - await self.clients[replica_id].abort_requests_async(replica_request_ids) + client = self.clients[replica_id] + if client is None: + continue + await client.abort_requests_async(replica_request_ids) # Clean up OutputProcessor state (e.g. mm_accumulated tensors) that # would otherwise leak — aborted requests never produce a final @@ -377,10 +691,15 @@ class StagePool: timeout: float | None = None, args: tuple[Any, ...] = (), kwargs: dict[str, Any] | None = None, - ) -> Any: + ) -> dict[str, Any] | Any: """Dispatch a stage-scoped control-plane RPC to one physical route.""" kwargs = dict(kwargs or {}) client = self.clients[replica_id] + if client is None: + return { + "supported": False, + "error": f"stage {self.stage_id} replica {replica_id} is not attached", + } try: return await client.collective_rpc_async( method=method, @@ -402,7 +721,11 @@ class StagePool: def shutdown_replica(self, replica_id: int) -> None: """Shutdown one backend handle in this stage pool.""" + if replica_id >= len(self.clients): + return client = self.clients[replica_id] + if client is None: + return try: client.shutdown() logger.info( diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index f2689d4b5..50b825d7c 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -9,6 +9,8 @@ import argparse import json import os import signal +import threading +from multiprocessing import connection from types import FrameType from typing import Any @@ -103,6 +105,62 @@ class OmniServeCommand(CLISubcommand): if args.stage_id is not None and (args.omni_master_address is None or args.omni_master_port is None): raise ValueError("--stage-id requires both --omni-master-address and --omni-master-port to be set") + # --omni-replica-address is only consulted in run_headless(); reject it + # on the head so a misconfigured launch fails loudly instead of being + # silently ignored. + if getattr(args, "omni_replica_address", None) is not None and not args.headless: + raise ValueError("--omni-replica-address requires --headless to be set") + + # --omni-dp-size-local is process-local. A value other than 1 only + # makes sense when this process owns a stage (head or headless). + omni_dp_size_local = getattr(args, "omni_dp_size_local", None) + if omni_dp_size_local is not None: + if omni_dp_size_local < 1: + raise ValueError(f"--omni-dp-size-local must be >= 1, got {omni_dp_size_local}") + if omni_dp_size_local != 1 and args.stage_id is None: + raise ValueError("--omni-dp-size-local != 1 requires --stage-id to be set") + + # vLLM CLI args that omni does not honor: parallelism comes from the + # per-stage YAML (parallel_config:, enable_expert_parallel:) and the + # process-local replica count from --omni-dp-size-local. Passing the + # vLLM equivalents on the command line would silently disagree with + # those sources of truth, so reject them at parse time. + if getattr(args, "omni", False): + explicit_cli_keys: set[str] = getattr(args, "_cli_explicit_keys", set()) or set() + prohibited_with_omni: dict[str, str] = { + "data_parallel_size": "--data-parallel-size", + "data_parallel_size_local": "--data-parallel-size-local", + "data_parallel_address": "--data-parallel-address", + "data_parallel_rpc_port": "--data-parallel-rpc-port", + "data_parallel_start_rank": "--data-parallel-start-rank", + "data_parallel_backend": "--data-parallel-backend", + "api_server_count": "--api-server-count", + "enable_expert_parallel": "--enable-expert-parallel", + } + offenders = sorted(flag for dest, flag in prohibited_with_omni.items() if dest in explicit_cli_keys) + if offenders: + raise ValueError( + "The following CLI args are not supported under --omni: " + f"{', '.join(offenders)}. Configure parallelism through the " + "per-stage YAML (`--deploy-config` / `--stage-configs-path`) " + "and replica count via `--omni-dp-size-local`." + ) + + # --omni-lb-policy is validated against the LoadBalancingPolicy enum. + omni_lb_policy = getattr(args, "omni_lb_policy", None) + if omni_lb_policy is not None: + from vllm_omni.distributed.omni_coordinator import LoadBalancingPolicy + + try: + LoadBalancingPolicy(omni_lb_policy) + except ValueError as exc: + valid = ", ".join(p.value for p in LoadBalancingPolicy) + raise ValueError(f"--omni-lb-policy={omni_lb_policy!r} is not one of: {valid}") from exc + + omni_heartbeat_timeout = getattr(args, "omni_heartbeat_timeout", None) + if omni_heartbeat_timeout is not None and omni_heartbeat_timeout <= 0: + raise ValueError(f"--omni-heartbeat-timeout must be > 0, got {omni_heartbeat_timeout}") + # Skip validation for diffusion models as they have different requirements from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model @@ -190,8 +248,12 @@ class OmniServeCommand(CLISubcommand): omni_config_group.add_argument( "--replica-id", type=int, - default=0, - help="Replica id to register when launching a single headless stage.", + default=None, + help=( + "Deprecated and ignored — replica ids are auto-assigned by the " + "master server. Specifying this flag prints a warning and has " + "no effect." + ), ) omni_config_group.add_argument( "--stage-init-timeout", @@ -253,6 +315,50 @@ class OmniServeCommand(CLISubcommand): type=int, help="Port of the Omni orchestrator (master).", ) + omni_config_group.add_argument( + "--omni-replica-address", + "-ora", + type=str, + default=None, + help=( + "Local bind address (this host's IP) that the headless stage " + "advertises to the Omni master for its handshake/input/output " + "ZMQ sockets. If unset, auto-detected via a UDP-connect " + "routing probe against --omni-master-address. Override only " + "when the auto-detected IP is wrong (e.g. multi-NIC host " + "where the master is reachable on the wrong interface)." + ), + ) + omni_config_group.add_argument( + "--omni-dp-size-local", + type=int, + default=1, + help=( + "Number of stage replicas this runtime launches locally for its " + "own --stage-id. Process-local: head and every headless invocation " + "read their own copy; values may differ across invocations. " + "Requires --stage-id to be set when not equal to 1." + ), + ) + omni_config_group.add_argument( + "--omni-lb-policy", + type=str, + default="random", + choices=["random", "round-robin", "least-queue-length"], + help=( + "Per-stage load-balancing policy used by the head's StagePool to " + "route requests across UP replicas. Only consulted on the head runtime." + ), + ) + omni_config_group.add_argument( + "--omni-heartbeat-timeout", + type=float, + default=30.0, + help=( + "Seconds before an unreporting replica is marked ERROR in the " + "OmniCoordinator. Only consulted on the head runtime." + ), + ) # Diffusion model specific arguments omni_config_group.add_argument( @@ -564,9 +670,15 @@ def _create_default_diffusion_stage_cfg(args: argparse.Namespace) -> list[dict[s def run_headless(args: argparse.Namespace) -> None: - """Run a single stage in headless mode.""" + """Run a single stage in headless mode. + + Honors ``--omni-dp-size-local``: launches that many replicas locally for + ``--stage-id``. Each replica registers with the head's OmniMasterServer + (auto-assigned replica id when ``--omni-dp-size-local > 1`` so multiple + headless invocations can coexist) and reports heartbeats to the head's + OmniCoordinator. + """ from vllm.v1.engine.coordinator import DPCoordinator - from vllm.v1.engine.utils import CoreEngineProcManager from vllm.v1.executor.multiproc_executor import MultiprocExecutor from vllm.version import __version__ as VLLM_VERSION @@ -575,6 +687,7 @@ def run_headless(args: argparse.Namespace) -> None: spawn_diffusion_proc, ) from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_omni_kv_config_for_stage + from vllm_omni.engine.omni_core_engine_proc_manager import OmniCoreEngineProcManager from vllm_omni.engine.stage_engine_startup import register_stage_with_omni_master from vllm_omni.engine.stage_init_utils import ( build_diffusion_config, @@ -582,23 +695,40 @@ def run_headless(args: argparse.Namespace) -> None: build_vllm_config, extract_stage_metadata, get_stage_connector_spec, + get_stage_devices_per_replica, inject_kv_stage_info, load_omni_transfer_config_for_model, prepare_engine_environment, + setup_stage_devices, + split_devices_for_replicas, terminate_alive_proc, ) from vllm_omni.entrypoints.utils import inject_omni_kv_config, load_and_resolve_stage_configs + from vllm_omni.platforms import current_omni_platform model = args.model stage_id: int | None = args.stage_id - replica_id: int = args.replica_id omni_master_address: str | None = args.omni_master_address omni_master_port: int | None = args.omni_master_port + omni_replica_address: str | None = getattr(args, "omni_replica_address", None) + omni_dp_size_local: int = max(1, int(getattr(args, "omni_dp_size_local", 1) or 1)) if stage_id is None: raise ValueError("--stage-id is required in headless mode") - if replica_id < 0: - raise ValueError("--replica-id must be >= 0 in headless mode") + + # ``--replica-id`` is deprecated and ignored — replica ids are + # auto-assigned by ``OmniMasterServer`` so headless processes carry + # no knowledge of their per-replica id at launch time. Warn (don't + # error) when the operator still supplies it so existing launchers + # keep working with a single log line. + explicit_cli_keys: set[str] = getattr(args, "_cli_explicit_keys", set()) or set() + if "replica_id" in explicit_cli_keys: + logger.warning( + "[Headless] --replica-id is deprecated and ignored " + "(supplied value: %s). Replica ids are auto-assigned by the " + "master server.", + args.replica_id, + ) if omni_master_address is None or omni_master_port is None: raise ValueError("--omni-master-address and --omni-master-port are required in headless mode") api_server_count = args.api_server_count or 0 @@ -609,6 +739,11 @@ def run_headless(args: argparse.Namespace) -> None: args_dict = vars(args).copy() args_dict.pop("_cli_explicit_keys", None) + # Forward ``--deploy-config`` so the headless reads the same YAML the + # head was launched with — otherwise ``load_and_resolve_stage_configs`` + # falls back to ``vllm_omni/deploy/.yaml`` and the headless's + # view of ``stage.runtime.devices`` diverges from the head's, breaking + # the per-replica device split. config_path, stage_configs = load_and_resolve_stage_configs( model, args_dict.get("stage_configs_path"), @@ -631,6 +766,39 @@ def run_headless(args: argparse.Namespace) -> None: omni_transfer_config = load_omni_transfer_config_for_model(model, config_path) omni_conn_cfg, omni_from, omni_to = resolve_omni_kv_config_for_stage(omni_transfer_config, stage_id) + # When ``--omni-dp-size-local > 1``, slice the YAML's ``devices:`` field + # into per-replica subsets so each subprocess we spawn below sees a + # narrowed ``CUDA_VISIBLE_DEVICES`` and doesn't stack on cuda:0. Mirrors + # the head-side per-replica device application at + # ``async_omni_engine.py`` (setup_stage_devices around each launch). + runtime_cfg = getattr(stage_cfg, "runtime", None) + devices_str: str | None = None + if runtime_cfg is not None: + devices_str = ( + runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None) + ) + devices_per_replica = get_stage_devices_per_replica(stage_cfg) + if devices_str: + # Always remap YAML's logical devices through setup_stage_devices, + # even for omni_dp_size_local==1. The launcher's CUDA_VISIBLE_DEVICES + # is dropped from the engine-subprocess env between vllm-serve and + # OmniCoreEngineProcManager.Process, so the worker would otherwise + # default cuda:0 to physical GPU 0 and collide with a co-located + # head on the same host (see hyi3_multi_host_1 reproducer). + per_replica_devices: list[str | None] = split_devices_for_replicas( + devices_str, omni_dp_size_local, devices_per_replica, stage_id + ) + logger.info( + "[Headless] Stage %d: %d local replicas, devices_per_replica=%d, per-replica devices: %s", + stage_id, + omni_dp_size_local, + devices_per_replica, + per_replica_devices, + ) + else: + per_replica_devices = [None] * omni_dp_size_local + device_control_env = current_omni_platform.device_control_env_var + if stage_cfg.stage_type == "diffusion": metadata = extract_stage_metadata(stage_cfg) if omni_conn_cfg: @@ -639,39 +807,107 @@ def run_headless(args: argparse.Namespace) -> None: od_config = build_diffusion_config(model, stage_cfg, metadata) logger.info( - "[Headless] Launching diffusion stage %d replica %d via OmniMasterServer at %s:%d", + "[Headless] Launching %d diffusion replica(s) for stage %d via OmniMasterServer at %s:%d", + omni_dp_size_local, stage_id, - replica_id, omni_master_address, omni_master_port, ) - proc = None + procs: list[Any] = [] try: - handshake_address, request_address, response_address = register_stage_with_omni_master( - omni_master_address=omni_master_address, - omni_master_port=omni_master_port, - omni_stage_id=stage_id, - omni_stage_config=stage_cfg, - return_addresses=True, - replica_id=replica_id, + for _rep_idx in range(omni_dp_size_local): + # Always auto-assign: headless processes carry no knowledge + # of their per-replica id and the master server is the sole + # authority on the per-stage id namespace. + response = register_stage_with_omni_master( + omni_master_address=omni_master_address, + omni_master_port=omni_master_port, + omni_stage_id=stage_id, + omni_stage_config=stage_cfg, + replica_id=None, + return_full_response=True, + replica_bind_address=omni_replica_address, + ) + # Apply this replica's CUDA_VISIBLE_DEVICES (only when + # ``--omni-dp-size-local > 1`` and the YAML's stage devices + # field is set). The spawned subprocess inherits the env at + # spawn time; we restore the parent env afterwards so the + # next replica's setup sees the same baseline. + previous_visible_devices = os.environ.get(device_control_env) + try: + if per_replica_devices[_rep_idx] is not None: + setup_stage_devices(stage_id, {"devices": per_replica_devices[_rep_idx]}) + # Each StageDiffusionProc starts its own + # torch.distributed group bound to + # ``od_config.master_port``. Without an explicit + # per-replica override all spawned subprocesses + # share the value ``OmniDiffusionConfig.__post_init__`` + # picked once (and the second binder hits EADDRINUSE + # on ``init_process_group``). We can't use + # kernel-ephemeral allocation either, because the + # master server's pre-allocated ZMQ ports (returned + # by ``register_stage_with_omni_master``) also live + # in the ephemeral range and are not actually bound + # until the headless ``_perform_diffusion_handshake`` + # runs — so picking an ephemeral port here can steal + # a port the master server already promised to a + # sibling headless. Use ``settle_port`` from a base + # above the Linux default ephemeral range + # (32768-60999) so torch.distributed master ports + # never overlap with ZMQ allocations. + if omni_dp_size_local > 1: + od_config.master_port = od_config.settle_port( + 61000 + _rep_idx * 100, + port_inc=37, + ) + proc, _, _, _ = spawn_diffusion_proc( + model, + od_config, + handshake_address=response.handshake_address, + request_address=response.input_address, + response_address=response.output_address, + omni_coordinator_address=response.coordinator_router_address, + omni_stage_id=stage_id, + omni_replica_id=response.replica_id, + ) + finally: + if previous_visible_devices is None: + current_omni_platform.unset_device_control_env_var() + else: + current_omni_platform.set_device_control_env_var(previous_visible_devices) + complete_diffusion_handshake(proc, response.handshake_address, args.stage_init_timeout) + procs.append(proc) + logger.info( + "[Headless] Diffusion replica id=%d for stage %d is up (coord=%s)", + response.replica_id, + stage_id, + response.coordinator_router_address, + ) + + # Block on the sentinel set so any replica crash is detected + # immediately (the previous per-proc join loop only noticed + # crashes in registration order). Any exit triggers fleet + # shutdown via the finally block; non-zero exits propagate. + sentinel_to_proc = {p.sentinel: p for p in procs} + died = connection.wait(list(sentinel_to_proc.keys())) + first = sentinel_to_proc[died[0]] + logger.info( + "[Headless] Diffusion replica %s exited (code=%s); shutting down stage %d.", + first.name, + first.exitcode, + stage_id, ) - proc, _, _, _ = spawn_diffusion_proc( - model, - od_config, - handshake_address=handshake_address, - request_address=request_address, - response_address=response_address, - ) - complete_diffusion_handshake(proc, handshake_address, args.stage_init_timeout) - proc.join() - if proc.exitcode not in (None, 0): - raise RuntimeError(f"Diffusion stage {stage_id} replica {replica_id} exited with code {proc.exitcode}") + if first.exitcode not in (None, 0): + raise RuntimeError( + f"Diffusion stage {stage_id} replica {first.name!r} exited with code {first.exitcode}" + ) return finally: - logger.info("[Headless] Shutting down stage %d replica %d.", stage_id, replica_id) - if proc is not None and proc.is_alive(): - terminate_alive_proc(proc) + logger.info("[Headless] Shutting down %d diffusion replica(s) for stage %d.", len(procs), stage_id) + for p in procs: + if p.is_alive(): + terminate_alive_proc(p) stage_connector_spec = get_stage_connector_spec( omni_transfer_config=omni_transfer_config, @@ -679,8 +915,10 @@ def run_headless(args: argparse.Namespace) -> None: async_chunk=False, ) - # Device assignment is managed externally (e.g. CUDA_VISIBLE_DEVICES); - # runtime_cfg is intentionally ignored in headless mode. + # ``runtime_cfg`` is mostly inherited from the parent's + # CUDA_VISIBLE_DEVICES; when ``--omni-dp-size-local > 1`` we additionally + # bracket each replica's spawn below with setup_stage_devices so they + # don't all stack on cuda:0 (see ``per_replica_devices`` above). engine_args_dict = build_engine_args_dict( stage_cfg, model, @@ -746,57 +984,116 @@ def run_headless(args: argparse.Namespace) -> None: enable_wave_coordination=vllm_config.model_config.is_moe, ) logger.info( - "[Headless] Started DP Coordinator process for stage %d replica %d (PID: %d)", + "[Headless] Started DP Coordinator process for stage %d (PID: %d)", stage_id, - replica_id, coordinator.proc.pid, ) logger.info( - "[Headless] Launching %d engine core(s) for stage %d replica %d via OmniMasterServer at %s:%d", + "[Headless] Launching %d omni replica(s) (vLLM dp_size_local=%d each) for stage %d " + "via OmniMasterServer at %s:%d", + omni_dp_size_local, local_engine_count, stage_id, - replica_id, omni_master_address, omni_master_port, ) - # Headless mode launches all local engine cores for a single stage. - # The OmniMasterServer allocates one handshake endpoint per stage, so we - # register the stage once here and let every local engine core reuse the - # returned handshake address directly. - handshake_address = register_stage_with_omni_master( - omni_master_address=omni_master_address, - omni_master_port=omni_master_port, - omni_stage_id=stage_id, - omni_stage_config=stage_cfg, - coordinator=coordinator, - replica_id=replica_id, - ) - - engine_manager = None + # One OmniMasterServer registration per omni replica; each registration + # yields its own (handshake, input, output) allocation and the head's + # OmniCoordinator ROUTER address. We then spawn one + # OmniCoreEngineProcManager per replica so its subprocess gets the + # right replica id wired into its OmniCoordClientForStage. log_stats = bool(args.log_stats) if args.disable_log_stats: log_stats = False + engine_managers: list[Any] = [] + monitor_threads: list[threading.Thread] = [] + + def _monitor_target(mgr: Any) -> None: + try: + mgr.monitor_engine_liveness() + except Exception: + logger.exception("[Headless] monitor_engine_liveness raised") + try: - engine_manager = CoreEngineProcManager( - local_engine_count=local_engine_count, - start_index=dp_rank, - local_start_index=0, - vllm_config=vllm_config, - local_client=False, - handshake_address=handshake_address, - executor_class=executor_class, - log_stats=log_stats, - ) - # vllm>=0.19 renamed CoreEngineProcManager.join_first() to - # monitor_engine_liveness() (see upstream PR #35862). - engine_manager.monitor_engine_liveness() + for _rep_idx in range(omni_dp_size_local): + # Always auto-assign: see the diffusion branch comment above + # for the rationale (headless owns no replica-id namespace). + response = register_stage_with_omni_master( + omni_master_address=omni_master_address, + omni_master_port=omni_master_port, + omni_stage_id=stage_id, + omni_stage_config=stage_cfg, + coordinator=coordinator, + replica_id=None, + return_full_response=True, + replica_bind_address=omni_replica_address, + # LLM headless: the head binds *all* three sockets — + # handshake ROUTER (``connect_remote_engine_cores``), + # input ROUTER and output PULL (``CoreClient`` — + # ``make_zmq_socket`` defaults bind=True for PULL). + # The remote LLM worker is purely a connector: it + # opens 3 outbound TCP connections to the master's + # host. So the master must keep every address on + # its own host; rewriting any of them to this + # replica's NIC makes the head's ``bind`` go + # EADDRNOTAVAIL on a cross-host launch. + replica_binds_sockets=False, + ) + # Per-replica CUDA_VISIBLE_DEVICES, same pattern as the diffusion + # branch above. OmniCoreEngineProcManager.__init__ spawns its + # subprocesses via context.Process inside the constructor, so we + # must set the env *before* instantiation and restore after. + previous_visible_devices = os.environ.get(device_control_env) + try: + if per_replica_devices[_rep_idx] is not None: + setup_stage_devices(stage_id, {"devices": per_replica_devices[_rep_idx]}) + mgr = OmniCoreEngineProcManager( + local_engine_count=local_engine_count, + start_index=dp_rank, + local_start_index=0, + vllm_config=vllm_config, + local_client=False, + handshake_address=response.handshake_address, + executor_class=executor_class, + log_stats=log_stats, + omni_stage_id=stage_id, + omni_coordinator_address=response.coordinator_router_address, + omni_replica_base_id=response.replica_id, + ) + finally: + if previous_visible_devices is None: + current_omni_platform.unset_device_control_env_var() + else: + current_omni_platform.set_device_control_env_var(previous_visible_devices) + engine_managers.append(mgr) + logger.info( + "[Headless] Stage %d replica id=%d up (coord=%s)", + stage_id, + response.replica_id, + response.coordinator_router_address, + ) + + # Run all managers' liveness monitors in parallel. Each blocks + # until its own subprocesses exit (or fail). + if len(engine_managers) == 1: + engine_managers[0].monitor_engine_liveness() + else: + for mgr in engine_managers: + t = threading.Thread(target=_monitor_target, args=(mgr,), name=f"omni-replica-monitor-{id(mgr):x}") + t.start() + monitor_threads.append(t) + for t in monitor_threads: + t.join() finally: - logger.info("[Headless] Shutting down stage %d.", stage_id) - if engine_manager is not None: - engine_manager.shutdown() + logger.info("[Headless] Shutting down stage %d (%d managers).", stage_id, len(engine_managers)) + for mgr in engine_managers: + try: + mgr.shutdown() + except Exception: + logger.exception("[Headless] engine manager shutdown failed") if coordinator is not None: coordinator.shutdown() diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index 7b725f469..b86c59fcb 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -95,6 +95,24 @@ def _map_device_list(stage_id: int, device_list: list[str], visible_device_list: if not all(device.isdigit() for device in device_list): raise ValueError("Logical devices must be non-negative integers") + # Idempotency: if every requested id already names an entry in the visible + # pool (by value, not by index), treat the YAML field as physical and skip + # the remap. This avoids a double-mapping crash when a parent harness has + # already narrowed CUDA_VISIBLE_DEVICES to specific physical ids before + # spawning the subprocess (the test runtime's OmniServerStageCli does this + # — it sets each stage subprocess's env via ``_set_stage_device_env``). + # Without this guard, a YAML ``devices: "1"`` against a narrowed visible + # list ``["1"]`` would try ``visible[1]`` (out of range) and raise. + visible_set = set(visible_device_list) + if all(device in visible_set for device in device_list): + logger.info( + "Stage %s requested devices %s are already in the visible set %s; using as-is (no remap).", + stage_id, + device_list, + visible_device_list, + ) + return list(device_list) + logical_ids = [int(device) for device in device_list] mapped_devices = [visible_device_list[idx] for idx in logical_ids if idx < num_visible] mapping_pairs = [ From bd83ac9b4b6f7f3a64d13a1695d5a51e73164075 Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Tue, 19 May 2026 14:28:47 +0800 Subject: [PATCH 11/38] [CI] invalid_param reliability suite and weekly http_invalid jobs (#3652) Signed-off-by: wangyu <410167048@qq.com> --- .buildkite/test-weekly.yml | 62 ++ tests/conftest.py | 2 + .../invalid_param_test/conftest.py | 23 + .../test_invalid_audio_diffusion.py | 118 +++ .../test_invalid_audio_speech.py | 862 ++++++++++++++++++ .../test_invalid_image_editing.py | 90 ++ .../test_invalid_image_generation.py | 142 +++ .../test_invalid_omni_chat.py | 216 +++++ .../test_invalid_server_control.py | 95 ++ .../test_invalid_video_generation.py | 171 ++++ ..._qwen2_5_omni_autoround_w4a16_expansion.py | 10 +- ...st_qwen3_omni_autoround_w4a16_expansion.py | 12 +- .../test_omnivoice_expansion.py | 13 - tests/e2e/online_serving/test_voxtral_tts.py | 17 - .../test_qwen3_tts_speaker_embedding.py | 66 +- .../openai_api/test_serving_speech.py | 6 +- tests/helpers/assertions.py | 177 ++-- tests/helpers/runtime.py | 749 ++++++++++++++- 18 files changed, 2647 insertions(+), 184 deletions(-) create mode 100644 tests/dfx/reliability/invalid_param_test/conftest.py create mode 100644 tests/dfx/reliability/invalid_param_test/test_invalid_audio_diffusion.py create mode 100644 tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py create mode 100644 tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py create mode 100644 tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py create mode 100644 tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py create mode 100644 tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py create mode 100644 tests/dfx/reliability/invalid_param_test/test_invalid_video_generation.py diff --git a/.buildkite/test-weekly.yml b/.buildkite/test-weekly.yml index 00c9b25f2..57eadeafd 100644 --- a/.buildkite/test-weekly.yml +++ b/.buildkite/test-weekly.yml @@ -84,3 +84,65 @@ steps: hostPath: path: /mnt/hf-cache type: DirectoryOrCreate + + # Invalid HTTP / WS payloads vs live Omni servers (Issue #3649 tracks known skips). + - group: "Reliability Test - Invalid parameters Test" + depends_on: upload-weekly-pipeline + if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test" + steps: + - label: "Invalid parameters Test · H100" + timeout_in_minutes: 120 + commands: + - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 2 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + + # Stable Audio Open, Qwen3-TTS speech/batch/voices paths marked ``L4``. + - label: "Invalid parameters Test · L4" + timeout_in_minutes: 120 + commands: + - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and L4" + agents: + queue: "gpu_1_queue" + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" diff --git a/tests/conftest.py b/tests/conftest.py index 760404870..163309b8c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -36,6 +36,7 @@ def pytest_terminal_summary(terminalreporter, exitstatus, config): _ASSERTION_EXPORT_NAMES = ( "assert_audio_speech_response", "assert_diffusion_response", + "assert_http_error", "assert_image_diffusion_response", "assert_image_valid", "assert_omni_response", @@ -54,6 +55,7 @@ _MEDIA_EXPORT_NAMES = ( _STAGE_CONFIG_EXPORT_NAMES = ("modify_stage_config",) _RUNTIME_EXPORT_NAMES = ( "DiffusionResponse", + "HttpResponse", "OmniResponse", "OmniRunner", "OmniRunnerHandler", diff --git a/tests/dfx/reliability/invalid_param_test/conftest.py b/tests/dfx/reliability/invalid_param_test/conftest.py new file mode 100644 index 000000000..13fab0bc8 --- /dev/null +++ b/tests/dfx/reliability/invalid_param_test/conftest.py @@ -0,0 +1,23 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared env for ``http_invalid`` (real :func:`omni_server`).""" + +from __future__ import annotations + +import io +import os + +import pytest +from PIL import Image + +# Match ``tests/e2e/online_serving/*`` module-level env for subprocess serve. +os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") +os.environ.setdefault("VLLM_TEST_CLEAN_GPU_MEMORY", "0") + + +@pytest.fixture +def tiny_png_bytes() -> bytes: + img = Image.new("RGB", (32, 32), color="gray") + buf = io.BytesIO() + img.save(buf, format="PNG") + return buf.getvalue() diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_diffusion.py b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_diffusion.py new file mode 100644 index 000000000..ee1515fa6 --- /dev/null +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_diffusion.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""``POST /v1/audio/generate`` validation (live Stable Audio Open server).""" + +from __future__ import annotations + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler + +pytestmark = [pytest.mark.slow, pytest.mark.diffusion] + +# https://github.com/vllm-project/vllm-omni/issues/3649 — invalid inputs still HTTP 200 / timeout / missing validation. +_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649") + +_PARAMS = [ + pytest.param( + OmniServerParams( + model="stabilityai/stable-audio-open-1.0", + server_args=[ + "--trust-remote-code", + "--enforce-eager", + "--model-class-name", + "StableAudioPipeline", + ], + ), + id="stable_audio_open", + marks=hardware_marks(res={"cuda": "L4"}), + ), +] + + +@pytest.mark.parametrize( + "overrides, err_message", + [ + pytest.param({"input": ""}, ("cannot be empty", "input"), id="input_empty", marks=_SKIP_ISSUE_3649), + pytest.param( + {"input": " "}, ("cannot be empty", "input"), id="input_whitespace_only", marks=_SKIP_ISSUE_3649 + ), + pytest.param({"input": 123}, ("input", "a valid string"), id="input_wrong_type"), + pytest.param( + {"response_format": "mpeg"}, + ("response_format", "literal_error", "Input should be"), + id="response_format_invalid", + ), + pytest.param({"speed": 0.1}, ("speed", "greater_than_equal", "0.25"), id="speed_too_low"), + pytest.param({"speed": 99.0}, ("speed", "less_than_equal", "4"), id="speed_too_high"), + pytest.param({"stream_format": 1}, ("stream_format", "literal_error", "audio"), id="stream_format_wrong_type"), + pytest.param( + {"stream_format": "sse"}, + ("sse", "stream_format", "value_error", "not a supported stream_format"), + id="stream_format_sse_blocked", + ), + pytest.param( + {"audio_length": -1.0}, ("audio_end_in_s", "higher than", "audio_start_in_s"), id="audio_length_negative" + ), + pytest.param({"audio_length": 0}, "audio_length", id="audio_length_zero", marks=_SKIP_ISSUE_3649), + # ``OmniOpenAIServingAudioGenerate`` only forwards ``audio_start`` when ``audio_length`` is set. + pytest.param( + {"audio_start": -0.5, "audio_length": 5.0}, + ("audio_start_in_s", "between", "0", "512"), + id="audio_start_negative", + ), + pytest.param( + {"negative_prompt": ["noise"]}, ("negative_prompt", "a valid string"), id="negative_prompt_wrong_type" + ), + pytest.param({"guidance_scale": -1.0}, "guidance_scale", id="guidance_scale_negative", marks=_SKIP_ISSUE_3649), + pytest.param({"guidance_scale": 0}, "guidance_scale", id="guidance_scale_zero", marks=_SKIP_ISSUE_3649), + pytest.param( + {"num_inference_steps": 0}, "num_inference_steps", id="num_inference_steps_zero", marks=_SKIP_ISSUE_3649 + ), + pytest.param( + {"num_inference_steps": -1}, ("number of steps", "non-negative"), id="num_inference_steps_negative" + ), + pytest.param( + {"num_inference_steps": 6000}, + "num_inference_steps", + id="num_inference_steps_above_max", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + {"guidance_scale": 1001.0}, "guidance_scale", id="guidance_scale_above_max", marks=_SKIP_ISSUE_3649 + ), + pytest.param( + {"audio_length": 86401.0}, ("Requested audio length", "exceeds maximum"), id="audio_length_above_max" + ), + pytest.param( + {"audio_start": 86401.0, "audio_length": 1.0}, + ("audio_start_in_s", "between", "0", "512"), + id="audio_start_above_max", + ), + ], +) +@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True) +def test_audio_generate_invalid_field_values( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, + overrides: dict[str, object], + err_message: str | tuple[str, ...], +) -> None: + body = {"model": omni_server.model, "input": "ambient electronic pad"} + body.update(overrides) + openai_client.send_audio_generate_http_request( + {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message} + ) + + +@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True) +def test_audio_generate_missing_input(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + openai_client.send_audio_generate_http_request( + { + "json": {"model": omni_server.model}, + "timeout": 120, + "err_code": 400, + "err_message": ("missing", "input", "field required"), + } + ) diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py new file mode 100644 index 000000000..51a0cbe2b --- /dev/null +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py @@ -0,0 +1,862 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HTTP + WebSocket validation for Qwen3-TTS Base: ``/v1/audio/speech``, ``/v1/audio/speech/stream``, batch, voices.""" + +from __future__ import annotations + +import io +import json +import uuid +import wave +from typing import Any + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.media import load_test_audio_data_url +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +from tests.helpers.stage_config import get_deploy_config_path + +# Speech / voice-upload numeric caps for these tests only (not exported from vllm_omni). +# Align where possible with serving_speech defaults: _TTS_MAX_INSTRUCTIONS_LENGTH (500), +# _TTS_MAX_NEW_TOKENS_MAX (4096). Voice form max lengths are not public constants in prod; +# values below keep tests self-contained and deterministic. +_SPEECH_API_MAX_INSTRUCTIONS_CHARS = 500 +_SPEECH_API_MAX_NEW_TOKENS = 4096 +_VOICE_UPLOAD_MAX_CONSENT_LEN = 1024 +_VOICE_UPLOAD_MAX_REF_TEXT_CHARS = 8192 +_VOICE_UPLOAD_MAX_SPEAKER_DESCRIPTION_CHARS = 2048 +REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav") +REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you." +# Qwen3-TTS 0.6B expects 1024-dim speaker embeddings (see serving_speech / talker hidden_size). +_QWEN3_TTS_06B_SPEAKER_EMBEDDING_DIM = 1024 +# Pop-only marker for ``test_speech_invalid_field_values``: omit ``ref_audio`` / ``ref_text`` so Base +# ``speaker_embedding`` cases mirror standalone embedding-only JSON (no ref_audio clash). +_SPEECH_INVALID_EMBEDDING_ONLY_BODY = "__speech_invalid_embedding_only_body__" +# Batch-only: merged items inherit batch ``ref_audio`` / ``ref_text`` because ``_pick`` treats item ``None`` +# as "unset". Strip clone fields from the batch envelope before asserting CustomVoice preset errors. +_SPEECH_BATCH_NO_CLONE_FIELDS = "__speech_batch_no_clone_fields__" + +_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649") + + +pytestmark = [pytest.mark.slow, pytest.mark.tts] + +_L4_SPEECH_HW = hardware_marks(res={"cuda": "L4"}) +_SPEECH_SERVER_ARGS = ["--trust-remote-code", "--disable-log-stats"] + +_QWEN3_TTS_SPEECH = [ + pytest.param( + OmniServerParams( + model="Qwen/Qwen3-TTS-12Hz-0.6B-Base", + stage_config_path=get_deploy_config_path("qwen3_tts.yaml"), + server_args=_SPEECH_SERVER_ARGS, + ), + id="qwen3_tts", + marks=_L4_SPEECH_HW, + ), +] + + +def _apply_batch_overrides(body: dict[str, object], *, loc: str, overrides: dict[str, object]) -> None: + if loc == "batch": + body.update(overrides) + return + items = body["items"] + assert isinstance(items, list) and len(items) >= 1 + first = items[0] + assert isinstance(first, dict) + body["items"] = [{**first, **overrides}] + + +def _pcm_wav_mono_bytes(*, duration_s: float = 1.1, sample_rate: int = 44100) -> bytes: + buf = io.BytesIO() + with wave.open(buf, "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(sample_rate) + nframes = max(1, int(sample_rate * duration_s)) + wf.writeframes(b"\x00\x00" * nframes) + return buf.getvalue() + + +# ───────────────────────────────────────────────────────────────────────────── +# POST /v1/audio/speech (JSON) +# ───────────────────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True) +def test_speech_malformed_json(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + openai_client.send_audio_speech_http_request( + { + "raw_body": "{", + "timeout": 120, + "err_code": 400, + "err_message": ("json", "decode", "invalid"), + } + ) + + +@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True) +def test_speech_missing_required_fields(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + openai_client.send_audio_speech_http_request( + { + "json": {"model": omni_server.model}, + "timeout": 120, + "err_code": 400, + "err_message": ("input", "field required", "missing"), + } + ) + + +@pytest.mark.parametrize( + "overrides, err_message", + [ + pytest.param({"input": ""}, ("input", "empty"), id="input_empty"), + pytest.param({"input": " "}, ("input", "empty"), id="input_whitespace_only"), + pytest.param({"voice": ""}, "Invalid voice", id="voice_empty", marks=_SKIP_ISSUE_3649), + pytest.param( + {"instructions": 123}, ("instructions", "string_type", "valid string"), id="instructions_wrong_type" + ), + pytest.param( + {"response_format": "mpeg"}, ("response_format", "literal_error", "wav"), id="response_format_invalid" + ), + pytest.param({"stream_format": 1}, ("stream_format", "literal_error", "audio"), id="stream_format_wrong_type"), + pytest.param( + {"stream_format": "sse"}, ("stream_format", "value_error", "audio"), id="stream_format_sse_blocked" + ), + pytest.param({"stream": "wrong_type"}, ("stream", "bool_parsing", "validation error"), id="stream_wrong_type"), + pytest.param( + { + "voice": "vivian", + "task_type": "CustomVoice", + "ref_audio": None, + "ref_text": None, + }, + ("CustomVoice", "does not support"), + id="customvoice_task_not_supported", + ), + pytest.param( + {"task_type": "InvalidEnum"}, ("task_type", "literal_error", "CustomVoice"), id="task_type_invalid" + ), + pytest.param({"language": ""}, ("language", "invalid language", "chinese"), id="language_empty"), + pytest.param({"ref_audio": "ftp://example.com/a.wav"}, ("ref_audio", "URL"), id="ref_audio_bad_scheme"), + pytest.param( + {"ref_audio": "not_a_valid_uri"}, + ("ref_audio", "url"), + id="ref_audio_invalid_uri", + ), + pytest.param({"ref_text": ["x"]}, ("ref_text", "string_type", "valid string"), id="ref_text_wrong_type"), + pytest.param( + {"x_vector_only_mode": "wrong"}, + ("x_vector_only_mode", "valid boolean"), + id="x_vector_only_mode_string_not_boolean", + ), + pytest.param( + {"speaker_embedding": [0.01] * _QWEN3_TTS_06B_SPEAKER_EMBEDDING_DIM}, + ("mutually exclusive", "speaker_embedding", "ref_audio"), + id="speaker_embedding_with_ref_audio_mutually_exclusive", + ), + pytest.param({"max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min"), + pytest.param( + {"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1}, + ("max_new_tokens", "exceed 4096"), + id="max_new_tokens_above_max", + ), + pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative"), + pytest.param({"seed": 2**63}, ("seed", "less_than_equal", "9223372036854775807"), id="seed_above_max"), + pytest.param( + {"initial_codec_chunk_frames": -1}, + ("initial_codec_chunk_frames", "greater_than_equal", "0"), + id="initial_codec_chunk_frames_negative", + ), + pytest.param( + {"instructions": "x" * (_SPEECH_API_MAX_INSTRUCTIONS_CHARS + 1)}, + ("instructions", "too long"), + id="instructions_exceed_api_limit", + ), + # Base clone: whitespace-only ref_text fails serving validation (not pydantic). + pytest.param( + { + "task_type": "Base", + "ref_audio": "data:audio/wav;base64,UklGRiQAAABXQVZFZm10IBAAAAABAAEARKwAAIhYAQACABAAZGF0YQAAAAA=", + "ref_text": " ", + }, + ("ref_text", "base task", "non-empty"), + id="ref_text_whitespace_only_base_clone", + ), + pytest.param( + { + _SPEECH_INVALID_EMBEDDING_ONLY_BODY: True, + "task_type": "Base", + "speaker_embedding": [], + }, + ("speaker_embedding", "non-empty"), + id="base_speaker_embedding_empty", + ), + pytest.param( + { + _SPEECH_INVALID_EMBEDDING_ONLY_BODY: True, + "task_type": "Base", + "x_vector_only_mode": True, + "speaker_embedding": [0.01] * 512, + }, + ("speaker_embedding", "dimensions", str(_QWEN3_TTS_06B_SPEAKER_EMBEDDING_DIM)), + id="base_speaker_embedding_wrong_dim_short", + ), + pytest.param( + { + _SPEECH_INVALID_EMBEDDING_ONLY_BODY: True, + "task_type": "Base", + "x_vector_only_mode": True, + "speaker_embedding": [0.01] * 2048, + }, + ("speaker_embedding", "2048", "dimensions", str(_QWEN3_TTS_06B_SPEAKER_EMBEDDING_DIM)), + id="base_speaker_embedding_wrong_dim_long", + ), + ], +) +@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True) +def test_speech_invalid_field_values( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, + overrides: dict[str, object], + err_message: str | tuple[str, ...], +) -> None: + ov = dict(overrides) + embedding_only = bool(ov.pop(_SPEECH_INVALID_EMBEDDING_ONLY_BODY, False)) + body = { + "model": omni_server.model, + "input": "Hello.", + "voice": "clone", + "language": "English", + "response_format": "wav", + "ref_audio": REF_AUDIO_URL, + "ref_text": REF_TEXT, + } + if embedding_only: + body.pop("ref_audio", None) + body.pop("ref_text", None) + body.update(ov) + openai_client.send_audio_speech_http_request( + {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message} + ) + + +# ───────────────────────────────────────────────────────────────────────────── +# POST /v1/audio/speech/batch (JSON) +# ───────────────────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True) +def test_speech_batch_empty_items(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + openai_client.send_audio_speech_batch_http_request( + {"json": {"items": []}, "timeout": 120, "err_code": (400, 422), "err_message": ("items", "least 1 item")} + ) + + +@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True) +def test_speech_batch_missing_items(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + openai_client.send_audio_speech_batch_http_request( + { + "json": {"model": omni_server.model}, + "timeout": 120, + "err_code": (400, 422), + "err_message": ("items", "field required", "missing"), + } + ) + + +# Align field-level checks with ``POST /v1/audio/speech`` where ``BatchSpeechRequest`` / +# ``SpeechBatchItem`` expose the same knobs. Not on batch schema (single-speech only): +# ``stream``, ``stream_format``, ``seed``, ``speaker_embedding`` — see dedicated speech tests. +@pytest.mark.parametrize( + "loc, overrides, err_message", + [ + pytest.param("item", {"input": ""}, ("input", "empty"), id="item_input_empty", marks=_SKIP_ISSUE_3649), + pytest.param( + "item", + {"input": " "}, + ("input", "empty"), + id="item_input_whitespace_only", + marks=_SKIP_ISSUE_3649, + ), + pytest.param("batch", {"voice": ""}, "Invalid voice", id="batch_voice_empty", marks=_SKIP_ISSUE_3649), + pytest.param("item", {"voice": ""}, "Invalid voice", id="item_voice_empty", marks=_SKIP_ISSUE_3649), + pytest.param( + "batch", + { + "voice": "nonexistent_voice_xyz", + "task_type": "CustomVoice", + "ref_audio": None, + "ref_text": None, + }, + "Invalid voice", + id="batch_voice_unknown_preset", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "item", + { + _SPEECH_BATCH_NO_CLONE_FIELDS: True, + "voice": "nonexistent_voice_xyz", + "task_type": "CustomVoice", + }, + "Invalid voice", + id="item_voice_unknown_preset", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "batch", + {"instructions": 123}, + ("instructions", "string_type", "valid string"), + id="batch_instructions_wrong_type", + ), + pytest.param( + "item", + {"instructions": 123}, + ("instructions", "string_type", "valid string"), + id="item_instructions_wrong_type", + ), + pytest.param( + "batch", + {"instructions": "x" * (_SPEECH_API_MAX_INSTRUCTIONS_CHARS + 1)}, + ("instructions", "too long"), + id="batch_instructions_exceed_api_limit", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "item", + {"instructions": "x" * (_SPEECH_API_MAX_INSTRUCTIONS_CHARS + 1)}, + ("instructions", "too long"), + id="item_instructions_exceed_api_limit", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "batch", + {"response_format": "mpeg"}, + ("response_format", "literal_error", "wav"), + id="batch_response_format_invalid", + ), + pytest.param( + "item", + {"response_format": "mpeg"}, + ("response_format", "literal_error", "wav"), + id="item_response_format_invalid", + ), + pytest.param("batch", {"speed": 0.1}, ("speed", "greater_than_equal", "0.25"), id="batch_speed_too_low"), + pytest.param("item", {"speed": 10.0}, ("speed", "less_than_equal", "4"), id="item_speed_too_high"), + pytest.param( + "batch", + {"task_type": "InvalidEnum"}, + ("task_type", "literal_error", "CustomVoice"), + id="batch_task_type_invalid", + ), + pytest.param( + "item", + {"task_type": "InvalidEnum"}, + ("task_type", "literal_error", "CustomVoice"), + id="item_task_type_invalid", + ), + pytest.param( + "batch", + {"language": ""}, + ("language", "invalid language", "chinese"), + id="batch_language_empty", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "item", + {"language": ""}, + ("language", "invalid language", "chinese"), + id="item_language_empty", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "batch", + {"ref_audio": "ftp://example.com/a.wav"}, + ("ref_audio", "URL"), + id="batch_ref_audio_bad_scheme", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "item", + {"ref_audio": "ftp://example.com/a.wav"}, + ("ref_audio", "URL"), + id="item_ref_audio_bad_scheme", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "batch", + {"ref_audio": "not_a_valid_uri"}, + ("ref_audio", "url"), + id="batch_ref_audio_invalid_uri", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "item", + {"ref_audio": "not_a_valid_uri"}, + ("ref_audio", "url"), + id="item_ref_audio_invalid_uri", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "batch", + {"ref_text": ["x"]}, + ("ref_text", "string_type", "valid string"), + id="batch_ref_text_wrong_type", + ), + pytest.param( + "item", + {"ref_text": ["x"]}, + ("ref_text", "string_type", "valid string"), + id="item_ref_text_wrong_type", + ), + pytest.param( + "batch", + {"x_vector_only_mode": "wrong_type"}, + ("x_vector_only_mode", "bool_parsing", "validation error"), + id="batch_x_vector_only_mode_wrong_type", + ), + pytest.param( + "item", + {"x_vector_only_mode": "wrong_type"}, + ("x_vector_only_mode", "bool_parsing", "validation error"), + id="item_x_vector_only_mode_wrong_type", + ), + pytest.param( + "batch", + {"max_new_tokens": 0}, + ("max_new_tokens", "least 1"), + id="batch_max_new_tokens_below_min", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "item", + {"max_new_tokens": 0}, + ("max_new_tokens", "least 1"), + id="item_max_new_tokens_below_min", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "batch", + {"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1}, + ("max_new_tokens", "exceed 4096"), + id="batch_max_new_tokens_above_max", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "item", + {"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1}, + ("max_new_tokens", "exceed 4096"), + id="item_max_new_tokens_above_max", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "batch", + {"initial_codec_chunk_frames": -1}, + ("initial_codec_chunk_frames", "greater_than_equal", "0"), + id="batch_initial_codec_chunk_frames_negative", + ), + pytest.param( + "item", + {"initial_codec_chunk_frames": -1}, + ("initial_codec_chunk_frames", "greater_than_equal", "0"), + id="item_initial_codec_chunk_frames_negative", + ), + # Base clone: whitespace-only ref_text (same serving check as single speech). + pytest.param( + "batch", + { + "task_type": "Base", + "ref_audio": "data:audio/wav;base64,UklGRiQAAABXQVZFZm10IBAAAAABAAEARKwAAIhYAQACABAAZGF0YQAAAAA=", + "ref_text": " ", + }, + ("ref_text", "base task", "non-empty"), + id="batch_ref_text_whitespace_only_base_clone", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "item", + { + "task_type": "Base", + "ref_audio": "data:audio/wav;base64,UklGRiQAAABXQVZFZm10IBAAAAABAAEARKwAAIhYAQACABAAZGF0YQAAAAA=", + "ref_text": " ", + }, + ("ref_text", "base task", "non-empty"), + id="item_ref_text_whitespace_only_base_clone", + marks=_SKIP_ISSUE_3649, + ), + ], +) +@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True) +def test_speech_batch_invalid_field_values( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, + loc: str, + overrides: dict[str, object], + err_message: str | tuple[str, ...], +) -> None: + ov = dict(overrides) + strip_clone = bool(ov.pop(_SPEECH_BATCH_NO_CLONE_FIELDS, False)) + body = { + "model": omni_server.model, + "voice": "clone", + "ref_audio": REF_AUDIO_URL, + "ref_text": REF_TEXT, + "language": "English", + "response_format": "wav", + "items": [{"input": "Hello."}], + } + if strip_clone: + body.pop("ref_audio", None) + body.pop("ref_text", None) + _apply_batch_overrides(body, loc=loc, overrides=ov) + openai_client.send_audio_speech_batch_http_request( + {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message} + ) + + +# ───────────────────────────────────────────────────────────────────────────── +# WebSocket /v1/audio/speech/stream +# ───────────────────────────────────────────────────────────────────────────── + +# ``test_speech_stream_invalid_requests``: build ``session.config`` JSON with ``omni_server.model``. +_SPEECH_STREAM_WS_SESSION_STREAM_AUDIO_WAV_FRAMES = object() + + +@pytest.mark.parametrize( + "send_frames_spec, err_message", + [ + pytest.param("{not-json", "Invalid JSON", id="invalid_json"), + pytest.param( + json.dumps({"type": "input.text", "text": "hello"}), + "Expected session.config", + id="wrong_first_message_type", + ), + pytest.param( + _SPEECH_STREAM_WS_SESSION_STREAM_AUDIO_WAV_FRAMES, + ("invalid session config", "response_format", "pcm"), + id="session_config_stream_audio_requires_pcm", + ), + ], +) +@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True) +def test_speech_stream_invalid_requests( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, + send_frames_spec: Any, + err_message: str | tuple[str, ...], +) -> None: + """Bad first WebSocket text frames: non-JSON, wrong ``type``, or invalid ``session.config``.""" + if send_frames_spec is _SPEECH_STREAM_WS_SESSION_STREAM_AUDIO_WAV_FRAMES: + send_frames = json.dumps( + { + "type": "session.config", + "model": omni_server.model, + "response_format": "wav", + "stream_audio": True, + } + ) + else: + send_frames = send_frames_spec + assert isinstance(send_frames, str) + openai_client.send_audio_speech_stream_ws_request( + { + "send_frames": send_frames, + "timeout": 120, + "ws_max_size": None, + "ws_json_type": "error", + "err_message": err_message, + } + ) + + +# ───────────────────────────────────────────────────────────────────────────── +# POST /v1/audio/voices (multipart + embedding-only) and DELETE /v1/audio/voices/{name} +# ───────────────────────────────────────────────────────────────────────────── + +# Templates for ``test_voices_create_invalid_requests``: resolved in ``_finalize_voices_form_data``. +_VF_LONG_CONSENT = object() +_VF_LONG_REF_TEXT = object() +_VF_LONG_SPEAKER_DESC = object() +_VF_EMBEDDING_NAN_JSON = object() + + +def _finalize_voices_form_data(template: dict[str, Any], uuid_hex: str) -> dict[str, Any]: + out: dict[str, Any] = {} + for key, val in template.items(): + if val is _VF_LONG_CONSENT: + out[key] = "x" * (_VOICE_UPLOAD_MAX_CONSENT_LEN + 1) + elif val is _VF_LONG_REF_TEXT: + out[key] = "y" * (_VOICE_UPLOAD_MAX_REF_TEXT_CHARS + 1) + elif val is _VF_LONG_SPEAKER_DESC: + out[key] = "z" * (_VOICE_UPLOAD_MAX_SPEAKER_DESCRIPTION_CHARS + 1) + elif val is _VF_EMBEDDING_NAN_JSON: + out[key] = json.dumps([0.1] * 1023 + [float("nan")]) + elif isinstance(val, str): + out[key] = val.replace("{uuid}", uuid_hex) + else: + out[key] = val + return out + + +def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None: + """``None`` => no multipart file (embedding-only or neither).""" + if kind is None: + return None + if kind == "wav_ok": + return {"audio_sample": ("clip.wav", _pcm_wav_mono_bytes(), "audio/wav")} + if kind == "wav_short": + return {"audio_sample": ("clip.wav", _pcm_wav_mono_bytes(duration_s=0.3), "audio/wav")} + if kind == "wav_bad_body": + return {"audio_sample": ("clip.wav", b"not valid wav content", "audio/wav")} + if kind == "wav_pdf_type": + return {"audio_sample": ("clip.wav", _pcm_wav_mono_bytes(), "application/pdf")} + raise AssertionError(f"unknown multipart kind {kind!r}") + + +@pytest.mark.parametrize( + "multipart_kind, form_template, err_message", + [ + pytest.param( + None, + {"consent": "c1", "name": "v1"}, + ("audio_sample", "speaker_embedding", "provided"), + id="neither_audio_nor_embedding", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "wav_ok", + {"name": "v_miss_consent_{uuid}"}, + ("missing", "consent", "field required"), + id="missing_consent", + ), + pytest.param( + "wav_ok", + {"consent": "consent_ok"}, + ("missing", "name", "field required"), + id="missing_name", + ), + pytest.param( + "wav_ok", + {"consent": "", "name": "v_empty_consent_{uuid}"}, + ("missing", "consent", "field required"), + id="empty_consent", + ), + pytest.param( + "wav_ok", + {"consent": " ", "name": "v_ws_consent_{uuid}"}, + "consent", + id="whitespace_consent", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "wav_ok", + {"consent": "bad/consent", "name": "v_bad_cons_{uuid}"}, + "consent", + id="consent_path_sep", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "wav_ok", + {"consent": _VF_LONG_CONSENT, "name": "v_long_cons_{uuid}"}, + ("consent", "too long", "Failed to save"), + id="consent_too_long", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "wav_ok", + {"consent": "consent_ok", "name": ""}, + ("missing", "name", "field required"), + id="empty_name", + ), + pytest.param( + "wav_ok", + {"consent": "consent_ok", "name": "evil/name"}, + ("voice name", "invalid voice name"), + id="invalid_name_path_sep", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "wav_ok", + {"consent": "consent_ok", "name": "v_long_ref_{uuid}", "ref_text": _VF_LONG_REF_TEXT}, + "ref_text", + id="ref_text_too_long", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "wav_ok", + { + "consent": "consent_ok", + "name": "v_long_desc_{uuid}", + "speaker_description": _VF_LONG_SPEAKER_DESC, + }, + "speaker_description", + id="speaker_description_too_long", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "wav_pdf_type", + {"consent": "consent_ok", "name": "v_bad_mime_{uuid}"}, + ("mime", "unsupported", "audio/mp4"), + id="audio_unsupported_mime", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "wav_short", + {"consent": "consent_ok", "name": "v_short_{uuid}"}, + ("too short", "reference audio"), + id="audio_too_short", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "wav_bad_body", + {"consent": "consent_ok", "name": "v_bad_wav_{uuid}"}, + ("decode", "audio", "Format not recognised"), + id="audio_decode_error", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + None, + { + "speaker_embedding": "not valid json [[[", + "consent": "consent_emb", + "name": "v_bad_json_{uuid}", + }, + ("speaker_embedding", "valid JSON"), + id="speaker_embedding_invalid_json", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + None, + { + "speaker_embedding": json.dumps([]), + "consent": "consent_emb", + "name": "v_empty_emb_{uuid}", + }, + ("speaker_embedding", "non-empty"), + id="speaker_embedding_empty", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + None, + { + "speaker_embedding": _VF_EMBEDDING_NAN_JSON, + "consent": "consent_emb", + "name": "v_nan_emb_{uuid}", + }, + ("speaker_embedding", "finite"), + id="speaker_embedding_nan", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + None, + { + "speaker_embedding": json.dumps([0.01] * 512), + "consent": "consent_emb", + "name": "v_wrong_dim_{uuid}", + }, + ("expected", "dimensions", "1024"), + id="speaker_embedding_wrong_dim", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "wav_ok", + { + "consent": "consent_emb", + "name": "v_both_{uuid}", + "speaker_embedding": json.dumps([0.01] * 1024), + }, + ("mutually exclusive", "audio_sample", "speaker_embedding"), + id="audio_and_embedding_mutually_exclusive", + marks=_SKIP_ISSUE_3649, + ), + ], +) +@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True) +def test_voices_create_invalid_requests( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, + multipart_kind: str | None, + form_template: dict[str, Any], + err_message: str | tuple[str, ...], +) -> None: + """Invalid ``POST /v1/audio/voices`` bodies (multipart or embedding-only).""" + uid = uuid.uuid4().hex[:8] + cfg: dict[str, Any] = { + "data": _finalize_voices_form_data(form_template, uid), + "timeout": 120, + "err_code": 400, + "err_message": err_message, + } + files = _voices_upload_multipart_files(multipart_kind) + if files is not None: + cfg["files"] = files + openai_client.send_audio_voices_create_http_request(cfg) + + +@pytest.mark.parametrize( + "voice_name, err_code, err_message", + [ + pytest.param( + "missing-voice-404", + 404, + "not found", + id="not_found", + ), + pytest.param( + " ", + 404, + "not found", + id="whitespace_only", + ), + pytest.param( + "evil/name", + 404, + "not found", + id="path_separator", + ), + pytest.param( + ".", + 405, + ("Method Not Allowed", "detail"), + id="dot_single", + ), + pytest.param( + "..", + 404, + "not found", + id="dot_dot", + ), + pytest.param( + "bad\x00voice", + 404, + "not found", + id="nul_byte", + ), + pytest.param( + "", + 405, + "Method Not Allowed", + id="empty_name_path", + ), + ], +) +@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True) +def test_voices_delete_invalid_requests( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, + voice_name: str, + err_code: int | tuple[int, ...], + err_message: str | tuple[str, ...], +) -> None: + """DELETE ``/v1/audio/voices/{name}``: missing voice (404), odd segments (often 404), ``.``paths (405).""" + openai_client.send_audio_voices_delete_http_request( + { + "name": voice_name, + "timeout": 120, + "err_code": err_code, + "err_message": err_message, + } + ) diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py b/tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py new file mode 100644 index 000000000..fa7f0f905 --- /dev/null +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""``POST /v1/images/edits`` multipart validation (live ``Qwen/Qwen-Image-Edit``).""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler + +pytestmark = [pytest.mark.slow, pytest.mark.diffusion] + +_PARAMS = [ + pytest.param( + OmniServerParams(model="Qwen/Qwen-Image-Edit"), + id="qwen_image_edit", + marks=hardware_marks(res={"cuda": "H100"}), + ), +] + +# Populated in ``test_images_edits_invalid_requests`` from ``omni_server.model``. +_IMAGE_EDITS_SERVER_MODEL = object() + + +def _finalize_edits_form_data(template: dict[str, Any], server_model: str) -> dict[str, Any]: + return {k: (server_model if v is _IMAGE_EDITS_SERVER_MODEL else v) for k, v in template.items()} + + +@pytest.mark.parametrize( + "include_image, data_template, err_message, err_code", + [ + pytest.param( + False, + {"prompt": "make it brighter", "model": _IMAGE_EDITS_SERVER_MODEL}, + ("image", "required", "Unprocessable Entity"), + 422, + id="missing_image", + ), + pytest.param( + True, + {"prompt": "edit", "model": "wrong-model-id"}, + ("model", "mismatch"), + 400, + id="model_mismatch", + ), + pytest.param( + True, + {"prompt": "edit", "response_format": "url"}, + ("response_format", "b64_json", "supported"), + 400, + id="unsupported_response_format", + ), + pytest.param( + True, + {"prompt": "edit", "model": _IMAGE_EDITS_SERVER_MODEL, "output_compression": "101"}, + ("less_than_equal", "output_compression", "100"), + 400, + id="output_compression_above_max", + ), + pytest.param( + True, + {"prompt": "edit", "model": _IMAGE_EDITS_SERVER_MODEL, "layers": "2"}, + ("Invalid layers", "3", "1"), + 400, + id="invalid_layers", + ), + ], +) +@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True) +def test_images_edits_invalid_requests( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, + tiny_png_bytes: bytes, + include_image: bool, + data_template: dict[str, Any], + err_message: str | tuple[str, ...], + err_code: int | tuple[int, ...], +) -> None: + cfg: dict[str, Any] = { + "data": _finalize_edits_form_data(data_template, omni_server.model), + "timeout": 300, + "err_code": err_code, + "err_message": err_message, + } + if include_image: + cfg["files"] = {"image": ("x.png", tiny_png_bytes, "image/png")} + openai_client.send_images_edits_http_request(cfg) diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py b/tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py new file mode 100644 index 000000000..6f7cf984d --- /dev/null +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""``POST /v1/images/generations`` validation (live ``Qwen/Qwen-Image``).""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler + +pytestmark = [pytest.mark.slow, pytest.mark.diffusion] + +_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649") + +_PARAMS = [ + pytest.param( + OmniServerParams(model="Qwen/Qwen-Image"), + id="qwen_image", + marks=hardware_marks(res={"cuda": "H100"}), + ), +] + +# Full replacement body for ``test_images_generations_invalid_requests`` (no minimal prompt/model merge). +_IMAGES_GEN_MISSING_PROMPT_JSON = object() + + +def _minimal_images_gen_json(omni_server: OmniServer) -> dict[str, object]: + return {"prompt": "a simple red apple icon", "model": omni_server.model} + + +# ───────────────────────────────────────────────────────────────────────────── +# POST /v1/images/generations (JSON) +# ───────────────────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "body_spec, err_message", + [ + pytest.param(_IMAGES_GEN_MISSING_PROMPT_JSON, ("prompt", "Field required", "Missing"), id="missing_prompt"), + pytest.param( + {"prompt": "cat", "model": "not-the-loaded-checkpoint"}, + ("model", "mismatch"), + id="model_mismatch", + ), + pytest.param({"prompt": ""}, ("prompt", "empty"), id="prompt_empty", marks=_SKIP_ISSUE_3649), + pytest.param({"prompt": " "}, ("prompt", "empty"), id="prompt_whitespace_only", marks=_SKIP_ISSUE_3649), + pytest.param({"prompt": 123}, ("prompt", "string_type", "valid string"), id="prompt_wrong_type"), + pytest.param({"n": 0}, ("n", "greater_than_equal", "1"), id="n_below_min"), + pytest.param({"n": 11}, ("n", "less_than_equal", "10"), id="n_above_max"), + pytest.param({"size": "1024"}, ("size", "value_error", "WIDTHxHEIGHT"), id="size_missing_height_token"), + pytest.param({"size": "not-a-size"}, ("size", "value_error", "WIDTHxHEIGHT"), id="size_missing_separator"), + pytest.param({"size": "1024xabc"}, ("Invalid size format", "integers"), id="size_non_integer_height"), + pytest.param({"size": "1024x"}, ("Invalid size format", "integers"), id="size_incomplete_height"), + pytest.param({"size": "0x1024"}, ("Invalid size", "positive integers"), id="size_non_positive_width"), + pytest.param( + {"response_format": "url"}, + ("response_format", "value_error", "b64_json"), + id="response_format_not_b64_json", + ), + pytest.param( + {"use_system_prompt": "bogus"}, + ("use_system_prompt", "value_error", "dynamic"), + id="use_system_prompt_invalid", + ), + pytest.param( + {"num_inference_steps": 0}, + ("num_inference_steps", "'greater_than_equal", "1"), + id="num_inference_steps_below_min", + ), + pytest.param( + {"num_inference_steps": 201}, + ("num_inference_steps", "less_than_equal", "200"), + id="num_inference_steps_above_max", + ), + pytest.param( + {"guidance_scale": -1.0}, ("guidance_scale", "greater_than_equal", "0"), id="guidance_scale_negative" + ), + pytest.param( + {"guidance_scale": 21.0}, ("guidance_scale", "less_than_equal", "20"), id="guidance_scale_above_max" + ), + pytest.param( + {"true_cfg_scale": -1.0}, ("true_cfg_scale", "greater_than_equal", "0"), id="true_cfg_scale_negative" + ), + pytest.param( + {"true_cfg_scale": 21.0}, ("true_cfg_scale", "less_than_equal", "20"), id="true_cfg_scale_above_max" + ), + pytest.param({"layers": 2}, ("layers", "value_error", "3", "10"), id="layers_below_min"), + pytest.param({"layers": 11}, ("layers", "value_error", "3", "10"), id="layers_above_max"), + pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative", marks=_SKIP_ISSUE_3649), + pytest.param( + {"seed": 2**32}, + ("seed", "less_than_equal", "4294967295"), + id="seed_above_uint32", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + {"output_format": "gif"}, + ("output_format", "value_error", "b64_json"), + id="output_format_invalid", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + {"vae_use_slicing": "wrong_type"}, + ("vae_use_slicing", "bool_parsing", "validation error"), + id="vae_use_slicing_wrong_type", + ), + pytest.param({"user": 123}, ("user", "string_type", "valid string"), id="user_wrong_type"), + pytest.param( + {"negative_prompt": ["noise"]}, + ("negative_prompt", "string_type", "valid string"), + id="negative_prompt_wrong_type", + ), + pytest.param( + {"generator_device": 1}, + ("generator_device", "string_type", "valid string"), + id="generator_device_wrong_type", + ), + pytest.param( + {"lora": {"foo": "bar"}}, ("lora", "both name and path", "required"), id="lora_missing_required_fields" + ), + ], +) +@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True) +def test_images_generations_invalid_requests( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, + body_spec: Any, + err_message: str | tuple[str, ...], +) -> None: + """Invalid JSON bodies for ``POST /v1/images/generations`` (missing prompt, model mismatch, bad fields).""" + if body_spec is _IMAGES_GEN_MISSING_PROMPT_JSON: + body: dict[str, object] = {"size": "1024x1024"} + else: + assert isinstance(body_spec, dict) + body = _minimal_images_gen_json(omni_server) + body.update(body_spec) + openai_client.send_images_generations_http_request( + {"json": body, "timeout": 300, "err_code": 400, "err_message": err_message} + ) diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py b/tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py new file mode 100644 index 000000000..9d1782e76 --- /dev/null +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Invalid inputs on Qwen3-Omni: ``POST /v1/chat/completions``, ``WS /v1/video/chat/stream``, ``WS /v1/realtime``.""" + +from __future__ import annotations + +import json +from typing import Any + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +from tests.helpers.stage_config import get_deploy_config_path + +pytestmark = [pytest.mark.slow, pytest.mark.omni] + +_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649") + + +def _minimal_chat_json(omni_server: OmniServer) -> dict[str, object]: + """Minimal valid chat body; individual tests override one offending field.""" + return { + "model": omni_server.model, + "messages": [{"role": "user", "content": "hello"}], + "stream": False, + } + + +_QWEN3_OMNI_SERVER = [ + pytest.param( + OmniServerParams( + model="Qwen/Qwen3-Omni-30B-A3B-Instruct", + stage_config_path=get_deploy_config_path("qwen3_omni_moe.yaml"), + server_args=["--async-chunk"], + ), + id="qwen3_omni", + ), +] + + +def _chat_completions_request_without_expectations(omni_server: OmniServer, case_id: str) -> dict[str, Any]: + """Build ``send_chat_completions_http_request`` kwargs excluding ``err_code`` / ``err_message``.""" + if case_id == "malformed_json": + return {"raw_body": "{not-json", "timeout": 120} + if case_id == "missing_messages": + return {"json": {"model": omni_server.model, "stream": False}, "timeout": 120} + if case_id == "model_mismatch": + return { + "json": { + "model": "this-model-is-not-served-on-this-instance", + "messages": [{"role": "user", "content": "hello"}], + "stream": False, + }, + "timeout": 120, + } + body = _minimal_chat_json(omni_server) + if case_id == "stream_wrong_type": + body["stream"] = "not-a-boolean" # type: ignore[assignment] + elif case_id == "temperature_wrong_type": + body["temperature"] = "hot" # type: ignore[assignment] + elif case_id == "max_tokens_wrong_type": + body["max_tokens"] = "not-an-int" # type: ignore[assignment] # type: ignore[assignment] + elif case_id == "modalities_list_bad": + body["modalities"] = [123] # type: ignore[assignment] + elif case_id == "response_format_json_schema_incomplete": + body["response_format"] = {"type": "json_schema"} + elif case_id == "logprobs_wrong_type": + body["logprobs"] = "yes" # type: ignore[assignment] + elif case_id == "logprobs_top_without_enabled": + body["logprobs"] = False + body["top_logprobs"] = 5 + elif case_id == "speaker_unknown": + body["speaker"] = "zz_invalid_qwen3_omni_chat_speaker_xyz" + else: + raise AssertionError(f"unknown chat completions invalid case_id {case_id!r}") + return {"json": body, "timeout": 120} + + +# ───────────────────────────────────────────────────────────────────────────── +# POST /v1/chat/completions +# ───────────────────────────────────────────────────────────────────────────── + + +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize( + "case_id, err_code, err_message", + [ + pytest.param( + "malformed_json", + 400, + ("json_invalid", "decode error"), + id="malformed_json", + ), + pytest.param("missing_messages", 400, ("messages", "missing", "Field required"), id="missing_messages"), + pytest.param( + "model_mismatch", + 404, + ("model", "does not exist", "NotFoundError"), + id="model_mismatch", + ), + pytest.param("stream_wrong_type", 400, ("stream", "bool_parsing", "not-a-boolean"), id="invalid_stream_type"), + pytest.param( + "temperature_wrong_type", + 400, + ("temperature", "float_parsing", "a valid number"), + id="invalid_temperature_type", + ), + pytest.param( + "max_tokens_wrong_type", 400, ("max_tokens", "int_parsing", "integer"), id="invalid_max_tokens_type" + ), + pytest.param( + "modalities_list_bad", + 400, + ("modalities", "value_error", ""), + id="modalities_list_bad_element", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "response_format_json_schema_incomplete", + 400, + ("response_format", "value_error", "json_schema"), + id="invalid_response_format_json_schema", + ), + pytest.param("logprobs_wrong_type", 400, "logprobs", id="logprobs_wrong_type", marks=_SKIP_ISSUE_3649), + pytest.param( + "logprobs_top_without_enabled", + 400, + ("top_logprobs", "logprobs", "True"), + id="logprobs_top_without_logprobs_true", + ), + pytest.param( + "speaker_unknown", + 400, + ("Invalid speaker", "Supported"), + id="speaker_unknown_preset", + ), + ], +) +@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True) +def test_chat_completions_invalid_requests( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, + case_id: str, + err_code: int | tuple[int, ...], + err_message: str | tuple[str, ...], +) -> None: + cfg = _chat_completions_request_without_expectations(omni_server, case_id) + cfg["err_code"] = err_code + cfg["err_message"] = err_message + openai_client.send_chat_completions_http_request(cfg)[0] + + +# ───────────────────────────────────────────────────────────────────────────── +# WS /v1/video/chat/stream +# ───────────────────────────────────────────────────────────────────────────── + +# First-frame JSON with ``omni_server.model`` for ``test_video_chat_stream_invalid_requests``. +_VIDEO_CHAT_WS_SESSION_MODALITIES_SCALAR = object() + + +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize( + "send_frames_spec, err_message", + [ + pytest.param("{not-json", "Invalid JSON", id="invalid_json"), + pytest.param( + _VIDEO_CHAT_WS_SESSION_MODALITIES_SCALAR, + "Invalid session config", + id="session_modalities_scalar", + ), + ], +) +@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True) +def test_video_chat_stream_invalid_requests( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, + send_frames_spec: Any, + err_message: str, +) -> None: + if send_frames_spec is _VIDEO_CHAT_WS_SESSION_MODALITIES_SCALAR: + send_frames = json.dumps({"type": "session.config", "model": omni_server.model, "modalities": "text"}) + else: + send_frames = send_frames_spec + assert isinstance(send_frames, str) + openai_client.send_video_chat_stream_ws_request( + { + "send_frames": send_frames, + "timeout": 120, + "ws_max_size": None, + "ws_json_type": "error", + "err_message": err_message, + } + ) + + +# ───────────────────────────────────────────────────────────────────────────── +# WS /v1/realtime +# ───────────────────────────────────────────────────────────────────────────── + + +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True) +def test_realtime_rejected_when_async_chunk_enabled( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, +) -> None: + openai_client.send_realtime_ws_request( + { + "timeout": 120, + "ws_max_size": None, + "ws_json_type": "error", + "ws_error_code": "unsupported", + "err_message": "async_chunk", + } + ) diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py b/tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py new file mode 100644 index 000000000..0a8240c4b --- /dev/null +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Omni sleep/wakeup JSON validation (live ``Qwen/Qwen-Image`` process).""" + +from __future__ import annotations + +from typing import Any, Literal + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler + +pytestmark = [pytest.mark.slow, pytest.mark.diffusion, pytest.mark.omni] + +_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649") + +_QWEN_IMAGE = [ + pytest.param( + OmniServerParams(model="Qwen/Qwen-Image"), + id="qwen_image", + marks=hardware_marks(res={"cuda": "H100"}), + ), +] + + +# ───────────────────────────────────────────────────────────────────────────── +# POST omni sleep / omni wakeup (invalid JSON) +# ───────────────────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "endpoint, json_body, err_message", + [ + pytest.param("sleep", {"level": 2}, ("stage_ids", "Field required", "Missing"), id="sleep_missing_stage_ids"), + pytest.param( + "sleep", + {"stage_ids": [], "level": 2}, + ("stage_ids", "Field required", "Missing"), + id="sleep_empty_stage_ids", + marks=_SKIP_ISSUE_3649, + ), + pytest.param( + "sleep", + {"stage_ids": "not-a-list", "level": 2}, + ("stage_ids", "list_type", "a valid list"), + id="sleep_invalid_stage_ids_type", + ), + pytest.param( + "sleep", + {"stage_ids": [0], "level": "high"}, + ("level", "int_parsing", "a valid integer"), + id="sleep_invalid_level_type", + ), + pytest.param( + "sleep", + {"stage_ids": [0], "level": -1}, + ("level", "greater"), + id="sleep_negative_level", + marks=_SKIP_ISSUE_3649, + ), + pytest.param("wakeup", {}, ("stage_ids", "missing", "Field required"), id="wakeup_missing_stage_ids"), + pytest.param( + "wakeup", + {"stage_ids": None}, + ("stage_ids", "list_type", "a valid list"), + id="wakeup_invalid_stage_ids_type", + ), + pytest.param( + "wakeup", + {"stage_ids": []}, + ("stage_ids", "length"), + id="wakeup_empty_stage_ids", + marks=_SKIP_ISSUE_3649, + ), + ], +) +@pytest.mark.parametrize("omni_server", _QWEN_IMAGE, indirect=True) +def test_omni_sleep_wakeup_invalid_requests( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, + endpoint: Literal["sleep", "wakeup"], + json_body: dict[str, Any], + err_message: str | tuple[str, ...], +) -> None: + cfg: dict[str, Any] = { + "json": json_body, + "timeout": 120, + "err_code": 400, + "err_message": err_message, + } + if endpoint == "sleep": + openai_client.send_omni_sleep_http_request(cfg) + else: + openai_client.send_omni_wakeup_http_request(cfg) diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_video_generation.py b/tests/dfx/reliability/invalid_param_test/test_invalid_video_generation.py new file mode 100644 index 000000000..c0dd133c9 --- /dev/null +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_video_generation.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HTTP video job routes (live ``Wan-AI/Wan2.2-T2V-A14B-Diffusers``). + +WebSocket ``/v1/video/chat/stream`` and ``/v1/realtime`` invalid-input checks live in +``test_invalid_omni_chat.py`` with the same Qwen3-Omni server module fixture (Wan cannot serve those paths). +""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler + +pytestmark = [pytest.mark.slow, pytest.mark.diffusion] + +_WAN_T2V = [ + pytest.param( + OmniServerParams(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers"), + id="wan22_t2v", + marks=hardware_marks(res={"cuda": "H100"}), + ), +] + + +def _minimal_video_form_data() -> dict[str, str]: + return {"prompt": "a cat walking on grass"} + + +# ───────────────────────────────────────────────────────────────────────────── +# POST /v1/videos (async) +# ───────────────────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "form_extra, err_message", + [ + pytest.param(None, ("prompt", "Field required", "Missing"), id="missing_prompt"), + pytest.param({"seconds": "0"}, ("seconds", "string_pattern_mismatch", "1"), id="seconds_zero"), + pytest.param({"size": "bad"}, ("size", "string_pattern_mismatch"), id="size_bad"), + pytest.param( + {"num_inference_steps": "0"}, + ("num_inference_steps", "greater_than_equal", "1"), + id="num_inference_steps_zero", + ), + pytest.param( + {"guidance_scale": "25"}, ("guidance_scale", "less_than_equal", "20"), id="guidance_scale_above_max" + ), + ], +) +@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True) +def test_post_videos_invalid_requests( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, + form_extra: dict[str, str] | None, + err_message: str | tuple[str, ...], +) -> None: + data: dict[str, str] = {} if form_extra is None else {**_minimal_video_form_data(), **form_extra} + openai_client.send_videos_create_http_request( + {"data": data, "timeout": 120, "err_code": 400, "err_message": err_message} + ) + + +# ───────────────────────────────────────────────────────────────────────────── +# POST /v1/videos/sync +# ───────────────────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "form_extra, err_message", + [ + pytest.param(None, ("prompt", "Field required", "Missing"), id="missing_prompt"), + pytest.param({"seconds": "0"}, ("seconds", "string_pattern_mismatch", "1"), id="seconds_zero"), + pytest.param({"size": "bad"}, ("size", "string_pattern_mismatch"), id="size_bad"), + pytest.param( + {"num_inference_steps": "201"}, + ("num_inference_steps", "less_than_equal", "200"), + id="num_inference_steps_above_max", + ), + ], +) +@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True) +def test_post_videos_sync_invalid_requests( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, + form_extra: dict[str, str] | None, + err_message: str | tuple[str, ...], +) -> None: + data: dict[str, str] = {} if form_extra is None else {**_minimal_video_form_data(), **form_extra} + openai_client.send_videos_sync_http_request( + {"data": data, "timeout": 120, "err_code": 400, "err_message": err_message} + ) + + +# ───────────────────────────────────────────────────────────────────────────── +# GET /v1/videos (list) +# ───────────────────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "params, err_message", + [ + pytest.param({"limit": -1}, ("limit", "greater_than_equal", "0"), id="invalid_limit_negative"), + pytest.param({"limit": 101}, ("limit", "less_than_equal", "100"), id="limit_above_max"), + pytest.param({"order": "newest"}, ("order", "literal_error", "asc", "desc"), id="invalid_order"), + ], +) +@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True) +def test_get_videos_list_invalid_requests( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, + params: dict[str, Any], + err_message: str | tuple[str, ...], +) -> None: + openai_client.send_videos_list_http_request( + {"params": params, "timeout": 120, "err_code": 400, "err_message": err_message} + ) + + +# ───────────────────────────────────────────────────────────────────────────── +# GET /v1/videos/{video_id} +# ───────────────────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True) +def test_get_video_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + openai_client.send_video_retrieve_http_request( + { + "video_id": "does-not-exist", + "timeout": 120, + "err_code": 404, + "err_message": ("not found", "video"), + } + ) + + +# ───────────────────────────────────────────────────────────────────────────── +# DELETE /v1/videos/{video_id} +# ───────────────────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True) +def test_delete_video_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + openai_client.send_video_delete_http_request( + { + "video_id": "does-not-exist", + "timeout": 120, + "err_code": 404, + "err_message": ("not found", "video"), + } + ) + + +# ───────────────────────────────────────────────────────────────────────────── +# GET /v1/videos/{video_id}/content +# ───────────────────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True) +def test_get_video_content_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + openai_client.send_video_content_http_request( + { + "video_id": "does-not-exist", + "timeout": 120, + "err_code": 404, + "err_message": ("not found", "video"), + } + ) diff --git a/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16_expansion.py index 250f37b94..190857e3b 100644 --- a/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16_expansion.py +++ b/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16_expansion.py @@ -70,7 +70,7 @@ def test_text_to_text(omni_runner, omni_runner_handler): "modalities": ["text"], } response = omni_runner_handler.send_omni_request(request_config) - assert response.success, f"Request failed: {response.error_message}" + assert response.success, "Request failed" assert response.text_content and len(response.text_content.strip()) > 0 @@ -93,7 +93,7 @@ def test_audio_to_text(omni_runner, omni_runner_handler): "modalities": ["text"], } response = omni_runner_handler.send_omni_request(request_config) - assert response.success, f"Request failed: {response.error_message}" + assert response.success, "Request failed" assert response.text_content and len(response.text_content.strip()) > 0 @@ -114,7 +114,7 @@ def test_image_to_text(omni_runner, omni_runner_handler): "modalities": ["text"], } response = omni_runner_handler.send_omni_request(request_config) - assert response.success, f"Request failed: {response.error_message}" + assert response.success, "Request failed" assert response.text_content and len(response.text_content.strip()) > 0 @@ -135,7 +135,7 @@ def test_video_to_text(omni_runner, omni_runner_handler): "modalities": ["text"], } response = omni_runner_handler.send_omni_request(request_config) - assert response.success, f"Request failed: {response.error_message}" + assert response.success, "Request failed" assert response.text_content and len(response.text_content.strip()) > 0 @@ -162,4 +162,4 @@ def test_mix_to_audio(omni_runner, omni_runner_handler): "modalities": ["audio"], } response = omni_runner_handler.send_omni_request(request_config) - assert response.success, f"Request failed: {response.error_message}" + assert response.success, "Request failed" diff --git a/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16_expansion.py index 6cd630e77..9111740ea 100644 --- a/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16_expansion.py +++ b/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16_expansion.py @@ -87,7 +87,7 @@ def test_text_to_text(omni_runner, omni_runner_handler): "modalities": ["text"], } response = omni_runner_handler.send_request(request_config) - assert response.success, f"Request failed: {response.error_message}" + assert response.success, "Request failed" assert response.text_content and len(response.text_content.strip()) > 0 @@ -111,7 +111,7 @@ def test_audio_to_text(omni_runner, omni_runner_handler): "modalities": ["text"], } response = omni_runner_handler.send_request(request_config) - assert response.success, f"Request failed: {response.error_message}" + assert response.success, "Request failed" assert response.text_content and len(response.text_content.strip()) > 0 @@ -133,7 +133,7 @@ def test_image_to_text(omni_runner, omni_runner_handler): "modalities": ["text"], } response = omni_runner_handler.send_request(request_config) - assert response.success, f"Request failed: {response.error_message}" + assert response.success, "Request failed" assert response.text_content and len(response.text_content.strip()) > 0 @@ -155,7 +155,7 @@ def test_video_to_text(omni_runner, omni_runner_handler): "modalities": ["text"], } response = omni_runner_handler.send_request(request_config) - assert response.success, f"Request failed: {response.error_message}" + assert response.success, "Request failed" assert response.text_content and len(response.text_content.strip()) > 0 @@ -177,7 +177,7 @@ def test_video_to_audio(omni_runner, omni_runner_handler): "modalities": ["audio"], } response = omni_runner_handler.send_request(request_config) - assert response.success, f"Request failed: {response.error_message}" + assert response.success, "Request failed" # ------------------------------------------------------------------ @@ -204,4 +204,4 @@ def test_mix_to_audio(omni_runner, omni_runner_handler): "modalities": ["audio"], } response = omni_runner_handler.send_request(request_config) - assert response.success, f"Request failed: {response.error_message}" + assert response.success, "Request failed" diff --git a/tests/e2e/online_serving/test_omnivoice_expansion.py b/tests/e2e/online_serving/test_omnivoice_expansion.py index 7ad3b9b25..631c799cd 100644 --- a/tests/e2e/online_serving/test_omnivoice_expansion.py +++ b/tests/e2e/online_serving/test_omnivoice_expansion.py @@ -105,16 +105,3 @@ class TestOmniVoiceVoiceCloning: "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES, } openai_client.send_audio_speech_request(request_config) - - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_voice_clone_invalid_ref_audio_format(self, omni_server, openai_client) -> None: - """Test that invalid ref_audio format returns a clear error.""" - request_config = { - "model": omni_server.model, - "input": get_prompt("text"), - "ref_audio": "not_a_valid_uri", - "response_format": "wav", - "timeout": 180.0, - "status_code": (400, 422), - } - openai_client.send_audio_speech_request(request_config) diff --git a/tests/e2e/online_serving/test_voxtral_tts.py b/tests/e2e/online_serving/test_voxtral_tts.py index dff53d539..a1b5fdbb8 100644 --- a/tests/e2e/online_serving/test_voxtral_tts.py +++ b/tests/e2e/online_serving/test_voxtral_tts.py @@ -91,20 +91,3 @@ class TestVoxtralTTSFixedVoice: "timeout": 120.0, } ) - - @pytest.mark.advanced_model - @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) - def test_speech_invalid_voice_rejected(self, omni_server, openai_client) -> None: - """Request with a non-existent voice should be rejected (HTTP 200 with JSON or 4xx, depending on server).""" - openai_client.send_audio_speech_request( - { - "model": omni_server.model, - "input": "This voice does not exist.", - "voice": "nonexistent_voice_xyz", - "language": "English", - "response_format": "wav", - "timeout": 120.0, - "err_message": "Invalid speaker", - } - ) diff --git a/tests/entrypoints/openai_api/test_qwen3_tts_speaker_embedding.py b/tests/entrypoints/openai_api/test_qwen3_tts_speaker_embedding.py index 48e3670b4..1f71cf8fb 100644 --- a/tests/entrypoints/openai_api/test_qwen3_tts_speaker_embedding.py +++ b/tests/entrypoints/openai_api/test_qwen3_tts_speaker_embedding.py @@ -5,6 +5,10 @@ E2E tests for speaker_embedding API parameter. Tests Base models (0.6B-Base and 1.7B-Base) which support voice cloning via pre-computed ECAPA-TDNN embeddings, bypassing ref_audio extraction. + +Invalid combinations (mutually exclusive with ref_audio, empty embedding, +wrong embedding dimensions) live in +``tests/dfx/reliability/http_invalid/test_invalid_audio_speech.py``. """ import os @@ -147,45 +151,6 @@ class TestSpeakerEmbeddingBase: assert len(response.content) > MIN_AUDIO_BYTES assert_not_silence(response.content) - @pytest.mark.advanced_model - @pytest.mark.tts - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_speaker_embedding_mutually_exclusive_with_ref_audio(self, omni_server) -> None: - """Sending both speaker_embedding and ref_audio returns 400.""" - url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech" - payload = { - "model": omni_server.model, - "input": SYN_TEXT, - "task_type": "Base", - "speaker_embedding": DUMMY_EMBEDDING_1024, - "ref_audio": "https://example.com/audio.wav", - "response_format": "wav", - } - with httpx.Client(timeout=30.0) as client: - response = client.post(url, json=payload) - - assert response.status_code == 400, f"Expected 400, got {response.status_code}: {response.text}" - assert "mutually exclusive" in response.text - - @pytest.mark.advanced_model - @pytest.mark.tts - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_speaker_embedding_empty_rejected(self, omni_server) -> None: - """Empty speaker_embedding list returns 400.""" - url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech" - payload = { - "model": omni_server.model, - "input": SYN_TEXT, - "task_type": "Base", - "speaker_embedding": [], - "response_format": "wav", - } - with httpx.Client(timeout=30.0) as client: - response = client.post(url, json=payload) - - assert response.status_code == 400, f"Expected 400, got {response.status_code}: {response.text}" - assert "non-empty" in response.text - # ── 1.7B-Base model tests (2048-dim embeddings) ── @@ -216,26 +181,3 @@ class TestSpeakerEmbedding1_7B: assert response.headers.get("content-type") == "audio/wav" assert verify_wav_audio(response.content), "Response is not valid WAV" assert len(response.content) > MIN_AUDIO_BYTES, f"Audio too small: {len(response.content)} bytes" - - @pytest.mark.advanced_model - @pytest.mark.tts - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_1024_dim_on_1_7b_model_rejected(self, omni_server) -> None: - """1024-dim speaker_embedding on a 1.7B model (expects 2048) should fail gracefully.""" - url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech" - payload = { - "model": omni_server.model, - "input": SYN_TEXT, - "task_type": "Base", - "speaker_embedding": DUMMY_EMBEDDING_1024, - "x_vector_only_mode": True, - "response_format": "wav", - "max_new_tokens": MAX_NEW_TOKENS, - } - with httpx.Client(timeout=120.0) as client: - response = client.post(url, json=payload) - - # Wrong dimensions should produce an error, not silent garbage. - assert response.status_code == 400, f"Expected 400, got {response.status_code}: {response.text}" - assert "speaker_embedding has 1024 dimensions" in response.text - assert "expected 2048" in response.text diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 72767ec02..2440fe586 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -1,5 +1,6 @@ # tests/entrypoints/openai/test_serving_speech.py import asyncio +import json import logging import os import struct @@ -564,7 +565,6 @@ class TestSpeechAPI: def test_upload_voice_embedding_success(self, client): """Upload a voice via speaker_embedding JSON.""" - import json emb = [0.1] * 1024 data = { @@ -585,7 +585,6 @@ class TestSpeechAPI: def test_upload_voice_embedding_appears_in_listing(self, client): """Embedding-uploaded voice appears in list with correct source.""" - import json emb = [0.2] * 2048 data = { @@ -606,7 +605,6 @@ class TestSpeechAPI: def test_upload_voice_embedding_and_audio_mutually_exclusive(self, client): """Providing both audio_sample and speaker_embedding returns 400.""" - import json emb = [0.1] * 1024 files = {"audio_sample": ("test.wav", b"fake", "audio/wav")} @@ -632,7 +630,6 @@ class TestSpeechAPI: def test_upload_voice_embedding_nan_rejected(self, client): """NaN values in speaker_embedding return 400.""" - import json data = { "speaker_embedding": json.dumps([0.1] * 1023 + [float("nan")]), @@ -930,7 +927,6 @@ class TestTTSMethods: def test_upload_voice_embedding_wrong_dims_rejected(self, speech_server): """Embedding uploads must match the loaded Qwen3-TTS model before being stored.""" - import json speech_server._tts_model_type = "qwen3_tts" speech_server.engine_client.model_config = SimpleNamespace( diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py index 65a7c1404..5898b624b 100644 --- a/tests/helpers/assertions.py +++ b/tests/helpers/assertions.py @@ -1,12 +1,13 @@ """Assertion and response validation helpers for tests.""" import io +import json import tempfile import threading import wave from io import BytesIO from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal if TYPE_CHECKING: from tests.helpers.runtime import DiffusionResponse @@ -495,62 +496,12 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve ) -def _speech_combined_error_text_for_assert(response: Any) -> str: - """Join ``error_message`` and UTF-8 body (e.g. HTTP 200 + JSON) for substring checks.""" - parts: list[str] = [] - em = getattr(response, "error_message", None) - if em: - parts.append(str(em)) - ab = getattr(response, "audio_bytes", None) - if ab: - parts.append(ab.decode("utf-8", errors="replace")) - return "\n".join(parts) - - -def _assert_speech_error_substrings(haystack: str, err_expected: Any) -> None: - if isinstance(err_expected, (list, tuple)): - assert any(str(s) in haystack for s in err_expected), ( - f"Expected error text to contain one of {err_expected!r}, got: {haystack!r}" - ) - else: - assert str(err_expected) in haystack, f"Expected error text to contain {str(err_expected)!r}, got: {haystack!r}" - - def assert_audio_speech_response(response: Any, request_config: dict[str, Any], run_level: str) -> None: - """Validate speech API results. + """Validate speech API results from :class:`~tests.helpers.runtime.OmniResponse`. - If ``status_code`` is set: assert HTTP status, then optional ``err_message`` (4xx: ``error_message`` only; - HTTP 200: ``error_message`` and/or body, e.g. JSON with HTTP 200). - - If only ``err_message`` is set: assert substring(s) on combined error text. Otherwise success/min_audio/etc. + Success path only. Use :meth:`OpenAIClientHandler.send_audio_speech_http_request` with ``err_code`` / + ``err_message`` for expected HTTP errors. """ - # When present, only compare HTTP status and skip all other checks (e.g. expected 4xx). - # ``status_code`` may be a single int or a collection of allowed codes (e.g. ``(400, 422)``). - expected_status = request_config.get("status_code") - err_expected = request_config.get("err_message") - if expected_status is not None: - actual = getattr(response, "status_code", None) - assert actual is not None, ( - "request_config has status_code but response has no .status_code " - "(set OmniResponse.status_code or pass a response object with status_code)" - ) - if isinstance(expected_status, (list, tuple, set, frozenset)): - assert actual in expected_status, f"Expected HTTP status in {expected_status!r}, got {actual}" - else: - assert actual == int(expected_status), f"Expected HTTP status {int(expected_status)}, got {actual}" - - if err_expected is not None: - if actual is not None and actual != 200: - msg = (getattr(response, "error_message", None) or "") + "" - _assert_speech_error_substrings(msg, err_expected) - else: - _assert_speech_error_substrings(_speech_combined_error_text_for_assert(response), err_expected) - return - - if err_expected is not None: - _assert_speech_error_substrings(_speech_combined_error_text_for_assert(response), err_expected) - return - assert response.success, "The request failed." # Optional floor on decoded audio size (models with very short clips may use a lower value). @@ -599,10 +550,128 @@ def assert_diffusion_response(response: "DiffusionResponse", request_config: dic assert_audio_diffusion_response(response=response, request_config=request_config, run_level=run_level) +def _http_response_body_materialize(resp: Any) -> tuple[bytes, dict[str, Any] | None]: + """Serialize ``HttpResponse``-like body to UTF-8 bytes and parse a JSON object when possible.""" + jb = getattr(resp, "json_body", None) + if jb is not None: + raw = json.dumps(jb, ensure_ascii=False).encode("utf-8") + if isinstance(jb, dict): + return raw, jb + try: + parsed = json.loads(raw.decode("utf-8", errors="replace")) + except json.JSONDecodeError: + return raw, None + return raw, parsed if isinstance(parsed, dict) else None + err = getattr(resp, "error_message", None) + raw = (err or "").encode("utf-8", errors="replace") + try: + parsed = json.loads(raw.decode("utf-8", errors="replace")) + except json.JSONDecodeError: + return raw, None + return raw, parsed if isinstance(parsed, dict) else None + + +def assert_err_message_in_text( + haystack: str, + err_message: str | tuple[str, ...] | list[str] | set[str] | frozenset[str], + *, + sequence_match: Literal["all", "any"] = "all", +) -> None: + """Assert ``err_message`` appears in ``haystack`` (case-insensitive). + + For non-string sequences: ``sequence_match='all'`` requires every substring (HTTP-style); + ``sequence_match='any'`` requires at least one (WebSocket first-frame JSON helpers). + """ + hl = haystack.lower() + if isinstance(err_message, (list, tuple, set, frozenset)): + if sequence_match == "all": + missing = [s for s in err_message if str(s).lower() not in hl] + assert not missing, ( + f"Expected error text to contain all of {err_message!r}; missing {missing!r}. haystack={haystack!r}" + ) + else: + assert any(str(s).lower() in hl for s in err_message), ( + f"Expected error text to contain one of {err_message!r}. haystack={haystack!r}" + ) + else: + assert str(err_message).lower() in hl, f"Expected error text to contain {str(err_message)!r}, got: {haystack!r}" + + +def assert_http_error( + resp: Any, + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + websocket_json_message: bool = False, +) -> dict[str, Any] | None: + """Validate a raw-HTTP :class:`~tests.helpers.runtime.HttpResponse`-like object. + + Used by :class:`~tests.helpers.runtime.OpenAIClientHandler` ``send_*_http_request`` helpers when + ``request_config`` contains optional ``err_code`` and/or ``err_message``. + + When ``websocket_json_message=True``, only ``json_body`` is checked (first JSON WebSocket text frame). + Tuple/list ``err_message`` then uses **any** substring match; HTTP mode still requires **all** pieces. + + - ``err_code``: exact HTTP ``int``, or membership if a non-string sequence (e.g. ``(400, 422)``). + When ``err_code`` is set and the actual status is a client error in ``400..499`` other than ``404``, + the JSON body must include FastAPI ``detail`` and/or OpenAI-style ``error`` (2xx skips this). + - ``err_message``: substring match (case-insensitive) against serialized ``json_body`` and ``error_message``. + If ``err_message`` is a non-string sequence (``list`` / ``tuple`` / ``set`` / ``frozenset``), **every** + element must appear as a substring; a plain ``str`` still requires that single substring. + """ + if websocket_json_message: + if err_code is not None: + raise ValueError("assert_http_error: err_code is incompatible with websocket_json_message=True") + jb_ws = getattr(resp, "json_body", None) + assert jb_ws is not None, resp + if err_message is None: + return jb_ws if isinstance(jb_ws, dict) else None + assert_err_message_in_text( + json.dumps(jb_ws, ensure_ascii=False), + err_message, + sequence_match="any", + ) + return jb_ws if isinstance(jb_ws, dict) else None + + if err_code is None and err_message is None: + return None + + actual = getattr(resp, "status_code", None) + assert actual is not None, "response missing status_code" + + if err_code is not None: + if isinstance(err_code, int): + assert actual == err_code, (resp, err_code) + else: + allowed = tuple(err_code) + assert actual in allowed, (resp, allowed) + + body_bytes, payload = _http_response_body_materialize(resp) + + if err_code is not None and actual is not None and 400 <= actual < 500 and actual != 404: + assert payload is not None, getattr(resp, "error_message", resp) + assert "detail" in payload or "error" in payload, payload + + if err_message is not None: + pieces: list[str] = [] + jb = getattr(resp, "json_body", None) + if jb is not None: + pieces.append(json.dumps(jb, ensure_ascii=False)) + em = getattr(resp, "error_message", None) + if em: + pieces.append(str(em)) + haystack = "\n".join(pieces) if pieces else body_bytes.decode("utf-8", errors="replace") + assert_err_message_in_text(haystack, err_message, sequence_match="all") + + return payload + + __all__ = [ "assert_audio_diffusion_response", "assert_audio_speech_response", "assert_diffusion_response", + "assert_err_message_in_text", + "assert_http_error", "assert_image_diffusion_response", "assert_image_valid", "assert_omni_response", diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 66ded055d..8b007dc16 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -1,5 +1,6 @@ """Server/client/runner runtime primitives for tests.""" +import asyncio import base64 import concurrent.futures import copy @@ -19,6 +20,7 @@ from dataclasses import asdict, dataclass from io import BytesIO from pathlib import Path from typing import Any, NamedTuple +from urllib.parse import quote import psutil import requests @@ -37,6 +39,7 @@ from vllm.logger import init_logger from tests.helpers.assertions import ( assert_audio_speech_response, assert_diffusion_response, + assert_http_error, assert_omni_response, ) from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup @@ -86,6 +89,16 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False): logger.warning("torch._C._host_emptyCache() only available in Pytorch >=2.5") +def _parse_response_json(r: requests.Response) -> dict[str, Any] | list[Any] | None: + try: + data = r.json() + if isinstance(data, (dict, list)): + return data + except Exception: + pass + return None + + def _split_request_config_by_per_output_sizes(cfg: dict[str, Any]) -> list[dict[str, Any]] | None: """If ``extra_body`` has list ``height``/``width``, return one config per index (scalar h/w, ``num_outputs_per_prompt=1``).""" eb = cfg.get("extra_body") @@ -647,6 +660,8 @@ class OmniServerStageCli(OmniServer): @dataclass class OmniResponse: + """Decoded multimodal / chat output from the OpenAI SDK or offline runner (not raw ``requests``).""" + text_content: str | None = None audio_data: list[str] | None = None audio_content: str | None = None @@ -656,16 +671,15 @@ class OmniResponse: #: OpenAI client call through response parsing and local post-process (e.g. audio decode). e2e_latency: float | None = None success: bool = False - error_message: str | None = None prompt_tokens: int | None = None cached_tokens: int | None = None - # Set for HTTP-level assertions (e.g. ``assert_audio_speech_response`` with ``status_code`` in config). - status_code: int | None = None logprobs: list | None = None @dataclass class DiffusionResponse: + """Decoded diffusion output from chat completions or offline runner (not raw ``requests``).""" + text_content: str | None = None images: list[Image.Image] | None = None audios: list[Any] | None = None @@ -674,7 +688,67 @@ class DiffusionResponse: #: ``chat.completions.create`` through local image / audio decode. e2e_latency: float | None = None success: bool = False + + +@dataclass +class HttpResponse: + """Normalized view of a ``requests`` response from :class:`OpenAIClientHandler` HTTP helpers.""" + + status_code: int + success: bool error_message: str | None = None + json_body: dict[str, Any] | list[Any] | None = None + + +@dataclass +class WebSocketJsonResponse: + """First JSON object delivered as a text WebSocket frame (streaming endpoints).""" + + json_body: dict[str, Any] + + +def _merge_http_expectation_kwargs( + base: dict[str, Any] | None, + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, +) -> dict[str, Any]: + cfg = dict(base or {}) + if err_code is not None: + cfg["err_code"] = err_code + if err_message is not None: + cfg["err_message"] = err_message + return cfg + + +def _merge_ws_expectation_kwargs( + base: dict[str, Any] | None, + *, + err_message: str | tuple[str, ...] | list[str] | None = None, + ws_json_type: str | None = None, + ws_error_code: str | None = None, +) -> dict[str, Any]: + cfg = dict(base or {}) + if err_message is not None: + cfg["err_message"] = err_message + if ws_json_type is not None: + cfg["ws_json_type"] = ws_json_type + if ws_error_code is not None: + cfg["ws_error_code"] = ws_error_code + return cfg + + +def _run_ws_expectations_from_request_config(cfg: dict[str, Any], resp: WebSocketJsonResponse) -> None: + jb = resp.json_body + want_type = cfg.get("ws_json_type") + if want_type is not None: + assert jb.get("type") == want_type, (jb, want_type) + want_code = cfg.get("ws_error_code") + if want_code is not None: + assert jb.get("code") == want_code, (jb, want_code) + err_message = cfg.get("err_message") + if err_message is not None: + assert_http_error(resp, err_message=err_message, websocket_json_message=True) def _merge_diffusion_responses(parts: list[DiffusionResponse]) -> DiffusionResponse: @@ -739,8 +813,8 @@ class OpenAIClientHandler: result.e2e_latency = time.perf_counter() - wall_start result.success = True except Exception as e: - result.error_message = f"Stream processing error: {str(e)}" - print(f"Error: {result.error_message}") + msg = f"Stream processing error: {str(e)}" + print(f"Error: {msg}") return result def _process_non_stream_omni_response(self, chat_completion, *, wall_start: float) -> OmniResponse: @@ -771,8 +845,8 @@ class OpenAIClientHandler: result.logprobs = chat_completion.choices[0].logprobs.content result.success = True except Exception as e: - result.error_message = f"Non-stream processing error: {str(e)}" - print(f"Error: {result.error_message}") + msg = f"Non-stream processing error: {str(e)}" + print(f"Error: {msg}") return result def _process_diffusion_response(self, chat_completion, *, wall_start: float) -> DiffusionResponse: @@ -811,11 +885,592 @@ class OpenAIClientHandler: result.e2e_latency = time.perf_counter() - wall_start result.success = True except Exception as e: - result.error_message = f"Diffusion response processing error: {str(e)}" - print(f"Error: {result.error_message}") + msg = f"Diffusion response processing error: {str(e)}" + print(f"Error: {msg}") return result + def _http_response_from_requests(self, r: requests.Response) -> HttpResponse: + payload = _parse_response_json(r) + ok = 200 <= r.status_code < 300 + return HttpResponse( + status_code=r.status_code, + success=ok, + error_message=None if ok else (r.text[:8000] if r.text else None), + json_body=payload, + ) + + def send_health_http_request( + self, + request_config: dict[str, Any] | None = None, + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """GET ``/health`` (raw ``requests``). + + ``request_config``: optional ``timeout`` plus optional ``err_code`` / ``err_message`` for + :func:`~tests.helpers.assertions.assert_http_error` (also as keyword-only args). + """ + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = requests.get(self._build_url("/health"), timeout=float(cfg.get("timeout", 120.0))) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_models_http_request( + self, + request_config: dict[str, Any] | None = None, + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """GET ``/v1/models``. Optional ``timeout`` and HTTP assertions (see :func:`~tests.helpers.assertions.assert_http_error`).""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = requests.get( + self._build_url("/v1/models"), + headers={"Accept": "application/json"}, + timeout=float(cfg.get("timeout", 120.0)), + ) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_chat_completions_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/chat/completions`` with ``json`` or ``raw_body`` (malformed-body / contract tests).""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_json_endpoint("/v1/chat/completions", cfg, default_timeout=120.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_omni_sleep_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/omni/sleep`` — ``json`` or ``raw_body``, ``timeout``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_json_endpoint("/v1/omni/sleep", cfg, default_timeout=120.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_omni_wakeup_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/omni/wakeup``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_json_endpoint("/v1/omni/wakeup", cfg, default_timeout=120.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_audio_voices_list_http_request( + self, + request_config: dict[str, Any] | None = None, + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """GET ``/v1/audio/voices``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = requests.get( + self._build_url("/v1/audio/voices"), + headers={"Accept": "application/json"}, + timeout=float(cfg.get("timeout", 120.0)), + ) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_audio_voices_create_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/audio/voices`` (multipart): ``data`` / ``files`` / ``timeout``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_form_endpoint("/v1/audio/voices", cfg, default_timeout=120.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_audio_voices_delete_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """DELETE ``/v1/audio/voices/{name}`` — requires ``name``, optional ``timeout``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + name = cfg["name"] + timeout = float(cfg.get("timeout", 120.0)) + path = f"/v1/audio/voices/{quote(str(name), safe='')}" + r = requests.delete( + self._build_url(path), + headers={"Accept": "application/json"}, + timeout=timeout, + ) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_audio_speech_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/audio/speech`` with ``json`` or ``raw_body``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_json_endpoint("/v1/audio/speech", cfg, default_timeout=120.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_audio_speech_batch_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/audio/speech/batch``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_json_endpoint("/v1/audio/speech/batch", cfg, default_timeout=120.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_audio_generate_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/audio/generate``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_json_endpoint("/v1/audio/generate", cfg, default_timeout=120.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_images_generations_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/images/generations`` — ``json`` or ``raw_body``, ``timeout``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_json_endpoint("/v1/images/generations", cfg, default_timeout=300.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_images_edits_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/images/edits`` — ``data`` / ``files`` / ``timeout``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_form_endpoint("/v1/images/edits", cfg, default_timeout=300.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_videos_create_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/videos`` (async job) — multipart ``data`` / ``files``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_form_endpoint("/v1/videos", cfg, default_timeout=120.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_videos_sync_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/videos/sync``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_form_endpoint("/v1/videos/sync", cfg, default_timeout=120.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_videos_list_http_request( + self, + request_config: dict[str, Any] | None = None, + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """GET ``/v1/videos`` — optional ``params``, ``timeout``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = requests.get( + self._build_url("/v1/videos"), + params=cfg.get("params"), + headers={"Accept": "application/json"}, + timeout=float(cfg.get("timeout", 120.0)), + ) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_video_retrieve_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """GET ``/v1/videos/{video_id}``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + video_id = cfg["video_id"] + timeout = float(cfg.get("timeout", 120.0)) + r = requests.get( + self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}"), + headers={"Accept": "application/json"}, + timeout=timeout, + ) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_video_delete_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """DELETE ``/v1/videos/{video_id}``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + video_id = cfg["video_id"] + timeout = float(cfg.get("timeout", 120.0)) + r = requests.delete( + self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}"), + headers={"Accept": "application/json"}, + timeout=timeout, + ) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_video_content_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """GET ``/v1/videos/{video_id}/content``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + video_id = cfg["video_id"] + timeout = float(cfg.get("timeout", 120.0)) + r = requests.get( + self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}/content"), + timeout=timeout, + ) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def _build_ws_url(self, path: str) -> str: + """Turn HTTP ``base_url`` into ``ws`` / ``wss`` for WebSocket helpers.""" + base = self.base_url.rstrip("/") + suffix = "/" + path.lstrip("/") + if base.startswith("http://"): + return "ws://" + base.removeprefix("http://") + suffix + if base.startswith("https://"): + return "wss://" + base.removeprefix("https://") + suffix + raise ValueError(f"Unsupported base_url for WebSocket: {base!r}") + + def _send_websocket_first_json_request( + self, + path: str, + cfg: dict[str, Any], + ) -> list[WebSocketJsonResponse]: + """Connect, optionally send text frames, return first JSON text frame as :class:`WebSocketJsonResponse`. + + ``request_config`` keys: + + - ``send_frames``: optional ``str`` or sequence of ``str`` raw WebSocket text frames (omit when the server + speaks first, e.g. ``/v1/realtime`` rejection path). + - ``timeout``: seconds to wait for the first inbound text frame (default ``120``). + - ``ws_max_size``: passed through as ``max_size`` to :func:`websockets.connect` when the key is present. + """ + send_frames_raw = cfg.get("send_frames") + if send_frames_raw is None: + frames: list[str] = [] + elif isinstance(send_frames_raw, str): + frames = [send_frames_raw] + else: + frames = list(send_frames_raw) + + timeout = float(cfg.get("timeout", 120.0)) + uri = self._build_ws_url(path) + + connect_kw: dict[str, Any] = {} + if "ws_max_size" in cfg: + connect_kw["max_size"] = cfg["ws_max_size"] + + async def _recv_first_json_object() -> WebSocketJsonResponse: + import websockets + + async with websockets.connect(uri, **connect_kw) as ws: + for frame in frames: + await ws.send(frame) + raw = await asyncio.wait_for(ws.recv(), timeout=timeout) + if not isinstance(raw, str): + raise AssertionError(f"Expected JSON text frame from {uri}, got {type(raw).__name__}") + try: + data = json.loads(raw) + except json.JSONDecodeError as exc: + raise AssertionError(f"Expected JSON text frame from {uri}, body={raw[:500]!r}") from exc + if not isinstance(data, dict): + raise AssertionError(f"Expected JSON object from {uri}, got {type(data).__name__}") + return WebSocketJsonResponse(json_body=data) + + resp = asyncio.run(_recv_first_json_object()) + _run_ws_expectations_from_request_config(cfg, resp) + return [resp] + + def send_audio_speech_stream_ws_request( + self, + request_config: dict[str, Any], + *, + err_message: str | tuple[str, ...] | list[str] | None = None, + ws_json_type: str | None = None, + ws_error_code: str | None = None, + ) -> list[WebSocketJsonResponse]: + """WebSocket ``/v1/audio/speech/stream`` — send ``send_frames`` then read first JSON text frame.""" + cfg = _merge_ws_expectation_kwargs( + request_config, + err_message=err_message, + ws_json_type=ws_json_type, + ws_error_code=ws_error_code, + ) + return self._send_websocket_first_json_request("/v1/audio/speech/stream", cfg) + + def send_video_chat_stream_ws_request( + self, + request_config: dict[str, Any], + *, + err_message: str | tuple[str, ...] | list[str] | None = None, + ws_json_type: str | None = None, + ws_error_code: str | None = None, + ) -> list[WebSocketJsonResponse]: + """WebSocket ``/v1/video/chat/stream`` — send ``send_frames`` then read first JSON text frame.""" + cfg = _merge_ws_expectation_kwargs( + request_config, + err_message=err_message, + ws_json_type=ws_json_type, + ws_error_code=ws_error_code, + ) + return self._send_websocket_first_json_request("/v1/video/chat/stream", cfg) + + def send_realtime_ws_request( + self, + request_config: dict[str, Any] | None = None, + *, + err_message: str | tuple[str, ...] | list[str] | None = None, + ws_json_type: str | None = None, + ws_error_code: str | None = None, + ) -> list[WebSocketJsonResponse]: + """WebSocket ``/v1/realtime`` — optional outbound frames, then first JSON text frame (often server-initiated).""" + cfg = _merge_ws_expectation_kwargs( + request_config, + err_message=err_message, + ws_json_type=ws_json_type, + ws_error_code=ws_error_code, + ) + return self._send_websocket_first_json_request("/v1/realtime", cfg) + def send_omni_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]: + """Chat completions via the OpenAI Python SDK (not raw HTTP).""" responses: list[OmniResponse] = [] stream = request_config.get("stream", False) modalities = request_config.get("modalities", ["text", "audio"]) @@ -939,8 +1594,8 @@ class OpenAIClientHandler: result.audio_format = result.audio_format.headers.get("content-type", "") except Exception as e: - result.error_message = f"Audio speech stream processing error: {str(e)}" - print(f"Error: {result.error_message}") + msg = f"Audio speech stream processing error: {str(e)}" + print(f"Error: {msg}") return result @@ -978,8 +1633,8 @@ class OpenAIClientHandler: result.audio_format = result.audio_format.headers.get("content-type", "") except Exception as e: - result.error_message = f"Audio speech non-stream processing error: {str(e)}" - print(f"Error: {result.error_message}") + msg = f"Audio speech non-stream processing error: {str(e)}" + print(f"Error: {msg}") return result @@ -1201,6 +1856,7 @@ class OpenAIClientHandler: Send OpenAI requests for diffusion models. If ``extra_body`` has list ``height``/``width``, sends one chat completion per index in parallel (scalar h/w, ``num_outputs_per_prompt=1`` each) and merges images in list order. + Args: request_config: A single request configuration dict, or a list of request configuration dicts (one request per element) @@ -1300,10 +1956,13 @@ class OpenAIClientHandler: self, request_config: dict[str, Any], request_num: int = 1 ) -> list[DiffusionResponse]: """ - Send native /v1/videos requests. + Send native /v1/videos requests: multipart ``form_data`` job create, poll until done, download content. + + For raw HTTP to video routes without polling, use ``send_videos_create_http_request``, etc. """ if request_num != 1: raise NotImplementedError("Concurrent video diffusion requests are not currently implemented") + form_data = request_config.get("form_data") if not isinstance(form_data, dict): raise ValueError("Video request_config must contain 'form_data'") @@ -1346,6 +2005,54 @@ class OpenAIClientHandler: self._print_client_stat("[diffusion] request#1 completed") return [result] + def _post_json_endpoint( + self, + path: str, + request_config: dict[str, Any], + *, + default_timeout: float, + ) -> requests.Response: + url = self._build_url(path) + timeout = float(request_config.get("timeout", default_timeout)) + if "raw_body" in request_config: + raw = request_config["raw_body"] + payload = raw.encode("utf-8") if isinstance(raw, str) else raw + return requests.post( + url, + data=payload, + headers={"Content-Type": "application/json", "Accept": "application/json"}, + timeout=timeout, + ) + if "json" not in request_config: + raise ValueError(f"{path} request_config must include 'json' or 'raw_body'") + return requests.post( + url, + json=request_config["json"], + headers={"Accept": "application/json", "Content-Type": "application/json"}, + timeout=timeout, + ) + + def _post_form_endpoint( + self, + path: str, + request_config: dict[str, Any], + *, + default_timeout: float = 120.0, + ) -> requests.Response: + url = self._build_url(path) + timeout = float(request_config.get("timeout", default_timeout)) + data = request_config.get("data") + files = request_config.get("files") + if data is None and not files: + data = {} + return requests.post( + url, + data=data, + files=files, + headers={"Accept": "application/json"} if not files else {"Accept": "application/json"}, + timeout=timeout, + ) + def _wait_until_video_completed( self, video_id: str, poll_interval_seconds: int = 2, timeout_seconds: int = 300 ) -> None: @@ -1693,9 +2400,9 @@ class OmniRunnerHandler: result.text_content = text_content result.success = True except Exception as e: - result.error_message = f"Output processing error: {str(e)}" + msg = f"Output processing error: {str(e)}" result.success = False - print(f"Error: {result.error_message}") + print(f"Error: {msg}") return result def _process_diffusion_output(self, outputs: list[OmniRequestOutput]) -> DiffusionResponse: @@ -1831,15 +2538,11 @@ class OmniRunnerHandler: mm_out = stage_out.request_output.outputs[0].multimodal_output break if mm_out is None: - result = OmniResponse(success=False, error_message="No audio output from pipeline") - assert result.success, result.error_message - return result + raise AssertionError("No audio output from pipeline") audio_data = mm_out.get("audio") if audio_data is None: - result = OmniResponse(success=False, error_message="No audio tensor in multimodal output") - assert result.success, result.error_message - return result + raise AssertionError("No audio tensor in multimodal output") sr_raw = mm_out.get("sr") sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw @@ -2004,6 +2707,8 @@ def iter_omni_runner( __all__ = [ "DiffusionResponse", + "HttpResponse", + "WebSocketJsonResponse", "OmniResponse", "OmniRunner", "OmniRunnerHandler", From 2af2a50e0e2981ec2eef32e704f5a66c3d451c95 Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Tue, 19 May 2026 15:22:02 +0800 Subject: [PATCH 12/38] [CI] improve Buildkite testcase statistics reports (#3543) Signed-off-by: wangyu <410167048@qq.com> --- .../nightly/buildkite_testcase_statistics.py | 424 ++++++++++++++---- 1 file changed, 337 insertions(+), 87 deletions(-) diff --git a/tools/nightly/buildkite_testcase_statistics.py b/tools/nightly/buildkite_testcase_statistics.py index 5f3217981..7a4f9bff3 100644 --- a/tools/nightly/buildkite_testcase_statistics.py +++ b/tools/nightly/buildkite_testcase_statistics.py @@ -1,10 +1,17 @@ #!/usr/bin/env python3 """ -Parse pytest commands from Buildkite test-ready.yml, test-merge.yml, test-nightly.yml; -collect test cases (including parametrized) via pytest --collect-only -q and produce an HTML report. +Parse pytest commands from Buildkite test-ready.yml, test-merge.yml, test-nightly.yml, +and test-weekly.yml; collect test cases (including parametrized) via pytest --collect-only -q +and produce an HTML report. + +Leaf steps may live under ``group:`` blocks (nested ``steps``); those are flattened with the +group title carried into the report. Usage (run from repo root): - python scripts/buildkite_testcase_statistics.py -o buildkite_testcase_statistics.html + python tools/nightly/buildkite_testcase_statistics.py -o buildkite_testcase_statistics.html + python tools/nightly/buildkite_testcase_statistics.py --sample-preview -o buildkite_testcase_statistics_preview.html + +``--sample-preview`` emits placeholder rows only (no pytest / no real counts). Requires: PyYAML (pip install pyyaml) """ @@ -24,7 +31,7 @@ import yaml # Repo root (parent of the directory containing this script) REPO_ROOT = Path(__file__).resolve().parent.parent.parent BUILDKITE_DIR = REPO_ROOT / ".buildkite" -PIPELINE_FILES = ["test-ready.yml", "test-merge.yml", "test-nightly.yml"] +PIPELINE_FILES = ["test-ready.yml", "test-merge.yml", "test-nightly.yml", "test-weekly.yml"] def load_yaml(path: Path) -> dict: @@ -32,6 +39,30 @@ def load_yaml(path: Path) -> dict: return yaml.safe_load(f) +def iter_pytest_leaf_steps(steps: list, *, group: str | None = None) -> list[tuple[dict, str | None]]: + """ + Flatten Buildkite YAML steps, expanding ``group`` blocks that contain nested ``steps``. + + Returns ``(leaf_step_dict, group_title)`` for each leaf step that has non-empty ``commands``. + ``group_title`` is ``None`` for top-level steps outside any ``group:`` block. + """ + out: list[tuple[dict, str | None]] = [] + for raw in steps or []: + if not isinstance(raw, dict): + continue + nested = raw.get("steps") + if isinstance(nested, list) and nested: + next_group = group + if "group" in raw: + g = raw.get("group") + next_group = g if isinstance(g, str) else (str(g) if g is not None else group) + out.extend(iter_pytest_leaf_steps(nested, group=next_group)) + continue + if raw.get("commands"): + out.append((raw, group)) + return out + + PYTEST_CMD_RE = re.compile(r"(?:timeout\s+\S+\s+)?(?:python3? -m\s+)?pytest\s+[^\n&|;]*") @@ -55,10 +86,10 @@ def normalize_commands(step: dict) -> list[str]: def extract_pytest_targets_from_line(line: str) -> list[str]: """ - Parse pytest test targets (file path or marker expression) from a single command line. - Returns e.g. ['tests/xxx.py']. - When there is no explicit test file but only -m/--run-level, returns - ['-m expr'] or ['-m expr --run-level level'] with the full marker args. + Parse pytest test targets (file path, directory under tests/, or marker expression) from a single command line. + Returns e.g. ['tests/xxx.py'] or ['tests/e2e']. + When there is no explicit path but markers/options apply to the whole tree, returns + a synthetic token built from sidecar args (``--ignore``, ``-m``, ``-k``, ``--run-level``). """ if "pytest" not in line: return [] @@ -74,17 +105,33 @@ def extract_pytest_targets_from_line(line: str) -> list[str]: args = parts[1].strip() if len(parts) > 1 else "" if not args: return [] - targets = [] + targets: list[str] = [] + # Omit paths that only appear as --ignore / --test-config-file values when discovering roots. + sanitized_for_roots = _strip_test_config_file_clauses_from_args_fragment( + _strip_ignore_clauses_from_args_fragment(args) + ) # Match tests/.../*.py - for m in re.finditer(r"tests/[^\s'\"#]+\.py", args): + for m in re.finditer(r"tests/[^\s'\"#]+\.py", sanitized_for_roots): path = m.group(0) if path not in targets: targets.append(path) + # Directory / namespace roots: tests/e2e, tests/e2e/, tests/foo/bar (no .py suffix) + for m in re.finditer(r"\s(tests/[a-zA-Z0-9_./\-]+)(?=\s|$|[\"'])", " " + sanitized_for_roots): + tok = m.group(1).rstrip("/") + if tok.endswith(".py"): + continue + if tok not in targets: + targets.append(tok) + for m in re.finditer(r"['\"](tests/[a-zA-Z0-9_./\-]+)['\"]", sanitized_for_roots): + tok = m.group(1).rstrip("/") + if tok.endswith(".py"): + continue + if tok not in targets: + targets.append(tok) if not targets: - # May be -m/--run-level only; record full marker arguments instead of a generic marker. - extra = _parse_extra_args_from_line(line) + # Marker-only / option-driven selection relative to tests/ + extra = _pytest_collect_sidecar_args(line) if extra: - # Join into a single string, e.g. "-m core and cpu --run-level nightly" targets.append(" ".join(extra)) return targets @@ -93,40 +140,77 @@ def get_pytest_targets_from_step(step: dict) -> list[tuple[str, str]]: """ Extract all pytest-related targets from a step. Returns [(target, raw_line), ...] where target is tests/xxx.py or '(marker/run-level)'. + + Deduplication uses ``(target, raw_line)`` so the same script path with different options + (e.g. ``--test-config-file …``) stays distinct across Buildkite jobs. """ lines = normalize_commands(step) out = [] seen = set() for line in lines: for t in extract_pytest_targets_from_line(line): - if t not in seen: - seen.add(t) + key = (t, line) + if key not in seen: + seen.add(key) out.append((t, line)) return out -def _parse_extra_args_from_line(raw_line: str) -> list[str]: - """Extract -m and --run-level args from a raw pytest command line. Strips extra spaces and quotes. - Correctly handles quoted expressions with spaces, e.g. -m 'core_model and cpu'. +def _strip_ignore_clauses_from_args_fragment(fragment: str) -> str: + """Replace ``--ignore=…`` / ``--ignore …`` spans so path discovery does not treat ignores as roots.""" + return re.sub(r"--ignore(?:=\S+|\s+\S+)", " ", fragment) + + +def _strip_test_config_file_clauses_from_args_fragment(fragment: str) -> str: + """Replace ``--test-config-file`` spans so JSON paths under ``tests/`` are not mistaken for collection roots.""" + return re.sub(r"--test-config-file(?:=\S+|\s+\S+)", " ", fragment) + + +def _pytest_collect_sidecar_args(raw_line: str) -> list[str]: + """Extract pytest args that affect collection: ``--ignore``, ``--test-config-file``, ``-m``, + ``--run-level``, ``-k``. + + Order matches typical CI lines (ignores / perf driver config, then markers / run-level / keyword). + Correctly handles quoted ``-m`` / ``-k`` / ``--run-level`` values with spaces. """ extra: list[str] = [] - # -m: match -m followed by '...', "..." or a single token (so -m expression is one argv) - m = re.search(r"-m\s+(?:'([^']*)'|\"([^\"]*)\"|(\S+))", raw_line) - if m: - value = (m.group(1) or m.group(2) or m.group(3) or "").strip() + for m in re.finditer(r"--ignore(?:=(\S+)|\s+(\S+))", raw_line): + path = (m.group(1) or m.group(2) or "").strip() + if path: + extra.extend(["--ignore", path]) + + for m in re.finditer(r"--test-config-file(?:=(\S+)|\s+(\S+))", raw_line): + path = (m.group(1) or m.group(2) or "").strip() + if path: + extra.extend(["--test-config-file", path]) + + marker = re.search(r"-m\s+(?:'([^']*)'|\"([^\"]*)\"|(\S+))", raw_line) + if marker: + value = (marker.group(1) or marker.group(2) or marker.group(3) or "").strip() if value: extra.extend(["-m", value]) - r = re.search(r"--run-level\s+(?:'([^']*)'|\"([^\"]*)\"|(\S+))", raw_line) - if r: - value = (r.group(1) or r.group(2) or r.group(3) or "").strip() + runlvl = re.search(r"--run-level\s+(?:'([^']*)'|\"([^\"]*)\"|(\S+))", raw_line) + if runlvl: + value = (runlvl.group(1) or runlvl.group(2) or runlvl.group(3) or "").strip() if value: extra.extend(["--run-level", value]) + kw = re.search(r"-k\s+(?:'([^']*)'|\"([^\"]*)\"|(\S+))", raw_line) + if kw: + value = (kw.group(1) or kw.group(2) or kw.group(3) or "").strip() + if value: + extra.extend(["-k", value]) + return extra +def _parse_extra_args_from_line(raw_line: str) -> list[str]: + """Backward-compatible alias for :func:`_pytest_collect_sidecar_args`.""" + return _pytest_collect_sidecar_args(raw_line) + + def _parse_collect_only_stdout(stdout: str, *, raise_on_empty: bool = True, stderr: str = "") -> list[str]: """ Parse pytest --collect-only -q stdout into a list of node ids. @@ -167,8 +251,15 @@ def _resolve_pytest_target(target: str, repo_root: Path, extra: list[str], raw_l if not path.exists(): raise FileNotFoundError(f"Pytest target path does not exist: {path}") return [str(path)], target.replace("\\", "/"), 60 + + if " " not in target and not target.startswith("-") and target.startswith("tests/"): + rel_dir = Path(target.rstrip("/")) + dir_path = (repo_root / rel_dir).resolve() + if dir_path.exists() and dir_path.is_dir(): + return [str(dir_path)], target.replace("\\", "/").rstrip("/"), 120 + if not extra: - raise RuntimeError(f"Failed to parse -m/--run-level from pytest line: {raw_line!r}") + raise RuntimeError(f"Failed to parse pytest sidecar args from line: {raw_line!r}") if not (repo_root / "tests").exists(): raise FileNotFoundError("tests/ directory not found under repo root") return ["tests/"], "tests/", 120 @@ -177,9 +268,13 @@ def _resolve_pytest_target(target: str, repo_root: Path, extra: list[str], raw_l def collect_test_names(target: str, repo_root: Path, raw_line: str) -> list[str]: """ Run pytest --collect-only -q for a target and return test node id list (incl. parametrized). - target: path like tests/foo.py or marker string (e.g. "-m expr"). raw_line used for -m/--run-level. + + ``path_args`` come from ``target`` (file, dir under ``tests/``, or fallback ``tests/``). Sidecar + flags ``--ignore``, ``--test-config-file``, ``-m``, ``-k``, ``--run-level`` are taken from the + full ``raw_line`` and appended to the same invocation so collection matches CI (perf scripts + parametrize from ``--test-config-file`` at import time). """ - extra = _parse_extra_args_from_line(raw_line) + extra = _pytest_collect_sidecar_args(raw_line) path_args, _fallback, timeout_quiet = _resolve_pytest_target(target, repo_root, extra, raw_line) cmd_quiet = [sys.executable, "-m", "pytest", *path_args, "--collect-only", "-q"] @@ -210,6 +305,15 @@ def _extra_with_skip_marker(extra: list[str]) -> list[str]: elif extra[i] == "--run-level" and i + 1 < len(extra): out.extend(["--run-level", extra[i + 1]]) i += 2 + elif extra[i] == "-k" and i + 1 < len(extra): + out.extend(["-k", extra[i + 1]]) + i += 2 + elif extra[i] == "--ignore" and i + 1 < len(extra): + out.extend(["--ignore", extra[i + 1]]) + i += 2 + elif extra[i] == "--test-config-file" and i + 1 < len(extra): + out.extend(["--test-config-file", extra[i + 1]]) + i += 2 else: i += 1 if not has_m: @@ -222,7 +326,7 @@ def get_skip_status(target: str, raw_line: str, repo_root: Path) -> set[str]: Use pytest -m skip --collect-only -q to collect tests marked with skip; return their node ids. Does not run tests; only collects. Supports both pytest xxx.py and pytest -m "marker". """ - extra = _parse_extra_args_from_line(raw_line) + extra = _pytest_collect_sidecar_args(raw_line) try: path_args, _, timeout_quiet = _resolve_pytest_target(target, repo_root, extra, raw_line) except (FileNotFoundError, RuntimeError): @@ -308,6 +412,7 @@ def write_html(all_stats: list[tuple], out_path: Path, total: int, repo_root: Pa "test-ready": "ready", "test-merge": "merge", "test-nightly": "nightly", + "test-weekly": "weekly", } def esc(s: str) -> str: @@ -318,7 +423,7 @@ def write_html(all_stats: list[tuple], out_path: Path, total: int, repo_root: Pa pipeline_totals: dict[str, int] = {} for item in all_stats: - _, pipeline_name, _, _, _, _, count, _, _, _ = item + _, pipeline_name, _, _, _, _, _, count, _, _, _ = item pipeline_totals[pipeline_name] = pipeline_totals.get(pipeline_name, 0) + count pipeline_summary_rows = [] for idx, (pipeline_name, pipeline_count) in enumerate(sorted(pipeline_totals.items())): @@ -331,7 +436,8 @@ def write_html(all_stats: list[tuple], out_path: Path, total: int, repo_root: Pa # Summary rows by module (item: ..., names, docstrings, skipped_ids) summary_rows = [] for i, item in enumerate(all_stats): - _, pipeline_name, label, test_files, _, marker_targets, count, _, _, _ = item + _, pipeline_name, group, label, test_files, _, marker_targets, count, _, _, _ = item + group_cell = esc(group) if group else "—" if test_files: files_html = "
".join(esc(f) for f in test_files) elif marker_targets: @@ -343,14 +449,15 @@ def write_html(all_stats: list[tuple], out_path: Path, total: int, repo_root: Pa badge_span = f'{esc(pipeline_name)}' if badge else esc(pipeline_name) summary_rows.append( f'' - f'{badge_span}{esc(label)}{count}' + f'{badge_span}{group_cell}{esc(label)}{count}' f'{files_html}' ) - # Detail rows grouped by Pipeline, then by Test Suite (two-level collapsible) + # Detail rows grouped by Pipeline -> Buildkite group -> Test Suite -> files case_details_thead = """ Pipeline + Group Test Suite Test File No. @@ -358,18 +465,19 @@ def write_html(all_stats: list[tuple], out_path: Path, total: int, repo_root: Pa Status Description """ - # Per-item: (pipeline_name, label, count, badge_span, suite_rows_html) - pipeline_suites: dict[str, list[tuple[str, int, str, str]]] = {} # pipeline -> [(label, count, summary, body)] + # pipeline -> group_title -> [(label, count, suite_summary, body)] + pipeline_group_suites: dict[str, dict[str | None, list[tuple[str, int, str, str]]]] = {} for item in all_stats: - _, pipeline_name, label, _, _, _, count, names, docstrings, skipped_ids = item + _, pipeline_name, group, label, _, _, _, count, names, docstrings, skipped_ids = item + group_cell = esc(group) if group else "—" badge = pipeline_badges.get(pipeline_name, "") badge_span = f'{esc(pipeline_name)}' if badge else esc(pipeline_name) suite_summary = f"{esc(label)} ({count} cases)" file_sections: list[str] = [] if not names: suite_rows = [ - f'' - f"{badge_span}{esc(label)}" + f'' + f"{badge_span}{group_cell}{esc(label)}" f'(marker-only / no collection)0' f'No collection or only -m/--run-level' f'' @@ -394,9 +502,11 @@ def write_html(all_stats: list[tuple], out_path: Path, total: int, repo_root: Pa for i, (name, desc, is_skipped) in enumerate(file_groups[file_name], 1): desc_html = esc(desc).replace("\n", "
") if desc else "" status_html = 'Skipped' if is_skipped else "—" + skip_attr = ' data-case-skipped="1"' if is_skipped else ' data-case-skipped="0"' suite_rows.append( - f'' - f'{badge_span}{esc(label)}{esc(file_name)}' + f'' + f"{badge_span}{group_cell}{esc(label)}" + f'{esc(file_name)}' f'{i}' f'{esc(name)}' f'{status_html}{desc_html}' @@ -408,44 +518,40 @@ def write_html(all_stats: list[tuple], out_path: Path, total: int, repo_root: Pa f"{''.join(suite_rows)}
" ) body = "".join(file_sections) - if pipeline_name not in pipeline_suites: - pipeline_suites[pipeline_name] = [] - pipeline_suites[pipeline_name].append((label, count, suite_summary, body)) + pipeline_group_suites.setdefault(pipeline_name, {}).setdefault(group, []).append( + (label, count, suite_summary, body) + ) - # Build two-level HTML: Pipeline (outer) -> Test Suite (inner) -> table - detail_sections_html_parts = [] - for pipeline_name in sorted(pipeline_badges.keys()): - if pipeline_name not in pipeline_suites: + def _group_sort_key(k: str | None) -> tuple[bool, str]: + return (k is None, (k or "").lower()) + + detail_sections_html_parts: list[str] = [] + all_pipeline_keys = sorted(set(pipeline_badges.keys()) | set(pipeline_group_suites.keys())) + for pipeline_name in all_pipeline_keys: + if pipeline_name not in pipeline_group_suites: continue - suites = pipeline_suites[pipeline_name] - pipeline_total = sum(c for _, c, _, _ in suites) + groups_dict = pipeline_group_suites[pipeline_name] + pipeline_total = sum(c for suites in groups_dict.values() for (_, c, _, _) in suites) badge = pipeline_badges.get(pipeline_name, "") badge_span = f'{esc(pipeline_name)}' if badge else esc(pipeline_name) pipeline_summary = f"{badge_span} ({pipeline_total} cases)" - inner_html = "".join( - f'
{suite_summary}' - f'
{body}
' - for _, _, suite_summary, body in suites - ) + group_sections: list[str] = [] + for g_key in sorted(groups_dict.keys(), key=_group_sort_key): + suites = groups_dict[g_key] + group_total = sum(c for _, c, _, _ in suites) + display_g = "—" if g_key is None else g_key + inner_html = "".join( + f'
{suite_summary}' + f'
{body}
' + for _, _, suite_summary, body in suites + ) + group_sections.append( + f'
{esc(display_g)} ({group_total} cases)' + f'
{inner_html}
' + ) detail_sections_html_parts.append( f'
{pipeline_summary}' - f'
{inner_html}
' - ) - # Include pipelines that might be in all_stats but not in pipeline_badges (e.g. custom names) - for pipeline_name in sorted(pipeline_suites.keys()): - if pipeline_name in pipeline_badges: - continue - suites = pipeline_suites[pipeline_name] - pipeline_total = sum(c for _, c, _, _ in suites) - pipeline_summary = f"{esc(pipeline_name)} ({pipeline_total} cases)" - inner_html = "".join( - f'
{suite_summary}' - f'
{body}
' - for _, _, suite_summary, body in suites - ) - detail_sections_html_parts.append( - f'
{pipeline_summary}' - f'
{inner_html}
' + f'
{"".join(group_sections)}
' ) detail_sections_html = "".join(detail_sections_html_parts) @@ -466,6 +572,7 @@ def write_html(all_stats: list[tuple], out_path: Path, total: int, repo_root: Pa --ready: #3fb950; --merge: #a371f7; --nightly: #f0883e; + --weekly: #79c0ff; }} * {{ box-sizing: border-box; }} body {{ @@ -524,6 +631,20 @@ def write_html(all_stats: list[tuple], out_path: Path, total: int, repo_root: Pa border-color: var(--accent); box-shadow: 0 0 0 2px rgba(88, 166, 255, 0.18); }} + .filter-select {{ + min-width: 180px; + padding: 0.65rem 0.85rem; + border-radius: 8px; + border: 1px solid var(--border); + background: rgba(255,255,255,0.04); + color: var(--text); + outline: none; + font-size: 0.9rem; + }} + .filter-select:focus {{ + border-color: var(--accent); + box-shadow: 0 0 0 2px rgba(88, 166, 255, 0.18); + }} .filter-hint {{ color: var(--muted); font-size: 0.82rem; @@ -566,6 +687,7 @@ def write_html(all_stats: list[tuple], out_path: Path, total: int, repo_root: Pa .badge.ready {{ background: rgba(63, 185, 80, 0.2); color: var(--ready); }} .badge.merge {{ background: rgba(163, 113, 247, 0.2); color: var(--merge); }} .badge.nightly {{ background: rgba(240, 136, 62, 0.2); color: var(--nightly); }} + .badge.weekly {{ background: rgba(121, 192, 255, 0.2); color: var(--weekly); }} .badge.skip {{ background: rgba(139, 148, 158, 0.25); color: var(--muted); }} td.status {{ color: var(--muted); font-size: 0.85rem; }} .case-details details {{ @@ -595,6 +717,12 @@ def write_html(all_stats: list[tuple], out_path: Path, total: int, repo_root: Pa .case-details details th {{ border-top: 0; }} .case-details .pipeline-details {{ margin-bottom: 0.5rem; }} .case-details .pipeline-details .case-details-inner {{ padding-left: 0.5rem; margin-top: 0.25rem; }} + .case-details .group-details {{ margin-bottom: 0.35rem; margin-left: 0.35rem; }} + .case-details .group-details summary::before {{ + content: "▶ "; font-size: 0.7rem; color: var(--muted); margin-right: 0.35rem; + }} + .case-details .group-details[open] summary::before {{ content: "▼ "; }} + .case-details .group-details .group-details-inner {{ padding-left: 0.35rem; margin-top: 0.25rem; }} .case-details .suite-details {{ margin-bottom: 0.35rem; margin-left: 0.5rem; }} .case-details .suite-details summary::before {{ content: "▶ "; font-size: 0.7rem; color: var(--muted); margin-right: 0.35rem; @@ -663,7 +791,7 @@ def write_html(all_stats: list[tuple], out_path: Path, total: int, repo_root: Pa
- + {"".join(summary_rows)} @@ -674,7 +802,7 @@ def write_html(all_stats: list[tuple], out_path: Path, total: int, repo_root: Pa

Case Details

-

Expand by Pipeline, then by Test Suite, then by Test File.

+

Expand by Pipeline, then Buildkite group, then Test Suite, then by Test File.

- Filter applies to the Test File fold sections below. + + + Path filter applies to file sections; status hides non-matching rows.
{detail_sections_html}
@@ -692,28 +825,48 @@ def write_html(all_stats: list[tuple], out_path: Path, total: int, repo_root: Pa @@ -723,6 +876,86 @@ def write_html(all_stats: list[tuple], out_path: Path, total: int, repo_root: Pa out_path.write_text(html_content, encoding="utf-8") +def _sample_preview_stats() -> list[tuple]: + """Synthetic stats tuples matching ``main()`` layout — UI demo only.""" + names_a = [ + "tests/e2e/demo/test_alpha.py::test_smoke", + "tests/e2e/demo/test_alpha.py::test_params[x0]", + "tests/e2e/demo/test_beta.py::test_widgets", + ] + skipped_a = {"tests/e2e/demo/test_alpha.py::test_smoke"} + return [ + ( + "test-ready / :card_index_dividers: Simple Test / Simple · Diffusion", + "test-ready", + ":card_index_dividers: Simple Test", + "Simple · Diffusion & Model Executor Test", + ["tests/diffusion/test_dummy.py", "tests/model_executor/test_dummy.py"], + False, + [], + len(names_a) - len(skipped_a), + names_a, + ["Alpha smoke.", "", "Widgets."], + skipped_a, + ), + ( + "test-ready / Custom Pipeline Test", + "test-ready", + None, + "Custom Pipeline Test", + ["tests/e2e/offline_inference/custom_pipeline/test_cp.py"], + False, + [], + 1, + ["tests/e2e/offline_inference/custom_pipeline/test_cp.py::test_cp_one"], + [""], + set(), + ), + ( + "test-merge / :card_index_dividers: Diffusion Test / LoRA", + "test-merge", + ":card_index_dividers: Diffusion Test", + "Diffusion · LoRA Test", + ["tests/diffusion/lora/test_lora.py"], + False, + [], + 2, + [ + "tests/diffusion/lora/test_lora.py::test_load", + "tests/diffusion/lora/test_lora.py::test_infer[a]", + ], + ["", ""], + set(), + ), + ( + "test-nightly / Omni marker-only (sample)", + "test-nightly", + ":card_index_dividers: Omni Model Test", + "Omni · marker-only step (sample)", + [], + True, + ["-m full_model and omni and H100", "--run-level full_model"], + 0, + [], + [], + set(), + ), + ( + "test-weekly / Reliability (sample)", + "test-weekly", + None, + "Reliability Test - qwen3-omni", + ["tests/dfx/reliability/test_reliability_qwen3_omni.py"], + False, + [], + 1, + ["tests/dfx/reliability/test_reliability_qwen3_omni.py::test_placeholder"], + ["Weekly reliability placeholder."], + set(), + ), + ] + + def main() -> None: parser = argparse.ArgumentParser(description="Count pytest cases from Buildkite pipelines") parser.add_argument( @@ -739,8 +972,23 @@ def main() -> None: default=default_out, help=f"Output HTML file path (default: {default_out})", ) + parser.add_argument( + "--sample-preview", + action="store_true", + help="Write HTML using placeholder pipelines/groups/cases only (no pytest, no YAML-driven counts).", + ) args = parser.parse_args() + if args.sample_preview: + all_stats = _sample_preview_stats() + out_path = args.output.resolve() + out_path.parent.mkdir(parents=True, exist_ok=True) + total = sum(item[7] for item in all_stats) + write_html(all_stats, out_path, total, REPO_ROOT) + print(f"Written (sample preview): {out_path}") + print(f"Test Suites: {len(all_stats)}, total placeholder cases: {total}") + return + all_stats = [] for filename in PIPELINE_FILES: @@ -752,11 +1000,10 @@ def main() -> None: steps = data.get("steps") or [] pipeline_name = filename.replace(".yml", "") - for step in steps: + for step, group in iter_pytest_leaf_steps(steps): if not isinstance(step, dict): continue label = step.get("label") or "(no label)" - # Skip commented-out steps (usually no label or commands) if not step.get("commands"): continue targets = get_pytest_targets_from_step(step) @@ -764,6 +1011,8 @@ def main() -> None: continue module_key = f"{pipeline_name} / {label}" + if group: + module_key = f"{pipeline_name} / {group} / {label}" # test_files: only real .py files; marker-only steps will have an empty list. test_files = [t for t, _ in targets if t.endswith(".py")] # marker_targets: synthetic marker targets (e.g. "-m expr", "-m expr --run-level level"). @@ -785,6 +1034,7 @@ def main() -> None: ( module_key, pipeline_name, + group, label, test_files, marker_only, @@ -798,7 +1048,7 @@ def main() -> None: out_path = args.output.resolve() out_path.parent.mkdir(parents=True, exist_ok=True) - total = sum(item[6] for item in all_stats) + total = sum(item[7] for item in all_stats) write_html(all_stats, out_path, total, REPO_ROOT) print(f"Written: {out_path}") print(f"Test Suites: {len(all_stats)}, total cases (sum of steps): {total}") From 14e5baceaf240e78d1a0c5dcc883563db23eb703 Mon Sep 17 00:00:00 2001 From: Lu Date: Tue, 19 May 2026 19:19:58 +0800 Subject: [PATCH 13/38] [Qwen-Image] Drop unused vision tower from text encoder (#3608) Signed-off-by: lulugoodcoder Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com> --- .../models/qwen_image/pipeline_qwen_image.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index 8f05e6d0f..6f7a6bb58 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py @@ -290,7 +290,22 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP ) self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( model, subfolder="text_encoder", local_files_only=local_files_only - ).to(self.device) + ) + # Qwen2.5-VL ships a vision tower that text-to-image does not use. + # Drop it while the model is still on CPU, before moving to GPU, so + # the vision tower never consumes GPU memory. Handle both transformers + # layouts: newer puts visual under .model, older puts it directly on + # the model. + visual_owner = None + if hasattr(self.text_encoder, "model") and hasattr(self.text_encoder.model, "visual"): + visual_owner = self.text_encoder.model + elif hasattr(self.text_encoder, "visual"): + visual_owner = self.text_encoder + if visual_owner is not None: + del visual_owner.visual + else: + logger.warning("Qwen-Image: vision tower not found on text encoder; skipping drop") + self.text_encoder = self.text_encoder.to(self.device) self.vae = DistributedAutoencoderKLQwenImage.from_pretrained( model, subfolder="vae", local_files_only=local_files_only ).to(self.device) From 18186db216319684e3e0d2c268d6a0409525fc2e Mon Sep 17 00:00:00 2001 From: Schatten <3192396192@qq.com> Date: Tue, 19 May 2026 19:23:45 +0800 Subject: [PATCH 14/38] [Cleanup] Remove unused build_base_engine_args after #1115 (#3720) Signed-off-by: Schatten --- vllm_omni/entrypoints/utils.py | 48 ---------------------------------- 1 file changed, 48 deletions(-) diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index d728e7641..de7eb5f4c 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -746,54 +746,6 @@ def filter_dataclass_kwargs(cls: Any, kwargs: dict) -> dict: return filtered_kwargs -# TODO(wuhang): Remove after PR #1115. -def build_base_engine_args(source: Any) -> dict[str, Any] | None: - """Build base engine args with tokenizer and parallel configuration. - - Automatically detects whether source is a dict-like object or namespace object. - - Args: - source: Source object (args namespace or kwargs dict) containing configuration. - - Returns: - Dictionary containing tokenizer and parallel configuration overrides, - or None if no configuration is present. - """ - # Auto-detect source type: dict-like objects have 'get' method - is_dict_like = hasattr(source, "get") and callable(getattr(source, "get")) - - # Extract tokenizer - if is_dict_like: - tokenizer = source.get("tokenizer", None) - else: - tokenizer = getattr(source, "tokenizer", None) - - base_engine_args = {"tokenizer": tokenizer} if tokenizer is not None else None - - # Extract parallel configuration - parallel_keys = [ - "tensor_parallel_size", - "pipeline_parallel_size", - "data_parallel_size", - "data_parallel_size_local", - "data_parallel_backend", - "distributed_executor_backend", - ] - - if is_dict_like: - parallel_overrides = {k: source[k] for k in parallel_keys if k in source and source[k] is not None} - else: - parallel_overrides = { - k: getattr(source, k) for k in parallel_keys if hasattr(source, k) and getattr(source, k) is not None - } - - if parallel_overrides: - base_engine_args = base_engine_args or {} - base_engine_args.update(parallel_overrides) - - return base_engine_args - - # The following code detects if the process is running in a container and if # PID host is available. If so, we can use process-scoped memory tracking; # otherwise we need sequential init locks. From da5361879395d45d5017fb575a7446cb36774bf4 Mon Sep 17 00:00:00 2001 From: Shin Date: Tue, 19 May 2026 19:56:38 +0800 Subject: [PATCH 15/38] [Recipe] Qwen/Qwen-Image-Edit (#3684) Signed-off-by: yixiaoer --- recipes/Qwen/Qwen-Image-Edit.md | 171 ++++++++++++++++++++++++++++++++ recipes/README.md | 1 + 2 files changed, 172 insertions(+) create mode 100644 recipes/Qwen/Qwen-Image-Edit.md diff --git a/recipes/Qwen/Qwen-Image-Edit.md b/recipes/Qwen/Qwen-Image-Edit.md new file mode 100644 index 000000000..88e42c60e --- /dev/null +++ b/recipes/Qwen/Qwen-Image-Edit.md @@ -0,0 +1,171 @@ +# Qwen-Image-Edit + +> Image editing serving + +## Summary + +- Vendor: Qwen +- Model: `Qwen/Qwen-Image-Edit` +- Task: Text-guided single-image editing +- Mode: Online serving with the OpenAI-compatible multimodal endpoint +- Maintainer: Community + +## When to use this recipe + +Use this recipe as a baseline for serving `Qwen/Qwen-Image-Edit` for single-image editing on H200 141GB. Two configurations are validated below, sharing one software environment: + +| Configuration | Parallelism | Notes | +| --- | --- | --- | +| `1x H200` | none | Single-card baseline; ~60 GiB peak VRAM on a 141 GB card | +| `2x H200` | TP=2 | Lower single-request latency (≈ 1.60× speedup) and lower per-GPU VRAM | + +## References + +- Model: +- Example guide: [`examples/online_serving/image_to_image/README.md`](../../examples/online_serving/image_to_image/README.md) +- User guide: [`docs/user_guide/examples/online_serving/image_to_image.md`](../../docs/user_guide/examples/online_serving/image_to_image.md) + +## Hardware Support + +### GPU + +The two hardware configurations below share the same software environment. + +#### Environment + +- OS: Linux +- Python: 3.12+ +- Driver / runtime: CUDA-capable runtime matching the repository build +- vLLM: use versions from your current checkout, >=0.20.0 +- vLLM-Omni: use versions from your current checkout + +#### 1x H200 141GB + +##### Command + +```bash +CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen-Image-Edit --omni \ + --port 8092 +``` + +The bundled launcher is equivalent: + +```bash +bash examples/online_serving/image_to_image/run_server.sh +``` + +The server is ready when the log shows `INFO: Application startup complete.` (preceded by `Started server process [PID]` and `Waiting for application startup.`). + +##### Verification + +Run a single-image edit against the public example asset: + +```bash +wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/qwen-bear.png + +PROMPT="Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" +IMG_B64=$(base64 -w0 qwen-bear.png) + +curl -s http://localhost:8092/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d "$(jq -n --arg prompt "$PROMPT" --arg img "$IMG_B64" '{ + messages: [{ + role: "user", + content: [ + {type: "text", text: $prompt}, + {type: "image_url", image_url: {url: ("data:image/png;base64," + $img)}} + ] + }], + extra_body: { + height: 1024, + width: 1024, + num_inference_steps: 50, + guidance_scale: 1.0, + seed: 42 + } + }')" \ + | jq -r '.choices[0].message.content[0].image_url.url' \ + | cut -d',' -f2- | base64 -d > qwen_image_edit_1card.png +``` + +The decoded PNG is a 1024×1024 RGB image of the bear mascot dancing under stylized moonlight, surrounded by floating stars and a "Be Kind" text bubble. + +You can also run the offline CLI without standing up a server: + +```bash +python examples/offline_inference/image_to_image/image_edit.py \ + --image qwen-bear.png \ + --prompt "Let this mascot dance under the moon, ..." \ + --output qwen_image_edit_offline.png \ + --num-inference-steps 50 \ + --cfg-scale 4.0 +``` + +Validated reference numbers for 1024×1024, 50 denoising steps, batch size 1, `guidance_scale=1.0`, seed 42 (`Cold` = first request after startup, `Warm` = identical subsequent request): + +| Cold latency | Warm latency | Peak GPU memory | +| ---: | ---: | ---: | +| 17.80 s | 17.78 s | 61,091 MiB (~59.66 GiB) | + +##### Notes + +- Generation parameters live under `extra_body` for `/v1/chat/completions`, or as top-level form fields when using `/v1/images/edits`. See [`examples/online_serving/image_to_image/README.md`](../../examples/online_serving/image_to_image/README.md) for the full request / response schema. +- Memory: a 1024×1024 / 50-step edit uses ~60 GiB on a single H200, well within the 141 GB headroom. +- Multi-image variants of Qwen-Image-Edit (e.g. `Qwen-Image-Edit-2509`) are intentionally out of scope for this recipe and are not validated here. + +#### 2x H200 141GB (TP=2) + +##### Command + +```bash +CUDA_VISIBLE_DEVICES=0,1 vllm serve Qwen/Qwen-Image-Edit --omni \ + --port 8092 \ + --tensor-parallel-size 2 +``` + +In addition to the `Application startup complete.` line, the TP=2 server's `non-default args` log line echoes `'tensor_parallel_size': 2`, and a separate `Worker N: Initialized device and distributed environment` line is emitted for each rank — one per visible GPU. + +##### Verification + +Run the same `curl` command as in the **1x H200** Verification section above against the TP=2 server (same prompt, same seed, same image). Save the result to `qwen_image_edit_2card_tp2.png` so the two outputs sit side by side for comparison. During TP=2 inference, `nvidia-smi` shows sustained `~99 %` utilization on both GPUs, which is the expected signal that TP wiring is active. + +Same workload as the 1x table above (1024×1024, 50 steps, `guidance_scale=1.0`, seed 42): + +| Configuration | Cold latency | Warm latency | Peak per-GPU memory | +| --- | ---: | ---: | ---: | +| 1x H200 | 17.80 s | 17.78 s | 61,091 MiB (~59.66 GiB) | +| 2x H200 TP=2 | 11.16 s | 11.09 s | 49,377 MiB (~48.22 GiB) | + +TP=2 reduces single-request warm latency from `17.78 s` to `11.09 s` (≈ **1.60×**) and drops per-GPU peak VRAM by ~11.4 GiB. + +##### Notes + +- TP=2 shards model weights across both cards, which is what produces both the latency reduction and the per-GPU VRAM drop above. +- CFG-Parallel (`--cfg-parallel-size 2`) is supported on the Qwen-Image family but only when `cfg_scale > 1`. The verification command above uses `guidance_scale=1.0` to match the public example image, so it does not exercise CFG-parallel; a separate run with `cfg_scale > 1` is needed if you want to benchmark that path. +- For deeper acceleration knobs (Cache-DiT, sequence parallel, HSDP) see [`docs/user_guide/diffusion/parallelism/overview.md`](../../docs/user_guide/diffusion/parallelism/overview.md); this recipe intentionally documents only the validated TP=2 baseline. + +#### Profiling + +The numbers below are from a warm request at 1024×1024, 50 steps, and `guidance_scale=1.0`, measured with `--enable-diffusion-pipeline-profiler` and a 250 ms `nvidia-smi --query-gpu=memory.used` sampler. For TP=2, latency uses rank max, and VRAM is per GPU. + +Per-stage latency: + +| Stage | TP=1 | TP=2 | Notes | +| --- | ---: | ---: | --- | +| `text_encoder.forward` | 0.132 s | 0.135 s | ~1.0× | +| `vae.encode` | 0.040 s | 0.041 s | ~1.0× | +| `diffuse` | 17.105 s | 10.382 s | **1.65×** speedup | +| `vae.decode` | 0.064 s | 0.066 s | ~1.0× | +| `forward` | 17.372 s | 10.652 s | 1.63× speedup | + +Per-GPU VRAM: + +| Snapshot | TP=1 | TP=2 | Notes | +| --- | ---: | ---: | --- | +| Static VRAM | 58,777 MiB | 47,065 MiB | after model load, before request | +| Peak VRAM | 61,091 MiB | 49,377 MiB | during inference | +| Request-time increment | 2,314 MiB | 2,312 MiB | peak minus static | + +The latency gain mainly comes from the DiT denoising stage (`diffuse` in the table above). The other measured stages, such as text encoder and VAE encode/decode, stay essentially unchanged between TP=1 and TP=2, which is consistent with them not benefiting meaningfully from TP in this pipeline. The `forward` row's 1.63× speedup tracks `diffuse`, since `diffuse` accounts for about 98% of total pipeline time. + +For memory, TP=2 saves about 11.4 GiB per GPU, mostly from static model-state memory, likely dominated by tensor-parallel sharding of the DiT weights. The request-time increment is almost unchanged, so TP=2 does not materially reduce transient inference-time memory for this single-image, batch-1 workload. diff --git a/recipes/README.md b/recipes/README.md index 09b1aeaee..6a49e6449 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -34,6 +34,7 @@ recipes/ | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB | | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) | | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay | 1x A100 80GB | +| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB | | [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB | | [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB | | [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x H200 (144GB) | From 3c58868c9a4fb7f0b1754d07738d1f87d3af5dae Mon Sep 17 00:00:00 2001 From: dengyunyang <584797741@qq.com> Date: Tue, 19 May 2026 22:22:27 +0800 Subject: [PATCH 16/38] [BugFix] fix mult cli timeout with get kv (#3741) Signed-off-by: dengyunyang <584797741@qq.com> --- tests/entrypoints/test_serve.py | 5 ++++- vllm_omni/entrypoints/cli/serve.py | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 86c311a60..0915ddead 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -325,7 +325,10 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture) run_headless(_make_headless_args(stage_id=1)) - mock_inject.assert_called_once_with(stage_cfg, 1) + mock_inject.assert_called_once() + assert mock_inject.call_args.args[0] is stage_cfg + assert mock_inject.call_args.args[1] == 1 + assert mock_inject.call_args.args[2] == [stage_cfg] reg_kwargs = mock_register.call_args.kwargs assert reg_kwargs["omni_master_address"] == "127.0.0.1" diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 50b825d7c..e4fddc830 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -803,7 +803,10 @@ def run_headless(args: argparse.Namespace) -> None: metadata = extract_stage_metadata(stage_cfg) if omni_conn_cfg: inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to) - inject_kv_stage_info(stage_cfg, stage_id) + # Headless single-stage launch must still infer cross-stage TP topology + # from the loaded deploy config so heterogeneous KV routing keys match + # the head process (e.g. from_tp=2, to_tp=1). + inject_kv_stage_info(stage_cfg, stage_id, stage_configs) od_config = build_diffusion_config(model, stage_cfg, metadata) logger.info( From a3d4ed809d56977eb632e8a63aae1fc090a790e3 Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Wed, 20 May 2026 00:14:08 +0800 Subject: [PATCH 17/38] [Quantization][tools] Add diffusion quantization output comparison tool (#3175) Signed-off-by: david6666666 <530634352@qq.com> Signed-off-by: David Chen <530634352@qq.com> --- docs/user_guide/quantization/overview.md | 95 +++ .../test_trajectory_similarity_tool.py | 126 ++++ vllm_omni/quantization/tools/__init__.py | 0 ...compare_diffusion_trajectory_similarity.py | 637 ++++++++++++++++++ 4 files changed, 858 insertions(+) create mode 100644 tests/diffusion/quantization/test_trajectory_similarity_tool.py create mode 100644 vllm_omni/quantization/tools/__init__.py create mode 100644 vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py diff --git a/docs/user_guide/quantization/overview.md b/docs/user_guide/quantization/overview.md index c4605c8c2..39669bf5d 100644 --- a/docs/user_guide/quantization/overview.md +++ b/docs/user_guide/quantization/overview.md @@ -138,3 +138,98 @@ build_quant_config({"method": "gguf", "gguf_model": "/path/to/model.gguf"}) build_quant_config({"transformer": {"method": "fp8"}, "vae": None}) build_quant_config(None) ``` + +## Output Similarity Comparison Tool + +Use `vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity` +to compare a reference diffusion run with a quantized candidate run using the +same prompt, seed, resolution, scheduler settings, and inference steps. The +tool compares final decoded images or video frames, and also reports generation +latency and worker-reported peak memory when available. + +This is useful when validating whether online quantization, an offline +pre-quantized checkpoint, or a new `ignored_layers` choice keeps generation +quality close to the BF16 reference. + +### Online Quantization Example + +```bash +python -m vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity \ + --task t2i \ + --model Qwen/Qwen-Image \ + --candidate-quantization fp8 \ + --ignored-layers img_mlp \ + --prompt "a cup of coffee on the table" \ + --height 512 --width 512 \ + --num-inference-steps 20 \ + --seed 142 \ + --output-json /tmp/qwen_image_fp8_similarity/result.json \ + --save-output-dir /tmp/qwen_image_fp8_similarity/images \ + --enforce-eager +``` + +### Offline Checkpoint Example + +Use `--candidate-model` when the candidate is already quantized or lives at a +different model path: + +```bash +python -m vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity \ + --task t2i \ + --reference-model Qwen/Qwen-Image \ + --candidate-model /path/to/qwen-image-fp8-checkpoint \ + --prompt "a cup of coffee on the table" \ + --height 512 --width 512 \ + --num-inference-steps 20 \ + --seed 142 \ + --output-json /tmp/qwen_image_fp8_checkpoint_similarity/result.json +``` + +If the checkpoint does not include a loadable quantization config, pass one +explicitly: + +```bash +--candidate-quantization-config-json '{"method":"fp8"}' +``` + +### Output Metrics + +The output JSON includes `output_metrics`, `reference_generation`, and +`candidate_generation`. + +| Metric | Direction | Meaning | +|--------|-----------|---------| +| `cosine_similarity` | Higher is better | Vector direction similarity between output pixels or frames. Useful as a broad sanity check. | +| `mae` | Lower is better | Mean absolute pixel or frame error. For decoded outputs, values are in uint8 pixel units. | +| `mse` / `rmse` | Lower is better | Squared error and its square root. These penalize localized large differences more than `mae`. | +| `max_abs` | Lower is better | Worst single-element absolute error. Treat it as an outlier/debug signal, not as a release gate. | +| `l2` / `relative_l2` | Lower is better | Absolute and reference-normalized L2 distance. `relative_l2` is easier to compare across resolutions. | +| `psnr_db` | Higher is better | Pixel-space signal-to-noise ratio in dB for uint8 images or frames. | +| `avg_generation_time_s` | Lower is better | Average wall-clock generation time across measured runs. | +| `max_peak_memory_mb` | Lower is better | Maximum worker-reported peak device memory across measured runs, when the worker reports it. | + +Recommended starting thresholds for same-seed diffusion comparisons: + +| Metric | Smoke threshold | Stricter target | Notes | +|--------|-----------------|-----------------|-------| +| `psnr_db` | `>= 20.0` | `>= 25.0` | Good for quick image or frame regression checks. | +| `mae` | `<= 12.0` | `<= 6.0` | Interpreted in decoded uint8 pixel units. | +| `cosine_similarity` | `>= 0.98` | `>= 0.995` | Less sensitive to global scale than L2-style metrics. | +| `relative_l2` | `<= 0.20` | `<= 0.08` | Useful when comparing across prompts or resolutions. | + +These thresholds are heuristics. Tune them by model family, task, resolution, +quantization method, and deployment tolerance. For release gating, pair the +numeric report with visual inspection of saved reference and candidate outputs. + +The tool intentionally reports separate quality, latency, and memory metrics +instead of a single consolidated similarity score. A single score can hide +important tradeoffs, for example a candidate with good PSNR but a meaningful +memory regression, or a candidate with low average error but localized visual +artifacts. If you need a project-specific pass/fail gate, define it as an +explicit policy over the individual metrics. + +Pixel-level metrics do not measure semantic consistency. For higher-cost +evaluation, you can complement this report with a vision-language judge that +describes the reference and candidate outputs and compares those descriptions. +Keep that semantic check separate from this lightweight tool so users can +choose whether the additional model cost and latency are appropriate. diff --git a/tests/diffusion/quantization/test_trajectory_similarity_tool.py b/tests/diffusion/quantization/test_trajectory_similarity_tool.py new file mode 100644 index 000000000..45d51d356 --- /dev/null +++ b/tests/diffusion/quantization/test_trajectory_similarity_tool.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math +import sys +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from PIL import Image + +from vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity import ( + VariantRun, + _build_variant_config, + _request_peak_memory_mb, + _run_summary, + compute_tensor_metrics, + compute_uint8_image_metrics, + metric_guidance, + parse_args, + summarize_output_image_metrics, +) + + +def test_compute_tensor_metrics_identical_tensors(): + metrics = compute_tensor_metrics(torch.ones(2, 3), torch.ones(2, 3)) + + assert metrics["cosine_similarity"] == pytest.approx(1.0) + assert metrics["mae"] == 0.0 + assert metrics["mse"] == 0.0 + assert metrics["rmse"] == 0.0 + assert metrics["max_abs"] == 0.0 + assert metrics["l2"] == 0.0 + assert metrics["relative_l2"] == 0.0 + + +def test_compute_uint8_image_metrics_adds_psnr(): + lhs = np.zeros((2, 2, 3), dtype=np.uint8) + rhs = np.zeros((2, 2, 3), dtype=np.uint8) + + metrics = compute_uint8_image_metrics(lhs, rhs) + + assert math.isinf(metrics["psnr_db"]) + + +def test_summarize_output_image_metrics_stacks_pil_images(): + reference = [Image.fromarray(np.zeros((2, 2, 3), dtype=np.uint8))] + candidate = [Image.fromarray(np.ones((2, 2, 3), dtype=np.uint8))] + + summary = summarize_output_image_metrics(reference, candidate) + + assert summary["num_images"] == 1 + assert summary["image0_metrics"]["mae"] == 1.0 + assert summary["all_images_metrics"]["mse"] == 1.0 + + +def test_run_summary_reports_worker_peak_memory(): + summary = _run_summary( + VariantRun( + label="candidate", + result=object(), + generation_times_s=[1.0, 3.0], + peak_memory_mb=[100.0, 150.0], + ) + ) + + assert summary["peak_memory_mb"] == 150.0 + assert summary["avg_peak_memory_mb"] == 125.0 + assert summary["max_peak_memory_mb"] == 150.0 + assert summary["per_run_peak_memory_mb"] == [100.0, 150.0] + + +def test_request_peak_memory_prefers_inner_diffusion_output(): + inner = type("InnerOutput", (), {"peak_memory_mb": 321.0})() + outer = type("OuterOutput", (), {"request_output": inner, "peak_memory_mb": 123.0})() + + assert _request_peak_memory_mb(outer) == 321.0 + + +def test_metric_guidance_describes_thresholds(): + guidance = metric_guidance() + + assert "cosine_similarity" in guidance["descriptions"] + assert guidance["recommended_thresholds"]["output_images_or_frames_uint8"]["psnr_db"]["recommended_min"] == 20.0 + assert ( + guidance["recommended_thresholds"]["performance"]["max_peak_memory_ratio_candidate_over_reference"][ + "recommended_max" + ] + == 1.00 + ) + + +def test_candidate_model_can_point_to_offline_checkpoint_without_online_quantization(): + args = SimpleNamespace( + model="Qwen/Qwen-Image", + candidate_model="Qwen/Qwen-Image-FP8", + candidate_quantization=None, + candidate_quantization_config_json=None, + candidate_ignored_layers=None, + ignored_layers=None, + candidate_gguf_model=None, + candidate_load_format="default", + ) + + config = _build_variant_config(args, "candidate") + + assert config.model == "Qwen/Qwen-Image-FP8" + assert config.quantization is None + assert config.quantization_config is None + + +def test_step_execution_defaults_to_false(monkeypatch): + monkeypatch.setattr( + sys, + "argv", + [ + "compare_diffusion_trajectory_similarity.py", + "--output-json", + "result.json", + ], + ) + + args = parse_args() + + assert args.step_execution is False diff --git a/vllm_omni/quantization/tools/__init__.py b/vllm_omni/quantization/tools/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py b/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py new file mode 100644 index 000000000..3d4f68d23 --- /dev/null +++ b/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py @@ -0,0 +1,637 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compare diffusion quantization quality with final output metrics. + +This tool runs a reference diffusion configuration and a candidate +configuration with the same prompt, seed, resolution, and step count. It reports +final image or video frame similarity metrics plus generation latency and memory. + +The main metrics are: + +* output images: the same pixel-space metrics plus PSNR +* runtime: per-run and average generation latency, and worker-reported peak + device memory when available + +Online FP8 Qwen-Image example: + + python -m vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity \\ + --model Qwen/Qwen-Image \\ + --candidate-quantization fp8 \\ + --ignored-layers img_mlp \\ + --prompt "a cup of coffee on the table" \\ + --height 1024 --width 1024 --num-inference-steps 20 --seed 42 \\ + --output-json /mnt/data4/cwq/tmp/qwen_image_fp8_similarity.json + +Offline checkpoint example: + + python -m vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity \\ + --reference-model Qwen/Qwen-Image \\ + --candidate-model /path/to/quantized/qwen-image-checkpoint \\ + --candidate-quantization-config-json '{"method":"inc","weight_bits":4}' \\ + --prompt "a cup of coffee on the table" \\ + --output-json /mnt/data4/cwq/tmp/qwen_image_checkpoint_similarity.json + +Offline FP8 model-id example: + + python -m vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity \\ + --model Qwen/Qwen-Image \\ + --candidate-model Qwen/Qwen-Image-FP8 \\ + --prompt "a cup of coffee on the table" \\ + --height 1024 --width 1024 --num-inference-steps 20 --seed 42 \\ + --output-json /mnt/data4/cwq/tmp/qwen_image_fp8_checkpoint_similarity.json + +If the checkpoint does not carry a loadable quantization config, add +--candidate-quantization-config-json, for example '{"method":"fp8"}'. +""" + +from __future__ import annotations + +import argparse +import gc +import json +import math +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import torch.nn.functional as F + + +@dataclass +class VariantConfig: + label: str + model: str + quantization: str | None + quantization_config: dict[str, Any] | None + diffusion_load_format: str + + +@dataclass +class VariantRun: + label: str + result: Any + generation_times_s: list[float] + peak_memory_mb: list[float | None] + + +def metric_guidance() -> dict[str, Any]: + """Human-readable metric notes and default heuristic thresholds.""" + return { + "scope": ( + "Heuristic defaults for same-model comparisons with identical prompt, seed, resolution, scheduler, " + "and inference steps. Tune thresholds per model family, task, and quantization method; visual review is " + "still recommended for release gating." + ), + "descriptions": { + "output_metrics": ( + "Pixel-space metrics computed on generated RGB images or video frames after decode. These are often " + "the most actionable signal for image/video quality regressions." + ), + "cosine_similarity": ( + "Higher is better. Measures vector direction similarity and is less sensitive to scale." + ), + "mae": ( + "Lower is better. Mean absolute error in the tensor's native scale; uint8 output MAE is in " + "pixel values." + ), + "mse": "Lower is better. Mean squared error; more sensitive to large localized differences than MAE.", + "rmse": "Lower is better. Square root of MSE, in the same unit as the compared tensor.", + "max_abs": ( + "Lower is better. Worst absolute element error; useful for debugging outliers, not stable as a gate." + ), + "l2": "Lower is better. Absolute L2 distance; shape- and scale-dependent.", + "relative_l2": "Lower is better. L2 distance normalized by the reference L2 norm.", + "psnr_db": "Higher is better. Pixel-space signal-to-noise ratio in dB for uint8 images/frames.", + "generation_time_s": "Lower is better. Wall-clock time for the last measured generation run.", + "avg_generation_time_s": "Lower is better. Mean wall-clock time across measured runs.", + "peak_memory_mb": ( + "Lower is better. Worker-reported peak device memory for the last measured generation run." + ), + "max_peak_memory_mb": "Lower is better. Maximum worker-reported peak memory across measured runs.", + }, + "recommended_thresholds": { + "output_images_or_frames_uint8": { + "psnr_db": { + "recommended_min": 20.0, + "good_min": 25.0, + "strict_min": 30.0, + "note": ( + "For stochastic diffusion outputs, 20 dB is a practical smoke threshold; 25+ dB is usually " + "close." + ), + }, + "mae": { + "recommended_max": 12.0, + "good_max": 6.0, + "strict_max": 3.0, + "note": "Measured on uint8 RGB pixels or frames.", + }, + "cosine_similarity": { + "recommended_min": 0.98, + "strict_min": 0.995, + "note": "Useful as a broad image/frame similarity sanity check.", + }, + "relative_l2": { + "recommended_max": 0.20, + "strict_max": 0.08, + "note": "Pixel-space relative L2 is generally more interpretable than latent-space relative L2.", + }, + }, + "performance": { + "avg_generation_time_ratio_candidate_over_reference": { + "recommended_max": 1.10, + "note": ( + "A value above 1.10 suggests a latency regression unless expected for the quantization mode." + ), + }, + "max_peak_memory_ratio_candidate_over_reference": { + "recommended_max": 1.00, + "target_max": 0.95, + "note": ( + "Quantization should usually not increase peak memory; <=0.95 indicates a meaningful saving." + ), + }, + }, + }, + } + + +def _split_csv(value: str | None) -> list[str]: + if not value: + return [] + return [item.strip() for item in value.split(",") if item.strip()] + + +def _load_json_arg(value: str | None) -> dict[str, Any] | None: + if value is None: + return None + path = Path(value).expanduser() + if path.exists(): + value = path.read_text(encoding="utf-8") + parsed = json.loads(value) + if not isinstance(parsed, dict): + raise ValueError("quantization config JSON must decode to an object.") + return parsed + + +def _cosine_similarity(flat_a: torch.Tensor, flat_b: torch.Tensor) -> float: + norm_a = torch.linalg.vector_norm(flat_a).item() + norm_b = torch.linalg.vector_norm(flat_b).item() + if norm_a == 0.0 and norm_b == 0.0: + return 1.0 + if norm_a == 0.0 or norm_b == 0.0: + return 0.0 + cosine = float(F.cosine_similarity(flat_a, flat_b, dim=0).item()) + return max(-1.0, min(1.0, cosine)) + + +def compute_tensor_metrics(lhs: Any, rhs: Any) -> dict[str, float]: + lhs_tensor = torch.as_tensor(lhs).detach().cpu().double() + rhs_tensor = torch.as_tensor(rhs).detach().cpu().double() + if lhs_tensor.shape != rhs_tensor.shape: + raise ValueError(f"Metric shape mismatch: {tuple(lhs_tensor.shape)} vs {tuple(rhs_tensor.shape)}") + + diff = lhs_tensor - rhs_tensor + mse = float(diff.square().mean().item()) + rmse = float(math.sqrt(mse)) + mae = float(diff.abs().mean().item()) + max_abs = float(diff.abs().max().item()) + l2 = float(torch.linalg.vector_norm(diff).item()) + ref_l2 = float(torch.linalg.vector_norm(lhs_tensor).item()) + relative_l2 = 0.0 if ref_l2 == 0.0 and l2 == 0.0 else float("inf") if ref_l2 == 0.0 else l2 / ref_l2 + cosine = _cosine_similarity(lhs_tensor.reshape(-1), rhs_tensor.reshape(-1)) + return { + "cosine_similarity": cosine, + "mae": mae, + "mse": mse, + "rmse": rmse, + "max_abs": max_abs, + "l2": l2, + "relative_l2": relative_l2, + } + + +def compute_uint8_image_metrics(lhs: Any, rhs: Any) -> dict[str, float]: + metrics = compute_tensor_metrics(lhs, rhs) + mse = metrics["mse"] + metrics["psnr_db"] = float("inf") if mse == 0.0 else 20 * math.log10(255.0) - 10 * math.log10(mse) + return metrics + + +def summarize_output_frame_metrics(reference_frames: list[Any], candidate_frames: list[Any]) -> dict[str, Any]: + if len(reference_frames) != len(candidate_frames): + raise ValueError(f"Output frame count mismatch: {len(reference_frames)} vs {len(candidate_frames)}") + if not reference_frames: + raise ValueError("No output frames available for comparison.") + + ref_stack = np.stack([np.asarray(frame.convert("RGB")) for frame in reference_frames], axis=0) + cand_stack = np.stack([np.asarray(frame.convert("RGB")) for frame in candidate_frames], axis=0) + + frame0_metrics = compute_uint8_image_metrics(ref_stack[0], cand_stack[0]) + mid_index = len(reference_frames) // 2 + mid_metrics = compute_uint8_image_metrics(ref_stack[mid_index], cand_stack[mid_index]) + all_metrics = compute_uint8_image_metrics(ref_stack, cand_stack) + + return { + "num_frames": len(reference_frames), + "frame0_metrics": frame0_metrics, + "mid_frame_index": mid_index, + "mid_frame_metrics": mid_metrics, + "all_frames_metrics": all_metrics, + # Backward-compatible aliases for existing image callers. + "num_images": len(reference_frames), + "image0_metrics": frame0_metrics, + "mid_image_index": mid_index, + "mid_image_metrics": mid_metrics, + "all_images_metrics": all_metrics, + } + + +def summarize_output_image_metrics(reference_images: list[Any], candidate_images: list[Any]) -> dict[str, Any]: + return summarize_output_frame_metrics(reference_images, candidate_images) + + +def _extract_inner_output(outputs: Any) -> Any: + if isinstance(outputs, list): + if len(outputs) != 1: + raise ValueError(f"Expected one request output, got {len(outputs)}.") + outputs = outputs[0] + inner = getattr(outputs, "request_output", None) + if inner is not None: + return inner + return outputs + + +def _request_peak_memory_mb(result: Any) -> float | None: + diffusion_output = getattr(result, "request_output", None) + peak_memory_mb = getattr(diffusion_output, "peak_memory_mb", None) + if peak_memory_mb is None: + peak_memory_mb = getattr(result, "peak_memory_mb", None) + if peak_memory_mb is None: + return None + return float(peak_memory_mb) + + +def _build_quantization_config( + *, + quantization: str | None, + quantization_config_json: str | None, + ignored_layers: str | None, + gguf_model: str | None, +) -> dict[str, Any] | None: + explicit = _load_json_arg(quantization_config_json) + if explicit is not None: + return explicit + if quantization is None: + return None + + config: dict[str, Any] = {"method": quantization} + ignored = _split_csv(ignored_layers) + if ignored: + config["ignored_layers"] = ignored + if gguf_model is not None: + config["gguf_model"] = gguf_model + return config + + +def _build_variant_config(args: argparse.Namespace, variant: str) -> VariantConfig: + common_ignored_layers = args.ignored_layers if variant == "candidate" else None + ignored_layers = getattr(args, f"{variant}_ignored_layers") or common_ignored_layers + quantization = getattr(args, f"{variant}_quantization") + quantization_config = _build_quantization_config( + quantization=quantization, + quantization_config_json=getattr(args, f"{variant}_quantization_config_json"), + ignored_layers=ignored_layers, + gguf_model=getattr(args, f"{variant}_gguf_model"), + ) + return VariantConfig( + label=variant, + model=getattr(args, f"{variant}_model") or args.model, + quantization=quantization if quantization_config is None else None, + quantization_config=quantization_config, + diffusion_load_format=getattr(args, f"{variant}_load_format"), + ) + + +def _build_omni_kwargs(args: argparse.Namespace, config: VariantConfig) -> dict[str, Any]: + from vllm_omni.diffusion.data import DiffusionParallelConfig + + parallel_config = DiffusionParallelConfig( + ulysses_degree=args.ulysses_degree, + 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, + ) + kwargs: dict[str, Any] = { + "model": config.model, + "mode": "text-to-video" if args.task == "t2v" else "text-to-image", + "parallel_config": parallel_config, + "enforce_eager": args.enforce_eager, + "enable_cpu_offload": args.enable_cpu_offload, + "enable_layerwise_offload": args.enable_layerwise_offload, + "vae_use_slicing": args.vae_use_slicing, + "vae_use_tiling": args.vae_use_tiling, + "step_execution": args.step_execution if args.task == "t2i" else False, + "diffusion_load_format": config.diffusion_load_format, + } + if args.model_class_name: + kwargs["model_class_name"] = args.model_class_name + if args.flow_shift is not None: + kwargs["flow_shift"] = args.flow_shift + if config.quantization_config is not None: + kwargs["quantization_config"] = config.quantization_config + elif config.quantization is not None: + kwargs["quantization"] = config.quantization + return kwargs + + +def _build_sampling_params(args: argparse.Namespace, seed: int): + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + from vllm_omni.platforms import current_omni_platform + + generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed) + kwargs: dict[str, Any] = { + "height": args.height, + "width": args.width, + "generator": generator, + "seed": seed, + "guidance_scale": args.guidance_scale, + "guidance_scale_2": args.guidance_scale_2, + "num_inference_steps": args.num_inference_steps, + "return_frames": True, + "extra_args": { + "timesteps_shift": args.timesteps_shift, + "cfg_schedule": args.cfg_schedule, + "use_norm": args.use_norm, + "use_system_prompt": args.use_system_prompt, + "system_prompt": args.system_prompt, + }, + } + if args.task == "t2v": + kwargs["num_frames"] = args.num_frames + kwargs["fps"] = args.fps + kwargs["frame_rate"] = args.frame_rate + else: + kwargs["true_cfg_scale"] = args.cfg_scale + kwargs["num_outputs_per_prompt"] = args.num_images_per_prompt + return OmniDiffusionSamplingParams(**kwargs) + + +def _run_variant(args: argparse.Namespace, config: VariantConfig) -> VariantRun: + from vllm_omni.entrypoints.omni import Omni + from vllm_omni.platforms import current_omni_platform + + omni = Omni(**_build_omni_kwargs(args, config)) + measured_results: list[Any] = [] + generation_times: list[float] = [] + peak_memory_mb: list[float | None] = [] + try: + for _ in range(args.warmup_runs): + sampling_params = _build_sampling_params(args, args.seed) + omni.generate( + {"prompt": args.prompt, "negative_prompt": args.negative_prompt}, + sampling_params, + use_tqdm=False, + ) + + for _ in range(args.measure_runs): + sampling_params = _build_sampling_params(args, args.seed) + start = time.perf_counter() + outputs = omni.generate( + {"prompt": args.prompt, "negative_prompt": args.negative_prompt}, + sampling_params, + use_tqdm=False, + ) + current_omni_platform.synchronize() + generation_times.append(time.perf_counter() - start) + result = _extract_inner_output(outputs) + peak_memory_mb.append(_request_peak_memory_mb(result)) + measured_results.append(result) + finally: + omni.close() + del omni + gc.collect() + current_omni_platform.empty_cache() + + if not measured_results: + raise RuntimeError("No measured results were produced.") + return VariantRun( + label=config.label, + result=measured_results[-1], + generation_times_s=generation_times, + peak_memory_mb=peak_memory_mb, + ) + + +def _get_output_frames(result: Any) -> list[Any]: + images = list(getattr(result, "images", []) or []) + if len(images) == 1 and isinstance(images[0], list): + return list(images[0]) + return images + + +def _save_outputs(result: Any, output_dir: Path, label: str, task: str, fps: int) -> list[str]: + frames = _get_output_frames(result) + saved: list[str] = [] + if not frames: + return saved + variant_dir = output_dir / label + variant_dir.mkdir(parents=True, exist_ok=True) + if task == "t2v": + from diffusers.utils import export_to_video + + video_path = variant_dir / "video.mp4" + export_to_video(frames, str(video_path), fps=fps) + saved.append(str(video_path)) + for idx in (0, len(frames) // 2, len(frames) - 1): + path = variant_dir / f"frame_{idx}.png" + frames[idx].save(path) + saved.append(str(path)) + return saved + for idx, image in enumerate(frames): + path = variant_dir / f"image_{idx}.png" + image.save(path) + saved.append(str(path)) + return saved + + +def _run_summary(run: VariantRun) -> dict[str, Any]: + valid_peak = [value for value in run.peak_memory_mb if value is not None] + return { + "generation_time_s": run.generation_times_s[-1], + "avg_generation_time_s": sum(run.generation_times_s) / len(run.generation_times_s), + "per_run_generation_time_s": run.generation_times_s, + "peak_memory_mb": valid_peak[-1] if valid_peak else None, + "avg_peak_memory_mb": sum(valid_peak) / len(valid_peak) if valid_peak else None, + "max_peak_memory_mb": max(valid_peak) if valid_peak else None, + "per_run_peak_memory_mb": run.peak_memory_mb, + } + + +def _to_jsonable(result: dict[str, Any]) -> dict[str, Any]: + return json.loads(json.dumps(result, allow_nan=True)) + + +def run_comparison(args: argparse.Namespace) -> dict[str, Any]: + if args.warmup_runs < 0: + raise ValueError("--warmup-runs must be >= 0.") + if args.measure_runs <= 0: + raise ValueError("--measure-runs must be >= 1.") + + reference_config = _build_variant_config(args, "reference") + candidate_config = _build_variant_config(args, "candidate") + + reference_run = _run_variant(args, reference_config) + candidate_run = _run_variant(args, candidate_config) + + save_output_paths: dict[str, list[str]] = {} + if args.save_output_dir: + save_dir = Path(args.save_output_dir).expanduser().resolve() + save_output_paths["reference"] = _save_outputs(reference_run.result, save_dir, "reference", args.task, args.fps) + save_output_paths["candidate"] = _save_outputs(candidate_run.result, save_dir, "candidate", args.task, args.fps) + + result = { + "model": args.model, + "prompt": args.prompt, + "negative_prompt": args.negative_prompt, + "seed": args.seed, + "warmup_runs": args.warmup_runs, + "measure_runs": args.measure_runs, + "variant_configs": { + "reference": reference_config.__dict__, + "candidate": candidate_config.__dict__, + }, + "sampling_kwargs": { + "task": args.task, + "width": args.width, + "height": args.height, + "num_frames": args.num_frames if args.task == "t2v" else None, + "num_inference_steps": args.num_inference_steps, + "guidance_scale": args.guidance_scale, + "cfg_scale": args.cfg_scale, + "guidance_scale_2": args.guidance_scale_2, + "flow_shift": args.flow_shift, + "num_images_per_prompt": args.num_images_per_prompt, + "step_execution": args.step_execution if args.task == "t2i" else False, + }, + "metric_guidance": metric_guidance(), + "reference_generation": _run_summary(reference_run), + "candidate_generation": _run_summary(candidate_run), + "output_metrics": summarize_output_image_metrics( + _get_output_frames(reference_run.result), + _get_output_frames(candidate_run.result), + ), + } + if save_output_paths: + result["saved_outputs"] = save_output_paths + return _to_jsonable(result) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--task", choices=["t2i", "t2v"], default="t2i") + parser.add_argument("--model", default="Qwen/Qwen-Image", help="Default model for both variants.") + parser.add_argument("--model-class-name", help="Optional diffusion pipeline class override.") + parser.add_argument("--reference-model", help="Reference model path or HF ID. Defaults to --model.") + parser.add_argument("--candidate-model", help="Candidate model path or HF ID. Defaults to --model.") + parser.add_argument("--reference-quantization", help="Reference online quantization method.") + parser.add_argument("--candidate-quantization", help="Candidate online quantization method, e.g. fp8.") + parser.add_argument( + "--reference-quantization-config-json", + help="Reference quantization_config JSON string or file.", + ) + parser.add_argument( + "--candidate-quantization-config-json", + help="Candidate quantization_config JSON string or file.", + ) + parser.add_argument("--reference-ignored-layers", help="Comma-separated reference ignored layer patterns.") + parser.add_argument("--candidate-ignored-layers", help="Comma-separated candidate ignored layer patterns.") + parser.add_argument("--ignored-layers", help="Alias for --candidate-ignored-layers.") + parser.add_argument("--reference-gguf-model", help="Reference GGUF file or HF reference when using gguf.") + parser.add_argument("--candidate-gguf-model", help="Candidate GGUF file or HF reference when using gguf.") + parser.add_argument("--reference-load-format", default="default", help="Reference diffusion load format.") + parser.add_argument("--candidate-load-format", default="default", help="Candidate diffusion load format.") + parser.add_argument("--prompt", default="a cup of coffee on the table") + parser.add_argument("--negative-prompt") + parser.add_argument("--output-json", required=True) + parser.add_argument("--save-output-dir") + parser.add_argument("--width", type=int, default=1024) + parser.add_argument("--height", type=int, default=1024) + parser.add_argument("--num-frames", type=int, default=1, help="Number of frames for t2v.") + parser.add_argument("--fps", type=int, default=24, help="FPS for saved t2v videos.") + parser.add_argument("--frame-rate", type=float, help="Optional generation frame rate for t2v.") + parser.add_argument("--num-inference-steps", type=int, default=50) + parser.add_argument("--num-images-per-prompt", type=int, default=1) + parser.add_argument("--guidance-scale", type=float, default=4.0) + parser.add_argument("--cfg-scale", type=float, default=4.0, help="Qwen-Image true CFG scale.") + parser.add_argument("--guidance-scale-2", type=float) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--warmup-runs", type=int, default=0) + parser.add_argument("--measure-runs", type=int, default=1) + parser.add_argument("--ulysses-degree", type=int, default=1) + parser.add_argument("--ring-degree", type=int, default=1) + parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2]) + parser.add_argument("--tensor-parallel-size", type=int, default=1) + parser.add_argument("--vae-patch-parallel-size", type=int, default=1) + parser.add_argument("--enforce-eager", action="store_true") + parser.add_argument("--enable-cpu-offload", action="store_true") + parser.add_argument("--enable-layerwise-offload", action="store_true") + parser.add_argument("--vae-use-slicing", action="store_true") + parser.add_argument("--vae-use-tiling", action="store_true") + parser.add_argument("--flow-shift", type=float) + parser.add_argument( + "--step-execution", + action=argparse.BooleanOptionalAction, + default=False, + help=( + "Use the diffusion step scheduler for text-to-image models. Disabled by default so the tool follows the " + "normal offline inference path unless explicitly requested." + ), + ) + parser.add_argument("--timesteps-shift", type=float, default=1.0) + parser.add_argument("--cfg-schedule", choices=["constant", "linear"], default="constant") + parser.add_argument("--use-norm", action="store_true") + parser.add_argument( + "--use-system-prompt", + choices=[ + "None", + "dynamic", + "en_vanilla", + "en_recaption", + "en_think_recaption", + "en_unified", + "custom", + ], + ) + parser.add_argument("--system-prompt") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + output_json = Path(args.output_json).expanduser().resolve() + output_json.parent.mkdir(parents=True, exist_ok=True) + + result = run_comparison(args) + output_json.write_text(json.dumps(result, indent=2, sort_keys=True), encoding="utf-8") + + image0 = result["output_metrics"]["image0_metrics"] + summary = { + "output_json": str(output_json), + "reference_avg_generation_time_s": result["reference_generation"]["avg_generation_time_s"], + "candidate_avg_generation_time_s": result["candidate_generation"]["avg_generation_time_s"], + "image0_psnr_db": image0["psnr_db"], + "image0_mae": image0["mae"], + } + print(json.dumps(summary, indent=2)) + + +if __name__ == "__main__": + main() From 4e82192fd6c6bcbc7d40ac50e612b72883579ed3 Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Wed, 20 May 2026 09:15:53 +0800 Subject: [PATCH 18/38] [CI] optional --assert-baseline and update perf JSON baselines (#3695) Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: wangyu <53896905+yenuo26@users.noreply.github.com> Co-authored-by: Hongsheng Liu --- tests/dfx/conftest.py | 37 ++- tests/dfx/perf/scripts/run_benchmark.py | 95 ++----- .../perf/scripts/run_diffusion_benchmark.py | 21 +- .../tests/test_qwen3_omni_async_chunk.json | 196 +++++++++++--- .../tests/test_qwen3_omni_no_async_chunk.json | 104 ++++++-- .../perf/tests/test_qwen3_omni_vllm_text.json | 96 +++++-- .../test_qwen_image_edit_2509_vllm_omni.json | 44 ++-- .../tests/test_qwen_image_edit_vllm_omni.json | 48 ++-- .../test_qwen_image_layered_vllm_omni.json | 16 +- .../perf/tests/test_qwen_image_vllm_omni.json | 42 +-- tests/dfx/perf/tests/test_tts.json | 249 ++++++++++++++---- .../perf/tests/test_wan22_i2v_vllm_omni.json | 18 +- 12 files changed, 666 insertions(+), 300 deletions(-) diff --git a/tests/dfx/conftest.py b/tests/dfx/conftest.py index 9332a24f5..a22f760b2 100644 --- a/tests/dfx/conftest.py +++ b/tests/dfx/conftest.py @@ -279,6 +279,10 @@ def create_benchmark_indices( return indices +# Omni perf/stability: ``vllm bench serve --omni`` result template (used if bench exits without writing JSON). +OMNI_RESULT_TEMPLATE_PATH = Path(__file__).resolve().parent / "perf" / "scripts" / "result_omni_template.json" + + def _safe_filename_token(value: Any | None, *, default: str = "na") -> str: """Make a single path segment safe for result filenames on common filesystems.""" if value is None: @@ -289,7 +293,7 @@ def _safe_filename_token(value: Any | None, *, default: str = "na") -> str: return s if s else default -def _resolve_baseline_value( +def resolve_baseline_value( baseline_raw: Any, *, sweep_index: int | None, @@ -330,7 +334,7 @@ def _baseline_thresholds_for_step( ) -> dict[str, Any]: """Resolve baseline config to one threshold per metric for this iteration.""" return { - metric_name: _resolve_baseline_value( + metric_name: resolve_baseline_value( baseline_raw, sweep_index=sweep_index, max_concurrency=max_concurrency, @@ -354,7 +358,12 @@ def run_benchmark( random_input_len: Any | None = None, random_output_len: Any | None = None, ) -> dict[str, Any]: - """Run one ``vllm bench serve --omni`` iteration and return parsed metrics.""" + """Run one ``vllm bench serve --omni`` iteration and return parsed metrics. + + After ``vllm bench`` writes the JSON, ``result["baseline"]`` holds the resolved per-metric thresholds + (when ``baseline_config`` is provided). If the benchmark exits without writing a result file, + ``result_omni_template.json`` is used as a fallback. + """ current_dt = datetime.now().strftime("%Y%m%d-%H%M%S") ri = _safe_filename_token(random_input_len) ro = _safe_filename_token(random_output_len) @@ -403,8 +412,17 @@ def run_benchmark( result_dir = "./" result_path = os.path.join(result_dir, result_filename) - with open(result_path, encoding="utf-8") as f: - result = json.load(f) + if not os.path.exists(result_path): + with open(OMNI_RESULT_TEMPLATE_PATH, encoding="utf-8") as f: + template_result: dict[str, Any] = json.load(f) + Path(result_path).parent.mkdir(parents=True, exist_ok=True) + with open(result_path, "w", encoding="utf-8") as f: + json.dump(template_result, f, ensure_ascii=False, indent=2) + print(f"Benchmark result file not generated, fallback to template: {result_path}") + result = template_result + else: + with open(result_path, encoding="utf-8") as f: + result = json.load(f) if baseline_config: result["baseline"] = _baseline_thresholds_for_step( @@ -432,3 +450,12 @@ def pytest_addoption(parser: pytest.Parser) -> None: default=None, help=("Path to benchmark config JSON. Example: --test-config-file tests/dfx/perf/tests/test_tts.json"), ) + parser.addoption( + "--assert-baseline", + action="store_true", + default=False, + help=( + "When set, omni/diffusion perf runners compare metrics against the baseline block in the JSON config " + "(default: off)." + ), + ) diff --git a/tests/dfx/perf/scripts/run_benchmark.py b/tests/dfx/perf/scripts/run_benchmark.py index 74332889c..55ad2f8b2 100644 --- a/tests/dfx/perf/scripts/run_benchmark.py +++ b/tests/dfx/perf/scripts/run_benchmark.py @@ -12,12 +12,17 @@ from tests.dfx.conftest import ( create_unique_server_params, get_benchmark_params_for_server, load_configs, + resolve_baseline_value, run_benchmark, ) from tests.helpers.runtime import OmniServer pytestmark = [pytest.mark.full_model] +# Compare metrics to each test JSON ``baseline`` block only when pytest is run with ``--assert-baseline`` +# (registered in ``tests/dfx/conftest.py``; default: off). ``run_benchmark`` and ``_resolve_baseline_value`` are +# defined in the same module. + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" @@ -46,7 +51,6 @@ if CONFIG_FILE_PATH is None: CONFIG_FILE_PATH = _DEFAULT_CONFIG_FILE BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH) -OMNI_RESULT_TEMPLATE_PATH = Path(__file__).parent / "result_omni_template.json" DEPLOY_CONFIGS_DIR = Path(__file__).parent.parent / "deploy" @@ -87,16 +91,6 @@ def omni_server(request): print("OmniServer stopped") -def _safe_filename_token(value: Any | None, *, default: str = "na") -> str: - """Make a single path segment safe for result filenames on common filesystems.""" - if value is None: - return default - s = str(value).strip() - for bad in ("/", "\\", ":", "*", "?", '"', "<", ">", "|"): - s = s.replace(bad, "_") - return s if s else default - - benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping) @@ -126,76 +120,23 @@ def benchmark_params(request, omni_server): } -def _resolve_baseline_value( - baseline_raw: Any, - *, - sweep_index: int | None, - max_concurrency: Any = None, - request_rate: Any = None, -) -> Any: - """Pick the baseline threshold for this sweep step. - - Supported shapes per metric: - - **Scalar** — same threshold for every concurrency / QPS. - - **List** — aligned with ``max_concurrency`` / ``request_rate`` sweep order; use ``sweep_index``. - - **Dict** — keyed by concurrency or rate, e.g. ``{"1": 500, "4": 800}`` (keys are strings in JSON). - - For dict lookup, ``max_concurrency`` is preferred when both are set (concurrency sweep). - """ - if baseline_raw is None: - # If no baseline is set, the maximum value will be used. - return 100000 - if isinstance(baseline_raw, dict): - if max_concurrency is not None: - for key in (max_concurrency, str(max_concurrency)): - if key in baseline_raw: - return baseline_raw[key] - if request_rate is not None: - for key in (request_rate, str(request_rate)): - if key in baseline_raw: - return baseline_raw[key] - raise KeyError( - f"baseline dict has no key for max_concurrency={max_concurrency!r} " - f"or request_rate={request_rate!r}; keys={list(baseline_raw.keys())!r}" - ) - if isinstance(baseline_raw, (list, tuple)): - return baseline_raw[sweep_index] - return baseline_raw - - -def _baseline_thresholds_for_step( - baseline_data: dict[str, Any], - *, - sweep_index: int | None = None, - max_concurrency: Any = None, - request_rate: Any = None, -) -> dict[str, Any]: - """Resolve ``test.json`` ``baseline`` block to one threshold per metric (same as ``assert_result``).""" - return { - metric_name: _resolve_baseline_value( - baseline_raw, - sweep_index=sweep_index, - max_concurrency=max_concurrency, - request_rate=request_rate, - ) - for metric_name, baseline_raw in baseline_data.items() - } - - def assert_result( result, params, num_prompt, *, + assert_baseline: bool, sweep_index: int | None = None, - max_concurrency: Any = None, - request_rate: Any = None, + max_concurrency: Any | None = None, + request_rate: Any | None = None, ) -> None: assert result["completed"] == num_prompt, "Request failures exist" + if not assert_baseline: + return baseline_data = params.get("baseline", {}) for metric_name, baseline_raw in baseline_data.items(): current_value = result[metric_name] - baseline_value = _resolve_baseline_value( + baseline_value = resolve_baseline_value( baseline_raw, sweep_index=sweep_index, max_concurrency=max_concurrency, @@ -214,7 +155,7 @@ def assert_result( @pytest.mark.benchmark @pytest.mark.parametrize("omni_server", test_params, indirect=True) @pytest.mark.parametrize("benchmark_params", benchmark_indices, indirect=True) -def test_performance_benchmark(omni_server, benchmark_params): +def test_performance_benchmark(omni_server, benchmark_params, request): test_name = benchmark_params["test_name"] params = benchmark_params["params"] dataset_name = params.get("dataset_name", "") @@ -226,6 +167,8 @@ def test_performance_benchmark(omni_server, benchmark_params): print(f"Running benchmark for model: {model}") print(f"Benchmark parameters: {benchmark_params}") + assert_baseline = request.config.getoption("--assert-baseline", default=False) + def to_list(value, default=None): if value is None: return [] if default is None else [default] @@ -264,7 +207,7 @@ def test_performance_benchmark(omni_server, benchmark_params): elif not isinstance(value, bool): args.extend([arg_name, str(value)]) - # QPS test (sweep_index aligns with qps_list / num_prompt_list for this loop) + # QPS / request-rate sweep for i, (qps, num_prompt) in enumerate(zip(qps_list, num_prompt_list)): args = args + ["--request-rate", str(qps), "--num-prompts", str(num_prompt)] result = run_benchmark( @@ -283,12 +226,13 @@ def test_performance_benchmark(omni_server, benchmark_params): assert_result( result, params, - num_prompt=num_prompt, + num_prompt, + assert_baseline=assert_baseline, sweep_index=i, request_rate=qps, ) - # concurrency test (sweep_index aligns with max_concurrency_list for separate thresholds per concurrency) + # concurrency test for i, (concurrency, num_prompt) in enumerate(zip(max_concurrency_list, num_prompt_list)): args = args + ["--max-concurrency", str(concurrency), "--num-prompts", str(num_prompt), "--request-rate", "inf"] result = run_benchmark( @@ -307,7 +251,8 @@ def test_performance_benchmark(omni_server, benchmark_params): assert_result( result, params, - num_prompt=num_prompt, + num_prompt, + assert_baseline=assert_baseline, sweep_index=i, max_concurrency=concurrency, ) diff --git a/tests/dfx/perf/scripts/run_diffusion_benchmark.py b/tests/dfx/perf/scripts/run_diffusion_benchmark.py index 9059904e4..22e844427 100644 --- a/tests/dfx/perf/scripts/run_diffusion_benchmark.py +++ b/tests/dfx/perf/scripts/run_diffusion_benchmark.py @@ -11,6 +11,8 @@ This runner separates two concepts: A config JSON file is REQUIRED via --test-config-file: pytest run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json +Optional: ``--assert-baseline`` compares metrics to the ``baseline`` block in each benchmark entry (default: off). + All benchmark results for a session are consolidated into a single JSON file under BENCHMARK_RESULT_DIR (override via the DIFFUSION_BENCHMARK_DIR environment variable). Each entry in the file contains the test metadata (test_name, endpoint, benchmark_params, @@ -652,12 +654,14 @@ def run_benchmark( # --------------------------------------------------------------------------- -def assert_result(result: dict[str, Any], params: dict[str, Any]) -> None: - """Assert that benchmark metrics satisfy the configured baselines.""" +def assert_result(result: dict[str, Any], params: dict[str, Any], *, assert_baseline: bool) -> None: + """Assert request completion; optionally assert metrics against ``baseline`` in *params*.""" num_prompts = params.get("num-prompts", 10) completed = result.get("completed_requests", result.get("completed", 0)) assert completed == num_prompts, f"Expected {num_prompts} completed requests, got {completed}" + if not assert_baseline: + return if params.get("skip-performance-assertion", False): print("Skipping performance assertions.") return @@ -699,16 +703,14 @@ def _resolve_benchmark_endpoint(server_cfg: dict[str, Any], params: dict[str, An indirect=True, ) @pytest.mark.parametrize("benchmark_params", benchmark_indices, indirect=True) -def test_diffusion_performance_benchmark(diffusion_server, benchmark_params): - """Run the diffusion performance benchmark and assert against baselines. +def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, request): + """Run the diffusion performance benchmark and verify request completion. One server is started per unique parallel configuration (module scope). For each server, all benchmark parameter sets defined in the config JSON - are executed sequentially; results are asserted against the baselines. + are executed sequentially; metrics are recorded to the aggregated result file. - Tracked metrics: - - throughput_qps (higher is better) - - latency_p50, latency_p99 (lower is better) + Pass ``--assert-baseline`` to compare ``throughput_qps``, ``latency_mean``, etc. to the JSON ``baseline`` block. """ test_name = benchmark_params["test_name"] params = benchmark_params["params"] @@ -744,4 +746,5 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params): print(f"\n Aggregated results: {AGGREGATED_RESULT_FILE}") print("=" * 60) - assert_result(result, params) + assert_baseline = request.config.getoption("--assert-baseline", default=False) + assert_result(result, params, assert_baseline=assert_baseline) diff --git a/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json b/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json index 98e311748..bc954cee5 100644 --- a/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json +++ b/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json @@ -3,31 +3,67 @@ "test_name": "test_qwen3_omni_chunk", "server_params": { "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", - "extra_cli_args": ["--async-chunk"] + "extra_cli_args": [ + "--async-chunk" + ] }, "benchmark_params": [ { "dataset_name": "random", "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", - "num_prompts": [4, 16, 32, 64, 128], - "max_concurrency": [1, 4, 8, 16, 32], + "num_prompts": [ + 4, + 16, + 32, + 64, + 128 + ], + "max_concurrency": [ + 1, + 4, + 8, + 16, + 32 + ], "random_input_len": 2500, "random_output_len": 900, "ignore_eos": true, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000], - "mean_audio_ttfp_ms": [1000, 3000, 5000, 7000, 9000], - "mean_audio_rtf": [0.2, 0.35, 0.6, 0.85, 0.9] + "mean_ttft_ms": [ + 400, + 300, + 400, + 600, + 900 + ], + "mean_e2el_ms": [ + 26000, + 43000, + 52000, + 72000, + 109000 + ], + "mean_audio_rtf": [ + 0.15, + 0.30, + 0.35, + 0.50, + 0.75 + ] } }, { "dataset_name": "random-mm", "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", - "num_prompts": [10], - "request_rate": [0.1], + "num_prompts": [ + 10 + ], + "request_rate": [ + 0.1 + ], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, @@ -42,17 +78,27 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [2000], - "mean_audio_ttfp_ms": [2000], - "mean_audio_rtf": [0.25] + "mean_ttft_ms": [ + 600 + ], + "mean_e2el_ms": [ + 6000 + ], + "mean_audio_rtf": [ + 0.20 + ] } }, { "dataset_name": "random-mm", "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", - "num_prompts": [40], - "request_rate": [0.5], + "num_prompts": [ + 40 + ], + "request_rate": [ + 0.5 + ], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, @@ -69,17 +115,27 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [6000], - "mean_audio_ttfp_ms": [6000], - "mean_audio_rtf": [0.7] + "mean_ttft_ms": [ + 300 + ], + "mean_e2el_ms": [ + 7000 + ], + "mean_audio_rtf": [ + 0.25 + ] } }, { "dataset_name": "random-mm", "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", - "num_prompts": [100], - "request_rate": [1.0], + "num_prompts": [ + 100 + ], + "request_rate": [ + 1.0 + ], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, @@ -98,40 +154,79 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [12000], - "mean_audio_ttfp_ms": [12000], - "mean_audio_rtf": [1.0] + "mean_ttft_ms": [ + 700 + ], + "mean_e2el_ms": [ + 11000 + ], + "mean_audio_rtf": [ + 0.35 + ] } }, { "dataset_name": "random", "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", - "num_prompts": [4, 16, 32, 64, 128], - "max_concurrency": [1, 4, 8, 16, 32], + "num_prompts": [ + 4, + 16, + 32, + 64, + 128 + ], + "max_concurrency": [ + 1, + 4, + 8, + 16, + 32 + ], "random_input_len": 2500, "random_output_len": 900, "ignore_eos": true, "extra_body": { - "modalities": ["text"] + "modalities": [ + "text" + ] }, "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { - "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000] + "mean_ttft_ms": [ + 100, + 300, + 500, + 800, + 1500 + ], + "mean_e2el_ms": [ + 5000, + 8000, + 10000, + 14000, + 20000 + ] } }, { "dataset_name": "random-mm", "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", - "num_prompts": [10], - "request_rate": [0.1], + "num_prompts": [ + 10 + ], + "request_rate": [ + 0.1 + ], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, "ignore_eos": true, "extra_body": { - "modalities": ["text"] + "modalities": [ + "text" + ] }, "random_mm_base_items_per_request": 1, "random_mm_num_mm_items_range_ratio": 0.5, @@ -143,21 +238,32 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { - "mean_ttft_ms": [2000] + "mean_ttft_ms": [ + 500 + ], + "mean_e2el_ms": [ + 1000 + ] } }, { "dataset_name": "random-mm", "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", - "num_prompts": [40], - "request_rate": [0.5], + "num_prompts": [ + 40 + ], + "request_rate": [ + 0.5 + ], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, "ignore_eos": true, "extra_body": { - "modalities": ["text"] + "modalities": [ + "text" + ] }, "random_mm_base_items_per_request": 2, "random_mm_num_mm_items_range_ratio": 0.5, @@ -171,21 +277,32 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { - "mean_ttft_ms": [6000] + "mean_ttft_ms": [ + 200 + ], + "mean_e2el_ms": [ + 800 + ] } }, { "dataset_name": "random-mm", "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", - "num_prompts": [100], - "request_rate": [1.0], + "num_prompts": [ + 100 + ], + "request_rate": [ + 1.0 + ], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, "ignore_eos": true, "extra_body": { - "modalities": ["text"] + "modalities": [ + "text" + ] }, "random_mm_base_items_per_request": 3, "random_mm_num_mm_items_range_ratio": 0.5, @@ -201,7 +318,12 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { - "mean_ttft_ms": [6000] + "mean_ttft_ms": [ + 600 + ], + "mean_e2el_ms": [ + 1300 + ] } } ] diff --git a/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json b/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json index b91a8a89b..e57acecba 100644 --- a/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json +++ b/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json @@ -3,31 +3,67 @@ "test_name": "test_qwen3_omni", "server_params": { "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", - "extra_cli_args": ["--no-async-chunk"] + "extra_cli_args": [ + "--no-async-chunk" + ] }, "benchmark_params": [ { "dataset_name": "random", "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", - "num_prompts": [4, 16, 32, 64, 128], - "max_concurrency": [1, 4, 8, 16, 32], + "num_prompts": [ + 4, + 16, + 32, + 64, + 128 + ], + "max_concurrency": [ + 1, + 4, + 8, + 16, + 32 + ], "random_input_len": 2500, "random_output_len": 900, "ignore_eos": true, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000], - "mean_audio_ttfp_ms": [30000, 60000, 90000, 120000, 150000], - "mean_audio_rtf": [0.35, 0.45, 0.55, 0.65, 0.75] + "mean_ttft_ms": [ + 400, + 300, + 400, + 600, + 900 + ], + "mean_e2el_ms": [ + 32000, + 34000, + 45000, + 61000, + 90000 + ], + "mean_audio_rtf": [ + 0.20, + 0.25, + 0.30, + 0.45, + 0.60 + ] } }, { "dataset_name": "random-mm", "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", - "num_prompts": [10], - "request_rate": [0.1], + "num_prompts": [ + 10 + ], + "request_rate": [ + 0.1 + ], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, @@ -42,17 +78,27 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [2000], - "mean_audio_ttfp_ms": [10000], - "mean_audio_rtf": [0.25] + "mean_ttft_ms": [ + 500 + ], + "mean_e2el_ms": [ + 6000 + ], + "mean_audio_rtf": [ + 0.20 + ] } }, { "dataset_name": "random-mm", "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", - "num_prompts": [40], - "request_rate": [0.5], + "num_prompts": [ + 40 + ], + "request_rate": [ + 0.5 + ], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, @@ -69,17 +115,27 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [6000], - "mean_audio_ttfp_ms": [15000], - "mean_audio_rtf": [0.45] + "mean_ttft_ms": [ + 300 + ], + "mean_e2el_ms": [ + 7000 + ], + "mean_audio_rtf": [ + 0.20 + ] } }, { "dataset_name": "random-mm", "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", - "num_prompts": [100], - "request_rate": [1.0], + "num_prompts": [ + 100 + ], + "request_rate": [ + 1.0 + ], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, @@ -98,9 +154,15 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [12000], - "mean_audio_ttfp_ms": [18000], - "mean_audio_rtf": [0.9] + "mean_ttft_ms": [ + 700 + ], + "mean_e2el_ms": [ + 9000 + ], + "mean_audio_rtf": [ + 0.30 + ] } } ] diff --git a/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json b/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json index aa2f79033..2f580aac7 100644 --- a/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json +++ b/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json @@ -4,7 +4,9 @@ "server_params": { "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", "use_omni": false, - "extra_cli_args": ["--no-enable-prefix-caching"] + "extra_cli_args": [ + "--no-enable-prefix-caching" + ] }, "benchmark_params": [ { @@ -12,8 +14,20 @@ "backend": "vllm", "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", "endpoint": "/v1/completions", - "num_prompts": [4, 16, 32, 64, 128], - "max_concurrency": [1, 4, 8, 16, 32], + "num_prompts": [ + 4, + 16, + 32, + 64, + 128 + ], + "max_concurrency": [ + 1, + 4, + 8, + 16, + 32 + ], "random_input_len": 2500, "random_output_len": 900, "temperature": 0.4, @@ -23,11 +37,26 @@ "repetition_penalty": 1.05, "ignore_eos": true, "extra_body": { - "modalities": ["text"] + "modalities": [ + "text" + ] }, "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { - "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000] + "mean_ttft_ms": [ + 300, + 300, + 1000, + 600, + 800 + ], + "mean_e2el_ms": [ + 6000, + 8000, + 11000, + 13000, + 17000 + ] } }, { @@ -35,8 +64,12 @@ "backend": "vllm", "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", "endpoint": "/v1/completions", - "num_prompts": [10], - "request_rate": [0.1], + "num_prompts": [ + 10 + ], + "request_rate": [ + 0.1 + ], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, @@ -47,7 +80,9 @@ "repetition_penalty": 1.05, "ignore_eos": true, "extra_body": { - "modalities": ["text"] + "modalities": [ + "text" + ] }, "random_mm_base_items_per_request": 1, "random_mm_num_mm_items_range_ratio": 0.5, @@ -59,7 +94,12 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { - "mean_ttft_ms": [2000] + "mean_ttft_ms": [ + 30 + ], + "mean_e2el_ms": [ + 540 + ] } }, { @@ -67,8 +107,12 @@ "backend": "vllm", "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", "endpoint": "/v1/completions", - "num_prompts": [40], - "request_rate": [0.5], + "num_prompts": [ + 40 + ], + "request_rate": [ + 0.5 + ], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, @@ -79,7 +123,9 @@ "repetition_penalty": 1.05, "ignore_eos": true, "extra_body": { - "modalities": ["text"] + "modalities": [ + "text" + ] }, "random_mm_base_items_per_request": 2, "random_mm_num_mm_items_range_ratio": 0.5, @@ -93,7 +139,12 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { - "mean_ttft_ms": [6000] + "mean_ttft_ms": [ + 30 + ], + "mean_e2el_ms": [ + 600 + ] } }, { @@ -101,8 +152,12 @@ "backend": "vllm", "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", "endpoint": "/v1/completions", - "num_prompts": [100], - "request_rate": [1.0], + "num_prompts": [ + 100 + ], + "request_rate": [ + 1.0 + ], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, @@ -113,7 +168,9 @@ "repetition_penalty": 1.05, "ignore_eos": true, "extra_body": { - "modalities": ["text"] + "modalities": [ + "text" + ] }, "random_mm_base_items_per_request": 3, "random_mm_num_mm_items_range_ratio": 0.5, @@ -129,7 +186,12 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { - "mean_ttft_ms": [6000] + "mean_ttft_ms": [ + 40 + ], + "mean_e2el_ms": [ + 700 + ] } } ] diff --git a/tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json index 7d1fbbfa7..d13f37053 100644 --- a/tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json +++ b/tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json @@ -22,10 +22,10 @@ "num-input-images": 2, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.05, - "latency_mean": 18, - "peak_memory_mb_max": 78500, - "peak_memory_mb_mean": 78500 + "throughput_qps": 0.06, + "latency_mean": 20, + "peak_memory_mb_max": 57000, + "peak_memory_mb_mean": 57000 } }, { @@ -41,9 +41,9 @@ "enable-negative-prompt": true, "baseline": { "throughput_qps": 0.01, - "latency_mean": 70, - "peak_memory_mb_max": 81000, - "peak_memory_mb_mean": 81000 + "latency_mean": 60, + "peak_memory_mb_max": 67000, + "peak_memory_mb_mean": 67000 } } ] @@ -75,10 +75,10 @@ "num-input-images": 2, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.1, - "latency_mean": 12, - "peak_memory_mb_max": 69000, - "peak_memory_mb_mean": 69000 + "throughput_qps": 0.15, + "latency_mean": 10, + "peak_memory_mb_max": 55000, + "peak_memory_mb_mean": 55000 } }, { @@ -93,10 +93,10 @@ "num-input-images": 2, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.03, - "latency_mean": 28, - "peak_memory_mb_max": 69000, - "peak_memory_mb_mean": 69000 + "throughput_qps": 0.05, + "latency_mean": 20, + "peak_memory_mb_max": 55000, + "peak_memory_mb_mean": 55000 } } ] @@ -138,10 +138,10 @@ "num-input-images": 2, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.10, - "latency_mean": 12, - "peak_memory_mb_max": 73000, - "peak_memory_mb_mean": 73000 + "throughput_qps": 0.21, + "latency_mean": 10, + "peak_memory_mb_max": 57000, + "peak_memory_mb_mean": 57000 } }, { @@ -156,10 +156,10 @@ "num-input-images": 2, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.05, + "throughput_qps": 0.08, "latency_mean": 20, - "peak_memory_mb_max": 81000, - "peak_memory_mb_mean": 81000 + "peak_memory_mb_max": 66000, + "peak_memory_mb_mean": 66000 } } ] diff --git a/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json index f68201db5..8ae370df6 100644 --- a/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json +++ b/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json @@ -21,10 +21,10 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.05, - "latency_mean": 15.0, - "peak_memory_mb_max": 72500, - "peak_memory_mb_mean": 72500 + "throughput_qps": 0.1, + "latency_mean": 10, + "peak_memory_mb_max": 57000, + "peak_memory_mb_mean": 57000 } }, { @@ -38,10 +38,10 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.01, - "latency_mean": 65.6, - "peak_memory_mb_max": 80777, - "peak_memory_mb_mean": 80777 + "throughput_qps": 0.02, + "latency_mean": 50, + "peak_memory_mb_max": 65000, + "peak_memory_mb_mean": 65000 } } ] @@ -72,10 +72,10 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.10, - "latency_mean": 7.2, - "peak_memory_mb_max": 68100, - "peak_memory_mb_mean": 68100 + "throughput_qps": 0.21, + "latency_mean": 10, + "peak_memory_mb_max": 56000, + "peak_memory_mb_mean": 56000 } }, { @@ -89,10 +89,10 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.03, - "latency_mean": 24.0, - "peak_memory_mb_max": 68100, - "peak_memory_mb_mean": 68100 + "throughput_qps": 0.06, + "latency_mean": 20, + "peak_memory_mb_max": 56000, + "peak_memory_mb_mean": 56000 } } ] @@ -133,10 +133,10 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.1, - "latency_mean": 6.5, - "peak_memory_mb_max": 72600, - "peak_memory_mb_mean": 72600 + "throughput_qps": 0.26, + "latency_mean": 10, + "peak_memory_mb_max": 58000, + "peak_memory_mb_mean": 58000 } }, { @@ -150,10 +150,10 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.05, - "latency_mean": 16.0, - "peak_memory_mb_max": 81000, - "peak_memory_mb_mean": 81000 + "throughput_qps": 0.1, + "latency_mean": 10, + "peak_memory_mb_max": 67000, + "peak_memory_mb_mean": 67000 } } ] diff --git a/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json index 3cf13509c..c0cd3b7c0 100644 --- a/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json +++ b/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json @@ -21,10 +21,10 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.02, - "latency_mean": 40.0, - "peak_memory_mb_max": 70000, - "peak_memory_mb_mean": 70000 + "throughput_qps": 0.05, + "latency_mean": 20, + "peak_memory_mb_max": 63000, + "peak_memory_mb_mean": 63000 } }, { @@ -38,10 +38,10 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.005, - "latency_mean": 80.0, - "peak_memory_mb_max": 70000, - "peak_memory_mb_mean": 70000 + "throughput_qps": 0.03, + "latency_mean": 30, + "peak_memory_mb_max": 63000, + "peak_memory_mb_mean": 63000 } } ] diff --git a/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json index cdd0cac2c..e05f1b945 100644 --- a/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json +++ b/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json @@ -21,9 +21,9 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.30, - "latency_mean": 3.50, - "peak_memory_mb_mean": 67000 + "throughput_qps": 0.35, + "latency_mean": 3, + "peak_memory_mb_mean": 56000 } }, { @@ -37,9 +37,9 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.037, - "latency_mean": 27.0, - "peak_memory_mb_mean": 74000 + "throughput_qps": 0.03, + "latency_mean": 26, + "peak_memory_mb_mean": 65000 } } ] @@ -67,9 +67,9 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.30, - "latency_mean": 3.50, - "peak_memory_mb_mean": 67000 + "throughput_qps": 0.35, + "latency_mean": 3, + "peak_memory_mb_mean": 56000 } }, { @@ -83,9 +83,9 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.037, - "latency_mean": 27.0, - "peak_memory_mb_mean": 74000 + "throughput_qps": 0.03, + "latency_mean": 26, + "peak_memory_mb_mean": 65000 } } ] @@ -116,9 +116,9 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.1, - "latency_mean": 9.1, - "peak_memory_mb_mean": 61000 + "throughput_qps": 0.11, + "latency_mean": 9, + "peak_memory_mb_mean": 55000 } } ] @@ -159,9 +159,9 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.1, - "latency_mean": 3.50, - "peak_memory_mb_mean": 67000 + "throughput_qps": 0.46, + "latency_mean": 3, + "peak_memory_mb_mean": 56000 } }, { @@ -175,9 +175,9 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.1, - "latency_mean": 6.15, - "peak_memory_mb_mean": 74000 + "throughput_qps": 0.18, + "latency_mean": 6, + "peak_memory_mb_mean": 67000 } } ] diff --git a/tests/dfx/perf/tests/test_tts.json b/tests/dfx/perf/tests/test_tts.json index db197c34b..733c40963 100644 --- a/tests/dfx/perf/tests/test_tts.json +++ b/tests/dfx/perf/tests/test_tts.json @@ -4,7 +4,9 @@ "server_params": { "model": "Qwen/Qwen3-TTS-12Hz-1.7B-Base", "stage_overrides": { - "1": {"devices": "1"} + "1": { + "devices": "1" + } } }, "benchmark_params": [ @@ -15,13 +17,21 @@ "backend": "openai-audio-speech", "endpoint": "/v1/audio/speech", "dataset_path": "linyueqian/seed-tts-eval-subset", - "num_prompts": [20], - "max_concurrency": [1], + "num_prompts": [ + 20 + ], + "max_concurrency": [ + 1 + ], "seed_tts_locale": "en", "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [350], - "median_audio_rtf": [0.25] + "median_audio_ttfp_ms": [ + 3000 + ], + "median_audio_rtf": [ + 0.6 + ] } }, { @@ -31,14 +41,34 @@ "backend": "openai-audio-speech", "endpoint": "/v1/audio/speech", "dataset_path": "linyueqian/seed-tts-eval-subset", - "num_prompts": [80, 128, 128], - "max_concurrency": [8, 16, 64], + "num_prompts": [ + 80, + 128, + 128 + ], + "max_concurrency": [ + 8, + 16, + 64 + ], "seed_tts_locale": "en", "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [3500, 30000, 7000], - "median_audio_rtf": [0.75, 10.0, 1.5], - "audio_throughput": [10.0, 1.0, 40.0] + "median_audio_ttfp_ms": [ + 1000, + 3000, + 14000 + ], + "median_audio_rtf": [ + 0.70, + 1.15, + 3.90 + ], + "audio_throughput": [ + 12, + 14, + 14 + ] } }, { @@ -49,13 +79,19 @@ "backend": "openai-audio-speech", "endpoint": "/v1/audio/speech", "dataset_path": "linyueqian/seed-tts-eval-subset", - "num_prompts": [200], - "max_concurrency": [4], + "num_prompts": [ + 200 + ], + "max_concurrency": [ + 4 + ], "seed_tts_locale": "en", "seed_tts_wer_eval": true, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_audio_rtf": [0.45] + "mean_audio_rtf": [ + 0.45 + ] } } ] @@ -65,7 +101,9 @@ "server_params": { "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", "stage_overrides": { - "1": {"devices": "1"} + "1": { + "devices": "1" + } } }, "benchmark_params": [ @@ -76,14 +114,26 @@ "backend": "openai-audio-speech", "endpoint": "/v1/audio/speech", "dataset_path": "benchmarks/build_dataset/seed_tts_smoke", - "num_prompts": [20], - "max_concurrency": [1], + "num_prompts": [ + 20 + ], + "max_concurrency": [ + 1 + ], "seed_tts_locale": "en", - "extra_body": {"voice": "Vivian", "language": "English", "task_type": "CustomVoice"}, + "extra_body": { + "voice": "Vivian", + "language": "English", + "task_type": "CustomVoice" + }, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [150], - "median_audio_rtf": [0.15] + "median_audio_ttfp_ms": [ + 2000 + ], + "median_audio_rtf": [ + 0.40 + ] } }, { @@ -93,15 +143,39 @@ "backend": "openai-audio-speech", "endpoint": "/v1/audio/speech", "dataset_path": "benchmarks/build_dataset/seed_tts_smoke", - "num_prompts": [80, 128, 128], - "max_concurrency": [8, 16, 64], + "num_prompts": [ + 80, + 128, + 128 + ], + "max_concurrency": [ + 8, + 16, + 64 + ], "seed_tts_locale": "en", - "extra_body": {"voice": "Vivian", "language": "English", "task_type": "CustomVoice"}, + "extra_body": { + "voice": "Vivian", + "language": "English", + "task_type": "CustomVoice" + }, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [1500, 30000, 3000], - "median_audio_rtf": [0.30, 10.0, 0.60], - "audio_throughput": [30.0, 1.0, 100.0] + "median_audio_ttfp_ms": [ + 300, + 1000, + 5000 + ], + "median_audio_rtf": [ + 0.25, + 0.35, + 1.10 + ], + "audio_throughput": [ + 30, + 40, + 40 + ] } }, { @@ -112,14 +186,24 @@ "backend": "openai-audio-speech", "endpoint": "/v1/audio/speech", "dataset_path": "benchmarks/build_dataset/seed_tts_smoke", - "num_prompts": [200], - "max_concurrency": [4], + "num_prompts": [ + 200 + ], + "max_concurrency": [ + 4 + ], "seed_tts_locale": "en", - "extra_body": {"voice": "Vivian", "language": "English", "task_type": "CustomVoice"}, + "extra_body": { + "voice": "Vivian", + "language": "English", + "task_type": "CustomVoice" + }, "seed_tts_wer_eval": true, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_audio_rtf": [0.35] + "mean_audio_rtf": [ + 0.35 + ] } }, { @@ -129,15 +213,29 @@ "backend": "openai-audio-speech", "endpoint": "/v1/audio/speech", "dataset_path": "benchmarks/build_dataset/seed_tts_smoke", - "num_prompts": [100], - "request_rate": [2.0], + "num_prompts": [ + 100 + ], + "request_rate": [ + 2.0 + ], "seed_tts_locale": "en", - "extra_body": {"voice": "Vivian", "language": "English", "task_type": "CustomVoice"}, + "extra_body": { + "voice": "Vivian", + "language": "English", + "task_type": "CustomVoice" + }, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [3000], - "median_audio_rtf": [0.25], - "audio_throughput": [4.0] + "median_audio_ttfp_ms": [ + 100 + ], + "median_audio_rtf": [ + 0.20 + ], + "audio_throughput": [ + 10 + ] } }, { @@ -147,14 +245,25 @@ "backend": "openai-audio-speech", "endpoint": "/v1/audio/speech", "dataset_path": "benchmarks/build_dataset/seed_tts_design", - "num_prompts": [20], - "max_concurrency": [1], + "num_prompts": [ + 20 + ], + "max_concurrency": [ + 1 + ], "seed_tts_locale": "en", - "extra_body": {"task_type": "VoiceDesign", "language": "English"}, + "extra_body": { + "task_type": "VoiceDesign", + "language": "English" + }, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [150], - "median_audio_rtf": [0.15] + "median_audio_ttfp_ms": [ + 100 + ], + "median_audio_rtf": [ + 0.15 + ] } }, { @@ -164,15 +273,38 @@ "backend": "openai-audio-speech", "endpoint": "/v1/audio/speech", "dataset_path": "benchmarks/build_dataset/seed_tts_design", - "num_prompts": [80, 128, 128], - "max_concurrency": [8, 16, 64], + "num_prompts": [ + 80, + 128, + 128 + ], + "max_concurrency": [ + 8, + 16, + 64 + ], "seed_tts_locale": "en", - "extra_body": {"task_type": "VoiceDesign", "language": "English"}, + "extra_body": { + "task_type": "VoiceDesign", + "language": "English" + }, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [1500, 30000, 3500], - "median_audio_rtf": [0.35, 10.0, 0.70], - "audio_throughput": [25.0, 1.0, 80.0] + "median_audio_ttfp_ms": [ + 100, + 1000, + 5000 + ], + "median_audio_rtf": [ + 0.25, + 0.40, + 1.20 + ], + "audio_throughput": [ + 30, + 40, + 40 + ] } }, { @@ -182,15 +314,28 @@ "backend": "openai-audio-speech", "endpoint": "/v1/audio/speech", "dataset_path": "benchmarks/build_dataset/seed_tts_design", - "num_prompts": [100], - "request_rate": [2.0], + "num_prompts": [ + 100 + ], + "request_rate": [ + 2.0 + ], "seed_tts_locale": "en", - "extra_body": {"task_type": "VoiceDesign", "language": "English"}, + "extra_body": { + "task_type": "VoiceDesign", + "language": "English" + }, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [3500], - "median_audio_rtf": [0.30], - "audio_throughput": [4.0] + "median_audio_ttfp_ms": [ + 100 + ], + "median_audio_rtf": [ + 0.20 + ], + "audio_throughput": [ + 9 + ] } } ] diff --git a/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json b/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json index 58a17c980..d70a28efd 100644 --- a/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json +++ b/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json @@ -30,9 +30,9 @@ } ], "baseline": { - "throughput_qps": 0.034, - "latency_mean": 26.0, - "peak_memory_mb_mean": 80000 + "throughput_qps": 0.04, + "latency_mean": 30, + "peak_memory_mb_mean": 75000 } } ] @@ -72,9 +72,9 @@ } ], "baseline": { - "throughput_qps": 0.042, - "latency_mean": 21.6, - "peak_memory_mb_mean": 55300 + "throughput_qps": 0.06, + "latency_mean": 20, + "peak_memory_mb_mean": 48000 } }, { @@ -97,9 +97,9 @@ } ], "baseline": { - "throughput_qps": 0.0085, - "latency_mean": 101.6, - "peak_memory_mb_mean": 65200 + "throughput_qps": 0.01, + "latency_mean": 90, + "peak_memory_mb_mean": 56000 } } ] From 29179599fc464c6f527af0a1c044cd6135db1e5a Mon Sep 17 00:00:00 2001 From: "Y. Fisher" Date: Wed, 20 May 2026 09:55:33 +0800 Subject: [PATCH 19/38] [Feat] Enable VAE parallel in HunyuanImage3 (#3091) Signed-off-by: KexiongYu Signed-off-by: zzh <943967662@qq.com> Co-authored-by: zzh <943967662@qq.com> --- .../test_autoencoder_kl_hunyuan.py | 79 ++++++++++ .../autoencoders/autoencoder_kl_hunyuan.py | 145 ++++++++++++++++++ .../hunyuan_image3/pipeline_hunyuan_image3.py | 9 +- 3 files changed, 231 insertions(+), 2 deletions(-) create mode 100644 tests/diffusion/distributed/test_autoencoder_kl_hunyuan.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_hunyuan.py diff --git a/tests/diffusion/distributed/test_autoencoder_kl_hunyuan.py b/tests/diffusion/distributed/test_autoencoder_kl_hunyuan.py new file mode 100644 index 000000000..58ce14546 --- /dev/null +++ b/tests/diffusion/distributed/test_autoencoder_kl_hunyuan.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan import ( + DistributedAutoencoderKLHunyuan, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _DummyDistributedAutoencoderKLHunyuan(DistributedAutoencoderKLHunyuan): + def __init__(self): + torch.nn.Module.__init__(self) + self.tile_latent_min_size = 8 + self.tile_sample_min_size = 8 + self.tile_overlap_factor = 0.25 + self.use_spatial_tiling = False + + @property + def dtype(self) -> torch.dtype: + return torch.float32 + + def decoder(self, x: torch.Tensor) -> torch.Tensor: + return x + 10 + + def encoder(self, x: torch.Tensor) -> torch.Tensor: + return x + 20 + + def blend_v(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor: + return b + + def blend_h(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor: + return b + + +def test_hunyuan_vae_use_tiling_aliases_spatial_tiling(): + # Verify use_tiling property maps to use_spatial_tiling. + vae = _DummyDistributedAutoencoderKLHunyuan() + + assert not vae.use_tiling + + vae.use_tiling = True + + assert vae.use_spatial_tiling + + +def test_hunyuan_vae_decode_tiles_round_trip(): + # Validate decode tile split/exec/merge returns expected reconstructed tensor. + vae = _DummyDistributedAutoencoderKLHunyuan() + z = torch.arange(144, dtype=torch.float32).reshape(1, 1, 1, 12, 12) + + tile_tasks, grid_spec = vae.tile_split(z) + decoded_tiles = {task.grid_coord: vae.tile_exec(task) for task in tile_tasks} + output = vae.tile_merge(decoded_tiles, grid_spec) + + assert grid_spec.split_dims == (3, 4) + assert grid_spec.grid_shape == (2, 2) + assert grid_spec.tile_spec == {"blend_extent": 2, "row_limit": 6} + assert len(tile_tasks) == 4 + assert torch.equal(output, z + 10) + + +def test_hunyuan_vae_encode_tiles_round_trip(): + # Validate encode tile split/exec/merge returns expected latent tensor. + vae = _DummyDistributedAutoencoderKLHunyuan() + x = torch.arange(144, dtype=torch.float32).reshape(1, 1, 1, 12, 12) + + tile_tasks, grid_spec = vae.encode_tile_split(x) + encoded_tiles = {task.grid_coord: vae.encode_tile_exec(task) for task in tile_tasks} + output = vae.encode_tile_merge(encoded_tiles, grid_spec) + + assert grid_spec.split_dims == (3, 4) + assert grid_spec.grid_shape == (2, 2) + assert grid_spec.tile_spec == {"blend_extent": 2, "row_limit": 6} + assert len(tile_tasks) == 4 + assert torch.equal(output, x + 20) diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_hunyuan.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_hunyuan.py new file mode 100644 index 000000000..c0e3b4d6c --- /dev/null +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_hunyuan.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import torch +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.hunyuan_image3.autoencoder import AutoencoderKLConv3D + +logger = init_logger(__name__) + + +class DistributedAutoencoderKLHunyuan(AutoencoderKLConv3D, DistributedVaeMixin): + @classmethod + def from_config(cls, config: Any, **kwargs: Any): + model = super().from_config(config, **kwargs) + model.init_distributed() + return model + + @classmethod + def from_pretrained(cls, *args: Any, **kwargs: Any): + model = super().from_pretrained(*args, **kwargs) + model.init_distributed() + return model + + @property + def use_tiling(self) -> bool: + return self.use_spatial_tiling + + @use_tiling.setter + def use_tiling(self, use_tiling: bool) -> None: + self.use_spatial_tiling = use_tiling + + def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]: + _, _, _, height, width = z.shape + 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 = int(self.tile_sample_min_size - blend_extent) + + tiletask_list = [] + for i in range(0, height, overlap_size): + for j in range(0, width, 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[3] * tile.shape[4], + ) + ) + + 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={"blend_extent": blend_extent, "row_limit": row_limit}, + output_dtype=self.dtype, + ) + return tiletask_list, grid_spec + + def tile_exec(self, task: TileTask) -> torch.Tensor: + return self.decoder(task.tensor) + + def tile_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): + 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=-1)) + return torch.cat(result_rows, dim=-2) + + def encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]: + _, _, _, height, width = x.shape + overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) + row_limit = int(self.tile_latent_min_size - blend_extent) + + tiletask_list = [] + for i in range(0, height, overlap_size): + for j in range(0, width, overlap_size): + tile = x[:, :, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size] + tiletask_list.append( + TileTask( + len(tiletask_list), + (i // overlap_size, j // overlap_size), + tile, + workload=tile.shape[3] * tile.shape[4], + ) + ) + + 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={"blend_extent": blend_extent, "row_limit": row_limit}, + output_dtype=self.dtype, + ) + return tiletask_list, grid_spec + + def encode_tile_exec(self, task: TileTask) -> torch.Tensor: + return self.encoder(task.tensor) + + def encode_tile_merge( + self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec + ) -> torch.Tensor: + return self.tile_merge(coord_tensor_map, grid_spec) + + def spatial_tiled_encode(self, x: torch.Tensor): + if not self.is_distributed_enabled(): + return super().spatial_tiled_encode(x) + + logger.debug("Encode running with distributed executor") + return self.distributed_executor.execute( + x, + DistributedOperator( + split=self.encode_tile_split, + exec=self.encode_tile_exec, + merge=self.encode_tile_merge, + ), + broadcast_result=True, + ) + + def spatial_tiled_decode(self, z: torch.Tensor): + if not self.is_distributed_enabled(): + return super().spatial_tiled_decode(z) + + logger.debug("Decode running with distributed executor") + return self.distributed_executor.execute( + z, + DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge), + broadcast_result=True, + ) diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py index 73b89bb11..677d67796 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py @@ -30,7 +30,6 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.model_executor.models.hunyuan_image3.siglip2 import Siglip2VisionTransformer -from .autoencoder import AutoencoderKLConv3D from .hunyuan_image3_tokenizer import TokenizerWrapper from .hunyuan_image3_transformer import ( CausalMMOutputWithPast, @@ -343,7 +342,13 @@ class HunyuanImage3Pipeline( quant_config = od_config.quantization_config self.model = HunyuanImage3Model(self.hf_config, quant_config=quant_config) self.transformer = self.model - self.vae = AutoencoderKLConv3D.from_config(self.hf_config.vae) + # Lazy import to break circular dependency: + # autoencoder_kl_hunyuan -> hunyuan_image3/__init__ -> pipeline_hunyuan_image3 -> autoencoder_kl_hunyuan + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan import ( # noqa: PLC0415 + DistributedAutoencoderKLHunyuan, + ) + + self.vae = DistributedAutoencoderKLHunyuan.from_config(self.hf_config.vae) self.vae.use_spatial_tiling = self.od_config.vae_use_tiling self._pipeline = None self._tkwrapper = TokenizerWrapper(od_config.model) From d4e90bbba8bc9aea60a8b42650646cc7764c6700 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Wed, 20 May 2026 00:36:59 -0400 Subject: [PATCH 20/38] [Bugfix][TTS] Only populate voice_name for uploaded voices without inline ref_audio (#3523) Signed-off-by: Nick Cao Co-authored-by: Claude --- .../openai_api/test_serving_speech_voxcpm.py | 81 +++++++++++++++++++ .../entrypoints/openai/serving_speech.py | 43 ++++++---- 2 files changed, 109 insertions(+), 15 deletions(-) diff --git a/tests/entrypoints/openai_api/test_serving_speech_voxcpm.py b/tests/entrypoints/openai_api/test_serving_speech_voxcpm.py index 1f7279483..f49c1d9ce 100644 --- a/tests/entrypoints/openai_api/test_serving_speech_voxcpm.py +++ b/tests/entrypoints/openai_api/test_serving_speech_voxcpm.py @@ -192,3 +192,84 @@ class TestVoxCPM2Serving: Regression: previously returned 400 Invalid voice 'default'. """ assert "default" in voxcpm2_server.supported_speakers + + def test_voxcpm2_inline_ref_audio_does_not_set_voice_name(self, voxcpm2_server): + """Inline ref_audio must not populate voice_name in additional_information. + + Regression: the speaker cache keys on voice_name, so two requests with + the same voice but different inline ref_audio would collide — the second + request reused stale cached features from the first, causing a prefill + length mismatch and engine crash. + """ + mock_prompt = {"prompt_token_ids": [1, 1, 1], "additional_information": {}, "type": "token"} + voxcpm2_server._build_voxcpm2_prompt = AsyncMock(return_value=mock_prompt) + + request = OpenAICreateSpeechRequest( + input="test", + voice="default", + ref_audio="data:audio/wav;base64,QUJD", + ) + + asyncio.run(voxcpm2_server._prepare_speech_generation(request)) + + additional = mock_prompt["additional_information"] + assert "voice_name" not in additional, ( + "voice_name must not be set for inline ref_audio — " + "it causes speaker cache collisions across different audio clips" + ) + + def test_voxcpm2_uploaded_voice_sets_voice_name(self, voxcpm2_server): + """Uploaded voice without inline ref_audio must populate voice_name. + + This ensures the speaker cache still works for the intended use case: + a voice uploaded once and referenced by name in subsequent requests. + """ + mock_prompt = {"prompt_token_ids": [1, 1, 1], "additional_information": {}, "type": "token"} + voxcpm2_server._build_voxcpm2_prompt = AsyncMock(return_value=mock_prompt) + voxcpm2_server.uploaded_speakers["my_voice"] = { + "name": "my_voice", + "file_path": "/tmp/fake.wav", + "created_at": 12345, + } + voxcpm2_server.supported_speakers.add("my_voice") + + request = OpenAICreateSpeechRequest( + input="test", + voice="my_voice", + ) + + asyncio.run(voxcpm2_server._prepare_speech_generation(request)) + + additional = mock_prompt["additional_information"] + assert additional.get("voice_name") == "my_voice" + assert "voice_created_at" in additional + + def test_voxcpm2_uploaded_voice_with_inline_override_does_not_set_voice_name(self, voxcpm2_server): + """Uploaded voice with inline ref_audio override must not populate voice_name. + + Prevents cache poisoning: an attacker could send an uploaded voice name + with crafted ref_audio, caching malicious features under the uploaded + voice's name for subsequent legitimate requests. + """ + mock_prompt = {"prompt_token_ids": [1, 1, 1], "additional_information": {}, "type": "token"} + voxcpm2_server._build_voxcpm2_prompt = AsyncMock(return_value=mock_prompt) + voxcpm2_server.uploaded_speakers["my_voice"] = { + "name": "my_voice", + "file_path": "/tmp/fake.wav", + "created_at": 12345, + } + voxcpm2_server.supported_speakers.add("my_voice") + + request = OpenAICreateSpeechRequest( + input="test", + voice="my_voice", + ref_audio="data:audio/wav;base64,QUJD", + ) + + asyncio.run(voxcpm2_server._prepare_speech_generation(request)) + + additional = mock_prompt["additional_information"] + assert "voice_name" not in additional, ( + "voice_name must not be set when inline ref_audio overrides an " + "uploaded voice — it would poison the cache for future requests" + ) diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index e5bfb1329..70a491a94 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -1787,6 +1787,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): self, request: OpenAICreateSpeechRequest, ref_audio_data: tuple[list[float], int] | None = None, + *, + has_inline_ref_audio: bool = False, ) -> dict[str, Any]: """Build prompt for Fish Speech S2 Pro. @@ -1838,7 +1840,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Pass voice identity for model-side DAC code caching. if request.voice is not None: voice_lower = request.voice.lower() - if voice_lower in self.uploaded_speakers: + if voice_lower in self.uploaded_speakers and not has_inline_ref_audio: additional_information["voice_name"] = voice_lower additional_information["voice_created_at"] = self._voice_created_at(voice_lower) if request.max_new_tokens is not None: @@ -1852,6 +1854,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): async def _build_cosyvoice3_prompt( self, request: OpenAICreateSpeechRequest, + *, + has_inline_ref_audio: bool = False, ) -> dict[str, Any]: """Build prompt for CosyVoice3. @@ -1870,8 +1874,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Pass voice metadata for caching in the processor if request.voice: voice_lower = request.voice.lower() - mm_kwargs["voice_name"] = voice_lower - mm_kwargs["voice_created_at"] = self._voice_created_at(voice_lower) + if voice_lower in self.uploaded_speakers and not has_inline_ref_audio: + mm_kwargs["voice_name"] = voice_lower + mm_kwargs["voice_created_at"] = self._voice_created_at(voice_lower) return { "prompt": request.input, @@ -2020,6 +2025,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Resolve uploaded voice for non-Qwen3 models. # Qwen3 TTS has its own uploaded voice handling in _build_tts_params(). + has_inline_ref_audio = request.ref_audio is not None if self._tts_model_type in ("fish_tts", "cosyvoice3", "moss_tts_nano"): err = self._apply_uploaded_speaker(request) if err: @@ -2033,7 +2039,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if request.ref_audio is not None: wav_list, sr = await self._resolve_ref_audio(request.ref_audio) ref_audio_data = (wav_list, sr) - prompt = await self._build_fish_speech_prompt_async(request, ref_audio_data=ref_audio_data) + prompt = await self._build_fish_speech_prompt_async( + request, ref_audio_data=ref_audio_data, has_inline_ref_audio=has_inline_ref_audio + ) tts_params = {} elif self._tts_model_type == "omnivoice": if not request.input or not request.input.strip(): @@ -2050,8 +2058,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): prompt["ref_text"] = request.ref_text if request.voice: voice_lower = request.voice.lower() - prompt["voice_name"] = voice_lower - prompt["voice_created_at"] = self._voice_created_at(voice_lower) + if voice_lower in self.uploaded_speakers and not has_inline_ref_audio: + prompt["voice_name"] = voice_lower + prompt["voice_created_at"] = self._voice_created_at(voice_lower) if request.language: prompt["lang"] = request.language if request.instructions: @@ -2068,7 +2077,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if voice_lower not in self.uploaded_speakers and voice_lower not in self.supported_speakers: all_voices = sorted(self.uploaded_speakers.keys() | self.supported_speakers) raise ValueError(f"Invalid voice '{request.voice}'. Supported: {', '.join(all_voices) or 'none'}") - if voice_lower in self.uploaded_speakers: + if voice_lower in self.uploaded_speakers and not has_inline_ref_audio: if self.uploaded_speakers[voice_lower].get("embedding_source") == "direct": raise ValueError( f"Uploaded voice '{request.voice}' uses a speaker embedding (Qwen3-only). " @@ -2080,9 +2089,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): tts_params = {} if request.voice: voice_lower = request.voice.lower() - additional = prompt.setdefault("additional_information", {}) - additional["voice_name"] = voice_lower - additional["voice_created_at"] = self._voice_created_at(voice_lower) + if voice_lower in self.uploaded_speakers and not has_inline_ref_audio: + additional = prompt.setdefault("additional_information", {}) + additional["voice_name"] = voice_lower + additional["voice_created_at"] = self._voice_created_at(voice_lower) elif self._is_tts: validation_error = self._validate_tts_request(request) if validation_error: @@ -2092,7 +2102,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): prompt = await self._build_voxtral_prompt_async(request) tts_params = {} elif self._tts_model_type == "cosyvoice3": - prompt = await self._build_cosyvoice3_prompt(request) + prompt = await self._build_cosyvoice3_prompt(request, has_inline_ref_audio=has_inline_ref_audio) tts_params = {} elif self._tts_model_type == "ming_flash_omni_tts": prompt = self._build_ming_prompt(request) @@ -2101,8 +2111,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): tts_params = await self._build_moss_tts_params(request) if request.voice: voice_lower = request.voice.lower() - tts_params["voice_name"] = [voice_lower] - tts_params["voice_created_at"] = [self._voice_created_at(voice_lower)] + if voice_lower in self.uploaded_speakers and not has_inline_ref_audio: + tts_params["voice_name"] = [voice_lower] + tts_params["voice_created_at"] = [self._voice_created_at(voice_lower)] # Propagate seed from sampling params for deterministic generation. # MOSS-TTS-Nano uses its own internal sampling inside # inference_stream(), which reads seed from additional_information @@ -2375,6 +2386,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): all_voices = sorted(self.uploaded_speakers.keys() | self.supported_speakers) raise ValueError(f"Invalid voice '{request.voice}'. Supported: {', '.join(all_voices) or 'none'}") + has_inline_ref_audio = request.ref_audio is not None err = self._apply_uploaded_speaker(request) if err: raise ValueError(err) @@ -2388,8 +2400,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): prompt["ref_text"] = request.ref_text if request.voice: voice_lower = request.voice.lower() - prompt["voice_name"] = voice_lower - prompt["voice_created_at"] = self._voice_created_at(voice_lower) + if voice_lower in self.uploaded_speakers and not has_inline_ref_audio: + prompt["voice_name"] = voice_lower + prompt["voice_created_at"] = self._voice_created_at(voice_lower) if request.language: prompt["lang"] = request.language if request.instructions: From 12bcc32c3df050345234ce7b51dead6ce8389e3c Mon Sep 17 00:00:00 2001 From: "Chendi.Xue" Date: Tue, 19 May 2026 23:40:08 -0500 Subject: [PATCH 21/38] [XPU][CI] fix test_qwen2_5_omni_expansion.py::test_mix_to_audio (#3761) Signed-off-by: Chendi Xue --- docker/Dockerfile.xpu | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 2c49dded1..01459e813 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -133,6 +133,11 @@ ENTRYPOINT [] FROM ${VLLM_BASE} AS vllm-omni +RUN apt-get update && \ + apt-get install -y espeak-ng ffmpeg git jq && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* + WORKDIR /workspace/vllm-omni COPY . . From 06ab723afca4008fa7678a8e398398d2fc9bf269 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Wed, 20 May 2026 01:01:42 -0400 Subject: [PATCH 22/38] [Perf][VoxCPM2][Ming-Flash-Omni] Use global CUDA graph pool (#3361) Signed-off-by: Nick Cao --- .../models/ming_flash_omni/talker_module.py | 3 ++- .../model_executor/models/voxcpm2/voxcpm2_talker.py | 10 ++-------- 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py index 360a4f25a..6404b3fb8 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py @@ -30,6 +30,7 @@ import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedTokenizerBase, Qwen2Config, Qwen2Model, StaticCache from vllm.logger import init_logger +from vllm.platforms import current_platform from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding @@ -479,7 +480,7 @@ class CFMGraphExecutor: self.sde_rnd_placeholder = torch.empty_like(sde_rnd) self.graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(self.graph): + with torch.cuda.graph(self.graph, pool=current_platform.get_global_graph_pool()): self.gen_lat_placeholder = self.cfm.sample( self.last_hidden_state_placeholder, self.his_lat_placeholder, diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py index 635b570bc..69b2a29ea 100644 --- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py +++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py @@ -31,6 +31,7 @@ from vllm.model_executor.models.utils import ( maybe_prefix, ) from vllm.multimodal.audio import AudioResampler +from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.models.output_templates import OmniOutput @@ -456,7 +457,6 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): self._scaffold_graphs: dict[int, _CapturedGraph] = {} self._residual_graphs: dict[int, _CapturedGraph] = {} self._max_cached_graphs = self._max_batch_size - self._cuda_graph_pool: tuple | None = None self._cuda_graph_warmup_steps = 0 self._cuda_graph_warmup_threshold = 3 @@ -661,11 +661,6 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): tts = self.tts return tts.stop_head(tts.stop_actn(tts.stop_proj(lm_h))) - def _get_cuda_graph_pool(self) -> tuple: - if self._cuda_graph_pool is None: - self._cuda_graph_pool = torch.cuda.graph_pool_handle() - return self._cuda_graph_pool - @staticmethod def _nullify_volatile_metadata(ctx: Any) -> Any: """Set ``scheduler_metadata`` to None on all attention layers. @@ -698,7 +693,6 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): hidden_size = self.config.hidden_size dtype = self._side_dtype dev = torch.device(self._device) - pool = self._get_cuda_graph_pool() model.precompute_fused_qkv() @@ -721,7 +715,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): for _ in range(3): _ = model(**call_kwargs) - with torch.cuda.graph(g.graph, pool=pool): + with torch.cuda.graph(g.graph, pool=current_platform.get_global_graph_pool()): g.output = model(**call_kwargs) logger.info("CUDA Graph captured for %s (batch_size=%d)", label, batch_size) From f99e8b2d3ce00bf2c45358a29563953d1cecd317 Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Wed, 20 May 2026 01:08:56 -0400 Subject: [PATCH 23/38] [Bench] Add audio-streaming continuity metric for TTS (#3618) Signed-off-by: Yueqian Lin Co-authored-by: Hongsheng Liu --- benchmarks/tts/README.md | 24 ++++- benchmarks/tts/bench_tts.py | 2 +- tests/benchmarks/metrics/test_metrics.py | 32 ++++++ tests/benchmarks/test_audio_continuity.py | 123 ++++++++++++++++++++++ vllm_omni/benchmarks/audio_continuity.py | 91 ++++++++++++++++ vllm_omni/benchmarks/metrics/metrics.py | 26 ++++- vllm_omni/benchmarks/patch/patch.py | 62 ++++++++++- 7 files changed, 350 insertions(+), 10 deletions(-) create mode 100644 tests/benchmarks/test_audio_continuity.py create mode 100644 vllm_omni/benchmarks/audio_continuity.py diff --git a/benchmarks/tts/README.md b/benchmarks/tts/README.md index 9e2fd35b1..d8841f382 100644 --- a/benchmarks/tts/README.md +++ b/benchmarks/tts/README.md @@ -40,7 +40,7 @@ vllm bench serve --omni \ --num-prompts 20 --num-warmups 2 \ --extra-body '{"task_type":"Base"}' \ --max-concurrency 1 --request-rate inf \ - --percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \ + --percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun \ --save-result --result-dir ./results ``` @@ -58,7 +58,7 @@ vllm bench serve --omni \ --num-prompts 20 --num-warmups 2 \ --extra-body '{"voice":"Vivian","language":"English","task_type":"CustomVoice"}' \ --max-concurrency 1 --request-rate inf \ - --percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \ + --percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun \ --save-result --result-dir ./results ``` @@ -76,10 +76,28 @@ vllm bench serve --omni \ --num-prompts 20 --num-warmups 2 \ --extra-body '{"task_type":"VoiceDesign","language":"English"}' \ --max-concurrency 1 --request-rate inf \ - --percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \ + --percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun \ --save-result --result-dir ./results ``` +#### Streaming continuity (`audio_underrun`) + +`audio_underrun` (seconds) is the per-request worst-case buffer deficit +under a realtime-rate player simulation. The audio-speech backend records +chunk arrival times during the SSE stream and surfaces: + +- `Mean / Median / P{50,99} AUDIO_UNDERRUN (s)` per the standard percentile + output - any positive value means at least one inter-chunk gap was longer + than the chunk's audio duration. +- `Streaming continuity OK rate` - fraction of requests whose worst-case + underrun stayed under the threshold (default 100 ms, the commonly cited + "audible gap" budget). Override the threshold with + `VLLM_OMNI_BENCH_AUDIO_CONTINUITY_THRESHOLD_S=`. + +This captures the failure mode where `RTF_p50 < 1` (server keeps up in +aggregate) but per-stream chunk arrival is bursty enough that listeners +still hear gaps - common at high concurrency on streaming TTS pipelines. + #### Add WER / SIM / UTMOS to any of the above Append `--seed-tts-wer-eval` (and optionally `SEED_TTS_EVAL_DEVICE=cuda:0` diff --git a/benchmarks/tts/bench_tts.py b/benchmarks/tts/bench_tts.py index ba82b1c9b..8d2eb99c8 100644 --- a/benchmarks/tts/bench_tts.py +++ b/benchmarks/tts/bench_tts.py @@ -118,7 +118,7 @@ def build_bench_args( "--num-warmups", "2", "--percentile-metrics", - "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", + "ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun", ] if resolved_dataset_path: diff --git a/tests/benchmarks/metrics/test_metrics.py b/tests/benchmarks/metrics/test_metrics.py index 788e1227d..2688b224f 100644 --- a/tests/benchmarks/metrics/test_metrics.py +++ b/tests/benchmarks/metrics/test_metrics.py @@ -63,6 +63,38 @@ def test_total_input_aggregated_from_output_prompt_len(): ) +def test_audio_continuity_aggregation(): + """Continuity rate and underrun percentile must aggregate from per-output fields.""" + bad = _make_output(100) + bad.audio_underrun_s = 0.5 + bad.audio_continuity_ok = False + good_a = _make_output(100) + good_a.audio_underrun_s = 0.02 + good_a.audio_continuity_ok = True + good_b = _make_output(100) + good_b.audio_underrun_s = 0.0 + good_b.audio_continuity_ok = True + + metrics, _ = calculate_metrics( + input_requests=[], + outputs=[bad, good_a, good_b], + dur_s=10.0, + tokenizer=None, + selected_percentiles=[50.0, 99.0], + goodput_config_dict={}, + task_type=TaskType.GENERATION, + selected_percentile_metrics=["audio_underrun"], + max_concurrency=None, + request_rate=float("inf"), + benchmark_duration=10.0, + ) + + assert metrics.audio_continuity_ok_rate == pytest.approx(2 / 3, abs=1e-6) + # p99 of [0.5, 0.02, 0.0] is dominated by the 0.5 outlier. + p99 = dict(metrics.percentiles_audio_underrun_s).get(99.0) + assert p99 is not None and p99 > 0.4 + + # ============================================================================ # TTFT suppression for pure-audio (TTS) benchmarks # ============================================================================ diff --git a/tests/benchmarks/test_audio_continuity.py b/tests/benchmarks/test_audio_continuity.py new file mode 100644 index 000000000..7737986f0 --- /dev/null +++ b/tests/benchmarks/test_audio_continuity.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for the audio-streaming continuity helper. + +The helper is intentionally vllm-free: it operates on a flat timeline of +``(arrival_time_s, num_bytes)`` chunks plus PCM format params. These tests +exercise the underrun arithmetic against hand-crafted timelines so future +refactors of the streaming backend cannot silently regress the metric. +""" + +from __future__ import annotations + +import pytest + +from vllm_omni.benchmarks.audio_continuity import ( + ContinuityStats, + compute_continuity_stats, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu] + + +# PCM s16le mono at 24 kHz -> 48 000 bytes/sec, the Qwen3-TTS / VoxCPM2 default. +_SR = 24_000 +_BPS = _SR * 2 # 48 000 + + +def test_empty_timeline_is_continuous() -> None: + stats = compute_continuity_stats([], [], sample_rate=_SR) + assert stats == ContinuityStats( + max_underrun_s=0.0, + underrun_event_count=0, + is_continuous=True, + ) + + +def test_single_chunk_has_no_underrun() -> None: + # One large chunk, no later arrivals to check against. + stats = compute_continuity_stats( + chunk_arrival_times_s=[0.05], + chunk_bytes=[_BPS], # one second of audio in a single chunk + sample_rate=_SR, + ) + assert stats.max_underrun_s == 0.0 + assert stats.underrun_event_count == 0 + assert stats.is_continuous is True + + +def test_chunks_faster_than_realtime_have_no_underrun() -> None: + # 100ms of audio per chunk, but arrive every 50ms — buffer always builds. + stats = compute_continuity_stats( + chunk_arrival_times_s=[0.0, 0.05, 0.10, 0.15], + chunk_bytes=[_BPS // 10] * 4, + sample_rate=_SR, + ) + assert stats.max_underrun_s == 0.0 + assert stats.underrun_event_count == 0 + assert stats.is_continuous is True + + +def test_single_gap_above_threshold_breaks_continuity() -> None: + # First chunk carries 100ms of audio at t=0. Next chunk arrives at t=0.5s, + # so the player ran dry from t=0.1 to t=0.5 -> 400ms underrun. + stats = compute_continuity_stats( + chunk_arrival_times_s=[0.0, 0.5], + chunk_bytes=[_BPS // 10, _BPS // 10], # 100ms each + sample_rate=_SR, + threshold_s=0.1, + ) + assert stats.max_underrun_s == pytest.approx(0.4, abs=1e-6) + assert stats.underrun_event_count == 1 + assert stats.is_continuous is False + + +def test_underrun_is_max_across_gaps_not_sum() -> None: + # Two gaps: a 200ms shortfall and a 50ms shortfall. max_underrun_s + # must surface the 200ms one, not the sum. + stats = compute_continuity_stats( + chunk_arrival_times_s=[0.0, 0.3, 0.45], + chunk_bytes=[_BPS // 10, _BPS // 10, _BPS // 10], # 100ms chunks + sample_rate=_SR, + threshold_s=0.1, + ) + # At t=0.3 the player consumed 300ms but only 100ms received -> 200ms deficit + # At t=0.45 player consumed 450ms, received 200ms -> 250ms deficit. + assert stats.max_underrun_s == pytest.approx(0.25, abs=1e-6) + assert stats.underrun_event_count == 2 + assert stats.is_continuous is False + + +def test_continuity_threshold_boundary_inclusive() -> None: + # First chunk = 100ms audio at t=0. Next arrives at t=0.2 -> 100ms deficit + # exactly. With threshold=0.1, continuity_ok stays True (boundary). + stats = compute_continuity_stats( + chunk_arrival_times_s=[0.0, 0.2], + chunk_bytes=[_BPS // 10, _BPS // 10], + sample_rate=_SR, + threshold_s=0.1, + ) + assert stats.max_underrun_s == pytest.approx(0.1, abs=1e-6) + assert stats.is_continuous is True + + +def test_mismatched_lengths_returns_zero_stats() -> None: + # Defensive: never crash on malformed timeline. + stats = compute_continuity_stats( + chunk_arrival_times_s=[0.0, 0.1], + chunk_bytes=[100], + sample_rate=_SR, + ) + assert stats == ContinuityStats(0.0, 0, True) + + +def test_threshold_is_configurable() -> None: + # 200ms underrun: violates 100ms threshold, OK against 500ms threshold. + arrivals = [0.0, 0.3] + bytes_list = [_BPS // 10, _BPS // 10] + strict = compute_continuity_stats(arrivals, bytes_list, sample_rate=_SR, threshold_s=0.1) + lenient = compute_continuity_stats(arrivals, bytes_list, sample_rate=_SR, threshold_s=0.5) + assert strict.is_continuous is False + assert lenient.is_continuous is True + assert strict.max_underrun_s == lenient.max_underrun_s diff --git a/vllm_omni/benchmarks/audio_continuity.py b/vllm_omni/benchmarks/audio_continuity.py new file mode 100644 index 000000000..e92ed9a0a --- /dev/null +++ b/vllm_omni/benchmarks/audio_continuity.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Streaming-audio continuity metrics for the TTS bench. + +Captures the failure mode where ``RTF_p50 < 1`` (server keeps up with realtime +in aggregate) but per-stream chunk arrival is bursty enough that listeners +hear audible gaps. Companion to ``audio_ttfp`` / ``audio_rtf``. + +The math is purposefully simple so the helper has no ``vllm`` dependency and +can be unit-tested in isolation: given the wall-clock arrival timeline of +audio chunks and the PCM format, simulate a player that consumes at realtime +rate from the first chunk onward, and surface the worst-case deficit. + +A deficit > ``threshold_s`` at any point counts as an audible underrun. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class ContinuityStats: + """Result of a single-request continuity analysis. + + Attributes: + max_underrun_s: Worst-case wall-clock seconds the player was starved. + underrun_event_count: Inter-chunk intervals during which the buffer + went negative (one per gap, not per starved millisecond). + is_continuous: ``max_underrun_s <= threshold_s`` for the request. + """ + + max_underrun_s: float + underrun_event_count: int + is_continuous: bool + + +def compute_continuity_stats( + chunk_arrival_times_s: list[float], + chunk_bytes: list[int], + sample_rate: int, + sample_width: int = 2, + channels: int = 1, + threshold_s: float = 0.1, +) -> ContinuityStats: + """Compute the continuity stats for one streaming response. + + Args: + chunk_arrival_times_s: Wall-clock seconds (since request start) at + which each non-empty audio chunk's *last byte* arrived. Must be + monotonically non-decreasing. + chunk_bytes: Byte count of each chunk in the same order. + sample_rate: PCM sample rate (Hz). 24 000 for Qwen3-TTS / VoxCPM2. + sample_width: Bytes per sample (2 = s16le). + channels: PCM channel count (1 = mono). + threshold_s: Underrun budget. The default 0.1 s matches the + commonly-cited "audible gap" threshold for streaming TTS. + + Returns: + A :class:`ContinuityStats` summarising the worst-case deficit and + whether it stayed under the threshold. + """ + n = len(chunk_arrival_times_s) + if n == 0 or n != len(chunk_bytes): + return ContinuityStats(0.0, 0, True) + + bytes_per_s = sample_rate * sample_width * channels + if bytes_per_s <= 0: + return ContinuityStats(0.0, 0, True) + + t0 = chunk_arrival_times_s[0] + received_before = 0 + max_underrun_s = 0.0 + event_count = 0 + for i in range(n): + if i > 0: + played_bytes = (chunk_arrival_times_s[i] - t0) * bytes_per_s + deficit_bytes = played_bytes - received_before + if deficit_bytes > 0: + deficit_s = deficit_bytes / bytes_per_s + if deficit_s > max_underrun_s: + max_underrun_s = deficit_s + event_count += 1 + received_before += chunk_bytes[i] + + return ContinuityStats( + max_underrun_s=max_underrun_s, + underrun_event_count=event_count, + is_continuous=max_underrun_s <= threshold_s, + ) diff --git a/vllm_omni/benchmarks/metrics/metrics.py b/vllm_omni/benchmarks/metrics/metrics.py index 4b32662c5..56ebb20ea 100644 --- a/vllm_omni/benchmarks/metrics/metrics.py +++ b/vllm_omni/benchmarks/metrics/metrics.py @@ -25,6 +25,12 @@ class MultiModalsBenchmarkMetrics(BenchmarkMetrics): median_audio_duration_s: float = 0.0 std_audio_duration_s: float = 0.0 percentiles_audio_duration_s: list[tuple[float, float]] = None + # Streaming-audio continuity (see vllm_omni/benchmarks/audio_continuity.py). + mean_audio_underrun_s: float = 0.0 + median_audio_underrun_s: float = 0.0 + std_audio_underrun_s: float = 0.0 + percentiles_audio_underrun_s: list[tuple[float, float]] = None + audio_continuity_ok_rate: float = 1.0 def print_metrics( @@ -83,6 +89,12 @@ def print_audio_metrics(selected_percentile_metrics, metrics: MultiModalsBenchma print("{:<40} {:<10.2f}".format("Total audio duration generated(s):", metrics.total_audio_duration_s)) print("{:<40} {:<10}".format("Total audio frames generated:", metrics.total_audio_frames)) print("{:<40} {:<10.2f}".format("Audio throughput(audio duration/s):", metrics.audio_throughput)) + print( + "{:<40} {:<10.2%}".format( + "Streaming continuity OK rate:", + metrics.audio_continuity_ok_rate, + ) + ) for metric in selected_percentile_metrics: if metric.startswith("audio"): process_one_metric(metric, metrics) @@ -100,17 +112,18 @@ def process_one_metric( "audio_ttfp": "Time to First Packet", "audio_rtf": "Real Time Factor", "audio_duration": "Audio Duration", + "audio_underrun": "Streaming Audio Underrun", } header = metric_header_map.get(metric_attribute_name, metric_attribute_name) print("{s:{c}^{n}}".format(s=header, n=50, c="-")) is_audio_rtf = metric_attribute_name == "audio_rtf" - is_audio_duration = metric_attribute_name == "audio_duration" + is_audio_duration_or_underrun = metric_attribute_name in ("audio_duration", "audio_underrun") suffix = "_ms" unit_suffix = " (ms)" - if is_audio_duration: + if is_audio_duration_or_underrun: suffix = "_s" unit_suffix = " (s)" elif is_audio_rtf: @@ -172,6 +185,8 @@ def calculate_metrics( audio_rtfs: list[float] = [] audio_duration: list[float] = [] audio_frames: list[int] = [] + audio_underruns: list[float] = [] + audio_continuity_ok: list[bool] = [] input_audio_duration = 0.0 for i in range(len(outputs)): if outputs[i].success: @@ -205,6 +220,8 @@ def calculate_metrics( audio_rtfs.append(getattr(outputs[i], "audio_rtf", 0.0)) audio_duration.append(getattr(outputs[i], "audio_duration", 0.0)) audio_frames.append(getattr(outputs[i], "audio_frames", 0.0)) + audio_underruns.append(getattr(outputs[i], "audio_underrun_s", 0.0)) + audio_continuity_ok.append(bool(getattr(outputs[i], "audio_continuity_ok", True))) e2els.append(outputs[i].latency) input_audio_duration += outputs[i].input_audio_duration completed += 1 @@ -332,6 +349,11 @@ def calculate_metrics( std_audio_rtf=np.std(audio_rtfs or 0), median_audio_rtf=np.median(audio_rtfs or 0), percentiles_audio_rtf=[(p, np.percentile(audio_rtfs or 0, p)) for p in selected_percentiles], + mean_audio_underrun_s=np.mean(audio_underruns or 0), + std_audio_underrun_s=np.std(audio_underruns or 0), + median_audio_underrun_s=np.median(audio_underruns or 0), + percentiles_audio_underrun_s=[(p, np.percentile(audio_underruns or 0, p)) for p in selected_percentiles], + audio_continuity_ok_rate=(sum(audio_continuity_ok) / len(audio_continuity_ok)) if audio_continuity_ok else 1.0, mean_tpot_ms=np.mean(tpots or 0) * 1000, std_tpot_ms=np.std(tpots or 0) * 1000, median_tpot_ms=np.median(tpots or 0) * 1000, diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index ea3b6b4c9..12aeef7e4 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -36,6 +36,7 @@ from vllm.tokenizers import TokenizerLike logger = init_logger(__name__) +from vllm_omni.benchmarks.audio_continuity import compute_continuity_stats from vllm_omni.benchmarks.data_modules.daily_omni_dataset import DailyOmniDataset, DailyOmniSampleRequest from vllm_omni.benchmarks.data_modules.random_multi_modal_dataset import OmniRandomMultiModalDataset from vllm_omni.benchmarks.data_modules.seed_tts_dataset import ( @@ -46,6 +47,33 @@ from vllm_omni.benchmarks.data_modules.seed_tts_dataset import ( SeedTTSTextDataset, ) +_AUDIO_CONTINUITY_THRESHOLD_ENV = "VLLM_OMNI_BENCH_AUDIO_CONTINUITY_THRESHOLD_S" +_DEFAULT_AUDIO_CONTINUITY_THRESHOLD_S = 0.1 + + +def _audio_continuity_threshold_s() -> float: + """Return the per-request underrun budget (s). + + Read from ``VLLM_OMNI_BENCH_AUDIO_CONTINUITY_THRESHOLD_S`` so users can + re-aim the SLO without rebuilding. Defaults to 100 ms - the standard + "audible gap" budget for streaming TTS. + """ + raw = os.environ.get(_AUDIO_CONTINUITY_THRESHOLD_ENV) + if not raw: + return _DEFAULT_AUDIO_CONTINUITY_THRESHOLD_S + try: + value = float(raw) + except ValueError: + logger.warning( + "Invalid %s=%r; using default %.3fs", + _AUDIO_CONTINUITY_THRESHOLD_ENV, + raw, + _DEFAULT_AUDIO_CONTINUITY_THRESHOLD_S, + ) + return _DEFAULT_AUDIO_CONTINUITY_THRESHOLD_S + return max(value, 0.0) + + get_samples_old = datasets.get_samples _DEFAULT_DAILY_OMNI_REPO = "liarliar/Daily-Omni" @@ -315,6 +343,16 @@ class MixRequestFuncOutput(RequestFuncOutput): audio_frames: int = 0 audio_rtf: float = 0.0 text_latency: float = 0.0 + #: Worst-case streaming-audio underrun (wall-clock seconds the player + #: would have been starved). Populated by the audio-speech backend; ``0.0`` + #: for backends that do not run continuity analysis. + audio_underrun_s: float = 0.0 + #: Whether the request stayed under the continuity threshold (default + #: 100 ms). Mirrors ``audio_underrun_s <= threshold``. + audio_continuity_ok: bool = True + #: Number of inter-chunk intervals during which the player buffer went + #: negative. + audio_underrun_event_count: int = 0 #: Raw PCM s16le mono at 24 kHz for Seed-TTS WER: from ``/v1/audio/speech`` stream or #: resampled export after ``openai-chat-omni`` audio deltas. tts_output_pcm_bytes: bytes | None = None @@ -605,6 +643,8 @@ async def async_request_openai_audio_speech( total_pcm_bytes = 0 capture_wer_pcm = _seed_tts_capture_pcm_for_wer() and getattr(request_func_input, "seed_tts_row", False) pcm_capture = bytearray() if capture_wer_pcm else None + chunk_arrival_times_s: list[float] = [] + chunk_sizes: list[int] = [] try: async with session.post(url=api_url, json=payload, headers=headers) as response: if response.status == 200: @@ -617,6 +657,8 @@ async def async_request_openai_audio_speech( # not defined here; only audio TTFP is meaningful. output.audio_ttfp = timestamp - st total_pcm_bytes += len(chunk) + chunk_arrival_times_s.append(timestamp - st) + chunk_sizes.append(len(chunk)) if pcm_capture is not None: pcm_capture.extend(chunk) @@ -631,6 +673,18 @@ async def async_request_openai_audio_speech( else: output.audio_rtf = 0 logger.warning("Audio duration is zero") + + continuity = compute_continuity_stats( + chunk_arrival_times_s=chunk_arrival_times_s, + chunk_bytes=chunk_sizes, + sample_rate=sample_rate, + sample_width=sample_width, + channels=channels, + threshold_s=_audio_continuity_threshold_s(), + ) + output.audio_underrun_s = continuity.max_underrun_s + output.audio_continuity_ok = continuity.is_continuous + output.audio_underrun_event_count = continuity.underrun_event_count if pcm_capture is not None and pcm_capture: output.tts_output_pcm_bytes = bytes(pcm_capture) elif capture_wer_pcm: @@ -1020,14 +1074,14 @@ async def benchmark( return # No text tokens generated (e.g. pure TTS speech endpoint): per-token # latency metrics (ttft/tpot/itl) are undefined, so skip them. - is_text_token_metric = metric_attribute_name not in ("e2el", "audio_ttfp", "audio_rtf", "audio_duration") + is_text_token_metric = not (metric_attribute_name == "e2el" or metric_attribute_name.startswith("audio")) if is_text_token_metric and getattr(metrics, "total_output", 0) == 0: return is_audio_rtf = metric_attribute_name == "audio_rtf" - is_audio_duration = metric_attribute_name == "audio_duration" + is_audio_duration_or_underrun = metric_attribute_name in ("audio_duration", "audio_underrun") suffix = "_ms" - if is_audio_duration: + if is_audio_duration_or_underrun: suffix = "_s" elif is_audio_rtf: suffix = "" @@ -1038,7 +1092,7 @@ async def benchmark( median_attr_name = f"median_{metric_attribute_name}{suffix}" median_value = getattr(metrics, median_attr_name, 0.0) result[median_attr_name] = median_value - for p, value in getattr(metrics, f"percentiles_{metric_attribute_name}{suffix}"): + for p, value in getattr(metrics, f"percentiles_{metric_attribute_name}{suffix}", None) or []: p_word = str(int(p)) if int(p) == p else str(p) result[f"p{p_word}_{metric_attribute_name}{suffix}"] = value From 109bba8cdf322da5664a6b96bba8f74fb22f92b4 Mon Sep 17 00:00:00 2001 From: Lancer Date: Wed, 20 May 2026 13:18:57 +0800 Subject: [PATCH 24/38] [Bugfix] Treat kv_cache_dtype=auto as unset for ring attention (#3622) Signed-off-by: Lancer --- .../attention/test_attention_config.py | 99 +++++++++++++++++++ vllm_omni/diffusion/attention/layer.py | 12 ++- 2 files changed, 107 insertions(+), 4 deletions(-) diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py index 5dea43b91..a4c67851a 100644 --- a/tests/diffusion/attention/test_attention_config.py +++ b/tests/diffusion/attention/test_attention_config.py @@ -420,3 +420,102 @@ class TestAttentionInitUsesCurrentDiffusionConfig: assert attn.use_ring is True assert attn.ring_runner is not None assert attn.ring_runner.attn_backend_pref == "TORCH_SDPA" + + +class TestDiffusionKvCacheQuantization: + @staticmethod + def _install_attention_init_stubs(monkeypatch): + class _FakeAttentionImpl: + def __init__(self, **kwargs): + self.kwargs = kwargs + + def forward(self, query, key, value, attn_metadata=None): + return query + + class _FakeBackend: + @staticmethod + def get_name() -> str: + return "FAKE_BACKEND" + + @staticmethod + def get_impl_cls(): + return _FakeAttentionImpl + + @staticmethod + def supports_kv_cache_dtype(kv_cache_dtype, platform_key) -> bool: + return True + + class _FakeRingParallelAttention: + def __init__(self, sp_group, attn_backend_pref=None): + self.sp_group = sp_group + self.attn_backend_pref = attn_backend_pref + + monkeypatch.setattr( + layer_mod, + "get_attn_backend_for_role", + lambda role, head_size, attention_config=None, role_category=None: (_FakeBackend, None), + ) + monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl)) + monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object()) + monkeypatch.setattr(layer_mod, "get_sp_group", lambda: SimpleNamespace(ring_group="ring-group")) + monkeypatch.setattr(layer_mod, "RingParallelAttention", _FakeRingParallelAttention) + monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False) + + def test_diffusion_kv_cache_dtype_none_does_not_trigger_ring_quantization_error(self, monkeypatch): + self._install_attention_init_stubs(monkeypatch) + od_config = SimpleNamespace( + diffusion_attention_config=AttentionConfig(), + parallel_config=SimpleNamespace(ring_degree=2), + diffusion_kv_cache_dtype=None, + diffusion_kv_cache_skip_step_indices=None, + diffusion_kv_cache_skip_layer_indices=None, + ) + + with set_current_diffusion_config(od_config): + attn = Attention( + num_heads=4, + head_size=64, + causal=False, + softmax_scale=1.0, + ) + + assert attn._kv_cache_dtype is None + + def test_diffusion_kv_cache_dtype_auto_does_not_trigger_ring_quantization_error(self, monkeypatch): + self._install_attention_init_stubs(monkeypatch) + od_config = SimpleNamespace( + diffusion_attention_config=AttentionConfig(), + parallel_config=SimpleNamespace(ring_degree=2), + diffusion_kv_cache_dtype="auto", + diffusion_kv_cache_skip_step_indices=None, + diffusion_kv_cache_skip_layer_indices=None, + ) + + with set_current_diffusion_config(od_config): + attn = Attention( + num_heads=4, + head_size=64, + causal=False, + softmax_scale=1.0, + ) + + assert attn._kv_cache_dtype is None + + def test_diffusion_kv_cache_dtype_fp8_raises_with_ring_attention(self, monkeypatch): + self._install_attention_init_stubs(monkeypatch) + od_config = SimpleNamespace( + diffusion_attention_config=AttentionConfig(), + parallel_config=SimpleNamespace(ring_degree=2), + diffusion_kv_cache_dtype="fp8", + diffusion_kv_cache_skip_step_indices=None, + diffusion_kv_cache_skip_layer_indices=None, + ) + + with set_current_diffusion_config(od_config): + with pytest.raises(ValueError, match="KV quantization is not compatible with ring attention"): + Attention( + num_heads=4, + head_size=64, + causal=False, + softmax_scale=1.0, + ) diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 5cbc39eb1..e0fb8d264 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -171,9 +171,13 @@ class Attention(nn.Module): def _init_kv_cache_quantization(self, config) -> None: if config is None: return - dtype = config.diffusion_kv_cache_dtype + dtype = getattr(config, "diffusion_kv_cache_dtype", None) + if dtype == "auto": + dtype = None + parallel_config = getattr(config, "parallel_config", None) + ring_degree = getattr(parallel_config, "ring_degree", 1) if dtype: - if config.parallel_config.ring_degree > 1: + if ring_degree > 1: raise ValueError( "KV quantization is not compatible with ring attention " "(ring_degree > 1). Ring kernels do not propagate quantization descale " @@ -190,8 +194,8 @@ class Attention(nn.Module): ) dtype = None self._kv_cache_dtype = dtype - self._kv_cache_skip_steps = config.diffusion_kv_cache_skip_step_indices - self._kv_cache_skip_layers = config.diffusion_kv_cache_skip_layer_indices + self._kv_cache_skip_steps = getattr(config, "diffusion_kv_cache_skip_step_indices", None) + self._kv_cache_skip_layers = getattr(config, "diffusion_kv_cache_skip_layer_indices", None) def _should_apply_kv_cache_quant(self) -> bool: skip_steps = self._kv_cache_skip_steps From 9dd36e3afd0940752c2a96f9d39f8d211c2b9195 Mon Sep 17 00:00:00 2001 From: hxhhhlalala Date: Wed, 20 May 2026 14:21:29 +0800 Subject: [PATCH 25/38] [NPU][Quant] Add W4A4 MXFP4 online & MXFP4 dual-scale online/offline quantization support for Wan2.2 T2V / I2V inference on Ascend NPU (#3578) Signed-off-by: hyh_hh Co-authored-by: hyh_hh --- docs/user_guide/quantization/mxfp4.md | 479 ++++++++++++ docs/user_guide/quantization/overview.md | 25 +- .../image_to_video/image_to_video.py | 4 +- .../text_to_video/text_to_video.py | 4 +- .../wan2_2/test_create_transformer_quant.py | 383 ++++++++++ .../quantization/test_mxfp4_config.py | 539 +++++++++++++ .../quantization/test_mxfp4_key_remap.py | 399 ++++++++++ .../quantization/test_mxfp8_config.py | 308 ++++++++ .../quantization/test_mxfp8_key_remap.py | 214 ++++++ .../models/wan2_2/pipeline_wan2_2.py | 48 +- vllm_omni/quantization/factory.py | 102 +++ vllm_omni/quantization/mxfp4_config.py | 709 ++++++++++++++++++ vllm_omni/quantization/tools/__init__.py | 2 + .../tools/merge_mxfp4_dualscale_checkpoint.py | 491 ++++++++++++ 14 files changed, 3646 insertions(+), 61 deletions(-) create mode 100644 docs/user_guide/quantization/mxfp4.md create mode 100644 tests/diffusion/models/wan2_2/test_create_transformer_quant.py create mode 100644 tests/diffusion/quantization/test_mxfp4_config.py create mode 100644 tests/diffusion/quantization/test_mxfp4_key_remap.py create mode 100644 tests/diffusion/quantization/test_mxfp8_config.py create mode 100644 tests/diffusion/quantization/test_mxfp8_key_remap.py create mode 100644 vllm_omni/quantization/mxfp4_config.py create mode 100644 vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py diff --git a/docs/user_guide/quantization/mxfp4.md b/docs/user_guide/quantization/mxfp4.md new file mode 100644 index 000000000..2fd4a8395 --- /dev/null +++ b/docs/user_guide/quantization/mxfp4.md @@ -0,0 +1,479 @@ +# W4A4 MXFP4 Quantization + +## Overview + +W4A4 MXFP4 (Microscaling FP4) quantizes both weights and activations to FP4 +(`float4_e2m1fn_x2`, packed 2 values per byte) using the OCP MX format: groups +of 32 K-dimension elements share a single `float8_e8m0fnu` exponent scale. + +vLLM-Omni provides two quantization methods with different scale structures: + +| Method | Scale structure | Mode | Use case | +|--------|----------------|------|----------| +| `mxfp4` | Single-scale (per-32 fine only) | Online only | Quick accuracy baseline; no checkpoint prep needed | +| `mxfp4_dualscale` | Dual-scale (fine per-32 + coarse per-512 + per-channel `mul_scale`) | Online + Offline | Production; better accuracy; offline recommended | + +!!! tip "Recommended: `mxfp4_dualscale` offline" + For production deployments, use the `mxfp4_dualscale` offline mode with a + pre-quantized checkpoint produced by msModelSlim. Offline checkpoints load + calibrated `mul_scale` tensors from disk, providing measurably better accuracy + than any online method. The one-time preprocessing cost amortises across all + subsequent inference runs. + + Use `mxfp4` online only for quick experimentation where preprocessing time + is not acceptable and accuracy loss is tolerable. + +!!! warning "Online single-scale ≠ Offline dual-scale" + `mxfp4_dualscale` offline mode uses `NPUMxfp4DualScaleLinearMethod`: + fine scale (per-32 K), coarse scale (per-512 K), and per-input-channel + `mul_scale` from calibration — all loaded from the checkpoint. + `mxfp4_dualscale` online mode uses `NPUMxfp4DualScaleOnlineLinearMethod`: + dual-level scales computed on the fly from BF16 weights; no calibration + `mul_scale` is available. Loading an offline checkpoint with the online + method (or vice versa) will produce incorrect results or shape errors. + +## Hardware Support + +| Device | Support | +|--------|---------| +| NVIDIA Blackwell GPU (SM 100+) | ⭕ | +| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ | +| NVIDIA Ampere GPU (SM 80+) | ⭕ | +| AMD ROCm | ⭕ | +| Intel XPU | ⭕ | +| Ascend NPU (Atlas 950 A5) | ✅ | + +Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide. + +## Model Type Support + +### Diffusion Model (Wan2.2) + +| Model | Online | Offline | Notes | +|-------|--------|---------|-------| +| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade (`transformer` + `transformer_2`); both transformers quantized with the same config | +| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade; same scheme as T2V-A14B | +| Wan2.2-TI2V-5B | ❌ | ❌ | Parameter count too small; W4A4 causes unacceptable accuracy loss | + +The choice between `mxfp4` and `mxfp4_dualscale` in **online mode** is about +quantization quality, not model compatibility — both work on cascade (A14B) and +single-transformer models alike, the same as `mxfp8` online: + +- `mxfp4`: single-scale, lower overhead, simpler compute, online only +- `mxfp4_dualscale`: dual-scale + optional BF16 fallback, better accuracy, online **and** offline + +**Offline** checkpoints for A14B are always in `mxfp4_dualscale` format (produced +by the merge script); there is no offline `mxfp4` single-scale format. + +!!! note "Per-layer BF16 fallback in offline cascade models" + The A14B offline checkpoint uses `quant_method: mxfp4_dualscale`. Most + linear layers are stored as W4A4 MXFP4 DualScale; precision-sensitive layers + retain their original BF16 weights and are listed in `ignored_layers` inside + each transformer's `config.json`. The two transformers may have different + `ignored_layers` sets — the pipeline reads each transformer's own `config.json` + and rebuilds the config locally when they differ, so routing is always + per-transformer-accurate. + +!!! warning "TI2V-5B not supported" + Wan2.2-TI2V-5B is excluded from W4A4 quantization. Its smaller parameter + count makes it significantly more sensitive to 4-bit quantization noise, + resulting in unacceptable accuracy loss. Use [MXFP8](mxfp8.md) for TI2V-5B. + +## Configuration + +### `mxfp4` — Single-Scale Online Mode + +Online mode requires no pre-processing. vLLM-Omni quantizes BF16 weights to +MXFP4 at load time using `npu_dynamic_mx_quant`. A single block scale +(`float8_e8m0fnu`, one per 32 K elements) is computed on the fly; no +calibration `mul_scale` is available. Applies equally to single-transformer +and cascade (A14B) models — both transformers in a cascade receive the same +quantization config automatically. + +```python +from vllm_omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +omni = Omni(model="", quantization="mxfp4") +outputs = omni.generate( + "A cat sitting on a windowsill", + OmniDiffusionSamplingParams(num_inference_steps=50), +) +``` + +```bash +# Single-transformer or cascade model — same command +python text_to_video.py --model --quantization mxfp4 + +# Online serving +vllm serve --omni --quantization mxfp4 +``` + +### `mxfp4_dualscale` — DualScale Online Mode + +Online DualScale mode computes both fine and coarse scales on the fly from BF16 +weights using `npu_dynamic_dual_level_mx_quant`. Applies equally to +single-transformer and cascade (A14B) models. Compared to `mxfp4` online, +DualScale provides better quantization accuracy at higher compute cost. + +The default configuration keeps the leading 5 transformer blocks in BF16 +(`num_bf16_fallback_layers=5`). Accuracy evaluation on Wan2.2-A14B shows this +is sufficient to meet quality requirements and is the recommended setting. + +```python +omni = Omni(model="", quantization="mxfp4_dualscale") +``` + +```bash +python text_to_video.py --model --quantization mxfp4_dualscale +``` + +If accuracy debugging identifies additional precision-sensitive layers, they +can be pinned to BF16 via the Python API: + +```python +omni = Omni( + model="", + quantization={ + "method": "mxfp4_dualscale", + "ignored_layers": ["blocks.10.attn1.to_q"], # explicit per-layer override + }, +) +``` + +BF16 fallback routing in online mode applies two rules in priority order: + +1. **`ignored_layers`** (explicit per-layer override): any layer whose prefix + matches is kept in BF16 regardless of block index. +2. **`num_bf16_fallback_layers`** (coarse leading-block rule): the first N + transformer blocks (`blocks.0` … `blocks.N-1`) fall back to BF16. Defaults + to `5` (recommended). Layers outside `blocks.N.*` + (e.g. `condition_embedder`) are always quantized. + +### `mxfp4_dualscale` — DualScale Offline Mode (Recommended) + +Offline mode loads a pre-quantized DualScale checkpoint from msModelSlim. A +preprocessing step converts the raw quantized output to the diffusers format +expected by vLLM-Omni and injects the quantization config into each +`transformer/config.json` so that vLLM-Omni auto-detects the offline path +without a `--quantization` flag. + +BF16 fallback layers may be interleaved anywhere in the transformer — they are +not restricted to leading blocks. The merge script detects them from +`quant_model_description.json` and writes their prefixes into `ignored_layers` +inside `config.json`. At runtime, each layer's prefix is matched against +`ignored_layers` to decide BF16 vs. MXFP4 DualScale. + +#### Checkpoint tensor layout + +Each quantized linear layer stores four tensors: + +| Tensor | Shape | dtype | Description | +|--------|-------|-------|-------------| +| `weight` | `(N, K)` | float8_e4m3fn | FP4 packed (2 values per byte) | +| `weight_scale` | `(N, K//32)` | uint8 | Fine block scale (`float8_e8m0fnu` bit pattern) | +| `weight_dual_scale` | `(N, K//512, 1)` | float32 | Coarse block scale | +| `mul_scale` | `(K,)` | float32 | Per-input-channel smooth pre-scale (from calibration) | + +BF16 fallback layers have no quantization tensors; only the original `weight` +(and optional `bias`) are present, loaded directly from the base checkpoint. + +#### Step 1 — Quantize with msModelSlim + +```bash +msmodelslim quant \ + --model_path /path/to/Wan2.2-T2V-A14B-Diffusers \ + --save_path /path/to/wan2_2_t2v_quantized_raw \ + --device npu \ + --model_type Wan2_2 \ + --config_path /path/to/wan2_2_w4a4_mxfp4_dualscale.yaml \ + --trust_remote_code True +``` + +After this step, `--save_path` contains raw quantized safetensors files, +scale files, and a metadata JSON (`quant_model_description*.json`). + +For cascade MoE models (T2V-A14B, I2V-A14B), msModelSlim outputs two +subdirectories: `high_noise_model/` (transformer) and `low_noise_model/` +(transformer_2). + +#### Step 2 — Preprocess with merge_mxfp4_dualscale_checkpoint.py + +The script (`vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py`): + +1. Copies the original diffusers model to `--output-path` (VAE, text encoder, + scheduler, etc. are preserved). +2. Remaps tensor names from msModelSlim convention to diffusers convention and + strips `.linear.` / `.div.` wrappers added by the quantization tool. +3. Overlays MXFP4 tensors (weight, fine/coarse scales, `mul_scale`) onto the + BF16 base checkpoint. Non-quantized layers keep their original BF16 weights. +4. Detects all linear layers that remain in BF16 and writes their prefixes into + `ignored_layers` in `config.json`. +5. Injects `quantization_config` so vLLM-Omni auto-detects offline MXFP4 + DualScale. + +For cascade MoE models, steps 2–5 run separately for each transformer. + +```bash +python vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py \ + --model-type Wan2.2-T2V-A14B \ + --original-model /path/to/Wan2.2-T2V-A14B-Diffusers \ + --quant-path /path/to/wan2_2_t2v_quantized_raw \ + --output-path /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale +``` + +| Argument | Description | +|----------|-------------| +| `--model-type` | Model variant: `Wan2.2-T2V-A14B` or `Wan2.2-I2V-A14B` | +| `--original-model` | Root directory of the original BF16 diffusers model | +| `--quant-path` | Root directory of the msModelSlim quantized output | +| `--output-path` | Output directory for the merged model (created by the script) | + +The script outputs a complete diffusers model directory at `--output-path`, +with each transformer subfolder containing: + +- `diffusion_pytorch_model.safetensors` — MXFP4 weights + scale tensors, with BF16 fallback layers from the base checkpoint +- `config.json` — original transformer config with `quantization_config` injected +- `quant_model_description.json` — quantization metadata (reference only) + +The `quantization_config` injected into `config.json` for each transformer: + +```json +{ + "quant_method": "mxfp4_dualscale", + "is_checkpoint_serialized": true, + "ignored_layers": [ + "blocks.0.attn1.to_qkv", + "blocks.0.attn1.to_out", + "proj_out" + ] +} +``` + +`ignored_layers` lists every linear layer that retains its original BF16 weight, +using vllm-omni model parameter names (QKV-fused, FFN underscored, `to_out` +unindexed). The exact entries are determined by the quantization tool (msModelSlim) +and may differ between `transformer` and `transformer_2` in a cascade model. + +#### Step 3 — Serve + +```bash +python text_to_video.py --model /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale + +# Online serving +vllm serve /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale --omni +``` + +```python +omni = Omni(model="/path/to/Wan2.2-T2V-A14B-MXFP4-DualScale") +``` + +!!! note + No `--quantization` flag is needed for offline mode. The preprocessing + script injects `quantization_config` into each `transformer/config.json`, + which vLLM-Omni reads automatically to activate the correct offline path. + +## Parameters + +### `mxfp4` (single-scale, online only) + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `method` | str | — | `"mxfp4"` | +| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16 | + +### `mxfp4_dualscale` (dual-scale, online + offline) + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `method` | str | — | `"mxfp4_dualscale"` | +| `is_checkpoint_serialized` | bool | `False` | `True` for offline DualScale checkpoints; auto-set from `config.json` when using the preprocessing script | +| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16. **Works in both modes**: offline — populated by the merge script for interleaved sensitive layers; online — user-supplied for explicit per-layer precision override | +| `num_bf16_fallback_layers` | int | `5` | **Online mode only**: leading N transformer blocks (`blocks.0` … `blocks.N-1`) kept in BF16. Applied after `ignored_layers`; ignored in offline mode. Default of `5` is the evaluated recommended value for Wan2.2-A14B | + +#### BF16 fallback priority (online mode) + +``` +for each linear layer: + if prefix in ignored_layers → BF16 (explicit override, highest priority) + elif block_idx < num_bf16_fallback_layers → BF16 (coarse leading-block rule) + else → MXFP4 DualScale online +``` + +Layers outside `blocks.N.*` (e.g. `condition_embedder.*`) are always quantized +unless they appear in `ignored_layers`. + +## Validation and Notes + +1. **Online single-scale (`mxfp4`)** quantizes BF16 weights at load time using + `npu_dynamic_mx_quant` (single-scale). No calibration `mul_scale` is + available — all output partitions receive an identity pre-scale. No offline + checkpoint format exists for this method. + +2. **Online dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=False`)** + quantizes BF16 weights using `npu_dynamic_dual_level_mx_quant` (fine + coarse + scales computed on the fly). No calibration `mul_scale`; leading blocks or + explicit `ignored_layers` stay in BF16 for accuracy. + +3. **Offline dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=True`)** — + **recommended for production** — loads four tensors per quantized layer: FP4 + weight, fine scale (`uint8` reinterpreted as `float8_e8m0fnu`), coarse scale + (`float32`), and per-input-channel `mul_scale` (`float32`). BF16 fallback + layers have no quantization tensors and are routed via `ignored_layers`. + +4. **Scale dtype**: fine scales are stored as `uint8` in safetensors (same bit + layout as `float8_e8m0fnu`) and reinterpreted at load time without a lossy + float32 round-trip. + +5. **Cascade model config propagation**: in a cascade model (transformer + + transformer_2), vLLM-Omni reads each transformer's own `config.json` and + rebuilds the quant config locally when `ignored_layers` differs between + transformers, ensuring per-layer routing is accurate for each. The first + transformer's config is propagated to `od_config` so the second transformer + can reuse it as a starting point. + +6. **Self-attention QKV fusion**: Q, K, V projection weights are fused into a + single `QKVParallelLinear` layer at runtime. `ignored_layers` entries use the + fused name (`attn1.to_qkv`), written automatically by the merge script. + +7. W4A4 carries higher quantization noise than W8A8 (16 vs 256 levels). The + DualScale offline method mitigates this with calibrated `mul_scale` smooth + quantization. Use `ignored_layers` and `num_bf16_fallback_layers` to trade + off compression vs. accuracy for precision-sensitive layers. + +## Adapting MXFP4 for a New Model + +This section is aimed at developers who want to add MXFP4 support to a model +other than Wan2.2. The three integration points are: (1) discovering the correct +runtime layer names, (2) wiring `ignored_layers` into the model, and (3) writing +a merge script for offline checkpoints. + +### Step 1 — Discover runtime layer names + +`ignored_layers` entries must match the **runtime parameter names** used inside +vllm-omni, which may differ from the names stored in the diffusers checkpoint. +The canonical source of truth is the model's own `named_parameters()`. + +```python +from vllm_omni import Omni + +# Load the model without quantization to inspect parameter names. +omni = Omni(model="/path/to/your-model") # no --quantization flag +for name, _ in omni.pipeline.transformer.named_parameters(): + if "weight" in name and "scale" not in name: + print(name) +``` + +Compare the printed names against the diffusers checkpoint keys +(`safetensors.safe_open` or `torch.load`) to identify any renames your model +applies. Common patterns that differ in Wan2.2 (and may appear in other +models): + +| Diffusers checkpoint name | vllm-omni runtime name | Reason | +|---------------------------|------------------------|--------| +| `attn1.to_q`, `attn1.to_k`, `attn1.to_v` | `attn1.to_qkv` | Self-attention Q/K/V fused into `QKVParallelLinear` | +| `ffn.net.0.proj` | `ffn.net_0.proj` | Dots in sub-module names replaced with underscores | +| `ffn.net.2` | `ffn.net_2` | Same underscore rule | +| `to_out.0` | `to_out` | Sequential index stripped | + +If your model has different fusion patterns, inspect `packed_modules_mapping` +on the model class — this dict records how checkpoint keys are mapped to +fused runtime parameters. + +!!! warning "Partial QKV fallback is not allowed" + If your model fuses Q, K, V into a single layer, `ignored_layers` must + include **all three or none**. A partial fallback (e.g. `to_q` in BF16 but + `to_k`, `to_v` quantized) cannot be expressed at runtime because they share + one `QKVParallelLinear`. The merge script enforces this and raises an error + if only some of the trio appear as non-quantized. + +### Step 2 — Add ignored_layers to the model + +#### Online mode + +Pass `ignored_layers` directly in the quantization config using the **runtime +names** discovered in Step 1. No code changes to the model are required. + +```python +omni = Omni( + model="/path/to/your-model", + quantization={ + "method": "mxfp4_dualscale", + "ignored_layers": [ + "blocks.0.attn1.to_qkv", # runtime name, not diffusers name + "blocks.0.attn1.to_out", + "blocks.0.ffn.net_0.proj", + ], + }, +) +``` + +```bash +# CLI does not support list-typed ignored_layers directly. +# Use the Python API or set ignored_layers in config.json (offline). +python your_script.py --model /path/to/your-model --quantization mxfp4_dualscale +``` + +The `num_bf16_fallback_layers` coarse rule is an alternative to listing layers +individually: set it to N to keep all linear layers in blocks 0 … N-1 in BF16. +The right value depends on the model's sensitivity; evaluate on a validation +set and pick the smallest N that meets your accuracy target. + +#### Offline mode + +For offline checkpoints, `ignored_layers` is written into each transformer's +`config.json` by the merge script (see Step 3). No manual editing is needed if +the merge script is correct. The injected block: + +```json +{ + "quant_method": "mxfp4_dualscale", + "is_checkpoint_serialized": true, + "ignored_layers": [ + "blocks.0.attn1.to_qkv", + "blocks.0.attn1.to_out" + ] +} +``` + +To add a layer manually (e.g. to pin an additional layer to BF16 without +re-running the merge script), edit `config.json` inside the transformer +subfolder. Use runtime names, not diffusers checkpoint names. + +### Step 3 — Write a merge script for offline mode + +The merge script for a new model mirrors +`vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py`. The four +things it must do: + +1. **Remap tensor names** from the quantization tool convention to diffusers + convention (strip wrappers like `.linear.`, `.div.`; fix any prefix + differences). + +2. **Collect ignored_layers**: after loading, enumerate all `*.weight` keys that + have no corresponding `*.weight_scale` (i.e. layers the tool left in BF16). + Convert diffusers names to vllm-omni runtime names (fuse QKV, rename FFN + sub-modules, etc.). Write the result to `config.json`. + +3. **Inject `quantization_config`** into `config.json`: + ```python + config["quantization_config"] = { + "quant_method": "mxfp4_dualscale", + "is_checkpoint_serialized": True, + "ignored_layers": ignored_layers, # runtime names + } + ``` + +4. **Save** the merged safetensors and the updated `config.json`. + +The key helper to implement is the diffusers-to-runtime name translator +(equivalent to `_diffusers_to_vllm_ignored` in the Wan2.2 merge script). +For each non-quantized diffusers weight key, apply your model's specific +renaming rules and collect the results. + +!!! tip "Validate before serving" + After producing the offline checkpoint, load it without a `--quantization` + flag and verify that vLLM-Omni auto-detects the correct method. Check that + the layer count reported in the startup log matches expectations: quantized + layer count + `ignored_layers` count should equal total linear layer count. + Any mismatch indicates a name-mapping bug in the merge script. diff --git a/docs/user_guide/quantization/overview.md b/docs/user_guide/quantization/overview.md index 39669bf5d..38d027f12 100644 --- a/docs/user_guide/quantization/overview.md +++ b/docs/user_guide/quantization/overview.md @@ -9,20 +9,20 @@ type has a different quantization scope. | Mode | Guide | Description | Methods | |------|-------|-------------|---------| -| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, MXFP8 W8A8 | +| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, MXFP8 W8A8, MXFP4 W4A4 | | Runtime attention quantization | [Quantized KV Cache](quantized_kvcache.md) | vLLM-Omni dynamically quantizes eligible diffusion Flash Attention tensors during inference. | FP8 FA | -| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, GGUF, AutoRound, msModelSlim, serialized Int8, offline MXFP8 | +| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, GGUF, AutoRound, msModelSlim, serialized Int8, offline MXFP8, offline MXFP4 DualScale | ## Hardware Support -| Device | FP8 W8A8 | Int8 W8A8 | ModelOpt | MXFP8 W8A8 | GGUF | AutoRound | msModelSlim | -|--------|----------|-----------|----------|------------|------|-----------|-------------| -| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ⭕ | ✅ | ✅ | ❌ | -| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ⭕ | ✅ | ✅ | ❌ | -| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ | -| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ❌ | -| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ | -| Ascend NPU | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | +| Device | FP8 W8A8 | Int8 W8A8 | ModelOpt | MXFP8 W8A8 | MXFP4 W4A4 | GGUF | AutoRound | msModelSlim | +|--------|----------|-----------|----------|------------|------------|------|-----------|-------------| +| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ | +| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ | +| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ✅ | ❌ | +| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ❌ | +| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ | +| Ascend NPU | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide. FP8 on Ampere may use a weight-only path where available. @@ -42,6 +42,7 @@ otherwise. | Int8 W8A8 | [Int8](int8.md) | Online or serialized W8A8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image | | ModelOpt | [ModelOpt](modelopt.md) | Pre-quantized FP8 checkpoints | Qwen-Image, Z-Image, FLUX.2, HunyuanImage-3.0 | Validated for ModelOpt FP8 diffusion checkpoints | | MXFP8 W8A8 | [MXFP8](mxfp8.md) | Online W8A8 or offline pre-quantized | Wan2.2-T2V-A14B, I2V-A14B, TI2V-5B | Ascend NPU only; validated for Wan2.2 | +| MXFP4 W4A4 | [MXFP4](mxfp4.md) | `mxfp4`: online single-scale only; `mxfp4_dualscale`: online or offline dual-scale (offline recommended) | Wan2.2-T2V-A14B, I2V-A14B | Ascend NPU only; validated for Wan2.2 A14B cascade models; TI2V-5B not supported; offline `mxfp4_dualscale` uses calibrated `mul_scale` for best accuracy | | GGUF | [GGUF](gguf.md) | Pre-quantized transformer weights | Qwen-Image | Validated where a model-specific GGUF adapter exists | | AutoRound | [AutoRound](autoround.md) | Pre-quantized W4A16 checkpoints | FLUX.1-dev; Qwen-Image/Wan2.2 not validated | Checkpoint-driven | | msModelSlim | [msModelSlim](msmodelslim.md) | Pre-quantized Ascend checkpoints | Wan2.2 recipe; HunyuanImage-3.0 inference target | Ascend/NPU path | @@ -58,6 +59,7 @@ in BF16 unless the model guide explicitly adds support. | ModelOpt | [ModelOpt](modelopt.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path | | Int8 | [Int8](int8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | | MXFP8 | [MXFP8](mxfp8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | +| MXFP4 | [MXFP4](mxfp4.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | | GGUF | [GGUF](gguf.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | | AutoRound | [AutoRound](autoround.md) | Thinker or language-model checkpoint config | Qwen2.5-Omni, Qwen3-Omni | Supported through AutoRound checkpoints | | msModelSlim | [msModelSlim](msmodelslim.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | @@ -73,6 +75,7 @@ attached to the intended stage rather than applied globally. | Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation | | ModelOpt | [ModelOpt](modelopt.md) | Checkpoint-defined diffusion stage | BAGEL, GLM-Image | Requires model-specific validation | | MXFP8 | [MXFP8](mxfp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated | +| MXFP4 | [MXFP4](mxfp4.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated | | GGUF | [GGUF](gguf.md) | Stage-specific transformer weights | BAGEL, GLM-Image | No validated adapter listed | | AutoRound | [AutoRound](autoround.md) | Checkpoint-defined stage | BAGEL, GLM-Image | No validated checkpoint listed | | msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation | @@ -100,7 +103,7 @@ config = build_quant_config({ | Component | Default quantized? | Notes | |-----------|--------------------|-------| -| Diffusion transformer | Yes | Primary target for FP8, Int8, ModelOpt, GGUF, AutoRound, and msModelSlim | +| Diffusion transformer | Yes | Primary target for FP8, Int8, ModelOpt, MXFP8, MXFP4, GGUF, AutoRound, and msModelSlim | | Text encoder | No | Keep BF16 unless a method-specific guide documents support | | VAE | No | Keep BF16; storage-only paths are method-specific | | Scheduler/tokenizer | No | Loaded from the base model repository | diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index 9e2521233..e5bbacdca 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -173,8 +173,8 @@ def parse_args() -> argparse.Namespace: "--quantization", type=str, default=None, - choices=["fp8", "mxfp8", "int8", "gguf"], - help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).", + choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"], + help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).", ) # Distributed and parallel execution diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 9e1775dd1..8bbc871ea 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -163,8 +163,8 @@ def parse_args() -> argparse.Namespace: "--quantization", type=str, default=None, - choices=["fp8", "mxfp8", "int8", "gguf"], - help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).", + choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"], + help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).", ) # Distributed and parallel execution diff --git a/tests/diffusion/models/wan2_2/test_create_transformer_quant.py b/tests/diffusion/models/wan2_2/test_create_transformer_quant.py new file mode 100644 index 000000000..bb95363fd --- /dev/null +++ b/tests/diffusion/models/wan2_2/test_create_transformer_quant.py @@ -0,0 +1,383 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for transformer quant-config auto-detection. + +The loader path at pipeline_wan2_2.py carries two quantization contracts: + + create_transformer_from_config (~L137) + - Reads quantization_config from config.json (injected by the merge scripts) + - Auto-detects the quant method when no CLI quant_config is provided + - Rejects method mismatches (CLI vs disk) + - Upgrades online → offline when disk marks is_checkpoint_*_serialized=True + - Rebuilds when the active ignored_layers differs from the disk value + + Wan22Pipeline._create_transformer (~L456) + - Passes od_config.quantization_config (set by CLI or externally) to + create_transformer_from_config for each transformer independently. + - When od_config.quantization_config is None, each transformer auto-detects + from its own config.json; od_config is NOT modified by this call. + +All tests are pure-CPU and do not load model weights. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import create_transformer_from_config + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + +# Minimum config that create_transformer_from_config accepts without raising. +_MIN_CFG: dict = {"patch_size": [1, 2, 2]} + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_fake_transformer(): + """Return (FakeTransformer class, captured list). + + Each _FakeTransformer.__init__ call appends its **kwargs to captured, + letting tests inspect what quant_config was passed per transformer. + """ + captured: list[dict] = [] + + class _FakeTransformer: + def __init__(self, **kwargs): + captured.append(kwargs) + + return _FakeTransformer, captured + + +class _FakePipeline: + """Minimal stand-in exposing only what _create_transformer needs from self.""" + + def __init__(self, od_config: SimpleNamespace) -> None: + self.od_config = od_config + + # Bind the real unbound method so the tests exercise production code. + _create_transformer = wan22_module.Wan22Pipeline._create_transformer + + +# --------------------------------------------------------------------------- +# create_transformer_from_config — auto-detection +# --------------------------------------------------------------------------- + + +def test_create_transformer_detects_mxfp8_serialized_from_config_json(monkeypatch): + """When config.json carries MXFP8 quant and no CLI quant_config is provided, + the transformer must receive a DiffusionMXFP8Config with serialized=True.""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + config = { + **_MIN_CFG, + "quantization_config": { + "quant_method": "mxfp8", + "is_checkpoint_mxfp8_serialized": True, + }, + } + create_transformer_from_config(config) + + qc = captured[0].get("quant_config") + assert isinstance(qc, DiffusionMXFP8Config) + assert qc.is_checkpoint_mxfp8_serialized is True + + +def test_create_transformer_detects_mxfp4_dualscale_from_config_json(monkeypatch): + """config.json with mxfp4_dualscale + ignored_layers must produce + a DiffusionMXFP4DualScaleMixedConfig with the correct ignored_layers and + serialized flag.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + config = { + **_MIN_CFG, + "quantization_config": { + "quant_method": "mxfp4_dualscale", + "ignored_layers": ["blocks.0.attn1.to_q", "blocks.0.attn1.to_k"], + "is_checkpoint_serialized": True, + }, + } + create_transformer_from_config(config) + + qc = captured[0].get("quant_config") + assert isinstance(qc, DiffusionMXFP4DualScaleMixedConfig) + assert set(qc.ignored_layers) == {"blocks.0.attn1.to_q", "blocks.0.attn1.to_k"} + assert qc.is_checkpoint_serialized is True + + +def test_create_transformer_without_quantization_config_passes_no_quant(monkeypatch): + """A plain BF16 config.json (no quantization_config key) must result in no + quant_config being passed to WanTransformer3DModel.""" + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + create_transformer_from_config(_MIN_CFG) + + assert "quant_config" not in captured[0] + + +# --------------------------------------------------------------------------- +# create_transformer_from_config — method-mismatch guard +# --------------------------------------------------------------------------- + + +def test_create_transformer_rejects_method_mismatch(monkeypatch): + """Passing a CLI quant_config whose get_name() differs from the config.json + quant_method must raise ValueError immediately (prevents silent weight corruption). + + fp8 (vLLM built-in, get_name()=='fp8') vs disk 'mxfp8' triggers the guard. + These are distinct methods; using the same method for both would not trigger it. + """ + from vllm.model_executor.layers.quantization.fp8 import Fp8Config + + FakeTransformer, _ = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + fp8_cli = Fp8Config(is_checkpoint_fp8_serialized=True, activation_scheme="dynamic") + config = { + **_MIN_CFG, + "quantization_config": { + "quant_method": "mxfp8", + "is_checkpoint_mxfp8_serialized": True, + }, + } + with pytest.raises(ValueError, match="quant_method"): + create_transformer_from_config(config, quant_config=fp8_cli) + + +# --------------------------------------------------------------------------- +# create_transformer_from_config — online → offline upgrade +# --------------------------------------------------------------------------- + + +def test_create_transformer_upgrades_to_serialized_when_disk_marks_it(monkeypatch): + """CLI passes online (is_checkpoint_mxfp8_serialized=False) but config.json + marks is_checkpoint_mxfp8_serialized=True → must switch to offline (serialized) + so that pre-quantized FP8 tensors are loaded correctly.""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + online_cli = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=False) + config = { + **_MIN_CFG, + "quantization_config": { + "quant_method": "mxfp8", + "is_checkpoint_mxfp8_serialized": True, + }, + } + create_transformer_from_config(config, quant_config=online_cli) + + qc = captured[0].get("quant_config") + assert isinstance(qc, DiffusionMXFP8Config) + assert qc.is_checkpoint_mxfp8_serialized is True + + +# --------------------------------------------------------------------------- +# create_transformer_from_config — ignored_layers rebuild +# --------------------------------------------------------------------------- + + +def test_create_transformer_rebuilds_when_ignored_layers_differ(monkeypatch): + """When the active quant_config has different ignored_layers than config.json, + the config must be rebuilt from disk so per-layer BF16 routing is authoritative + for this specific transformer.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + stale = DiffusionMXFP4DualScaleMixedConfig( + is_checkpoint_serialized=True, + ignored_layers=["blocks.0.attn1.to_q"], + ) + config = { + **_MIN_CFG, + "quantization_config": { + "quant_method": "mxfp4_dualscale", + "ignored_layers": ["blocks.0.attn1.to_q", "blocks.1.attn1.to_q"], + "is_checkpoint_serialized": True, + }, + } + create_transformer_from_config(config, quant_config=stale) + + qc = captured[0].get("quant_config") + assert isinstance(qc, DiffusionMXFP4DualScaleMixedConfig) + assert set(qc.ignored_layers) == {"blocks.0.attn1.to_q", "blocks.1.attn1.to_q"} + + +def test_create_transformer_does_not_rebuild_when_ignored_layers_match(monkeypatch): + """When the active quant_config already has the same ignored_layers, + the same instance must be passed through unchanged (no unnecessary rebuild).""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + matching = DiffusionMXFP4DualScaleMixedConfig( + is_checkpoint_serialized=True, + ignored_layers=["blocks.0.attn1.to_q"], + ) + config = { + **_MIN_CFG, + "quantization_config": { + "quant_method": "mxfp4_dualscale", + "ignored_layers": ["blocks.0.attn1.to_q"], + "is_checkpoint_serialized": True, + }, + } + create_transformer_from_config(config, quant_config=matching) + + assert captured[0].get("quant_config") is matching + + +# --------------------------------------------------------------------------- +# Wan22Pipeline._create_transformer — od_config passthrough +# --------------------------------------------------------------------------- + + +def test_pipeline_create_transformer_auto_detects_from_config_json(monkeypatch): + """When od_config.quantization_config is None, _create_transformer must + auto-detect the quant method from config.json and pass it to the transformer. + od_config itself must remain unchanged (None).""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + od_config = SimpleNamespace(quantization_config=None) + pipeline = _FakePipeline(od_config) + + config = { + **_MIN_CFG, + "quantization_config": { + "quant_method": "mxfp8", + "is_checkpoint_mxfp8_serialized": True, + }, + } + pipeline._create_transformer(config) + + qc = captured[0].get("quant_config") + assert isinstance(qc, DiffusionMXFP8Config) + assert qc.is_checkpoint_mxfp8_serialized is True + # od_config must NOT be modified — each transformer auto-detects independently. + assert od_config.quantization_config is None + + +def test_pipeline_create_transformer_does_not_overwrite_existing_od_config(monkeypatch): + """If od_config.quantization_config is already set (e.g. via CLI), _create_transformer + must pass it through to the transformer unchanged and leave od_config unmodified.""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + FakeTransformer, _ = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + existing = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=True) + od_config = SimpleNamespace(quantization_config=existing) + pipeline = _FakePipeline(od_config) + + config = { + **_MIN_CFG, + "quantization_config": { + "quant_method": "mxfp8", + "is_checkpoint_mxfp8_serialized": True, + }, + } + pipeline._create_transformer(config) + + assert od_config.quantization_config is existing + + +# --------------------------------------------------------------------------- +# Wan22Pipeline._create_transformer — cascade contracts +# --------------------------------------------------------------------------- + + +def test_pipeline_cascade_both_transformers_get_mxfp8_serialized_config(monkeypatch): + """Cascade model (transformer + transformer_2) with MXFP8 checkpoint: + - Each transformer auto-detects from its own config.json independently. + Both must receive is_checkpoint_mxfp8_serialized=True.""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + od_config = SimpleNamespace(quantization_config=None) + pipeline = _FakePipeline(od_config) + + mxfp8_qc = {"quant_method": "mxfp8", "is_checkpoint_mxfp8_serialized": True} + pipeline._create_transformer({**_MIN_CFG, "quantization_config": mxfp8_qc}) + pipeline._create_transformer({**_MIN_CFG, "quantization_config": mxfp8_qc}) + + assert len(captured) == 2 + for i, kwargs in enumerate(captured): + qc = kwargs.get("quant_config") + assert isinstance(qc, DiffusionMXFP8Config), f"transformer[{i}]: expected DiffusionMXFP8Config, got {type(qc)}" + assert qc.is_checkpoint_mxfp8_serialized is True, f"transformer[{i}]: expected serialized=True" + + +def test_pipeline_cascade_mxfp4_dualscale_each_transformer_gets_correct_ignored_layers(monkeypatch): + """Cascade with mxfp4_dualscale where transformer and transformer_2 have + different ignored_layers in their config.json. + + Expected outcome: + transformer → ignored_layers=["blocks.0.attn1.to_q"] (auto-detected from its config.json) + transformer_2 → ignored_layers=["blocks.0.attn1.to_q", "blocks.1.attn1.to_q"] + (auto-detected from its own config.json independently) + od_config → quantization_config remains None (not modified by _create_transformer) + """ + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + od_config = SimpleNamespace(quantization_config=None) + pipeline = _FakePipeline(od_config) + + cfg1 = { + **_MIN_CFG, + "quantization_config": { + "quant_method": "mxfp4_dualscale", + "ignored_layers": ["blocks.0.attn1.to_q"], + "is_checkpoint_serialized": True, + }, + } + cfg2 = { + **_MIN_CFG, + "quantization_config": { + "quant_method": "mxfp4_dualscale", + "ignored_layers": ["blocks.0.attn1.to_q", "blocks.1.attn1.to_q"], + "is_checkpoint_serialized": True, + }, + } + + pipeline._create_transformer(cfg1) + pipeline._create_transformer(cfg2) + + assert len(captured) == 2 + qc1 = captured[0].get("quant_config") + qc2 = captured[1].get("quant_config") + + assert isinstance(qc1, DiffusionMXFP4DualScaleMixedConfig) + assert isinstance(qc2, DiffusionMXFP4DualScaleMixedConfig) + assert set(qc1.ignored_layers) == {"blocks.0.attn1.to_q"}, f"transformer expected 1 layer, got {qc1.ignored_layers}" + assert set(qc2.ignored_layers) == { + "blocks.0.attn1.to_q", + "blocks.1.attn1.to_q", + }, f"transformer_2 expected 2 layers, got {qc2.ignored_layers}" + + # od_config must remain unchanged — _create_transformer does not modify it. + assert od_config.quantization_config is None diff --git a/tests/diffusion/quantization/test_mxfp4_config.py b/tests/diffusion/quantization/test_mxfp4_config.py new file mode 100644 index 000000000..af3a9c9c4 --- /dev/null +++ b/tests/diffusion/quantization/test_mxfp4_config.py @@ -0,0 +1,539 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for MXFP4 quantization configs and the MXFP4 DualScale + BF16 mixed config.""" + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +@pytest.fixture(autouse=True) +def _patch_tp_state(monkeypatch): + """Patch TP rank/world_size so ModelWeightParameter can be instantiated on CPU + without an initialized distributed group. Returns TP=1 rank=0 for all tests.""" + monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1) + + +# --------------------------------------------------------------------------- +# DiffusionMXFP4Config +# --------------------------------------------------------------------------- + + +def test_mxfp4_config_get_name(): + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + assert DiffusionMXFP4Config.get_name() == "mxfp4" + + +def test_mxfp4_config_from_config_defaults(): + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + cfg = DiffusionMXFP4Config.from_config({}) + assert cfg.is_checkpoint_mxfp4_serialized is False + assert cfg.ignored_layers == [] + + +def test_mxfp4_config_from_config_serialized(): + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + cfg = DiffusionMXFP4Config.from_config({"is_checkpoint_mxfp4_serialized": True}) + assert cfg.is_checkpoint_mxfp4_serialized is True + + +def test_mxfp4_config_from_config_ignored_layers(): + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + cfg = DiffusionMXFP4Config.from_config({"ignored_layers": ["proj_out"]}) + assert cfg.ignored_layers == ["proj_out"] + + +def test_mxfp4_config_from_config_modules_to_not_convert_fallback(): + """modules_to_not_convert must be accepted as an alias for ignored_layers.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + cfg = DiffusionMXFP4Config.from_config({"modules_to_not_convert": ["proj_out"]}) + assert cfg.ignored_layers == ["proj_out"] + + +# --------------------------------------------------------------------------- +# build_quant_config integration +# --------------------------------------------------------------------------- + + +def test_build_quant_config_mxfp4_string(): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + cfg = build_quant_config("mxfp4") + assert isinstance(cfg, DiffusionMXFP4Config) + assert cfg.get_name() == "mxfp4" + assert cfg.is_checkpoint_mxfp4_serialized is False + + +def test_build_quant_config_mxfp4_dict(): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + cfg = build_quant_config({"method": "mxfp4", "is_checkpoint_mxfp4_serialized": True}) + assert isinstance(cfg, DiffusionMXFP4Config) + assert cfg.is_checkpoint_mxfp4_serialized is True + + +def test_build_quant_config_mxfp4_dualscale_string(): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = build_quant_config("mxfp4_dualscale") + assert isinstance(cfg, DiffusionMXFP4DualScaleMixedConfig) + assert cfg.is_checkpoint_serialized is False + assert cfg.num_bf16_fallback_layers == 5 + assert cfg.ignored_layers == [] + + +def test_build_quant_config_mxfp4_dualscale_dict_offline(): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = build_quant_config( + { + "method": "mxfp4_dualscale", + "is_checkpoint_serialized": True, + "ignored_layers": ["blocks.0.attn1.to_q", "blocks.0.attn1.to_k"], + } + ) + assert isinstance(cfg, DiffusionMXFP4DualScaleMixedConfig) + assert cfg.is_checkpoint_serialized is True + assert cfg.ignored_layers == ["blocks.0.attn1.to_q", "blocks.0.attn1.to_k"] + + +def test_build_quant_config_mxfp4_dualscale_dict_online_custom_fallback(): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = build_quant_config({"method": "mxfp4_dualscale", "num_bf16_fallback_layers": 10}) + assert isinstance(cfg, DiffusionMXFP4DualScaleMixedConfig) + assert cfg.num_bf16_fallback_layers == 10 + + +# --------------------------------------------------------------------------- +# Block-index dispatch (_parse_block_idx) +# --------------------------------------------------------------------------- + + +def test_parse_block_idx_valid(): + from vllm_omni.quantization.mxfp4_config import _parse_block_idx + + assert _parse_block_idx("blocks.0.attn1.to_q") == 0 + assert _parse_block_idx("blocks.5.ffn.net.0.proj") == 5 + assert _parse_block_idx("blocks.40.norm1.weight") == 40 + + +def test_parse_block_idx_non_block_prefixes(): + """Prefixes that do not start with 'blocks.N.' must return None.""" + from vllm_omni.quantization.mxfp4_config import _parse_block_idx + + assert _parse_block_idx("condition_embedder.time_embedder.linear_1") is None + assert _parse_block_idx("proj_out.weight") is None + assert _parse_block_idx("model.layers.0.self_attn.q_proj") is None + assert _parse_block_idx("scale_shift_table") is None + + +# --------------------------------------------------------------------------- +# SUPPORTED_QUANTIZATION_METHODS +# --------------------------------------------------------------------------- + + +def test_supported_methods_include_mxfp4_variants(): + from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS + + assert "mxfp4" in SUPPORTED_QUANTIZATION_METHODS + assert "mxfp8" in SUPPORTED_QUANTIZATION_METHODS + assert "mxfp4_dualscale" in SUPPORTED_QUANTIZATION_METHODS + + +# --------------------------------------------------------------------------- +# DiffusionMXFP4DualScaleMixedConfig — config roundtrips +# --------------------------------------------------------------------------- + + +def test_mixed_dualscale_config_get_name(): + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + assert DiffusionMXFP4DualScaleMixedConfig.get_name() == "mxfp4_dualscale" + + +def test_mixed_dualscale_config_no_args_defaults(): + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = DiffusionMXFP4DualScaleMixedConfig() + assert cfg.is_checkpoint_serialized is False + assert cfg.ignored_layers == [] + assert cfg.num_bf16_fallback_layers == 5 + + +def test_mixed_dualscale_config_from_config_offline(): + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = DiffusionMXFP4DualScaleMixedConfig.from_config( + { + "quant_method": "mxfp4_dualscale", + "is_checkpoint_serialized": True, + "ignored_layers": ["blocks.0.attn1.to_q", "proj_out"], + } + ) + assert cfg.is_checkpoint_serialized is True + assert cfg.ignored_layers == ["blocks.0.attn1.to_q", "proj_out"] + assert cfg.num_bf16_fallback_layers == 5 # default + + +def test_mixed_dualscale_config_from_config_online_custom_fallback(): + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = DiffusionMXFP4DualScaleMixedConfig.from_config({"num_bf16_fallback_layers": 10}) + assert cfg.is_checkpoint_serialized is False + assert cfg.num_bf16_fallback_layers == 10 + + +def test_mixed_dualscale_config_from_config_modules_to_not_convert_fallback(): + """modules_to_not_convert must be accepted as an alias for ignored_layers.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = DiffusionMXFP4DualScaleMixedConfig.from_config( + {"is_checkpoint_serialized": True, "modules_to_not_convert": ["proj_out"]} + ) + assert cfg.ignored_layers == ["proj_out"] + + +# --------------------------------------------------------------------------- +# DiffusionMXFP4DualScaleMixedConfig — get_quant_method dispatch +# --------------------------------------------------------------------------- + + +def test_mixed_dualscale_offline_ignored_layer_returns_unquantized( + mocker, + monkeypatch: pytest.MonkeyPatch, +): + """Offline: a prefix in ignored_layers must return UnquantizedLinearMethod.""" + from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + + from vllm_omni.platforms import current_omni_platform + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = DiffusionMXFP4DualScaleMixedConfig( + is_checkpoint_serialized=True, + ignored_layers=["blocks.0.attn1.to_q"], + ) + layer = mocker.Mock(spec=LinearBase) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + method = cfg.get_quant_method(layer, "blocks.0.attn1.to_q") + assert isinstance(method, UnquantizedLinearMethod) + + +def test_mixed_dualscale_offline_non_ignored_returns_mxfp4( + mocker, + monkeypatch: pytest.MonkeyPatch, +): + """Offline: a prefix NOT in ignored_layers must return NPUMxfp4DualScaleLinearMethod.""" + from vllm.model_executor.layers.linear import LinearBase + + from vllm_omni.platforms import current_omni_platform + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod + + cfg = DiffusionMXFP4DualScaleMixedConfig( + is_checkpoint_serialized=True, + ignored_layers=["blocks.0.attn1.to_q"], + ) + layer = mocker.Mock(spec=LinearBase) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + method = cfg.get_quant_method(layer, "blocks.1.attn1.to_q") + assert isinstance(method, NPUMxfp4DualScaleLinearMethod) + + +def test_mixed_dualscale_online_fallback_block_returns_unquantized( + mocker, + monkeypatch: pytest.MonkeyPatch, +): + """Online: blocks < num_bf16_fallback_layers must return UnquantizedLinearMethod.""" + from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + + from vllm_omni.platforms import current_omni_platform + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=False, num_bf16_fallback_layers=5) + layer = mocker.Mock(spec=LinearBase) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + assert isinstance(cfg.get_quant_method(layer, "blocks.0.attn1.to_q"), UnquantizedLinearMethod) + assert isinstance(cfg.get_quant_method(layer, "blocks.4.ffn.net.0.proj"), UnquantizedLinearMethod) + + +def test_mixed_dualscale_online_quantized_block_returns_mxfp4( + mocker, + monkeypatch: pytest.MonkeyPatch, +): + """Online: blocks >= num_bf16_fallback_layers must return NPUMxfp4DualScaleOnlineLinearMethod.""" + from vllm.model_executor.layers.linear import LinearBase + + from vllm_omni.platforms import current_omni_platform + from vllm_omni.quantization.mxfp4_config import ( + DiffusionMXFP4DualScaleMixedConfig, + NPUMxfp4DualScaleOnlineLinearMethod, + ) + + cfg = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=False, num_bf16_fallback_layers=5) + layer = mocker.Mock(spec=LinearBase) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + assert isinstance(cfg.get_quant_method(layer, "blocks.5.attn1.to_q"), NPUMxfp4DualScaleOnlineLinearMethod) + assert isinstance(cfg.get_quant_method(layer, "blocks.40.ffn.net.0.proj"), NPUMxfp4DualScaleOnlineLinearMethod) + + +def test_mixed_dualscale_online_non_block_prefix_returns_mxfp4( + mocker, + monkeypatch: pytest.MonkeyPatch, +): + """Online: layers outside 'blocks.N.*' (condition_embedder etc.) always use MXFP4 online.""" + from vllm.model_executor.layers.linear import LinearBase + + from vllm_omni.platforms import current_omni_platform + from vllm_omni.quantization.mxfp4_config import ( + DiffusionMXFP4DualScaleMixedConfig, + NPUMxfp4DualScaleOnlineLinearMethod, + ) + + cfg = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=False, num_bf16_fallback_layers=5) + layer = mocker.Mock(spec=LinearBase) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + method = cfg.get_quant_method(layer, "condition_embedder.time_embedder.linear_1") + assert isinstance(method, NPUMxfp4DualScaleOnlineLinearMethod) + + +def test_mixed_dualscale_online_ignored_layers_override( + mocker, + monkeypatch: pytest.MonkeyPatch, +): + """Online: explicit ignored_layers must return UnquantizedLinearMethod regardless of block index. + + A layer that is NOT in the leading-block range (block 10 >= num_bf16_fallback_layers=5) + but IS listed in ignored_layers must still fall back to BF16. This lets power users + pin specific interleaved layers to BF16 during online quantization without needing an + offline checkpoint. + """ + from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + + from vllm_omni.platforms import current_omni_platform + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = DiffusionMXFP4DualScaleMixedConfig( + is_checkpoint_serialized=False, + num_bf16_fallback_layers=5, + ignored_layers=["blocks.10.attn1.to_q"], + ) + layer = mocker.Mock(spec=LinearBase) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + # block 10 is above the leading-block threshold but is in ignored_layers → BF16 + assert isinstance(cfg.get_quant_method(layer, "blocks.10.attn1.to_q"), UnquantizedLinearMethod) + + +def test_mixed_dualscale_non_linear_returns_none(monkeypatch: pytest.MonkeyPatch): + """Non-LinearBase layers (norms, embeddings) must return None → no quantization.""" + from vllm_omni.platforms import current_omni_platform + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = DiffusionMXFP4DualScaleMixedConfig() + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + norm_layer = torch.nn.LayerNorm(64) + assert cfg.get_quant_method(norm_layer, "blocks.0.norm1") is None + + +# --------------------------------------------------------------------------- +# TP=2 create_weights: parameter shapes and input_dim/output_dim +# +# Two scenarios mirror real Wan2.2 A14B linear layer types: +# Column-parallel (to_q, ffn.net_0): output is sharded (N/TP), input is full (K). +# Row-parallel (to_out, ffn.net_2): input is sharded (K/TP), output is full (N). +# +# Tests verify: +# 1. Registered parameter shapes are correct for each partition configuration. +# 2. input_dim/output_dim attributes are set so RowParallelLinear.weight_loader +# can shard scale tensors correctly (the fix for the TP>1 shape-mismatch bug). +# 3. Simulated loader slicing: slicing the full checkpoint tensor along the +# declared input_dim produces the exact shape stored in the parameter — +# proving the dim declaration is consistent with the allocation. +# --------------------------------------------------------------------------- + +# K must be divisible by 32 (fine groups) and 512 (coarse groups). +_TP2_K, _TP2_N, _TP2 = 1024, 512, 2 + + +class _FakeLayer(torch.nn.Module): + """Bare nn.Module that accepts register_parameter without a real weight_loader.""" + + +def _create_weights(method, *, input_size_per_partition, output_partition_sizes): + layer = _FakeLayer() + method.create_weights( + layer=layer, + input_size_per_partition=input_size_per_partition, + output_partition_sizes=output_partition_sizes, + input_size=_TP2_K, + output_size=_TP2_N, + params_dtype=torch.bfloat16, + ) + return layer + + +def _shard(tensor, param, rank, tp, dim_attr): + """Slice `tensor` along the dimension given by `param.` for `rank`.""" + dim = getattr(param, dim_attr) + if dim is None: + return tensor # not sharded along this axis + shard_size = param.shape[dim] + slices = [slice(None)] * tensor.ndim + slices[dim] = slice(rank * shard_size, (rank + 1) * shard_size) + return tensor[tuple(slices)] + + +# ── DualScale method ───────────────────────────────────────────────────────── + + +def test_dualscale_column_parallel_tp2_shapes(): + """Column-parallel TP=2: output halved, fine/coarse groups stay full, mul_scale full.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod + + method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig()) + layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N // _TP2]) + + assert layer.weight.shape == (_TP2_N // _TP2, _TP2_K) + assert layer.weight_scale.shape == (_TP2_N // _TP2, _TP2_K // 32) + assert layer.weight_dual_scale.shape == (_TP2_N // _TP2, _TP2_K // 512, 1) + assert layer.mul_scale.shape == (_TP2_K,) + + +def test_dualscale_row_parallel_tp2_shapes(): + """Row-parallel TP=2: input halved, fine/coarse groups halved, mul_scale halved.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod + + method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig()) + layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N]) + + assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2) + assert layer.weight_scale.shape == (_TP2_N, (_TP2_K // _TP2) // 32) + assert layer.weight_dual_scale.shape == (_TP2_N, (_TP2_K // _TP2) // 512, 1) + assert layer.mul_scale.shape == (_TP2_K // _TP2,) + + +def test_dualscale_scale_parameter_input_dims(): + """weight_scale/weight_dual_scale must have input_dim=1; mul_scale must have input_dim=0. + + RowParallelLinear.weight_loader only shards a parameter when input_dim is set. + Without these, loading a full checkpoint tensor into a per-rank shape causes a + shape mismatch for TP>1 row-parallel layers (to_out, ffn.net_2). + """ + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod + + method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig()) + layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N]) + + assert layer.weight_scale.input_dim == 1 + assert layer.weight_scale.output_dim == 0 + assert layer.weight_dual_scale.input_dim == 1 + assert layer.weight_dual_scale.output_dim == 0 + assert layer.mul_scale.input_dim == 0 + assert layer.mul_scale.output_dim is None + + +def test_dualscale_row_parallel_tp2_loader_simulation(): + """Slicing full checkpoint tensors along input_dim must match row-parallel parameter shapes. + + Simulates what RowParallelLinear.weight_loader does: for each scale parameter, + take the slice at rank*shard_size:(rank+1)*shard_size along input_dim. + The resulting shape must equal the per-rank parameter shape allocated by create_weights. + """ + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod + + method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig()) + layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N]) + + # Full checkpoint tensors (what the loader reads from disk). + ckpt_weight_scale = torch.zeros(_TP2_N, _TP2_K // 32) + ckpt_weight_dual_scale = torch.zeros(_TP2_N, _TP2_K // 512, 1) + ckpt_mul_scale = torch.zeros(_TP2_K) + + for rank in range(_TP2): + assert _shard(ckpt_weight_scale, layer.weight_scale, rank, _TP2, "input_dim").shape == layer.weight_scale.shape + assert ( + _shard(ckpt_weight_dual_scale, layer.weight_dual_scale, rank, _TP2, "input_dim").shape + == layer.weight_dual_scale.shape + ) + assert _shard(ckpt_mul_scale, layer.mul_scale, rank, _TP2, "input_dim").shape == layer.mul_scale.shape + + +def test_dualscale_column_parallel_tp2_loader_simulation(): + """Slicing full checkpoint tensors along output_dim must match column-parallel parameter shapes. + + For column-parallel layers, the loader shards along output_dim (rows). + mul_scale has output_dim=None → not sharded (full tensor, same for all ranks). + """ + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod + + method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig()) + layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N // _TP2]) + + ckpt_weight_scale = torch.zeros(_TP2_N, _TP2_K // 32) + ckpt_weight_dual_scale = torch.zeros(_TP2_N, _TP2_K // 512, 1) + ckpt_mul_scale = torch.zeros(_TP2_K) + + for rank in range(_TP2): + assert _shard(ckpt_weight_scale, layer.weight_scale, rank, _TP2, "output_dim").shape == layer.weight_scale.shape + assert ( + _shard(ckpt_weight_dual_scale, layer.weight_dual_scale, rank, _TP2, "output_dim").shape + == layer.weight_dual_scale.shape + ) + # mul_scale: output_dim=None → no sharding → full tensor fits the column-parallel parameter + assert _shard(ckpt_mul_scale, layer.mul_scale, rank, _TP2, "output_dim").shape == layer.mul_scale.shape + + +# ── Single-scale method ─────────────────────────────────────────────────────── + + +def test_single_scale_row_parallel_tp2_shapes(): + """Row-parallel TP=2: input halved → weight_scale groups halved.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, NPUMxfp4LinearMethod + + method = NPUMxfp4LinearMethod(DiffusionMXFP4Config()) + layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N]) + + assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2) + assert layer.weight_scale.shape == (_TP2_N, (_TP2_K // _TP2) // 32) + + +def test_single_scale_scale_parameter_input_dims(): + """Single-scale weight_scale must have input_dim=1 for RowParallel TP sharding.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, NPUMxfp4LinearMethod + + method = NPUMxfp4LinearMethod(DiffusionMXFP4Config()) + layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N]) + + assert layer.weight_scale.input_dim == 1 + assert layer.weight_scale.output_dim == 0 + + +def test_single_scale_row_parallel_tp2_loader_simulation(): + """Slicing full checkpoint weight_scale along input_dim matches row-parallel parameter shape.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, NPUMxfp4LinearMethod + + method = NPUMxfp4LinearMethod(DiffusionMXFP4Config()) + layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N]) + + ckpt_weight_scale = torch.zeros(_TP2_N, _TP2_K // 32) + + for rank in range(_TP2): + assert _shard(ckpt_weight_scale, layer.weight_scale, rank, _TP2, "input_dim").shape == layer.weight_scale.shape diff --git a/tests/diffusion/quantization/test_mxfp4_key_remap.py b/tests/diffusion/quantization/test_mxfp4_key_remap.py new file mode 100644 index 000000000..e0cbc84df --- /dev/null +++ b/tests/diffusion/quantization/test_mxfp4_key_remap.py @@ -0,0 +1,399 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for merge_mxfp4_dualscale_checkpoint.py key-remapping helpers. + +These are pure-Python unit tests that exercise the transformation functions +without loading any actual checkpoint files or requiring NPU hardware. +""" + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + +# --------------------------------------------------------------------------- +# SUPPORTED_MODEL_TYPES +# --------------------------------------------------------------------------- + + +def test_supported_model_types_includes_a14b(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import SUPPORTED_MODEL_TYPES + + assert "Wan2.2-T2V-A14B" in SUPPORTED_MODEL_TYPES + assert "Wan2.2-I2V-A14B" in SUPPORTED_MODEL_TYPES + + +def test_supported_model_types_excludes_ti2v_5b(): + """TI2V-5B is explicitly NOT supported under W4A4 quantization.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import SUPPORTED_MODEL_TYPES + + assert "Wan2.2-TI2V-5B" not in SUPPORTED_MODEL_TYPES + + +# --------------------------------------------------------------------------- +# _apply_rename_dict +# --------------------------------------------------------------------------- + + +def test_apply_rename_dict_self_attn_qkvo(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict + + assert _apply_rename_dict("blocks.0.self_attn.q.weight") == "blocks.0.attn1.to_q.weight" + assert _apply_rename_dict("blocks.0.self_attn.k.weight") == "blocks.0.attn1.to_k.weight" + assert _apply_rename_dict("blocks.0.self_attn.v.weight") == "blocks.0.attn1.to_v.weight" + assert _apply_rename_dict("blocks.0.self_attn.o.weight") == "blocks.0.attn1.to_out.0.weight" + + +def test_apply_rename_dict_ffn(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict + + assert _apply_rename_dict("blocks.1.ffn.0.weight") == "blocks.1.ffn.net.0.proj.weight" + assert _apply_rename_dict("blocks.1.ffn.2.weight") == "blocks.1.ffn.net.2.weight" + + +def test_apply_rename_dict_norm_order_swap(): + """norm2↔norm3 swap: quant tool uses norm1/norm3/norm2 order, + Diffusers uses norm1/norm2/norm3.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict + + assert _apply_rename_dict("blocks.0.norm2.weight") == "blocks.0.norm3.weight" + assert _apply_rename_dict("blocks.0.norm3.weight") == "blocks.0.norm2.weight" + + +def test_apply_rename_dict_cross_attn(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict + + assert _apply_rename_dict("blocks.0.cross_attn.q.weight") == "blocks.0.attn2.to_q.weight" + assert _apply_rename_dict("blocks.0.cross_attn.k.weight") == "blocks.0.attn2.to_k.weight" + assert _apply_rename_dict("blocks.0.cross_attn.v.weight") == "blocks.0.attn2.to_v.weight" + assert _apply_rename_dict("blocks.0.cross_attn.o.weight") == "blocks.0.attn2.to_out.0.weight" + + +def test_apply_rename_dict_head_and_modulation(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict + + assert _apply_rename_dict("head.head.weight") == "proj_out.weight" + + +# --------------------------------------------------------------------------- +# _strip_mxfp4_wrapper +# --------------------------------------------------------------------------- + + +def test_strip_mxfp4_wrapper_linear_weight(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper + + assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.linear.weight") == "blocks.0.attn1.to_q.weight" + + +def test_strip_mxfp4_wrapper_linear_weight_scale(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper + + assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.linear.weight_scale") == "blocks.0.attn1.to_q.weight_scale" + + +def test_strip_mxfp4_wrapper_linear_weight_dual_scale(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper + + assert ( + _strip_mxfp4_wrapper("blocks.0.attn1.to_q.linear.weight_dual_scale") == "blocks.0.attn1.to_q.weight_dual_scale" + ) + + +def test_strip_mxfp4_wrapper_linear_bias(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper + + assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.linear.bias") == "blocks.0.attn1.to_q.bias" + + +def test_strip_mxfp4_wrapper_div_mul_scale(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper + + assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.div.mul_scale") == "blocks.0.attn1.to_q.mul_scale" + + +def test_strip_mxfp4_wrapper_noop_for_plain_weight(): + """MXFP8 / FLOAT tensors have no wrapper — must be returned unchanged.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper + + assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.weight") == "blocks.0.attn1.to_q.weight" + assert _strip_mxfp4_wrapper("blocks.0.norm_q.weight") == "blocks.0.norm_q.weight" + assert _strip_mxfp4_wrapper("condition_embedder.time_embedder.linear_1.weight") == ( + "condition_embedder.time_embedder.linear_1.weight" + ) + + +# --------------------------------------------------------------------------- +# _classify_blocks +# --------------------------------------------------------------------------- + + +def test_classify_blocks_mixed_mxfp8_and_mxfp4(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _classify_blocks + + quant_meta = { + "blocks.0.attn1.to_q.weight": "W8A8_MXFP8", + "blocks.1.attn1.to_q.weight": "W8A8_MXFP8", + "blocks.2.attn1.to_q.weight": "W4A4_MXFP4_DUALSCALE", + "blocks.3.attn1.to_q.weight": "W4A4_MXFP4_DUALSCALE", + "condition_embedder.time_embedder.linear_1.weight": "FLOAT", + } + block_types = _classify_blocks(quant_meta) + assert block_types[0] == "mxfp8" + assert block_types[1] == "mxfp8" + assert block_types[2] == "mxfp4_dualscale" + assert block_types[3] == "mxfp4_dualscale" + # Non-block key must not produce an entry + assert None not in block_types + + +def test_classify_blocks_float_entries_skipped(): + """FLOAT-typed tensors must not contribute to block classification.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _classify_blocks + + quant_meta = { + "blocks.0.bias": "FLOAT", + "blocks.1.attn1.to_q.weight": "W4A4_MXFP4_DUALSCALE", + } + block_types = _classify_blocks(quant_meta) + assert 0 not in block_types + assert block_types[1] == "mxfp4_dualscale" + + +def test_classify_blocks_empty(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _classify_blocks + + assert _classify_blocks({}) == {} + + +# --------------------------------------------------------------------------- +# _is_mxfp4_tensor +# --------------------------------------------------------------------------- + + +def test_is_mxfp4_tensor_quantized_weight(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _is_mxfp4_tensor + + assert _is_mxfp4_tensor("blocks.0.attn1.to_q.linear.weight", "W4A4_MXFP4_DUALSCALE") is True + assert _is_mxfp4_tensor("blocks.0.attn1.to_q.linear.weight_scale", "W4A4_MXFP4_DUALSCALE") is True + assert _is_mxfp4_tensor("blocks.0.attn1.to_q.linear.weight_dual_scale", "W4A4_MXFP4_DUALSCALE") is True + + +def test_is_mxfp4_tensor_companion_bias_and_mul_scale(): + """Companion tensors (bias, mul_scale) are FLOAT but belong to MXFP4 layers + and must be included in the merged checkpoint.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _is_mxfp4_tensor + + # bias inside .linear. wrapper → must be included + assert _is_mxfp4_tensor("blocks.0.attn1.to_q.linear.bias", "FLOAT") is True + # mul_scale inside .div. wrapper → must be included + assert _is_mxfp4_tensor("blocks.0.attn1.to_q.div.mul_scale", "FLOAT") is True + + +def test_is_mxfp4_tensor_bf16_fallback_weight(): + """BF16 fallback linear layers have plain .weight keys (no wrapper) → must NOT be included.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _is_mxfp4_tensor + + assert _is_mxfp4_tensor("blocks.0.attn1.to_q.weight", "FLOAT") is False + assert _is_mxfp4_tensor("condition_embedder.time_embedder.linear_1.weight", "FLOAT") is False + + +def test_is_mxfp4_tensor_norm_weight(): + """Norm layers are always FLOAT with plain .weight → not MXFP4 related.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _is_mxfp4_tensor + + assert _is_mxfp4_tensor("blocks.0.norm1.weight", "FLOAT") is False + + +# --------------------------------------------------------------------------- +# _collect_ignored_layers +# --------------------------------------------------------------------------- + + +def test_collect_ignored_layers_basic(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _collect_ignored_layers + + # block 0: all three of Q/K/V are BF16 → fused as to_qkv in ignored_layers + # block 1: to_qkv is MXFP4 → not in ignored_layers + merged = { + "blocks.0.attn1.to_q.weight": torch.zeros(4, 4), + "blocks.0.attn1.to_k.weight": torch.zeros(4, 4), + "blocks.0.attn1.to_v.weight": torch.zeros(4, 4), + "blocks.1.attn1.to_q.weight": torch.zeros(4, 4), + "blocks.1.attn1.to_q.weight_scale": torch.zeros(4, 1), + } + mxfp4_prefixes = {"blocks.1.attn1.to_q"} + + ignored = _collect_ignored_layers(merged, mxfp4_prefixes) + assert "blocks.0.attn1.to_qkv" in ignored # fused vllm-omni name + assert "blocks.0.attn1.to_q" not in ignored # diffusers name must not appear + assert "blocks.0.attn1.to_k" not in ignored + assert "blocks.1.attn1.to_qkv" not in ignored # MXFP4, not ignored + + +def test_collect_ignored_layers_empty_mxfp4(): + """When no layers are MXFP4 (all BF16), all .weight prefixes become ignored_layers.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _collect_ignored_layers + + # Use non-attn layers to avoid the Q/K/V completeness requirement. + merged = { + "blocks.0.ffn.net.0.proj.weight": torch.zeros(4, 4), + "proj_out.weight": torch.zeros(4, 4), + } + ignored = _collect_ignored_layers(merged, set()) + assert sorted(ignored) == ["blocks.0.ffn.net_0.proj", "proj_out"] + + +def test_collect_ignored_layers_returns_sorted(): + """ignored_layers must be sorted for deterministic config.json output.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _collect_ignored_layers + + # Use FFN layers (no Q/K/V fusion concern) to isolate the sort guarantee. + merged = { + "blocks.2.ffn.net.0.proj.weight": torch.zeros(1), + "blocks.0.ffn.net.0.proj.weight": torch.zeros(1), + "blocks.1.ffn.net.0.proj.weight": torch.zeros(1), + } + ignored = _collect_ignored_layers(merged, set()) + assert ignored == sorted(ignored) + + +# --------------------------------------------------------------------------- +# _build_quant_config (new mxfp4_dualscale format) +# --------------------------------------------------------------------------- + + +def test_build_quant_config_new_format(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _build_quant_config + + config = _build_quant_config(["blocks.0.attn1.to_q", "proj_out"]) + assert config["quant_method"] == "mxfp4_dualscale" + assert config["is_checkpoint_serialized"] is True + assert config["ignored_layers"] == ["blocks.0.attn1.to_q", "proj_out"] + assert "num_mxfp8_blocks" not in config + + +def test_build_quant_config_empty_ignored_layers(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _build_quant_config + + config = _build_quant_config([]) + assert config["ignored_layers"] == [] + + +# --------------------------------------------------------------------------- +# _diffusers_to_vllm_ignored +# --------------------------------------------------------------------------- + + +def test_remap_self_attn_qkv_fusion(): + """All three of attn1.to_q/k/v in BF16 → fused attn1.to_qkv.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored + + result = _diffusers_to_vllm_ignored(["blocks.0.attn1.to_k", "blocks.0.attn1.to_q", "blocks.0.attn1.to_v"]) + assert result == ["blocks.0.attn1.to_qkv"] + + +def test_remap_self_attn_qkv_partial_raises(): + """Partial Q/K/V in ignored_layers is invalid — must raise ValueError. + + Self-attention Q/K/V are fused into a single to_qkv layer at runtime; + partial precision (some BF16, some MXFP4) cannot be expressed and must + be caught early in the merge script. + """ + import pytest + + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored + + with pytest.raises(ValueError, match="Partial BF16 fallback"): + _diffusers_to_vllm_ignored(["blocks.0.attn1.to_q", "blocks.0.attn1.to_k"]) + + +def test_remap_cross_attn_kept_separate(): + """Cross-attention (attn2) to_q/k/v must NOT be fused.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored + + result = _diffusers_to_vllm_ignored(["blocks.0.attn2.to_k", "blocks.0.attn2.to_q", "blocks.0.attn2.to_v"]) + assert "blocks.0.attn2.to_qkv" not in result + assert "blocks.0.attn2.to_q" in result + assert "blocks.0.attn2.to_k" in result + assert "blocks.0.attn2.to_v" in result + + +def test_remap_ffn_dot_to_underscore(): + """ffn.net.0.proj → ffn.net_0.proj; ffn.net.2 → ffn.net_2.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored + + result = _diffusers_to_vllm_ignored(["blocks.0.ffn.net.0.proj", "blocks.0.ffn.net.2"]) + assert "blocks.0.ffn.net_0.proj" in result + assert "blocks.0.ffn.net_2" in result + assert "blocks.0.ffn.net.0.proj" not in result + assert "blocks.0.ffn.net.2" not in result + + +def test_remap_to_out_strips_index(): + """attn1.to_out.0 and attn2.to_out.0 must have the .0 suffix removed.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored + + result = _diffusers_to_vllm_ignored(["blocks.0.attn1.to_out.0", "blocks.0.attn2.to_out.0"]) + assert "blocks.0.attn1.to_out" in result + assert "blocks.0.attn2.to_out" in result + assert "blocks.0.attn1.to_out.0" not in result + assert "blocks.0.attn2.to_out.0" not in result + + +def test_remap_noop_for_non_block_layers(): + """Layers outside 'blocks.N.*' (condition_embedder, proj_out) pass through unchanged.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored + + inputs = [ + "condition_embedder.time_embedder.linear_1", + "condition_embedder.text_embedder.linear_1", + "proj_out", + ] + result = _diffusers_to_vllm_ignored(inputs) + assert result == sorted(inputs) + + +def test_remap_mixed_block_layers(): + """Typical mixed block: QKV fused, to_out stripped, FFN renamed, cross-attn separate.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored + + inputs = [ + "blocks.0.attn1.to_k", + "blocks.0.attn1.to_out.0", + "blocks.0.attn1.to_q", + "blocks.0.attn1.to_v", + "blocks.0.attn2.to_k", + "blocks.0.attn2.to_out.0", + "blocks.0.attn2.to_q", + "blocks.0.attn2.to_v", + "blocks.0.ffn.net.0.proj", + "blocks.0.ffn.net.2", + ] + result = _diffusers_to_vllm_ignored(inputs) + assert "blocks.0.attn1.to_qkv" in result + assert "blocks.0.attn1.to_out" in result + assert "blocks.0.attn2.to_q" in result + assert "blocks.0.attn2.to_k" in result + assert "blocks.0.attn2.to_v" in result + assert "blocks.0.attn2.to_out" in result + assert "blocks.0.ffn.net_0.proj" in result + assert "blocks.0.ffn.net_2" in result + # Old names must be gone + assert "blocks.0.attn1.to_q" not in result + assert "blocks.0.ffn.net.0.proj" not in result + + +def test_remap_returns_sorted(): + """Output of _diffusers_to_vllm_ignored must always be sorted.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored + + inputs = [ + "blocks.2.ffn.net.0.proj", + "blocks.0.attn1.to_q", + "blocks.0.attn1.to_k", + "blocks.0.attn1.to_v", + "blocks.1.attn2.to_out.0", + ] + result = _diffusers_to_vllm_ignored(inputs) + assert result == sorted(result) diff --git a/tests/diffusion/quantization/test_mxfp8_config.py b/tests/diffusion/quantization/test_mxfp8_config.py new file mode 100644 index 000000000..d831b5d47 --- /dev/null +++ b/tests/diffusion/quantization/test_mxfp8_config.py @@ -0,0 +1,308 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for MXFP8 quantization config and linear method dispatch. + +Coverage: +- DiffusionMXFP8Config.from_config roundtrips (CPU, no NPU required) +- get_quant_method dispatch (mocked platform) +- MXFPLinearMethodBase.apply() reshape skeleton (CPU) +- Weight / scale shape-transform arithmetic from process_weights_after_loading (CPU) +- build_quant_config integration +- MXFP8_QUANT_CONFIG structure as the auto-detection contract +""" + +import pytest +import torch +from pytest_mock import MockerFixture +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + +from vllm_omni.platforms import current_omni_platform +from vllm_omni.quantization import build_quant_config + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + +npu_available = pytest.mark.skipif(not current_omni_platform.is_npu(), reason="NPU platform not available") + + +# --------------------------------------------------------------------------- +# DiffusionMXFP8Config — from_config roundtrips +# --------------------------------------------------------------------------- + + +def test_mxfp8_config_get_name(): + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + assert DiffusionMXFP8Config.get_name() == "mxfp8" + + +def test_mxfp8_config_from_config_defaults(): + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + cfg = DiffusionMXFP8Config.from_config({}) + assert cfg.is_checkpoint_mxfp8_serialized is False + assert cfg.ignored_layers == [] + + +def test_mxfp8_config_from_config_serialized(): + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + cfg = DiffusionMXFP8Config.from_config({"is_checkpoint_mxfp8_serialized": True}) + assert cfg.is_checkpoint_mxfp8_serialized is True + + +def test_mxfp8_config_from_config_ignored_layers(): + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + cfg = DiffusionMXFP8Config.from_config({"ignored_layers": ["proj_out"]}) + assert cfg.ignored_layers == ["proj_out"] + + +def test_mxfp8_config_from_config_modules_to_not_convert_fallback(): + """modules_to_not_convert must be accepted as an alias for ignored_layers.""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + cfg = DiffusionMXFP8Config.from_config({"modules_to_not_convert": ["proj_out"]}) + assert cfg.ignored_layers == ["proj_out"] + + +# --------------------------------------------------------------------------- +# build_quant_config integration +# --------------------------------------------------------------------------- + + +def test_build_quant_config_mxfp8_string(): + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + cfg = build_quant_config("mxfp8") + assert isinstance(cfg, DiffusionMXFP8Config) + assert cfg.get_name() == "mxfp8" + assert cfg.is_checkpoint_mxfp8_serialized is False + + +def test_build_quant_config_mxfp8_dict(): + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + cfg = build_quant_config({"method": "mxfp8", "is_checkpoint_mxfp8_serialized": True}) + assert isinstance(cfg, DiffusionMXFP8Config) + assert cfg.is_checkpoint_mxfp8_serialized is True + + +def test_build_quant_config_mxfp8_config_json_format(): + """Verify that the exact quantization_config injected by merge_mxfp8_checkpoint.py + is accepted by build_quant_config and selects the offline (serialized) path. + + This is the critical auto-detection contract: TransformerConfig.from_dict() + reads quant_method + is_checkpoint_mxfp8_serialized to pick NPUMxfp8LinearMethod. + """ + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import MXFP8_QUANT_CONFIG + + cfg = build_quant_config(MXFP8_QUANT_CONFIG) + assert isinstance(cfg, DiffusionMXFP8Config) + assert cfg.is_checkpoint_mxfp8_serialized is True + + +def test_mxfp8_quant_config_structure(): + """MXFP8_QUANT_CONFIG must contain exactly the keys that auto-detection reads.""" + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import MXFP8_QUANT_CONFIG + + assert MXFP8_QUANT_CONFIG.get("quant_method") == "mxfp8" + assert MXFP8_QUANT_CONFIG.get("is_checkpoint_mxfp8_serialized") is True + + +# --------------------------------------------------------------------------- +# get_quant_method dispatch +# --------------------------------------------------------------------------- + + +def test_get_quant_method_npu_offline(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): + """Offline (serialized) path on NPU must return NPUMxfp8LinearMethod.""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config, NPUMxfp8LinearMethod + + config = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=True) + layer = mocker.Mock(spec=LinearBase) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + method = config.get_quant_method(layer, "blocks.0.attn1.to_q") + assert isinstance(method, NPUMxfp8LinearMethod) + + +def test_get_quant_method_npu_online(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): + """Online (BF16 checkpoint) path on NPU must return NPUMxfp8OnlineLinearMethod.""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config, NPUMxfp8OnlineLinearMethod + + config = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=False) + layer = mocker.Mock(spec=LinearBase) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + method = config.get_quant_method(layer, "blocks.0.attn1.to_q") + assert isinstance(method, NPUMxfp8OnlineLinearMethod) + + +def test_get_quant_method_non_npu_raises(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): + """Non-NPU platform must raise NotImplementedError (CUDA not yet supported).""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + config = DiffusionMXFP8Config() + layer = mocker.Mock(spec=LinearBase) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: False) + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False) + + with pytest.raises(NotImplementedError): + config.get_quant_method(layer, "blocks.0.attn1.to_q") + + +def test_get_quant_method_ignored_layer(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): + """A prefix in ignored_layers must bypass quantization → UnquantizedLinearMethod.""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + config = DiffusionMXFP8Config(ignored_layers=["proj_out"]) + layer = mocker.Mock(spec=LinearBase) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + method = config.get_quant_method(layer, "proj_out") + assert isinstance(method, UnquantizedLinearMethod) + + +def test_get_quant_method_non_linear_returns_none(monkeypatch: pytest.MonkeyPatch): + """Non-LinearBase layers (norms, embeddings) must get None → no quantization.""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + config = DiffusionMXFP8Config() + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + norm_layer = torch.nn.LayerNorm(64) + assert config.get_quant_method(norm_layer, "blocks.0.norm1") is None + + +# --------------------------------------------------------------------------- +# MXFPLinearMethodBase.apply() — reshape skeleton (CPU, no NPU) +# --------------------------------------------------------------------------- + + +def test_apply_reshape_skeleton(): + """apply() must flatten batch dims → _apply_inner → restore original leading dims.""" + from vllm_omni.quantization.mxfp8_config import MXFPLinearMethodBase + + OUT_FEATURES = 4 + + class _StubMethod(MXFPLinearMethodBase): + def create_weights( + self, + layer, + input_size_per_partition, + output_partition_sizes, + input_size, + output_size, + params_dtype, + **extra_weight_attrs, + ): + pass + + def _quantize_activation(self, x): + return x, None + + def _quant_matmul(self, x_q, x_scale, layer, bias, ori_dtype): + return torch.zeros(x_q.shape[0], OUT_FEATURES, dtype=ori_dtype) + + method = _StubMethod() + x = torch.randn(2, 3, 8) # (batch=2, seq=3, K=8) + out = method.apply(None, x) + assert out.shape == (2, 3, OUT_FEATURES) + + +def test_apply_reshape_with_bias(): + """apply() must pass bias through to _apply_inner unchanged.""" + from vllm_omni.quantization.mxfp8_config import MXFPLinearMethodBase + + received_bias = [] + + class _StubMethod(MXFPLinearMethodBase): + def create_weights( + self, + layer, + input_size_per_partition, + output_partition_sizes, + input_size, + output_size, + params_dtype, + **extra_weight_attrs, + ): + pass + + def _quantize_activation(self, x): + return x, None + + def _quant_matmul(self, x_q, x_scale, layer, bias, ori_dtype): + received_bias.append(bias) + return torch.zeros(x_q.shape[0], 4, dtype=ori_dtype) + + method = _StubMethod() + bias = torch.zeros(4) + method.apply(None, torch.randn(2, 8), bias=bias) + assert received_bias[0] is bias + + +# --------------------------------------------------------------------------- +# process_weights_after_loading shape arithmetic (pure torch, no NPU ops) +# +# These tests replicate the CPU-safe portions of process_weights_after_loading +# to guard the key layout contract: (N,K) weight → (K,N) and (N,S) scale → +# (S/2,N,2). They do NOT call NPU ops; they test only the math. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "n, k", + [(64, 128), (32, 64), (16, 96)], + ids=["64x128", "32x64", "16x96"], +) +def test_weight_transpose_contract(n: int, k: int): + """Weight must be transposed from (N, K) to (K, N) and be contiguous.""" + w = torch.zeros(n, k, dtype=torch.uint8) + w = w.transpose(0, 1).contiguous() + assert w.shape == (k, n) + assert w.is_contiguous() + + +@pytest.mark.parametrize( + "n, k_groups, expected_groups", + [ + (64, 4, 2), # even K_groups — no padding + (64, 3, 2), # odd K_groups — padded to 4 + (32, 1, 1), # odd K_groups — padded to 2 + ], + ids=["even", "odd-to-4", "odd-to-2"], +) +def test_weight_scale_reshape_contract(n: int, k_groups: int, expected_groups: int): + """Scale must be reshaped from (N, K_groups) to (K_groups_even//2, N, 2). + + Odd K_groups must be padded to even before the reshape. + """ + s = torch.zeros(n, k_groups, dtype=torch.uint8) + if k_groups % 2 == 1: + s = torch.cat([s, torch.zeros(n, 1, dtype=s.dtype)], dim=1) + k_groups += 1 + s = s.reshape(n, k_groups // 2, 2).transpose(0, 1).contiguous() + assert s.shape == (expected_groups, n, 2) + assert s.is_contiguous() + + +def test_num_groups_formula(): + """K_groups formula: ceil(K / 32) — spot-check boundary values.""" + assert (31 + 31) // 32 == 1 # K=31 → 1 group + assert (32 + 31) // 32 == 1 # K=32 → 1 group + assert (33 + 31) // 32 == 2 # K=33 → 2 groups + assert (128 + 31) // 32 == 4 # K=128 → 4 groups (even) + assert (96 + 31) // 32 == 3 # K=96 → 3 groups (odd → needs padding) + + +# --------------------------------------------------------------------------- +# SUPPORTED_QUANTIZATION_METHODS +# --------------------------------------------------------------------------- + + +def test_supported_methods_include_mxfp8(): + from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS + + assert "mxfp8" in SUPPORTED_QUANTIZATION_METHODS diff --git a/tests/diffusion/quantization/test_mxfp8_key_remap.py b/tests/diffusion/quantization/test_mxfp8_key_remap.py new file mode 100644 index 000000000..f6eea04a0 --- /dev/null +++ b/tests/diffusion/quantization/test_mxfp8_key_remap.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for merge_mxfp8_checkpoint.py key-remapping helpers and model metadata. + +These are pure-Python unit tests that exercise the transformation functions +without loading any actual checkpoint files or requiring NPU hardware. + +Key contracts verified: +- SUPPORTED_MODEL_TYPES includes Wan2.2-TI2V-5B (MXFP8 supports it; MXFP4 does not) +- MXFP8_QUANT_CONFIG structure matches what TransformerConfig.from_dict() reads +- _remap_keys correctly translates msModelSlim naming → Diffusers naming +- _get_transformer_dirs routes cascade vs single-transformer models +- _get_quant_subdir maps high/low noise subdirs for cascade models +""" + +import pathlib + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +# --------------------------------------------------------------------------- +# SUPPORTED_MODEL_TYPES and CASCADE_MODEL_TYPES +# --------------------------------------------------------------------------- + + +def test_supported_model_types_includes_all_wan22(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import SUPPORTED_MODEL_TYPES + + assert "Wan2.2-T2V-A14B" in SUPPORTED_MODEL_TYPES + assert "Wan2.2-I2V-A14B" in SUPPORTED_MODEL_TYPES + + +def test_supported_model_types_includes_ti2v_5b(): + """MXFP8 supports TI2V-5B (contrast: MXFP4 explicitly excludes it).""" + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import SUPPORTED_MODEL_TYPES + + assert "Wan2.2-TI2V-5B" in SUPPORTED_MODEL_TYPES + + +def test_cascade_model_types_excludes_ti2v(): + """TI2V-5B is a single-transformer model — must not be in CASCADE_MODEL_TYPES.""" + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import CASCADE_MODEL_TYPES + + assert "Wan2.2-TI2V-5B" not in CASCADE_MODEL_TYPES + assert "Wan2.2-T2V-A14B" in CASCADE_MODEL_TYPES + assert "Wan2.2-I2V-A14B" in CASCADE_MODEL_TYPES + + +# --------------------------------------------------------------------------- +# MXFP8_QUANT_CONFIG — auto-detection contract +# --------------------------------------------------------------------------- + + +def test_mxfp8_quant_config_has_required_keys(): + """MXFP8_QUANT_CONFIG must carry exactly the keys that auto-detection reads: + quant_method (selects DiffusionMXFP8Config) and is_checkpoint_mxfp8_serialized + (selects NPUMxfp8LinearMethod over the online path).""" + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import MXFP8_QUANT_CONFIG + + assert MXFP8_QUANT_CONFIG["quant_method"] == "mxfp8" + assert MXFP8_QUANT_CONFIG["is_checkpoint_mxfp8_serialized"] is True + + +# --------------------------------------------------------------------------- +# _get_transformer_dirs +# --------------------------------------------------------------------------- + + +def test_get_transformer_dirs_cascade(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_transformer_dirs + + assert _get_transformer_dirs("Wan2.2-T2V-A14B") == ["transformer", "transformer_2"] + assert _get_transformer_dirs("Wan2.2-I2V-A14B") == ["transformer", "transformer_2"] + + +def test_get_transformer_dirs_single(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_transformer_dirs + + assert _get_transformer_dirs("Wan2.2-TI2V-5B") == ["transformer"] + + +# --------------------------------------------------------------------------- +# _get_quant_subdir +# --------------------------------------------------------------------------- + + +def test_get_quant_subdir_cascade_high_noise(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_quant_subdir + + base = pathlib.Path("/quant") + result = _get_quant_subdir("Wan2.2-T2V-A14B", base, "transformer") + assert result == base / "high_noise_model" + + +def test_get_quant_subdir_cascade_low_noise(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_quant_subdir + + base = pathlib.Path("/quant") + result = _get_quant_subdir("Wan2.2-T2V-A14B", base, "transformer_2") + assert result == base / "low_noise_model" + + +def test_get_quant_subdir_non_cascade(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_quant_subdir + + base = pathlib.Path("/quant") + result = _get_quant_subdir("Wan2.2-TI2V-5B", base, "transformer") + assert result == base + + +# --------------------------------------------------------------------------- +# _remap_keys — msModelSlim naming → Diffusers naming +# --------------------------------------------------------------------------- + + +def test_remap_keys_self_attn_q(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys + + state = {"blocks.0.self_attn.q.weight": torch.zeros(1)} + meta = {"blocks.0.self_attn.q.weight": "W8A8_MXFP8"} + new_state, new_meta = _remap_keys(state, meta) + assert "blocks.0.attn1.to_q.weight" in new_state + assert new_meta.get("blocks.0.attn1.to_q.weight") == "W8A8_MXFP8" + + +def test_remap_keys_self_attn_all_heads(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys + + pairs = { + "self_attn.q": "attn1.to_q", + "self_attn.k": "attn1.to_k", + "self_attn.v": "attn1.to_v", + "self_attn.o": "attn1.to_out.0", + } + for src_part, dst_part in pairs.items(): + src_key = f"blocks.0.{src_part}.weight" + state = {src_key: torch.zeros(1)} + new_state, _ = _remap_keys(state, {}) + expected = f"blocks.0.{dst_part}.weight" + assert expected in new_state, f"{src_key} → expected {expected}, got {list(new_state)}" + + +def test_remap_keys_ffn(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys + + state = { + "blocks.1.ffn.0.weight": torch.zeros(1), + "blocks.1.ffn.2.weight": torch.zeros(1), + } + new_state, _ = _remap_keys(state, {}) + assert "blocks.1.ffn.net.0.proj.weight" in new_state + assert "blocks.1.ffn.net.2.weight" in new_state + + +def test_remap_keys_norm_order_swap(): + """norm2↔norm3 swap: msModelSlim uses norm1/norm3/norm2 order, + Diffusers uses norm1/norm2/norm3.""" + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys + + state = { + "blocks.0.norm2.weight": torch.zeros(1), + "blocks.0.norm3.weight": torch.zeros(1), + } + new_state, _ = _remap_keys(state, {}) + # norm2 → norm3 and norm3 → norm2 + assert "blocks.0.norm3.weight" in new_state + assert "blocks.0.norm2.weight" in new_state + # Both must be present (swap, not collapse) + assert len([k for k in new_state if "norm" in k and "norm_q" not in k and "norm_k" not in k]) == 2 + + +def test_remap_keys_head(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys + + state = {"head.head.weight": torch.zeros(1)} + new_state, _ = _remap_keys(state, {}) + assert "proj_out.weight" in new_state + + +def test_remap_keys_cross_attn(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys + + state = {"blocks.0.cross_attn.q.weight": torch.zeros(1)} + new_state, _ = _remap_keys(state, {}) + assert "blocks.0.attn2.to_q.weight" in new_state + + +def test_remap_keys_meta_only_mapped_for_existing_state_keys(): + """quant_meta entries are only emitted for keys present in state_dict.""" + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys + + state = {"blocks.0.self_attn.q.weight": torch.zeros(1)} + # meta has an extra key not in state_dict + meta = { + "blocks.0.self_attn.q.weight": "W8A8_MXFP8", + "blocks.0.self_attn.q.weight_scale": "W8A8_MXFP8", + } + _, new_meta = _remap_keys(state, meta) + assert "blocks.0.attn1.to_q.weight" in new_meta + # weight_scale was in meta but not state_dict → must be absent + assert "blocks.0.attn1.to_q.weight_scale" not in new_meta + + +def test_remap_keys_preserves_tensors(): + """Tensor values must survive the key rename unchanged.""" + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys + + t = torch.randn(4, 8) + state = {"blocks.0.self_attn.q.weight": t} + new_state, _ = _remap_keys(state, {}) + assert torch.equal(new_state["blocks.0.attn1.to_q.weight"], t) diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index b4438c795..3ee46ffb0 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -159,54 +159,10 @@ def create_transformer_from_config( if "pos_embed_seq_len" in config: kwargs["pos_embed_seq_len"] = config["pos_embed_seq_len"] - # Auto-detect quantization from transformer's config.json when not explicitly provided. - # merge_mxfp8_checkpoint.py injects quantization_config into config.json so that - # offline quantized checkpoints are recognized here without a CLI flag. if "quantization_config" in config: - from vllm_omni.quantization.factory import build_quant_config + from vllm_omni.quantization.factory import resolve_quant_config_from_disk - disk_qc = config["quantization_config"] - if isinstance(disk_qc, dict) and "quant_method" in disk_qc: - qc_method = disk_qc["quant_method"] - qc_kwargs = {k: v for k, v in disk_qc.items() if k != "quant_method"} - if quant_config is None: - # No CLI flag: full auto-detection. - quant_config = build_quant_config(qc_method, **qc_kwargs) - logger.info( - "Auto-detected quantization from transformer config.json: method=%s kwargs=%s", - qc_method, - qc_kwargs, - ) - elif quant_config.get_name() != qc_method: - # The caller supplied a quant_config of a different method than - # what the checkpoint was built with. Loading serialized tensors - # (e.g. MXFP8 weight scales) with the wrong linear method would - # produce corrupt output or a shape mismatch crash. Reject early - # so the user gets a clear message instead of a silent failure. - raise ValueError( - f"Checkpoint config.json declares quant_method={qc_method!r} but the " - f"active quantization config is {quant_config.get_name()!r}. " - "Pass a matching --quantization flag or omit it for auto-detection." - ) - elif ( - qc_kwargs.get("is_checkpoint_mxfp8_serialized", False) - and hasattr(quant_config, "is_checkpoint_mxfp8_serialized") - and not quant_config.is_checkpoint_mxfp8_serialized - ): - # Same method: CLI provided online mode but config.json marks this - # as a pre-quantized offline checkpoint. Switch to offline mode so - # users can pass --quantization mxfp8 without knowing the - # online/offline distinction. - quant_config = build_quant_config(qc_method, **qc_kwargs) - logger.info( - "config.json marks checkpoint as serialized; switching from online to offline MXFP8 mode.", - ) - elif isinstance(disk_qc, str) and quant_config is None: - quant_config = build_quant_config(disk_qc) - logger.info( - "Auto-detected quantization from transformer config.json: method=%s", - disk_qc, - ) + quant_config = resolve_quant_config_from_disk(quant_config, config["quantization_config"]) if quant_config is not None: kwargs["quant_config"] = quant_config diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py index ea7b8a7fe..955f97cef 100644 --- a/vllm_omni/quantization/factory.py +++ b/vllm_omni/quantization/factory.py @@ -48,6 +48,29 @@ def _build_mxfp8(**kw: Any) -> QuantizationConfig: return DiffusionMXFP8Config(**kw) +def _build_mxfp4(**kw: Any) -> QuantizationConfig: + """Lazy import for W4A4 MXFP4 diffusion config (NPU only).""" + from .mxfp4_config import DiffusionMXFP4Config + + return DiffusionMXFP4Config(**kw) + + +def _build_mxfp4_dualscale(**kw: Any) -> QuantizationConfig: + """Lazy import for MXFP4 DualScale + BF16 mixed diffusion config (NPU only). + + Offline mode (is_checkpoint_serialized=True): + ignored_layers from config.json marks interleaved BF16 fallback layers. + All other linear layers use W4A4 MXFP4 DualScale. + + Online mode (is_checkpoint_serialized=False): + num_bf16_fallback_layers leading transformer blocks use BF16 original weights + (default 5 when not specified). Remaining blocks use W4A4 MXFP4 DualScale online. + """ + from .mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + return DiffusionMXFP4DualScaleMixedConfig(**kw) + + def _build_inc(**kw: Any) -> QuantizationConfig: """Lazy import for INC/AutoRound config with checkpoint kwarg normalization.""" from .inc_config import OmniINCConfig @@ -66,6 +89,8 @@ _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = { "gguf": _build_gguf, "int8": _build_int8, "mxfp8": _build_mxfp8, + "mxfp4": _build_mxfp4, + "mxfp4_dualscale": _build_mxfp4_dualscale, "inc": _build_inc, "auto-round": _build_inc, "auto_round": _build_inc, @@ -255,3 +280,80 @@ def build_quant_config( return _build_from_method_and_config(method, merged) raise TypeError(f"quantization config must be str, dict, QuantizationConfig, or None, got {type(spec).__name__}") + + +def _disk_marks_serialized(qc_kwargs: dict[str, Any], quant_config: object) -> bool: + """Return True when config.json says serialized but the active quant_config does not. + + Matches any flag following the is_checkpoint_*_serialized naming convention, + so new quant methods don't require updating an explicit allowlist. + """ + for key, val in qc_kwargs.items(): + if key.startswith("is_checkpoint_") and key.endswith("_serialized"): + if val and hasattr(quant_config, key) and not getattr(quant_config, key): + return True + return False + + +def resolve_quant_config_from_disk( + quant_config: QuantizationConfig | None, + disk_qc: dict[str, Any] | str | None, +) -> QuantizationConfig | None: + """Reconcile an active quant_config against quantization_config from a transformer's config.json. + + Used when loading individual transformer blocks that each have their own config.json + (e.g. cascade models with separate transformer and transformer_2 directories). + + Rules: + - disk_qc is None: return quant_config unchanged. + - quant_config is None: auto-detect from disk_qc (full build). + - Methods mismatch: raise ValueError — prevents silent weight corruption. + - Disk marks serialized but quant_config is online: rebuild from disk. + - ignored_layers differ: rebuild from disk (per-transformer BF16 routing). + """ + if disk_qc is None: + return quant_config + + if isinstance(disk_qc, str): + if quant_config is None: + logger.info("Auto-detected quantization from config.json: method=%s", disk_qc) + return build_quant_config(disk_qc) + return quant_config + + if not isinstance(disk_qc, Mapping) or "quant_method" not in disk_qc: + return quant_config + + qc_method: str = disk_qc["quant_method"] + qc_kwargs: dict[str, Any] = {k: v for k, v in disk_qc.items() if k != "quant_method"} + + if quant_config is None: + logger.info( + "Auto-detected quantization from config.json: method=%s kwargs=%s", + qc_method, + qc_kwargs, + ) + return build_quant_config(qc_method, **qc_kwargs) + + if quant_config.get_name() != qc_method: + raise ValueError( + f"Checkpoint config.json declares quant_method={qc_method!r} but the " + f"active quantization config is {quant_config.get_name()!r}. " + "Pass a matching --quantization flag or omit it for auto-detection." + ) + + if _disk_marks_serialized(qc_kwargs, quant_config): + logger.info( + "config.json marks checkpoint as serialized; switching to offline %s mode.", + qc_method, + ) + return build_quant_config(qc_method, **qc_kwargs) + + if ( + "ignored_layers" in qc_kwargs + and hasattr(quant_config, "ignored_layers") + and set(qc_kwargs.get("ignored_layers") or []) != set(quant_config.ignored_layers or []) + ): + logger.info("config.json ignored_layers differs from active config; rebuilding quant_config.") + return build_quant_config(qc_method, **qc_kwargs) + + return quant_config diff --git a/vllm_omni/quantization/mxfp4_config.py b/vllm_omni/quantization/mxfp4_config.py new file mode 100644 index 000000000..6a02d5edb --- /dev/null +++ b/vllm_omni/quantization/mxfp4_config.py @@ -0,0 +1,709 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""W4A4 MXFP4 (Microscaling FP4) online/offline quantization for diffusion transformers. + +Architecture mirrors mxfp8_config.py: + + MXFPLinearMethodBase – platform-agnostic skeleton (imported from mxfp8_config) + NPUMxfp4LinearMethod – NPU single-scale offline (W4A4 MXFP4) + NPUMxfp4OnlineLinearMethod – NPU single-scale online (BF16 → FP4) + NPUMxfp4DualScaleLinearMethod – NPU dual-scale offline (W4A4 MXFP4 DualScale) + NPUMxfp4DualScaleOnlineLinearMethod – NPU dual-scale online (BF16 → FP4) + +Quantization configs: + + DiffusionMXFP4Config – single-scale online/offline (quant_method="mxfp4") + DiffusionMXFP4DualScaleMixedConfig – dual-scale + per-layer BF16 fallback (quant_method="mxfp4_dualscale") + Offline: ignored_layers from config.json routes interleaved BF16 layers + Online: num_bf16_fallback_layers leading blocks stay in BF16 (default 5) + +Key differences from MXFP8: + + 1. Precision: float4_e2m1fn_x2 (FP4 packed, 2 values per element). + npu_dynamic_mx_quant(x) without dst_type defaults to float4_e2m1fn_x2. + + 2. Weight layout: stored as (N, K) — NOT pre-transposed. + FP4 uses a packed format; transposing a packed tensor is not safely contiguous. + Transpose is done inline in _quant_matmul via layer.weight.transpose(0, 1). + + 3. GEMM signature: npu_quant_matmul requires explicit + x1_dtype=float4_e2m1fn_x2 and x2_dtype=float4_e2m1fn_x2. + + Scale layout: (N, S/2, 2) — same reshape as MXFP8, also NOT pre-transposed; + transposed inline in _quant_matmul. + +Dual-scale (W4A4_MXFP4_DUALSCALE): + + Two-level quantization: fine scale (per-32 K) + coarse scale (per-512 K) + per-channel + activation pre-scale (mul_scale from calibration). Uses npu_dynamic_dual_level_mx_quant + and npu_dual_level_quant_matmul. Weight stored in NZ hardware format via npu_format_cast. + + Checkpoint tensor shapes for dual-scale: + weight : (N, K) float8_e4m3fn – FP4 packed + weight_scale : (N, K//32) uint8 – fine scale (float8_e8m0fnu bits) + weight_dual_scale : (N, K//512, 1) float32 – coarse scale (extra dim avoids shape assert) + mul_scale : (K,) float32 – per-input-channel activation pre-scale + +Reference: MindIE-SD W4A4MXFP4QuantLinear / W4A4MXFP4DualQuantLinear (mindiesd/quantization/layer.py). +""" + +from __future__ import annotations + +import re +from typing import TYPE_CHECKING, Any + +import torch +from torch.nn import Module +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import ( + LinearBase, + UnquantizedLinearMethod, +) +from vllm.model_executor.layers.quantization import QuantizationMethods +from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from vllm.model_executor.layers.quantization.fp8 import _copy_missing_attrs +from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped +from vllm.model_executor.model_loader.weight_utils import initialize_single_dummy_weight +from vllm.model_executor.parameter import ModelWeightParameter +from vllm.model_executor.utils import replace_parameter + +from vllm_omni.platforms import current_omni_platform +from vllm_omni.quantization.mxfp8_config import ( + MXFPLinearMethodBase, + _LazyWeightMixin, +) + +if TYPE_CHECKING: + from vllm.model_executor.models.utils import WeightsMapper + +logger = init_logger(__name__) + + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- + + +class DiffusionMXFP4Config(QuantizationConfig): + """W4A4 MXFP4 quantization config for diffusion transformers. + + Supports both online (BF16 checkpoint → quantize at load time) and offline + (pre-quantized MXFP4 checkpoint) modes, mirroring DiffusionMXFP8Config. + + MX (microscaling) format: groups of 32 K-dimension elements share one + float8_e8m0fnu exponent scale. Weight and activation are float4_e2m1fn_x2. + """ + + def __init__( + self, + is_checkpoint_mxfp4_serialized: bool = False, + ignored_layers: list[str] | None = None, + ) -> None: + super().__init__() + self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized + self.ignored_layers = ignored_layers or [] + + @classmethod + def get_name(cls) -> QuantizationMethods: + return "mxfp4" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.bfloat16, torch.float16] + + @classmethod + def get_min_capability(cls) -> int: + return 80 + + @classmethod + def get_config_filenames(cls) -> list[str]: + return [] + + def apply_vllm_mapper(self, hf_to_vllm_mapper: WeightsMapper) -> None: + if self.ignored_layers: + self.ignored_layers = hf_to_vllm_mapper.apply_list(self.ignored_layers) + + @classmethod + def from_config(cls, config: dict[str, Any]) -> DiffusionMXFP4Config: + is_serialized = cls.get_from_keys_or(config, ["is_checkpoint_mxfp4_serialized"], False) + ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None) + if not ignored_layers: + ignored_layers = cls.get_from_keys_or(config, ["modules_to_not_convert"], None) + return cls( + is_checkpoint_mxfp4_serialized=is_serialized, + ignored_layers=ignored_layers, + ) + + def get_quant_method( + self, + layer: torch.nn.Module, + prefix: str, + ) -> QuantizeMethodBase | None: + if isinstance(layer, LinearBase): + if is_layer_skipped( + prefix=prefix, + ignored_layers=self.ignored_layers, + fused_mapping=self.packed_modules_mapping, + ): + return UnquantizedLinearMethod() + if current_omni_platform.is_npu(): + if self.is_checkpoint_mxfp4_serialized: + return NPUMxfp4LinearMethod(self) + return NPUMxfp4OnlineLinearMethod(self) + raise NotImplementedError( + "DiffusionMXFP4Config (W4A4 MXFP4) is currently only supported " + "on NPU (Ascend) platforms. CUDA support is not yet implemented." + ) + return None + + +# --------------------------------------------------------------------------- +# NPU MXFP4 single-scale offline method (pre-quantized checkpoint) +# --------------------------------------------------------------------------- + + +class NPUMxfp4LinearMethod(MXFPLinearMethodBase): + """NPU W4A4 MXFP4 offline linear method for pre-quantized checkpoints. + + Weight canonical layout after process_weights_after_loading: + weight : (N, K) in float4_e2m1fn_x2 — NOT pre-transposed (FP4 packed) + weight_scale: (N, S/2, 2) in float8_e8m0fnu — NOT pre-transposed + + Both are transposed inline in _quant_matmul, unlike MXFP8 which pre-transposes. + NPUMxfp4OnlineLinearMethod normalizes to the same layout so apply() is shared. + """ + + def __init__(self, quant_config: DiffusionMXFP4Config) -> None: + self.quant_config = quant_config + self.out_dtype = torch.get_default_dtype() + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.orig_dtype = params_dtype + layer.weight_block_size = None + + # BF16 placeholder; cast to float4_e2m1fn_x2 in process_weights. + layer.register_parameter( + "weight", + ModelWeightParameter( + data=torch.empty(output_size_per_partition, input_size_per_partition, dtype=params_dtype), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ), + ) + + # Scale stored as uint8 in safetensors (float8_e8m0fnu is same bit width). + # Using uint8 avoids a lossy float32 round-trip when loading the checkpoint. + num_groups = (input_size_per_partition + 31) // 32 + layer.register_parameter( + "weight_scale", + ModelWeightParameter( + data=torch.empty(output_size_per_partition, num_groups, dtype=torch.uint8), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ), + ) + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + import torch_npu + + # NPU: cast to float4_e2m1fn_x2. Weight stays (N, K) — no pre-transpose. + w = layer.weight + if w.dtype != torch_npu.float4_e2m1fn_x2: + w = torch_npu.npu_dtype_cast(w.npu(), torch_npu.float4_e2m1fn_x2) + + # Scale: checkpoint stores uint8 bytes that ARE float8_e8m0fnu bits. + # Only convert if neither uint8 nor the target NPU dtype already. + # (N, K_groups) → (N, K_groups/2, 2). Not pre-transposed; done inline in _quant_matmul. + s = layer.weight_scale.data + if s.dtype not in (torch.uint8, torch_npu.float8_e8m0fnu): + s = s.to(torch_npu.float8_e8m0fnu) + N, K_groups = s.shape + if K_groups % 2 == 1: + s = torch.cat([s, torch.zeros(N, 1, dtype=s.dtype, device=s.device)], dim=1) + s = s.reshape(N, -1, 2).contiguous() + + replace_parameter(layer, "weight", w) + replace_parameter(layer, "weight_scale", s) + layer._already_called_process_weights_after_loading = True + + # --- NPU MXFP4 ops — shared with online path via inheritance --- + + def _quantize_activation(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + import torch_npu + + # No dst_type: npu_dynamic_mx_quant defaults to float4_e2m1fn_x2. + return torch_npu.npu_dynamic_mx_quant(x) + + def _quant_matmul( + self, + x_q: torch.Tensor, + x_scale: torch.Tensor, + layer: torch.nn.Module, + bias: torch.Tensor | None, + ori_dtype: torch.dtype, + ) -> torch.Tensor: + import torch_npu + + if bias is not None and bias.dtype != torch.float32: + bias = bias.to(torch.float32) + # FP4 differences vs FP8: + # weight (N,K) transposed inline → (K,N); scale (N,S/2,2) transposed inline → (S/2,N,2). + # x1_dtype / x2_dtype required — FP4 dtype not inferred from tensor dtype. + return torch_npu.npu_quant_matmul( + x_q, + layer.weight.transpose(0, 1), # (K, N) inline + layer.weight_scale.transpose(0, 1), # (S/2, N, 2) inline + scale_dtype=torch_npu.float8_e8m0fnu, + x1_dtype=torch_npu.float4_e2m1fn_x2, + x2_dtype=torch_npu.float4_e2m1fn_x2, + pertoken_scale=x_scale, + pertoken_scale_dtype=torch_npu.float8_e8m0fnu, + bias=bias, + output_dtype=ori_dtype, + group_sizes=[1, 1, 32], + ) + + +# --------------------------------------------------------------------------- +# NPU MXFP4 single-scale online method (BF16 checkpoint → quantize at load time) +# --------------------------------------------------------------------------- + + +class NPUMxfp4OnlineLinearMethod(_LazyWeightMixin, NPUMxfp4LinearMethod): + """NPU W4A4 MXFP4 online linear method. + + MRO: NPUMxfp4OnlineLinearMethod → _LazyWeightMixin → NPUMxfp4LinearMethod + → MXFPLinearMethodBase → LinearMethodBase + + create_weights : _LazyWeightMixin (meta device + patched loader) + process_weights : NPUMxfp4OnlineLinearMethod (BF16 → FP4 + normalize) + apply / ops : NPUMxfp4LinearMethod / MXFPLinearMethodBase (shared) + """ + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + import torch_npu + + if layer.weight.device == torch.device("meta"): + weight = ModelWeightParameter( + data=torch.empty_like(layer.weight, device=layer._load_device), + input_dim=1, + output_dim=0, + weight_loader=layer.weight.weight_loader, + ) + _copy_missing_attrs(layer.weight, weight) + layer.register_parameter("weight", weight) + initialize_single_dummy_weight(layer.weight) + + # NPU: quantize BF16/FP16 (N, K) → FP4. No dst_type → float4_e2m1fn_x2. + weight_fp4, weight_scale_raw = torch_npu.npu_dynamic_mx_quant(layer.weight) + + # Weight stays (N, K) — no pre-transpose for FP4 packed format. + # Scale: (N, S) → (N, S/2, 2). Not pre-transposed; done inline in _quant_matmul. + weight_scale = weight_scale_raw.reshape(weight_scale_raw.shape[0], -1, 2).contiguous() + + replace_parameter(layer, "weight", weight_fp4) + replace_parameter(layer, "weight_scale", weight_scale) + layer._already_called_process_weights_after_loading = True + + +# --------------------------------------------------------------------------- +# NPU MXFP4 dual-scale offline method (W4A4_MXFP4_DUALSCALE checkpoint) +# --------------------------------------------------------------------------- + + +class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase): + """NPU W4A4 MXFP4 dual-scale offline method for pre-quantized checkpoints. + + Checkpoint tensors and their canonical post-load shapes: + weight : (N, K) float8_e4m3fn – FP4 packed (2 values per byte) + weight_scale : (N, K//32) uint8 – fine scale (float8_e8m0fnu bits); reshaped to (N, K//64, 2) + weight_dual_scale : (N, K//512, 1) float32 – coarse scale; transposed to (K//512, N) + mul_scale : (K,) float32 – per-input-channel activation pre-scale (from calibration) + + Forward pass: + x_q, l0, l1 = npu_dynamic_dual_level_mx_quant(x, smooth_scale=mul_scale) + out = npu_dual_level_quant_matmul(x_q, weight, l0, weight_dual_scale, l1, weight_scale) + + Reference: MindIE-SD W4A4MXFP4DualQuantLinear. + """ + + def __init__(self, quant_config: Any) -> None: + self.quant_config = quant_config + self.out_dtype = torch.get_default_dtype() + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.orig_dtype = params_dtype + layer.weight_block_size = None + + # FP4 packed: 2 values per float8_e4m3fn byte → checkpoint stores as float8_e4m3fn; register as BF16 + layer.register_parameter( + "weight", + ModelWeightParameter( + data=torch.empty(output_size_per_partition, input_size_per_partition, dtype=params_dtype), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ), + ) + + # Fine scale: one uint8 exponent (float8_e8m0fnu bit pattern) per group of 32 K elements. + # input_dim=1: RowParallelLinear weight_loader shards along the K-group dimension (dim 1) + # so each rank receives only its (K/TP)//32 groups; ColumnParallelLinear only uses + # output_dim=0 for sharding and leaves dim 1 intact. + num_groups_fine = (input_size_per_partition + 31) // 32 + layer.register_parameter( + "weight_scale", + ModelWeightParameter( + data=torch.empty(output_size_per_partition, num_groups_fine, dtype=torch.uint8), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ), + ) + + # Coarse scale: one float32 per group of 512 K elements. + # Shape (N, K_coarse, 1) matches checkpoint layout exactly, avoiding the + # shape-mismatch assert in linear.py:1344. + # input_dim=1: same TP sharding rationale as weight_scale above. + num_groups_coarse = (input_size_per_partition + 511) // 512 + layer.register_parameter( + "weight_dual_scale", + ModelWeightParameter( + data=torch.empty(output_size_per_partition, num_groups_coarse, 1, dtype=torch.float32), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ), + ) + + # mul_scale is a float32 calibration tensor; register as float32 + # to avoid precision loss from an implicit BF16 cast during weight loading. + # input_dim=0: RowParallelLinear shards the 1-D per-input-channel tensor along dim 0 + # so each rank receives only its K/TP channels; ColumnParallelLinear leaves it intact. + layer.register_parameter( + "mul_scale", + ModelWeightParameter( + data=torch.empty(input_size_per_partition, dtype=torch.float32), + input_dim=0, + output_dim=None, + weight_loader=weight_loader, + ), + ) + setattr(layer.mul_scale, "ignore_warning", True) + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + import torch_npu + + # float8_e4m3fn (FP4 packed) → float4_e2m1fn_x2 → NZ hardware format (format ID 29). + # NZ layout is required by npu_dual_level_quant_matmul for FP4 weight matrices. + w = torch_npu.npu_dtype_cast(layer.weight.data.npu(), torch_npu.float4_e2m1fn_x2) + w = torch_npu.npu_format_cast(w.view(torch.int8), 29, customize_dtype=torch.int8) + + # Fine scale: (N, K//32) uint8 → cast to float8_e8m0fnu → (N, K//64, 2). + # The cast reinterprets the stored uint8 bit-patterns as float8_e8m0fnu exponents, + # which is required for npu_dual_level_quant_matmul to compute correct scale values. + s = layer.weight_scale.data.npu() + s = s.reshape(s.shape[0], -1, 2).contiguous() + + # Coarse scale: (N, K//512, 1) → squeeze → (N, K//512) → transpose → (K//512, N). + ds = layer.weight_dual_scale.data.squeeze(-1).transpose(0, 1).npu().contiguous() + + ms = layer.mul_scale.to(torch.bfloat16).data.view(-1).npu().contiguous() + + replace_parameter(layer, "weight", w) + replace_parameter(layer, "weight_scale", s) + replace_parameter(layer, "weight_dual_scale", ds) + replace_parameter(layer, "mul_scale", ms) + layer._already_called_process_weights_after_loading = True + + def _apply_inner( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None, + ori_dtype: torch.dtype, + ) -> torch.Tensor: + """Dual-scale inner loop: dtype guard → 3-tuple quantize (mul_scale as smooth) → matmul. + + Overrides the default single-scale _apply_inner from MXFPLinearMethodBase. + apply() in the base class handles reshape/unreshape; this method is not + responsible for that. + """ + if ori_dtype not in (torch.bfloat16, torch.float16): + x = x.to(torch.bfloat16) + x_q, l0_scale, l1_scale = self._quantize_activation(x, layer.mul_scale) + return self._quant_matmul(x_q, l0_scale, l1_scale, layer, bias, ori_dtype) + + def _quantize_activation( # type: ignore[override] + self, + x: torch.Tensor, + smooth_scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + import torch_npu + + return torch_npu.npu_dynamic_dual_level_mx_quant(x, smooth_scale=smooth_scale) + + def _quant_matmul( # type: ignore[override] + self, + x_q: torch.Tensor, + l0_scale: torch.Tensor, + l1_scale: torch.Tensor, + layer: torch.nn.Module, + bias: torch.Tensor | None, + ori_dtype: torch.dtype, + ) -> torch.Tensor: + import torch_npu + + if bias is not None and bias.dtype != torch.float32: + bias = bias.to(torch.float32) + # weight_scale is (N, K//64, 2) float8_e8m0fnu — operator expects output-major layout. + # weight_dual_scale is (K//512, N) float32 — transposed to K-major in process_weights. + return torch_npu.npu_dual_level_quant_matmul( + x_q, + layer.weight, + l0_scale, + layer.weight_dual_scale, + l1_scale, + layer.weight_scale, + bias=bias, + output_dtype=ori_dtype, + ) + + +# --------------------------------------------------------------------------- +# NPU MXFP4 dual-scale online method (BF16 checkpoint → quantize at load time) +# --------------------------------------------------------------------------- + + +class NPUMxfp4DualScaleOnlineLinearMethod(_LazyWeightMixin, NPUMxfp4DualScaleLinearMethod): + """NPU W4A4 MXFP4 dual-scale online method: quantises BF16 weights at load time. + + MRO: NPUMxfp4DualScaleOnlineLinearMethod → _LazyWeightMixin + → NPUMxfp4DualScaleLinearMethod → MXFPLinearMethodBase + + create_weights : _LazyWeightMixin (meta device + patched loader) + process_weights : NPUMxfp4DualScaleOnlineLinearMethod (BF16 → FP4 + dual scales) + apply / _quant_matmul : NPUMxfp4DualScaleLinearMethod (shared with offline path) + """ + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + import torch_npu + + if layer.weight.device == torch.device("meta"): + weight = ModelWeightParameter( + data=torch.empty_like(layer.weight, device=layer._load_device), + input_dim=1, + output_dim=0, + weight_loader=layer.weight.weight_loader, + ) + _copy_missing_attrs(layer.weight, weight) + layer.register_parameter("weight", weight) + initialize_single_dummy_weight(layer.weight) + + # Quantize BF16 weight → FP4 + dual-level scales (no smooth pre-scale for online). + # Returns: (weight_fp4, l0_scale[coarse per-512], l1_scale[fine per-32]) + weight_fp4, w_l0_scale, w_l1_scale = torch_npu.npu_dynamic_dual_level_mx_quant( + layer.weight.data.npu(), smooth_scale=None + ) + + # NZ hardware format for the FP4 weight (same as offline path). + w = torch_npu.npu_format_cast(weight_fp4.view(torch.int8), 29, customize_dtype=torch.int8) + + # Fine scale (l1): (N, K//32) → (N, K//64, 2). Dtype from op output (float8_e8m0fnu). + s = w_l1_scale.reshape(w_l1_scale.shape[0], -1, 2).contiguous() + + # Coarse scale (l0): (N, K_coarse) → (K_coarse, N). Dtype from op output. + ds = w_l0_scale.reshape(w_l0_scale.shape[0], -1).transpose(0, 1).contiguous() + + # No calibration available: identity pre-scale (no smooth quantization effect). + ms = torch.ones(layer.input_size_per_partition, dtype=torch.bfloat16, device="npu") + + replace_parameter(layer, "weight", w) + replace_parameter(layer, "weight_scale", s) + replace_parameter(layer, "weight_dual_scale", ds) + replace_parameter(layer, "mul_scale", ms) + layer._already_called_process_weights_after_loading = True + + +# --------------------------------------------------------------------------- +# Block-index helper (shared by DiffusionMXFP4DualScaleMixedConfig) +# --------------------------------------------------------------------------- + +_BLOCK_IDX_RE = re.compile(r"^blocks\.(\d+)\.") + + +def _parse_block_idx(prefix: str) -> int | None: + """Extract block index from prefix like 'blocks.5.attn1.to_q'.""" + m = _BLOCK_IDX_RE.match(prefix) + return int(m.group(1)) if m else None + + +# --------------------------------------------------------------------------- +# Config: MXFP4 DualScale + per-layer BF16 fallback +# --------------------------------------------------------------------------- + + +class DiffusionMXFP4DualScaleMixedConfig(QuantizationConfig): + """W4A4 MXFP4 DualScale with per-layer BF16 fallback for diffusion transformers. + + Sensitive layers fall back to BF16 (original weights) while all other linear + layers use W4A4 MXFP4 DualScale. BF16 fallback layers may be interleaved + anywhere in the transformer. + + Offline mode (is_checkpoint_serialized=True): + Layers whose prefix appears in ignored_layers → UnquantizedLinearMethod (BF16) + All other linear layers → NPUMxfp4DualScaleLinearMethod + + ignored_layers is injected into transformer/config.json by the merge script + and contains the prefixes of all non-MXFP4 linear layers. + + Online mode (is_checkpoint_serialized=False): + Layer routing applies two rules in priority order: + 1. ignored_layers (explicit per-layer BF16 override, user-supplied) → BF16 + 2. Blocks 0 .. num_bf16_fallback_layers-1 (coarse leading-block rule) → BF16 + 3. All other linear layers → NPUMxfp4DualScaleOnlineLinearMethod + + num_bf16_fallback_layers defaults to 5 when not specified. + Set ignored_layers to pin arbitrary interleaved layers to BF16 without + needing an offline checkpoint (useful for accuracy debugging). + Layers outside "blocks.N.*" (condition_embedder etc.) always use online MXFP4 + unless they appear in ignored_layers. + + Config injected by merge_mxfp4_dualscale_checkpoint.py: + { + "quant_method": "mxfp4_dualscale", + "is_checkpoint_serialized": true, + "ignored_layers": ["blocks.0.attn1.to_q", ...] + } + """ + + def __init__( + self, + is_checkpoint_serialized: bool = False, + ignored_layers: list[str] | None = None, + num_bf16_fallback_layers: int = 5, + ) -> None: + super().__init__() + self.is_checkpoint_serialized = is_checkpoint_serialized + self.ignored_layers = ignored_layers or [] + self.num_bf16_fallback_layers = num_bf16_fallback_layers + + @classmethod + def get_name(cls) -> QuantizationMethods: + return "mxfp4_dualscale" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.bfloat16, torch.float16] + + @classmethod + def get_min_capability(cls) -> int: + return 80 + + @classmethod + def get_config_filenames(cls) -> list[str]: + return [] + + def apply_vllm_mapper(self, hf_to_vllm_mapper: WeightsMapper) -> None: + if self.ignored_layers: + self.ignored_layers = hf_to_vllm_mapper.apply_list(self.ignored_layers) + + @classmethod + def from_config(cls, config: dict[str, Any]) -> DiffusionMXFP4DualScaleMixedConfig: + is_serialized = cls.get_from_keys_or(config, ["is_checkpoint_serialized"], False) + ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None) + if not ignored_layers: + ignored_layers = cls.get_from_keys_or(config, ["modules_to_not_convert"], None) + num_bf16_fallback_layers = cls.get_from_keys_or(config, ["num_bf16_fallback_layers"], 5) + return cls( + is_checkpoint_serialized=is_serialized, + ignored_layers=ignored_layers, + num_bf16_fallback_layers=num_bf16_fallback_layers, + ) + + def get_quant_method( + self, + layer: torch.nn.Module, + prefix: str, + ) -> QuantizeMethodBase | None: + if not isinstance(layer, LinearBase): + return None + + if self.is_checkpoint_serialized: + # Offline: ignored_layers lists interleaved BF16 fallback layer prefixes. + if is_layer_skipped( + prefix=prefix, + ignored_layers=self.ignored_layers, + fused_mapping=self.packed_modules_mapping, + ): + return UnquantizedLinearMethod() + if not current_omni_platform.is_npu(): + raise NotImplementedError( + "DiffusionMXFP4DualScaleMixedConfig is currently only supported on NPU (Ascend) platforms." + ) + return NPUMxfp4DualScaleLinearMethod(self) + + # Online: explicit ignored_layers take priority (user-specified per-layer BF16 override), + # then fall back to the coarse leading-block rule. + if is_layer_skipped( + prefix=prefix, + ignored_layers=self.ignored_layers, + fused_mapping=self.packed_modules_mapping, + ): + return UnquantizedLinearMethod() + block_idx = _parse_block_idx(prefix) + if block_idx is not None and block_idx < self.num_bf16_fallback_layers: + return UnquantizedLinearMethod() + + if not current_omni_platform.is_npu(): + raise NotImplementedError( + "DiffusionMXFP4DualScaleMixedConfig is currently only supported on NPU (Ascend) platforms." + ) + return NPUMxfp4DualScaleOnlineLinearMethod(self) diff --git a/vllm_omni/quantization/tools/__init__.py b/vllm_omni/quantization/tools/__init__.py index e69de29bb..208f01a7c 100644 --- a/vllm_omni/quantization/tools/__init__.py +++ b/vllm_omni/quantization/tools/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py b/vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py new file mode 100644 index 000000000..0786d4607 --- /dev/null +++ b/vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py @@ -0,0 +1,491 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Merge W4A4_MXFP4_DUALSCALE quantized Wan2.2 weights into HF Diffusers format. + +msModelSlim produces a checkpoint where each linear layer is either: + - W4A4_MXFP4_DUALSCALE: quantized weights with fine/coarse scales and mul_scale + - FLOAT (BF16 fallback): kept as original BF16 weights for precision-sensitive layers + +BF16 fallback layers may be interleaved anywhere in the transformer (not just leading +blocks). The merge script detects them from quant_model_description.json and writes +their prefixes into config.json as ignored_layers so the runtime dynamically routes +each layer to the correct quantization method at weight-loading time. + +MXFP4_DUALSCALE key structure per linear layer (msModelSlim naming) +--------------------------------------------------------------------- + blocks.N.X.linear.weight W4A4_MXFP4_DUALSCALE – int8 (FP4 packed) + blocks.N.X.linear.weight_scale W4A4_MXFP4_DUALSCALE – uint8 (float8_e8m0fnu fine scale, per-32K) + blocks.N.X.linear.weight_dual_scale W4A4_MXFP4_DUALSCALE – float32 (coarse scale, per-512K) + blocks.N.X.linear.bias FLOAT – bias (if present) + blocks.N.X.div.mul_scale FLOAT – float32 per-input-channel activation pre-scale + +BF16 fallback key structure (no wrappers, plain linear weights): + blocks.N.X.weight FLOAT + condition_embedder.*.weight FLOAT (always BF16 — not quantized) + +Self-attention QKV notes +------------------------ +Self-attention Q/K/V weights are separate in the checkpoint (self_attn.q/k/v) but fused +into a single to_qkv layer in vllm-omni. The transformer's load_weights() handles this via +stacked_params_mapping. This script keeps Q/K/V keys separate — do NOT pre-fuse them here. + +For mul_scale specifically: even though Q/K/V process the same input (same mul_scale value), +they are kept separate. load_weights() routes all three to the same to_qkv.mul_scale parameter +and each overwrites the previous. Since Q=K=V for mul_scale, the final value is correct. + +NOTE: Pre-fusing as to_qkv.mul_scale would BREAK loading because ".attn1.to_q" is a +substring of ".attn1.to_qkv", causing load_weights() stacked_params_mapping to produce a +garbage key ("to_qkvkv") that is not in params_dict, triggering a break that skips the +direct-load else branch entirely. + +Supported model types: + - Wan2.2-T2V-A14B (MoE cascade: transformer + transformer_2) + - Wan2.2-I2V-A14B (MoE cascade: transformer + transformer_2) + +Note: Wan2.2-TI2V-5B is NOT supported. Its smaller parameter count causes +unacceptable accuracy loss under W4A4 quantization. Use MXFP8 for TI2V-5B. + +Usage: + python merge_mxfp4_dualscale_checkpoint.py \\ + --model-type Wan2.2-T2V-A14B \\ + --original-model /path/to/Wan2.2-T2V-A14B-Diffusers \\ + --quant-path /path/to/msmodelslim-output \\ + --output-path /path/to/merged-output +""" + +from __future__ import annotations + +import argparse +import json +import pathlib +import re +import shutil +from typing import Any + +import torch +from safetensors.torch import load_file, save_file + +# --------------------------------------------------------------------------- +# Key rename: msModelSlim naming → Diffusers / vllm-omni naming +# Identical to merge_mxfp8_checkpoint.py; applied to all blocks uniformly. +# --------------------------------------------------------------------------- + +TRANSFORMER_KEYS_RENAME_DICT: dict[str, str] = { + "time_embedding.0": "condition_embedder.time_embedder.linear_1", + "time_embedding.2": "condition_embedder.time_embedder.linear_2", + "text_embedding.0": "condition_embedder.text_embedder.linear_1", + "text_embedding.2": "condition_embedder.text_embedder.linear_2", + "time_projection.1": "condition_embedder.time_proj", + "head.modulation": "scale_shift_table", + "head.head": "proj_out", + "modulation": "scale_shift_table", + "ffn.0": "ffn.net.0.proj", + "ffn.2": "ffn.net.2", + # Norm order swap (quant tool: norm1, norm3, norm2 → diffusers: norm1, norm2, norm3) + "norm2": "norm__placeholder", + "norm3": "norm2", + "norm__placeholder": "norm3", + # Self-attention + "self_attn.q": "attn1.to_q", + "self_attn.k": "attn1.to_k", + "self_attn.v": "attn1.to_v", + "self_attn.o": "attn1.to_out.0", + "self_attn.norm_q": "attn1.norm_q", + "self_attn.norm_k": "attn1.norm_k", + # Cross-attention + "cross_attn.q": "attn2.to_q", + "cross_attn.k": "attn2.to_k", + "cross_attn.v": "attn2.to_v", + "cross_attn.o": "attn2.to_out.0", + "cross_attn.norm_q": "attn2.norm_q", + "cross_attn.norm_k": "attn2.norm_k", + "attn2.to_k_img": "attn2.add_k_proj", + "attn2.to_v_img": "attn2.add_v_proj", + "attn2.norm_k_img": "attn2.norm_added_k", + # I2V image embedder + "img_emb.proj.0": "condition_embedder.image_embedder.norm1", + "img_emb.proj.1": "condition_embedder.image_embedder.ff.net.0.proj", + "img_emb.proj.3": "condition_embedder.image_embedder.ff.net.2", + "img_emb.proj.4": "condition_embedder.image_embedder.norm2", + "img_emb.emb_pos": "condition_embedder.image_embedder.pos_embed", +} + +SUPPORTED_MODEL_TYPES = ["Wan2.2-T2V-A14B", "Wan2.2-I2V-A14B"] +CASCADE_MODEL_TYPES = {"Wan2.2-T2V-A14B", "Wan2.2-I2V-A14B"} + +# Suffixes that appear inside .linear.* wrapper for MXFP4 tensors. +_LINEAR_ATTRS = ("weight_dual_scale", "weight_scale", "weight", "bias") + +_BLOCK_IDX_RE = re.compile(r"^blocks\.(\d+)\.") + + +# --------------------------------------------------------------------------- +# Key transformation helpers +# --------------------------------------------------------------------------- + + +def _parse_block_idx(key: str) -> int | None: + """Extract block index from a key like 'blocks.5.attn1.to_q.weight'.""" + m = _BLOCK_IDX_RE.match(key) + return int(m.group(1)) if m else None + + +def _apply_rename_dict(key: str) -> str: + for src, dst in TRANSFORMER_KEYS_RENAME_DICT.items(): + key = key.replace(src, dst) + return key + + +def _strip_mxfp4_wrapper(key: str) -> str: + """Strip .linear.ATTR or .div.mul_scale wrappers added by msModelSlim. + + MXFP4 tensors are wrapped in sub-modules: + X.linear.weight / weight_scale / weight_dual_scale / bias + X.div.mul_scale + + MXFP8 and FLOAT tensors have no wrappers — this function is a no-op for them. + + Examples after apply_rename_dict: + attn1.to_q.linear.weight → attn1.to_q.weight + attn1.to_q.linear.weight_scale → attn1.to_q.weight_scale + attn1.to_q.div.mul_scale → attn1.to_q.mul_scale + attn1.norm_q.weight → attn1.norm_q.weight (unchanged) + """ + # Check longest attribute names first to avoid partial suffix matches. + for attr in _LINEAR_ATTRS: + suffix = f".linear.{attr}" + if key.endswith(suffix): + return key[: -len(suffix)] + f".{attr}" + if key.endswith(".div.mul_scale"): + return key[: -len(".div.mul_scale")] + ".mul_scale" + return key + + +# --------------------------------------------------------------------------- +# Quantization metadata helpers +# --------------------------------------------------------------------------- + +# Known quantized weight types (FLOAT tensors don't determine block type). +_MXFP8_TYPE = "mxfp8" +_MXFP4_DUALSCALE_TYPE = "mxfp4_dualscale" + + +def _classify_blocks(quant_meta: dict[str, str]) -> dict[int, str]: + """Classify each transformer block by quantization type from quant_meta. + + Returns a dict mapping block_idx → 'mxfp8' | 'mxfp4_dualscale'. + A block's type is determined by the first quantized (non-FLOAT) tensor found for it. + """ + block_types: dict[int, str] = {} + for key, qtype in quant_meta.items(): + idx = _parse_block_idx(key) + if idx is None or idx in block_types: + continue + if qtype.startswith("W8A8_MXFP8"): + block_types[idx] = _MXFP8_TYPE + elif qtype.startswith("W4A4_MXFP4_DUALSCALE"): + block_types[idx] = _MXFP4_DUALSCALE_TYPE + return block_types + + +def _print_block_summary(block_types: dict[int, str]) -> None: + """Print a compact run-length summary of the block layout.""" + if not block_types: + print(" Block layout: (empty)") + return + + sorted_indices = sorted(block_types) + runs: list[tuple[int, int, str]] = [] + run_start = sorted_indices[0] + run_type = block_types[run_start] + for idx in sorted_indices[1:]: + if block_types[idx] != run_type: + runs.append((run_start, idx - 1, run_type)) + run_start = idx + run_type = block_types[idx] + runs.append((run_start, sorted_indices[-1], run_type)) + + print(f" Block layout ({len(sorted_indices)} blocks classified):") + for start, end, btype in runs: + count = end - start + 1 + range_str = f"{start}" if start == end else f"{start}–{end}" + print(f" blocks {range_str:>8}: {btype} ({count} block{'s' if count > 1 else ''})") + + +# --------------------------------------------------------------------------- +# Safetensors I/O +# --------------------------------------------------------------------------- + + +def _load_safetensors_dir(directory: pathlib.Path, glob: str = "*.safetensors") -> dict[str, torch.Tensor]: + candidates = sorted(directory.glob(glob)) + if not candidates: + raise FileNotFoundError(f"No safetensors matching '{glob}' found in {directory}") + state: dict[str, torch.Tensor] = {} + for f in candidates: + state.update(load_file(str(f))) + return state + + +def _load_quant_safetensors(directory: pathlib.Path) -> dict[str, torch.Tensor]: + try: + return _load_safetensors_dir(directory, "quant_model_weight*.safetensors") + except FileNotFoundError: + return _load_safetensors_dir(directory) + + +def _load_quant_meta(directory: pathlib.Path) -> dict[str, str]: + candidates = sorted(directory.glob("quant_model_description*.json")) + if not candidates: + print(f" WARNING: No quant_model_description*.json in {directory}; treating all tensors as FLOAT.") + return {} + with open(candidates[0]) as f: + return json.load(f) + + +# --------------------------------------------------------------------------- +# Per-transformer conversion +# --------------------------------------------------------------------------- + + +def _is_mxfp4_tensor(key: str, qtype: str) -> bool: + """Return True for MXFP4 quantized tensors and their companion tensors. + + MXFP4 DualScale layers produce four tensor types: + .linear.weight / .linear.weight_scale / .linear.weight_dual_scale (W4A4_MXFP4_DUALSCALE) + .linear.bias / .div.mul_scale (FLOAT but companion to a quantized layer) + + BF16 fallback layers have plain .weight / .bias without .linear. / .div. wrappers. + """ + if qtype.startswith("W4A4_MXFP4_DUALSCALE"): + return True + # Companion tensors: bias and mul_scale that belong to an MXFP4 layer. + return ".linear." in key or ".div.mul_scale" in key + + +def _diffusers_to_vllm_ignored(diffusers_ignored: list[str]) -> list[str]: + """Translate diffusers checkpoint prefixes to vllm-omni model parameter names. + + Three transformations align the naming conventions: + + 1. Self-attention QKV fusion (attn1 only, not attn2): + attn1.to_q + attn1.to_k + attn1.to_v → attn1.to_qkv + Cross-attention (attn2) to_q/k/v are not fused and remain separate. + + 2. FFN naming: + ffn.net.0.proj → ffn.net_0.proj + ffn.net.2 → ffn.net_2 + + 3. to_out index: + attn1.to_out.0 / attn2.to_out.0 → attn1.to_out / attn2.to_out + """ + ignored_set = set(diffusers_ignored) + result: set[str] = set() + + for name in diffusers_ignored: + m = re.match(r"^(.*\.attn1)\.to_([qkv])$", name) + if m: + prefix = m.group(1) + if all(f"{prefix}.to_{c}" in ignored_set for c in ("q", "k", "v")): + result.add(f"{prefix}.to_qkv") + else: + present = [f"to_{c}" for c in ("q", "k", "v") if f"{prefix}.to_{c}" in ignored_set] + missing = [f"to_{c}" for c in ("q", "k", "v") if f"{prefix}.to_{c}" not in ignored_set] + raise ValueError( + f"Partial BF16 fallback for '{prefix}': " + f"{', '.join(present)} in ignored_layers but {', '.join(missing)} is not. " + f"Self-attention Q/K/V are fused into a single to_qkv layer at runtime; " + f"all three must share the same precision. " + f"Either quantize all of to_q/to_k/to_v or keep all three in BF16." + ) + continue + + name = re.sub(r"\.ffn\.net\.0\.proj$", ".ffn.net_0.proj", name) + name = re.sub(r"\.ffn\.net\.2$", ".ffn.net_2", name) + name = re.sub(r"\.to_out\.0$", ".to_out", name) + result.add(name) + + return sorted(result) + + +def _collect_ignored_layers(merged: dict[str, Any], mxfp4_layer_prefixes: set[str]) -> list[str]: + """Collect vllm-omni parameter name prefixes for BF16 fallback layers. + + A layer is BF16 if it has a .weight tensor but its prefix is not in + mxfp4_layer_prefixes. Prefixes are returned in vllm-omni parameter naming + (QKV-fused, FFN underscored, to_out unindexed) so the list can be written + directly into ignored_layers in config.json without further translation. + """ + all_weight_prefixes = {key[: -len(".weight")] for key in merged if key.endswith(".weight")} + return _diffusers_to_vllm_ignored(sorted(all_weight_prefixes - mxfp4_layer_prefixes)) + + +def _convert_transformer( + quant_subdir: pathlib.Path, + output_dir: pathlib.Path, + original_transformer_dir: pathlib.Path, +) -> None: + """Convert one transformer directory to the mxfp4_dualscale + BF16 mixed format.""" + output_dir.mkdir(parents=True, exist_ok=True) + + # BF16 base: provides the scaffold for non-MXFP4 tensors (norms, embeddings, + # BF16 fallback linear layers that msModelSlim may omit or keep as FLOAT). + print(f" Loading BF16 base from {original_transformer_dir} …") + base_state = _load_safetensors_dir(original_transformer_dir) + print(f" {len(base_state)} BF16 tensors loaded") + + print(f" Loading quantized weights from {quant_subdir} …") + quant_state = _load_quant_safetensors(quant_subdir) + quant_meta = _load_quant_meta(quant_subdir) + print(f" {len(quant_state)} quant tensors, {len(quant_meta)} meta entries") + + # Classify blocks for a compact summary (informational only in the new scheme). + block_types = _classify_blocks(quant_meta) + _print_block_summary(block_types) + + # Remap MXFP4 quantized tensors only. + # BF16 fallback tensors (FLOAT in quant_meta, no .linear./.div. wrapper) are + # skipped here; the base_state provides them unchanged. + remapped: dict[str, torch.Tensor] = {} + remapped_meta: dict[str, str] = {} + mxfp4_layer_prefixes: set[str] = set() + skipped: list[str] = [] + + for key, tensor in quant_state.items(): + qtype = quant_meta.get(key, "FLOAT") + + if not _is_mxfp4_tensor(key, qtype): + continue # BF16 fallback — base_state already covers this tensor + + renamed = _apply_rename_dict(key) + final_key = _strip_mxfp4_wrapper(renamed) + + if final_key.endswith(".quant_type"): + skipped.append(key) + continue + + remapped[final_key] = tensor + remapped_meta[final_key] = qtype + + # Track MXFP4 layer prefixes (via their .weight keys) for ignored_layers. + if final_key.endswith(".weight") and qtype.startswith("W4A4_MXFP4_DUALSCALE"): + mxfp4_layer_prefixes.add(final_key[: -len(".weight")]) + + if skipped: + print(f" Skipped {len(skipped)} metadata keys (quant_type markers): {skipped[:5]}") + print(f" {len(remapped)} MXFP4 tensors remapped, {len(mxfp4_layer_prefixes)} quantized layers") + + # Merge: base_state provides BF16 scaffold; MXFP4 tensors override their BF16 counterparts + # and add the new scale tensors (weight_scale, weight_dual_scale, mul_scale). + merged = {**base_state, **remapped} + + # Determine ignored_layers: BF16 fallback layer prefixes in vllm-omni parameter naming. + ignored_layers = _collect_ignored_layers(merged, mxfp4_layer_prefixes) + print(f" {len(ignored_layers)} BF16 fallback layers → ignored_layers in config.json") + + # Save weights. + out_weights = output_dir / "diffusion_pytorch_model.safetensors" + save_file(merged, str(out_weights)) + print(f" Saved {len(merged)} tensors → {out_weights}") + + # Save remapped quant metadata (for inspection / debugging). + out_meta_path = output_dir / "quant_model_description.json" + with open(out_meta_path, "w") as f: + json.dump(remapped_meta, f, indent=2) + + # Inject quantization_config into config.json. + src_config = original_transformer_dir / "config.json" + if src_config.is_file(): + with open(src_config) as f: + config = json.load(f) + config["quantization_config"] = _build_quant_config(ignored_layers) + out_config = output_dir / "config.json" + with open(out_config, "w") as f: + json.dump(config, f, indent=2) + print(f" Injected quantization_config (mxfp4_dualscale, {len(ignored_layers)} ignored_layers) → {out_config}") + else: + print(f" WARNING: No config.json at {src_config}; quantization_config not injected.") + + +def _build_quant_config(ignored_layers: list[str]) -> dict[str, Any]: + return { + "quant_method": "mxfp4_dualscale", + "is_checkpoint_serialized": True, + "ignored_layers": ignored_layers, + } + + +# --------------------------------------------------------------------------- +# Model-type helpers +# --------------------------------------------------------------------------- + + +def _get_transformer_dirs(model_type: str) -> list[str]: + return ["transformer", "transformer_2"] if model_type in CASCADE_MODEL_TYPES else ["transformer"] + + +def _get_quant_subdir(model_type: str, quant_path: pathlib.Path, transformer_dir: str) -> pathlib.Path: + if model_type in CASCADE_MODEL_TYPES: + sub = "high_noise_model" if transformer_dir == "transformer" else "low_noise_model" + return quant_path / sub + return quant_path + + +# --------------------------------------------------------------------------- +# Main repack +# --------------------------------------------------------------------------- + + +def repack( + model_type: str, + original_model_path: pathlib.Path, + quant_path: pathlib.Path, + output_path: pathlib.Path, +) -> None: + transformer_dirs = _get_transformer_dirs(model_type) + + print(f"Copying original model to {output_path} (skipping {transformer_dirs}) …") + shutil.copytree( + str(original_model_path), + str(output_path), + ignore=shutil.ignore_patterns(*transformer_dirs), + ) + + for tdir in transformer_dirs: + q_subdir = _get_quant_subdir(model_type, quant_path, tdir) + out_tdir = output_path / tdir + orig_tdir = original_model_path / tdir + print(f"\nConverting {tdir} (quant source: {q_subdir.name}) …") + _convert_transformer(q_subdir, out_tdir, orig_tdir) + + print(f"\nDone. Merged model → {output_path}") + print("\nRun inference (quantization auto-detected from config.json):") + print(f" python text_to_video.py --model {output_path}") + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--model-type", required=True, choices=SUPPORTED_MODEL_TYPES, help="Model variant.") + parser.add_argument("--original-model", required=True, help="Original HF Diffusers model directory (BF16).") + parser.add_argument("--quant-path", required=True, help="msModelSlim quantized weights directory.") + parser.add_argument("--output-path", required=True, help="Output directory for merged model.") + args = parser.parse_args() + + repack( + model_type=args.model_type, + original_model_path=pathlib.Path(args.original_model), + quant_path=pathlib.Path(args.quant_path), + output_path=pathlib.Path(args.output_path), + ) + + +if __name__ == "__main__": + main() From 1c530d984da729b75e7589797474ca8199bd24a0 Mon Sep 17 00:00:00 2001 From: yangsonglin13 Date: Wed, 20 May 2026 14:51:33 +0800 Subject: [PATCH 26/38] Yuanrong TransferEngine Connector for NPU (#3180) Signed-off-by: yangsonglin13 --- .../design/feature/disaggregated_inference.md | 2 + .../yuanrong_transfer_engine_connector.md | 211 +++++ .../test_mooncake_transfer_engine_buffer.py | 70 +- vllm_omni/deploy/hunyuan_image3.yaml | 10 + vllm_omni/distributed/__init__.py | 2 + .../distributed/omni_connectors/__init__.py | 8 + .../mooncake_transfer_engine_connector.py | 155 +--- .../yuanrong_transfer_engine_connector.py | 10 + .../distributed/omni_connectors/factory.py | 13 + .../omni_connectors/kv_transfer_manager.py | 58 +- .../omni_connectors/utils/config.py | 7 + .../distributed/omni_connectors/utils/env.py | 47 + .../omni_connectors/utils/initialization.py | 14 +- .../omni_connectors/utils/local_rank.py | 49 + .../omni_connectors/utils/memory_pool.py | 153 ++++ vllm_omni/engine/stage_engine_core_client.py | 8 +- .../platforms/npu/omni_connectors/__init__.py | 11 + .../yuanrong_transfer_engine_connector.py | 862 ++++++++++++++++++ 18 files changed, 1507 insertions(+), 183 deletions(-) create mode 100644 docs/design/feature/omni_connectors/yuanrong_transfer_engine_connector.md create mode 100644 vllm_omni/distributed/omni_connectors/connectors/yuanrong_transfer_engine_connector.py create mode 100644 vllm_omni/distributed/omni_connectors/utils/env.py create mode 100644 vllm_omni/distributed/omni_connectors/utils/local_rank.py create mode 100644 vllm_omni/distributed/omni_connectors/utils/memory_pool.py create mode 100644 vllm_omni/platforms/npu/omni_connectors/__init__.py create mode 100644 vllm_omni/platforms/npu/omni_connectors/yuanrong_transfer_engine_connector.py diff --git a/docs/design/feature/disaggregated_inference.md b/docs/design/feature/disaggregated_inference.md index 83c35ac01..755d6aeea 100644 --- a/docs/design/feature/disaggregated_inference.md +++ b/docs/design/feature/disaggregated_inference.md @@ -9,6 +9,7 @@ Backend-specific setup lives in separate docs: - [MooncakeStoreConnector](omni_connectors/mooncake_store_connector.md) - [MooncakeTransferEngineConnector](omni_connectors/mooncake_transfer_engine_connector.md) - [YuanrongConnector](omni_connectors/yuanrong_connector.md) +- [YuanrongTransferEngineConnector](omni_connectors/yuanrong_transfer_engine_connector.md) ## Overview @@ -23,6 +24,7 @@ Current connectors operate in D2H2D (device to host to device) mode. | Multi node (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. | | Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. | | Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. | +| Ascend NPU P2P (Yuanrong TE) | YuanrongTransferEngineConnector | Uses Yuanrong TransferEngine directly. Configure NPU device IPv4 and `memory_pool_device: "npu"`. | ## Core API diff --git a/docs/design/feature/omni_connectors/yuanrong_transfer_engine_connector.md b/docs/design/feature/omni_connectors/yuanrong_transfer_engine_connector.md new file mode 100644 index 000000000..21d5f1841 --- /dev/null +++ b/docs/design/feature/omni_connectors/yuanrong_transfer_engine_connector.md @@ -0,0 +1,211 @@ +# YuanrongTransferEngineConnector + +## When to Use + +Use `YuanrongTransferEngineConnector` for high-performance peer-to-peer KV cache +transfer on Ascend NPU. It is different from `YuanrongConnector`: + +- `YuanrongConnector` uses Yuanrong Datasystem as a distributed KV store and + requires Datasystem workers plus etcd. +- `YuanrongTransferEngineConnector` uses Yuanrong TransferEngine directly. It + registers local memory in each stage worker and lets the receiver pull data + from the sender with TransferEngine. + +The implementation currently lives under +`vllm_omni/platforms/npu/omni_connectors/` because the supported backend is +Ascend NPU only. The connector is still registered with the generic +OmniConnector factory under the name `YuanrongTransferEngineConnector`. + +For Ascend P2P transfer, the connector should use NPU memory as its transfer +pool. The Mooncake-style CPU pool is not supported for this connector. + +## Prerequisites + +Install Yuanrong Datasystem/TransferEngine Python bindings in the runtime +environment. The packaged install is: + +```bash +pip install openyuanrong-datasystem +``` + +To build from source, follow the upstream project instructions: + +- Source repository: [openeuler/yuanrong-datasystem](https://atomgit.com/openeuler/yuanrong-datasystem) +- Installation guide: [Yuanrong Datasystem Linux installation](https://pages.openeuler.openatom.cn/openyuanrong-datasystem/docs/zh-cn/latest/installation/installation_linux.html) + +The Ascend driver, CANN runtime, HCCL/HCCP, and NPU device network must also be +configured on every machine participating in the transfer. + +## Check NPU Device IPv4 + +Yuanrong TransferEngine with `protocol: "ascend"` uses the NPU device network, +not only the host management IP. Each participating NPU must have an IPv4 HCCN +device IP. + +Check the device IP, link, and health state: + +```bash +for i in {0..7}; do + echo "===== NPU $i =====" + hccn_tool -i "$i" -ip -g + hccn_tool -i "$i" -link -g + hccn_tool -i "$i" -net_health -g +done +``` + +Expected output for each NPU should include an IPv4 address and healthy link: + +```text +ipaddr:10.86.1.80 +netmask:255.255.255.0 +link status: UP +net health status: Success +``` + +If the output says `Get ipconf failed, because no ip was preset there!`, the NPU +device IP is not configured and Ascend P2P will fail. If the address is IPv6, +configure IPv4 for this connector path or validate that the underlying Yuanrong +TransferEngine build supports the IPv6 HCCN path before using it. + +## Check Device-to-Device Connectivity + +After collecting the target-stage NPU IPv4 addresses, ping them from each +source-stage NPU: + +```bash +# Example: AR uses NPU 0..3 and DiT uses NPU 4..7. +# Replace the destination IPs with the values reported by hccn_tool -i 4..7 -ip -g. +for src in 0 1 2 3; do + for dst_ip in 10.86.1.78 10.86.1.79 10.86.1.80 10.86.1.81; do + echo "===== NPU $src -> $dst_ip =====" + hccn_tool -i "$src" -ping -g address "$dst_ip" || true + done +done +``` + +The expected result is `0.00% packet loss` for every required source/target pair. +This verifies basic HCCN IPv4 connectivity. TransferEngine can still fail later +if RDMA/QP/HCCL setup is broken, but failed ping means the device network must be +fixed first. + +## YAML Configuration + +For Ascend P2P, configure the connector with `protocol: "ascend"` and an NPU +memory pool: + +```yaml +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 1 + connectors: + yuanrong_te_connector: + name: YuanrongTransferEngineConnector + extra: + host: "auto" + zmq_port: 50051 + rpc_port: "auto" + protocol: "ascend" + device_name: "auto" + memory_pool_size: 1073741824 + memory_pool_device: "npu" + edges: + - from: 0 + to: 1 + window_size: -1 +``` + +Important fields: + +| Parameter | Recommended Value | Notes | +|---|---|---| +| `host` | `"auto"` | Advertises a routable local host IP for ZMQ metadata exchange. | +| `zmq_port` | `50051` | Base ZMQ port. The runtime applies stage/rank offsets. | +| `rpc_port` | `"auto"` | Lets TransferEngine choose an available RPC port. | +| `protocol` | `"ascend"` | Uses the Ascend TransferEngine backend. | +| `device_name` | `"auto"` | Each TP worker resolves to its local logical NPU, for example `npu:0`, `npu:1`, etc. | +| `memory_pool_device` | `"npu"` | Required for the Ascend fast path. This registers NPU memory with TransferEngine. | +| `memory_pool_size` | `1073741824` | 1 GiB per worker is a practical starting point for KV-cache transfer validation. Increase if the transfer pool is exhausted. | + +`memory_pool_device: "cpu"` is rejected for +`YuanrongTransferEngineConnector` with `protocol: "ascend"`. CPU pool registration +can lead to Ascend P2P failures such as HCCL RA init errors, QP timeouts, or +FFTS/SDMA runtime errors. `memory_pool_device: "cpu"` is common for +`MooncakeTransferEngineConnector`, but Yuanrong TE on Ascend only supports +`"npu"`. + +## 4-Card AR-to-DiT TP Example + +For AR 4-way TP sending KV cache to DiT 4-way TP, set the stage TP topology in +both stage `omni_kv_config` blocks: + +```yaml +stage_args: + - stage_id: 0 + runtime: + devices: "0,1,2,3" + engine_args: + tensor_parallel_size: 4 + omni_kv_config: + need_send_cache: true + rank_mapping: + from_tp: 4 + to_tp: 4 + output_connectors: + to_stage_1: yuanrong_te_connector + + - stage_id: 1 + runtime: + devices: "4,5,6,7" + engine_args: + parallel_config: + tensor_parallel_size: 4 + omni_kv_config: + need_recv_cache: true + rank_mapping: + from_tp: 4 + to_tp: 4 + input_connectors: + from_stage_0: yuanrong_te_connector +``` + +In this topology each DiT rank receives the corresponding AR rank shard: + +```text +DiT rank 0 <- AR rank 0 +DiT rank 1 <- AR rank 1 +DiT rank 2 <- AR rank 2 +DiT rank 3 <- AR rank 3 +``` + +## Validate a Successful Run + +Search the inference log for Yuanrong TE receive-side messages: + +```bash +grep -E "\[YR GET\]|Successfully received KV cache|transfer_engine get failed" inference.log +``` + +A successful Ascend fast-path transfer contains lines like: + +```text +[YR GET] ... fast_path, zero-copy +Successfully received KV cache ... across 1 key(s) +Applied CFG KV caches +[KV Reuse] ... +``` + +Sender-side `KV transfer OK` only means that the sender registered the payload +and metadata. The actual P2P data transfer happens on the receiver side during +`[YR GET]`. + +## Troubleshooting + +| Symptom | Likely Cause | Action | +|---|---|---| +| `IPv4 device IP not found` | NPU HCCN IPv4 is missing. | Run `hccn_tool -i -ip -g` and configure IPv4 for every participating NPU. | +| `P2PCommInitRootInfo failed`, `ra init failed`, `RA_QP_STATUS_TIMEOUT` | Ascend P2P/RDMA/QP setup failed. | First verify device IPv4, link health, and device-to-device ping. Then test a minimal Yuanrong TE pair outside vLLM. | +| `fftsplus sdma error`, `context is abort` | Runtime SDMA/FFTS failure, often after an invalid or failed P2P transfer. | Use `memory_pool_device: "npu"` for Ascend TE, reduce `memory_pool_size` if OOM occurs, and isolate the failing NPU pair with a Yuanrong TE demo. | +| Sender logs `KV transfer OK` but receiver never logs `[YR GET]` | Sender `put()` succeeded, receiver P2P read failed or timed out. | Check receiver logs and search for `transfer_engine get failed`. | +| Pool allocation fails or OOM | Transfer pool is too small or too large for available NPU memory. | Start with 1 GiB (`1073741824`), reduce to 512 MiB for validation, or increase for larger KV payloads. | diff --git a/tests/distributed/omni_connectors/test_mooncake_transfer_engine_buffer.py b/tests/distributed/omni_connectors/test_mooncake_transfer_engine_buffer.py index 55351efcb..3e9b97d4e 100755 --- a/tests/distributed/omni_connectors/test_mooncake_transfer_engine_buffer.py +++ b/tests/distributed/omni_connectors/test_mooncake_transfer_engine_buffer.py @@ -15,8 +15,9 @@ from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_c BufferAllocator, ManagedBuffer, ) +from vllm_omni.distributed.omni_connectors.kv_transfer_manager import KVCacheTransferData -# All tests in this file are pure-CPU unit tests for the memory allocator. +# Most tests in this file are pure-CPU unit tests; CUDA smoke tests are skip-guarded. pytestmark = [pytest.mark.cpu, pytest.mark.parallel, pytest.mark.core_model] @@ -225,3 +226,70 @@ class TestManagedBuffer: # Space should be reusable new_offset = self.allocator.alloc(128) assert new_offset == offset + + def test_to_bytes_cpu_pool(self): + """Smoke-test Mooncake's ManagedBuffer byte extraction on CPU pools.""" + expected = bytes(range(64)) + offset = self.allocator.alloc(len(expected)) + buf = ManagedBuffer(self.allocator, offset, len(expected), self.pool) + self.pool[offset : offset + len(expected)] = torch.tensor(list(expected), dtype=torch.uint8) + + try: + assert buf.to_bytes() == expected + finally: + buf.release() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for GPU pool smoke test") + def test_to_bytes_cuda_pool(self): + """Smoke-test Mooncake's existing CUDA pool D2H byte path.""" + expected = bytes(range(64)) + allocator = BufferAllocator(total_size=4096, alignment=64) + pool = torch.zeros(4096, dtype=torch.uint8, device="cuda:0") + offset = allocator.alloc(len(expected)) + buf = ManagedBuffer(allocator, offset, len(expected), pool) + pool[offset : offset + len(expected)].copy_(torch.tensor(list(expected), dtype=torch.uint8, device="cuda:0")) + + try: + assert buf.to_bytes() == expected + finally: + buf.release() + + +def _make_kv_payload(device: str | torch.device = "cpu") -> tuple[KVCacheTransferData, torch.Tensor, torch.Tensor]: + key_tensor = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4) + value_tensor = key_tensor + 100 + payload = KVCacheTransferData( + request_id="moon-smoke", + layer_blocks={"key_cache": [key_tensor], "value_cache": [value_tensor]}, + block_ids=[1], + metadata={"seq_len": 3}, + ) + return payload, key_tensor, value_tensor + + +class TestMooncakePackedPayloadSmoke: + """Smoke tests for Mooncake raw payload helpers without real Mooncake/RDMA.""" + + def test_load_header_from_cpu_tensor(self): + payload, _, _ = _make_kv_payload() + raw = payload.to_bytes() + packed = torch.tensor(list(raw), dtype=torch.uint8) + header, data_start = KVCacheTransferData._load_header_from_tensor(packed) + + assert header["rid"] == "moon-smoke" + assert data_start == 4 + int.from_bytes(raw[:4], "big") + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for packed CUDA smoke test") + def test_from_bytes_device_cuda_round_trip(self): + payload, key_tensor, value_tensor = _make_kv_payload(device="cuda:0") + packed = payload.to_gpu_tensor() + + header, data_start = KVCacheTransferData._load_header_from_tensor(packed) + decoded = KVCacheTransferData.from_bytes_device(packed) + + assert header["rid"] == "moon-smoke" + assert data_start == 4 + int.from_bytes(payload.to_bytes()[:4], "big") + assert decoded["layer_blocks"]["key_cache"][0].is_cuda + assert decoded["layer_blocks"]["value_cache"][0].is_cuda + assert torch.equal(decoded["layer_blocks"]["key_cache"][0].cpu(), key_tensor.cpu()) + assert torch.equal(decoded["layer_blocks"]["value_cache"][0].cpu(), value_tensor.cpu()) diff --git a/vllm_omni/deploy/hunyuan_image3.yaml b/vllm_omni/deploy/hunyuan_image3.yaml index 20d72304b..a6e2add17 100644 --- a/vllm_omni/deploy/hunyuan_image3.yaml +++ b/vllm_omni/deploy/hunyuan_image3.yaml @@ -15,6 +15,16 @@ connectors: device_name: "" memory_pool_size: 4294967296 memory_pool_device: "cpu" + yuanrong_te_connector: + name: YuanrongTransferEngineConnector + extra: + host: "auto" + zmq_port: 50051 + rpc_port: "auto" + protocol: "ascend" + device_name: "auto" + memory_pool_size: 1073741824 + memory_pool_device: "npu" shared_memory_connector: name: SharedMemoryConnector extra: diff --git a/vllm_omni/distributed/__init__.py b/vllm_omni/distributed/__init__.py index f43a833f8..7af13a322 100644 --- a/vllm_omni/distributed/__init__.py +++ b/vllm_omni/distributed/__init__.py @@ -11,6 +11,7 @@ from .omni_connectors import ( OmniTransferConfig, SharedMemoryConnector, YuanrongConnector, + YuanrongTransferEngineConnector, load_omni_transfer_config, ) @@ -26,6 +27,7 @@ __all__ = [ "MooncakeTransferEngineConnector", "SharedMemoryConnector", "YuanrongConnector", + "YuanrongTransferEngineConnector", # Utilities "load_omni_transfer_config", ] diff --git a/vllm_omni/distributed/omni_connectors/__init__.py b/vllm_omni/distributed/omni_connectors/__init__.py index 383267928..98e3a5c36 100644 --- a/vllm_omni/distributed/omni_connectors/__init__.py +++ b/vllm_omni/distributed/omni_connectors/__init__.py @@ -6,6 +6,13 @@ from .connectors.mooncake_store_connector import MooncakeStoreConnector from .connectors.shm_connector import SharedMemoryConnector from .connectors.yuanrong_connector import YuanrongConnector +try: + from vllm_omni.platforms.npu.omni_connectors.yuanrong_transfer_engine_connector import ( + YuanrongTransferEngineConnector, + ) +except ImportError: + YuanrongTransferEngineConnector = None + try: from .connectors.mooncake_transfer_engine_connector import MooncakeTransferEngineConnector except ImportError: @@ -39,6 +46,7 @@ __all__ = [ "MooncakeTransferEngineConnector", "SharedMemoryConnector", "YuanrongConnector", + "YuanrongTransferEngineConnector", # Utilities "load_omni_transfer_config", "initialize_connectors_from_config", diff --git a/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py b/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py index 4b21ed0ec..1c058c7f0 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py @@ -16,6 +16,7 @@ import torch import zmq from ..utils.logging import get_connector_logger +from ..utils.memory_pool import BufferAllocator, ManagedBuffer from ..utils.serialization import OmniSerializer from .base import OmniConnectorBase @@ -66,160 +67,6 @@ class MooncakeAgentMetadata: lengths: list[int] -class BufferAllocator: - """ - Manages the allocation of memory segments within the registered pool. - Thread-safe implementation using a simple free list. - """ - - def __init__(self, total_size: int, alignment: int = 4096): - self.total_size = total_size - self.alignment = alignment - self.lock = threading.Lock() - # Free list: [(start, size), ...] sorted by start - self.free_blocks = [(0, total_size)] - - def alloc(self, size: int) -> int: - """ - Allocates a block of 'size' bytes. - Returns the starting offset. - """ - # Align size upwards - aligned_size = (size + self.alignment - 1) // self.alignment * self.alignment - - with self.lock: - for i, (start, block_size) in enumerate(self.free_blocks): - if block_size >= aligned_size: - # Found a block - new_start = start + aligned_size - new_size = block_size - aligned_size - - if new_size > 0: - self.free_blocks[i] = (new_start, new_size) - else: - self.free_blocks.pop(i) - return start - - raise MemoryError(f"Out of memory in buffer pool. Requested {size} bytes (aligned {aligned_size}).") - - def free(self, offset: int, size: int): - """ - Frees a previously allocated block. - """ - aligned_size = (size + self.alignment - 1) // self.alignment * self.alignment - - with self.lock: - # Check for double-free and corruption - for start, length in self.free_blocks: - # Case 1: Exact match = double free, safe to ignore - if offset == start and aligned_size == length: - logger.warning(f"Double free detected at offset {offset}, size {aligned_size}. Ignoring.") - return - # Case 2: Block is fully contained within an existing free block = also double free - # This happens when the block was freed and then merged with adjacent blocks - if offset >= start and offset + aligned_size <= start + length: - logger.warning( - f"Double free detected: block {offset}-{offset + aligned_size} " - f"is already within free block {start}-{start + length}. Ignoring." - ) - return - # Case 3: Partial overlap (but not fully contained) = memory corruption - if not (offset + aligned_size <= start or start + length <= offset): - raise RuntimeError( - f"Memory corruption detected: freeing {offset}-{offset + aligned_size} " - f"partially overlaps with free block {start}-{start + length}" - ) - - self.free_blocks.append((offset, aligned_size)) - self.free_blocks.sort() # Sort by offset - - # Merge adjacent blocks - i = 0 - while i < len(self.free_blocks) - 1: - curr_start, curr_size = self.free_blocks[i] - next_start, next_size = self.free_blocks[i + 1] - - if curr_start + curr_size == next_start: - self.free_blocks[i] = (curr_start, curr_size + next_size) - self.free_blocks.pop(i + 1) - else: - i += 1 - - -class ManagedBuffer: - """ - A temporary view into the global memory pool. - Must be kept alive while the data view is being used. - """ - - def __init__(self, allocator: BufferAllocator, offset: int, size: int, pool_tensor: torch.Tensor): - self.allocator = allocator - self.offset = offset - self.size = size - self.pool_tensor = pool_tensor - self._released = False - - def release(self): - """Explicitly release the buffer back to the pool.""" - if not self._released: - self.allocator.free(self.offset, self.size) - self._released = True - - def __del__(self): - self.release() - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - self.release() - - @property - def tensor(self) -> torch.Tensor: - """ - Returns a 1D uint8 zero-copy view of the buffer. - """ - return self.pool_tensor[self.offset : self.offset + self.size] - - def as_tensor(self, dtype: torch.dtype, shape: tuple) -> torch.Tensor: - """ - Returns a typed, shaped zero-copy view. - Validates size, shape, and alignment. - """ - itemsize = torch.tensor([], dtype=dtype).element_size() - - # Calculate expected size - expected_bytes = itemsize - for dim in shape: - if dim < 0: - raise ValueError("Dynamic dimension (-1) is not supported in as_tensor") - expected_bytes *= dim - - if expected_bytes != self.size: - raise ValueError( - f"Shape {shape} with dtype {dtype} requires {expected_bytes} bytes, but buffer size is {self.size}" - ) - - # Check alignment (offset must be divisible by itemsize) - if self.offset % itemsize != 0: - raise RuntimeError(f"Buffer offset {self.offset} is not aligned for dtype {dtype} (itemsize {itemsize})") - - raw_view = self.tensor - # view() requires contiguous memory, slice of contiguous tensor is contiguous - typed_view = raw_view.view(dtype) - return typed_view.reshape(shape) - - def to_bytes(self) -> bytes: - """ - Returns a copy of the data as python bytes. - Performs D2H copy if pool is on GPU. - """ - t = self.tensor - if t.is_cuda: - t = t.cpu() - return t.numpy().tobytes() - - class MooncakeTransferEngineConnector(OmniConnectorBase): """ OmniConnector implementation using Mooncake Transfer Engine with a managed memory pool. diff --git a/vllm_omni/distributed/omni_connectors/connectors/yuanrong_transfer_engine_connector.py b/vllm_omni/distributed/omni_connectors/connectors/yuanrong_transfer_engine_connector.py new file mode 100644 index 000000000..3254743c5 --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/connectors/yuanrong_transfer_engine_connector.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Compatibility import for the NPU-only Yuanrong TransferEngine connector.""" + +from vllm_omni.platforms.npu.omni_connectors.yuanrong_transfer_engine_connector import ( + YuanrongTransferEngineConnector, +) + +__all__ = ["YuanrongTransferEngineConnector"] diff --git a/vllm_omni/distributed/omni_connectors/factory.py b/vllm_omni/distributed/omni_connectors/factory.py index a680da953..7ee2dcde3 100644 --- a/vllm_omni/distributed/omni_connectors/factory.py +++ b/vllm_omni/distributed/omni_connectors/factory.py @@ -91,6 +91,18 @@ def _create_yuanrong_connector(config: dict[str, Any]) -> OmniConnectorBase: return YuanrongConnector(config) +def _create_yuanrong_transfer_engine_connector(config: dict[str, Any]) -> OmniConnectorBase: + try: + from vllm_omni.platforms.npu.omni_connectors import YuanrongTransferEngineConnector + except ImportError as exc: + raise ImportError( + "YuanrongTransferEngineConnector is only available in the NPU platform " + "environment. Install the Ascend/Yuanrong runtime dependencies before " + "using this connector." + ) from exc + return YuanrongTransferEngineConnector(config) + + def _create_mooncake_transfer_engine_connector(config: dict[str, Any]) -> OmniConnectorBase: try: from .connectors.mooncake_transfer_engine_connector import MooncakeTransferEngineConnector @@ -107,5 +119,6 @@ OmniConnectorFactory.register_connector("MooncakeStoreConnector", _create_moonca OmniConnectorFactory.register_connector("MooncakeTransferEngineConnector", _create_mooncake_transfer_engine_connector) OmniConnectorFactory.register_connector("SharedMemoryConnector", _create_shm_connector) OmniConnectorFactory.register_connector("YuanrongConnector", _create_yuanrong_connector) +OmniConnectorFactory.register_connector("YuanrongTransferEngineConnector", _create_yuanrong_transfer_engine_connector) # Backward-compatible aliases – will be removed in the future OmniConnectorFactory.register_connector("MooncakeConnector", _create_mooncake_store_connector) diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index 241a6070d..372c2cb3d 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -13,7 +13,8 @@ import torch from vllm.logger import init_logger from .factory import OmniConnectorFactory -from .utils.config import ConnectorSpec +from .utils.config import TRANSFER_ENGINE_CONNECTOR_NAMES, ConnectorSpec +from .utils.env import expand_env_int from .utils.initialization import KV_RANK_PORT_STRIDE from .utils.kv_utils import ( KVTPTopology, @@ -106,7 +107,7 @@ class KVCacheTransferData: t = tensor.detach().contiguous() if cpu: t = t.cpu() - elif device is None and t.is_cuda: + elif device is None and getattr(t.device, "type", "cpu") != "cpu": device = t.device nbytes = t.numel() * t.element_size() tensors_desc.append( @@ -143,10 +144,10 @@ class KVCacheTransferData: return b"".join([self._build_header_bytes(tensors_desc)] + chunks) def to_gpu_tensor(self) -> torch.Tensor: - """Convert to a packed GPU tensor for raw-data connectors.""" + """Convert to a packed device tensor for raw-data connectors.""" tensors_desc, chunks, data_offset, device = self._build_tensors_desc(cpu=False) if device is None: - raise RuntimeError("No CUDA tensors found, use to_bytes() instead") + raise RuntimeError("No device tensors found, use to_bytes() instead") header_prefix = self._build_header_bytes(tensors_desc) output = torch.empty(len(header_prefix) + data_offset, dtype=torch.uint8, device=device) output[: len(header_prefix)].copy_(torch.frombuffer(bytearray(header_prefix), dtype=torch.uint8)) @@ -169,19 +170,19 @@ class KVCacheTransferData: return json.loads(bytes(raw_mv[4 : 4 + header_len])), raw_mv[4 + header_len :] @staticmethod - def _load_header_from_tensor(gpu_tensor: torch.Tensor) -> tuple[dict[str, Any], int]: - if gpu_tensor.dtype != torch.uint8 or gpu_tensor.dim() != 1: - raise ValueError("Packed GPU KV payload must be a 1-D uint8 tensor") + def _load_header_from_tensor(tensor: torch.Tensor) -> tuple[dict[str, Any], int]: + if tensor.dtype != torch.uint8 or tensor.dim() != 1: + raise ValueError("Packed device KV payload must be a 1-D uint8 tensor") - total_bytes = int(gpu_tensor.numel()) + total_bytes = int(tensor.numel()) if total_bytes < 4: raise ValueError("Corrupted KV payload: missing 4-byte header length") - header_len = struct.unpack(">I", gpu_tensor[:4].cpu().numpy().tobytes())[0] + header_len = struct.unpack(">I", tensor[:4].cpu().numpy().tobytes())[0] if header_len > total_bytes - 4: raise ValueError(f"Corrupted KV payload: header_len={header_len} exceeds buffer size={total_bytes}") - header_bytes = gpu_tensor[4 : 4 + header_len].cpu().numpy().tobytes() + header_bytes = tensor[4 : 4 + header_len].cpu().numpy().tobytes() return json.loads(header_bytes), 4 + header_len @staticmethod @@ -260,17 +261,22 @@ class KVCacheTransferData: return KVCacheTransferData._populate_caches(header, _get) @staticmethod - def from_bytes_gpu(gpu_tensor: torch.Tensor) -> dict[str, Any]: - """Reconstruct KV cache data from a packed GPU tensor.""" - header, data_start = KVCacheTransferData._load_header_from_tensor(gpu_tensor) - data_len = int(gpu_tensor.numel()) - data_start + def from_bytes_device(tensor: torch.Tensor) -> dict[str, Any]: + """Reconstruct KV cache data from a packed device tensor.""" + header, data_start = KVCacheTransferData._load_header_from_tensor(tensor) + data_len = int(tensor.numel()) - data_start def _get(info: dict) -> torch.Tensor: offset, nbytes = KVCacheTransferData._validate_tensor_span(info["n"], info, data_len) - return gpu_tensor[data_start + offset : data_start + offset + nbytes].clone() + return tensor[data_start + offset : data_start + offset + nbytes].clone() return KVCacheTransferData._populate_caches(header, _get) + @staticmethod + def from_bytes_gpu(tensor: torch.Tensor) -> dict[str, Any]: + """Compatibility alias for callers using the old GPU-specific name.""" + return KVCacheTransferData.from_bytes_device(tensor) + class OmniKVTransferManager: """Unified management for OmniConnector and KV cache transfer. @@ -405,8 +411,8 @@ class OmniKVTransferManager: if cfg and (c_type := cfg.get("type")): try: c_extra = {k: v for k, v in cfg.items() if k != "type"} - if c_type == "MooncakeTransferEngineConnector": - base_port = c_extra.get("zmq_port", 50051) + if c_type in TRANSFER_ENGINE_CONNECTOR_NAMES: + base_port = expand_env_int(c_extra.get("zmq_port", 50051), "zmq_port") c_extra["from_stage"] = ( str(self.config.from_stage) if self.config.from_stage is not None else "0" ) @@ -429,8 +435,13 @@ class OmniKVTransferManager: c_extra["zmq_port"] = zmq_port elif self.config.need_recv_cache: c_extra["role"] = "receiver" - c_extra.setdefault("sender_host", c_extra.get("host", "127.0.0.1")) - c_extra.setdefault("sender_zmq_port", zmq_port) + # Receiver-side sender endpoints are request-scoped. + # They are attached by the orchestrator as + # kv_sender_info and applied in update_sender_info(). + # Do not derive sender_host from this receiver's + # local YAML host; on multi-node runs that points at + # the wrong process. Explicit sender_* YAML values + # are still preserved for standalone connector use. logger.info( "Initializing OmniConnector type=%s role=%s", @@ -1058,8 +1069,8 @@ class OmniKVTransferManager: managed_buffer = raw_data try: buf_tensor = raw_data.tensor - if buf_tensor.is_cuda: - data = KVCacheTransferData.from_bytes_gpu(buf_tensor) + if getattr(buf_tensor.device, "type", "cpu") != "cpu": + data = KVCacheTransferData.from_bytes_device(buf_tensor) raw_data.release() managed_buffer = None else: @@ -1075,7 +1086,10 @@ class OmniKVTransferManager: elif isinstance(raw_data, (bytes, bytearray)): data = KVCacheTransferData.from_bytes(raw_data) elif isinstance(raw_data, torch.Tensor) and raw_data.dtype == torch.uint8 and raw_data.dim() == 1: - data = KVCacheTransferData.from_bytes(raw_data.cpu().numpy().tobytes()) + if getattr(raw_data.device, "type", "cpu") != "cpu": + data = KVCacheTransferData.from_bytes_device(raw_data) + else: + data = KVCacheTransferData.from_bytes(raw_data.numpy().tobytes()) else: data = raw_data diff --git a/vllm_omni/distributed/omni_connectors/utils/config.py b/vllm_omni/distributed/omni_connectors/utils/config.py index f18e52c4d..abe6e6a20 100644 --- a/vllm_omni/distributed/omni_connectors/utils/config.py +++ b/vllm_omni/distributed/omni_connectors/utils/config.py @@ -8,6 +8,13 @@ from .logging import get_connector_logger logger = get_connector_logger(__name__) +TRANSFER_ENGINE_CONNECTOR_NAMES = frozenset( + { + "MooncakeTransferEngineConnector", + "YuanrongTransferEngineConnector", + } +) + @dataclass class ConnectorSpec: diff --git a/vllm_omni/distributed/omni_connectors/utils/env.py b/vllm_omni/distributed/omni_connectors/utils/env.py new file mode 100644 index 000000000..ee8b324ea --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/utils/env.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os +import re +from typing import Any + +AUTO_HOST_VALUES = {"", "auto", "*", "0.0.0.0", "::"} +AUTO_DEVICE_VALUES = {"", "auto"} + +_UNEXPANDED_ENV_VAR_PATTERN = re.compile(r"\$\{([^}]+)\}|\$([A-Za-z_][A-Za-z0-9_]*)|%([^%]+)%") + + +class EnvVarExpansionError(ValueError): + pass + + +def _unexpanded_env_vars(value: str) -> list[str]: + names: list[str] = [] + for match in _UNEXPANDED_ENV_VAR_PATTERN.finditer(value): + name = next(group for group in match.groups() if group) + if name not in names: + names.append(name) + return names + + +def expand_env_value(value: Any, *, strict: bool = False, field_name: str = "value") -> Any: + if not isinstance(value, str): + return value + + expanded = os.path.expandvars(value) + if strict: + unresolved = _unexpanded_env_vars(expanded) + if unresolved: + names = ", ".join(unresolved) + raise EnvVarExpansionError( + f"{field_name} references unset environment variable(s): {names}. Original value: {value!r}" + ) + return expanded + + +def expand_env_int(value: Any, field_name: str) -> int: + expanded = expand_env_value(value, strict=True, field_name=field_name) + try: + return int(expanded) + except (TypeError, ValueError) as exc: + raise ValueError(f"{field_name} must be an integer after environment expansion, got {expanded!r}") from exc diff --git a/vllm_omni/distributed/omni_connectors/utils/initialization.py b/vllm_omni/distributed/omni_connectors/utils/initialization.py index a8cafba61..f7a464809 100644 --- a/vllm_omni/distributed/omni_connectors/utils/initialization.py +++ b/vllm_omni/distributed/omni_connectors/utils/initialization.py @@ -10,7 +10,9 @@ from pathlib import Path from typing import TYPE_CHECKING, Any from ..factory import OmniConnectorFactory -from .config import ConnectorSpec, OmniTransferConfig +from .config import TRANSFER_ENGINE_CONNECTOR_NAMES, ConnectorSpec, OmniTransferConfig +from .env import expand_env_int +from .local_rank import get_connector_local_rank from .logging import get_connector_logger if TYPE_CHECKING: @@ -93,18 +95,20 @@ def create_connectors_from_config( for edge_key, connector_spec in connectors_config.items(): from_stage, to_stage = edge_key try: - if connector_spec.name == "MooncakeTransferEngineConnector": + if connector_spec.name in TRANSFER_ENGINE_CONNECTOR_NAMES: extra = dict(connector_spec.extra) if connector_spec.extra else {} - base_port = extra.get("zmq_port", 50051) + base_port = expand_env_int(extra.get("zmq_port", 50051), "zmq_port") try: stage_offset = int(from_stage) except (TypeError, ValueError): stage_offset = 0 + local_rank = 0 if str(caller_stage_id) == "orchestrator" else get_connector_local_rank() if str(caller_stage_id) == "orchestrator": - adjusted_port = base_port + orchestrator_port_offset + stage_offset + rank0_port = base_port + orchestrator_port_offset + stage_offset else: - adjusted_port = base_port + port_offset + stage_offset + rank0_port = base_port + port_offset + stage_offset + adjusted_port = rank0_port + local_rank * KV_RANK_PORT_STRIDE extra["zmq_port"] = adjusted_port if is_sender is not None: diff --git a/vllm_omni/distributed/omni_connectors/utils/local_rank.py b/vllm_omni/distributed/omni_connectors/utils/local_rank.py new file mode 100644 index 000000000..4c2b49e8f --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/utils/local_rank.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os + +_LOCAL_RANK_ENV_NAMES = ( + "LOCAL_RANK", + "VLLM_LOCAL_RANK", + "OMPI_COMM_WORLD_LOCAL_RANK", + "MV2_COMM_WORLD_LOCAL_RANK", +) +_GLOBAL_RANK_ENV_NAMES = ("RANK", "RANK_ID") + + +def _visible_device_count() -> int | None: + visible = os.environ.get("ASCEND_RT_VISIBLE_DEVICES") or os.environ.get("CUDA_VISIBLE_DEVICES") + if not visible: + return None + devices = [item.strip() for item in visible.split(",") if item.strip()] + return len(devices) or None + + +def _get_tensor_model_parallel_rank() -> int: + from vllm.distributed.parallel_state import get_tensor_model_parallel_rank + + return int(get_tensor_model_parallel_rank()) + + +def get_connector_local_rank() -> int: + try: + return _get_tensor_model_parallel_rank() + except Exception: + pass + + for env_name in _LOCAL_RANK_ENV_NAMES: + try: + return int(os.environ[env_name]) + except (KeyError, TypeError, ValueError): + pass + + visible_count = _visible_device_count() + for env_name in _GLOBAL_RANK_ENV_NAMES: + try: + rank = int(os.environ[env_name]) + except (KeyError, TypeError, ValueError): + continue + return rank % visible_count if visible_count else rank + + return 0 diff --git a/vllm_omni/distributed/omni_connectors/utils/memory_pool.py b/vllm_omni/distributed/omni_connectors/utils/memory_pool.py new file mode 100644 index 000000000..2e2e7ba00 --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/utils/memory_pool.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import threading + +import torch + +from .logging import get_connector_logger + +logger = get_connector_logger(__name__) + + +class BufferAllocator: + """ + Manages the allocation of memory segments within a registered pool. + Thread-safe implementation using a simple free list. + """ + + def __init__(self, total_size: int, alignment: int = 4096): + self.total_size = total_size + self.alignment = alignment + self.lock = threading.Lock() + self.free_blocks = [(0, total_size)] + + def alloc(self, size: int) -> int: + """ + Allocates a block of *size* bytes and returns the starting offset. + """ + aligned_size = (size + self.alignment - 1) // self.alignment * self.alignment + + with self.lock: + for i, (start, block_size) in enumerate(self.free_blocks): + if block_size >= aligned_size: + new_start = start + aligned_size + new_size = block_size - aligned_size + + if new_size > 0: + self.free_blocks[i] = (new_start, new_size) + else: + self.free_blocks.pop(i) + return start + + raise MemoryError(f"Out of memory in buffer pool. Requested {size} bytes (aligned {aligned_size}).") + + def free(self, offset: int, size: int) -> None: + """ + Frees a previously allocated block. + """ + aligned_size = (size + self.alignment - 1) // self.alignment * self.alignment + + with self.lock: + for start, length in self.free_blocks: + if offset == start and aligned_size == length: + logger.warning("Double free detected at offset %s, size %s. Ignoring.", offset, aligned_size) + return + if offset >= start and offset + aligned_size <= start + length: + logger.warning( + "Double free detected: block %s-%s is already within free block %s-%s. Ignoring.", + offset, + offset + aligned_size, + start, + start + length, + ) + return + if not (offset + aligned_size <= start or start + length <= offset): + raise RuntimeError( + f"Memory corruption detected: freeing {offset}-{offset + aligned_size} " + f"partially overlaps with free block {start}-{start + length}" + ) + + self.free_blocks.append((offset, aligned_size)) + self.free_blocks.sort() + + i = 0 + while i < len(self.free_blocks) - 1: + curr_start, curr_size = self.free_blocks[i] + next_start, next_size = self.free_blocks[i + 1] + + if curr_start + curr_size == next_start: + self.free_blocks[i] = (curr_start, curr_size + next_size) + self.free_blocks.pop(i + 1) + else: + i += 1 + + +class ManagedBuffer: + """ + A temporary view into a global memory pool. + Must be kept alive while the data view is being used. + """ + + def __init__(self, allocator: BufferAllocator, offset: int, size: int, pool_tensor: torch.Tensor): + self.allocator = allocator + self.offset = offset + self.size = size + self.pool_tensor = pool_tensor + self._released = False + + def release(self) -> None: + """Explicitly release the buffer back to the pool.""" + if not self._released: + self.allocator.free(self.offset, self.size) + self._released = True + + def __del__(self): + self.release() + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.release() + + @property + def tensor(self) -> torch.Tensor: + """ + Returns a 1D uint8 zero-copy view of the buffer. + """ + return self.pool_tensor[self.offset : self.offset + self.size] + + def as_tensor(self, dtype: torch.dtype, shape: tuple) -> torch.Tensor: + """ + Returns a typed, shaped zero-copy view. + Validates size, shape, and alignment. + """ + itemsize = torch.tensor([], dtype=dtype).element_size() + + expected_bytes = itemsize + for dim in shape: + if dim < 0: + raise ValueError("Dynamic dimension (-1) is not supported in as_tensor") + expected_bytes *= dim + + if expected_bytes != self.size: + raise ValueError( + f"Shape {shape} with dtype {dtype} requires {expected_bytes} bytes, but buffer size is {self.size}" + ) + + if self.offset % itemsize != 0: + raise RuntimeError(f"Buffer offset {self.offset} is not aligned for dtype {dtype} (itemsize {itemsize})") + + typed_view = self.tensor.view(dtype) + return typed_view.reshape(shape) + + def to_bytes(self) -> bytes: + """ + Returns a copy of the data as Python bytes. + Performs D2H copy if the pool is on a device. + """ + t = self.tensor + if getattr(t.device, "type", "cpu") != "cpu": + t = t.cpu() + return t.numpy().tobytes() diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index e9e2944a7..84bf72863 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -8,6 +8,7 @@ from __future__ import annotations import inspect import multiprocessing.connection +import os import socket import threading import weakref @@ -20,6 +21,9 @@ from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient from vllm.v1.engine.exceptions import EngineDeadError +from vllm_omni.distributed.omni_connectors.utils.config import ( + TRANSFER_ENGINE_CONNECTOR_NAMES, +) from vllm_omni.distributed.omni_connectors.utils.initialization import ( KV_TRANSFER_PORT_OFFSET, ) @@ -331,11 +335,13 @@ class StageEngineCoreClientBase(StageClientBase): if sender_host is not None: self._kv_sender_host = sender_host + connector_type = connector_config.get("type") sender_port = connector_config.get("sender_zmq_port") - if sender_port is None: + if connector_type in TRANSFER_ENGINE_CONNECTOR_NAMES or sender_port is None: base_port = connector_config.get("zmq_port") if base_port is None: return + base_port = os.path.expandvars(str(base_port)) omni_kv_config = getattr(self, "_omni_kv_config", None) from_stage = self.stage_id diff --git a/vllm_omni/platforms/npu/omni_connectors/__init__.py b/vllm_omni/platforms/npu/omni_connectors/__init__.py new file mode 100644 index 000000000..e411df850 --- /dev/null +++ b/vllm_omni/platforms/npu/omni_connectors/__init__.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +try: + from vllm_omni.platforms.npu.omni_connectors.yuanrong_transfer_engine_connector import ( + YuanrongTransferEngineConnector, + ) +except ImportError: + YuanrongTransferEngineConnector = None + +__all__ = ["YuanrongTransferEngineConnector"] diff --git a/vllm_omni/platforms/npu/omni_connectors/yuanrong_transfer_engine_connector.py b/vllm_omni/platforms/npu/omni_connectors/yuanrong_transfer_engine_connector.py new file mode 100644 index 000000000..649d4d45c --- /dev/null +++ b/vllm_omni/platforms/npu/omni_connectors/yuanrong_transfer_engine_connector.py @@ -0,0 +1,862 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import queue +import socket +import threading +import time as _time_mod +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from typing import Any + +import msgspec +import torch +import zmq + +from vllm_omni.distributed.omni_connectors.connectors.base import OmniConnectorBase +from vllm_omni.distributed.omni_connectors.utils.env import ( + AUTO_DEVICE_VALUES, + AUTO_HOST_VALUES, + expand_env_int, + expand_env_value, +) +from vllm_omni.distributed.omni_connectors.utils.local_rank import get_connector_local_rank +from vllm_omni.distributed.omni_connectors.utils.logging import get_connector_logger +from vllm_omni.distributed.omni_connectors.utils.memory_pool import BufferAllocator, ManagedBuffer +from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer + +logger = get_connector_logger(__name__) + +try: + from yr.datasystem import TransferEngine +except ImportError: + TransferEngine = None + +_BUFFER_TTL_SECONDS = 300 + +QUERY_INFO = b"query_info" +CLEANUP_KEY = b"cleanup_key" +ACK_OK = b"ack_ok" +ACK_ERROR = b"ack_error" +INFO_NOT_FOUND = b"info_not_found" + + +@dataclass +class QueryRequest: + """Request to query metadata for a specific key.""" + + request_id: str + + +@dataclass +class CleanupRequest: + """Response containing metadata for a request.""" + + request_id: str + + +@dataclass +class QueryResponse: + request_id: str + rpc_port: int + data_size: int + is_fast_path: bool + src_addrs: list[int] + lengths: list[int] + + +def _resolve_device_name(configured_device: Any, protocol: str) -> str: + device_name = str(expand_env_value(configured_device or "auto")).strip() + if device_name.lower() not in AUTO_DEVICE_VALUES: + return device_name + if protocol.lower() in {"ascend", "hccl", "npu"}: + return f"npu:{get_connector_local_rank()}" + return "" + + +def _resolve_pool_device(configured_device: Any) -> str: + pool_device = str(expand_env_value(configured_device or "npu")).strip().lower() + if pool_device in AUTO_DEVICE_VALUES: + return "npu" + if not (pool_device == "npu" or pool_device.startswith("npu:")): + raise ValueError( + f"YuanrongTransferEngineConnector requires an NPU memory pool. Got memory_pool_device={pool_device!r}." + ) + return pool_device + + +class YuanrongTransferEngineConnector(OmniConnectorBase): + """Pull-based connector backed by Yuanrong transfer_engine.""" + + supports_raw_data: bool = True + + def __init__(self, config: dict[str, Any]): + if TransferEngine is None: + raise ImportError("Yuanrong is not available.") + + self._closed = False + self._bind_error: Exception | None = None + self._stop_event = threading.Event() + self._listener_thread: threading.Thread | None = None + self._listener_ready = threading.Event() + self._control_executor: ThreadPoolExecutor | None = None + self._local_buffers: dict[str, Any] = {} + self._local_buffers_lock = threading.Lock() + self._req_local = threading.local() + self._worker_local = threading.local() + self._last_ttl_check: float = _time_mod.monotonic() + + self._metrics = { + "puts": 0, + "gets": 0, + "bytes_transferred": 0, + "errors": 0, + "timeouts": 0, + } + + self.config = {key: expand_env_value(value) for key, value in config.items()} + config = self.config + + raw_host_config = config.get("host") + host_config = "auto" if raw_host_config is None else str(raw_host_config).strip() + self.host = self._get_local_ip() if host_config.lower() in AUTO_HOST_VALUES else host_config + self.zmq_port = self._resolve_port(config.get("zmq_port"), self.host, "zmq_port") + self.rpc_port = self._resolve_port(config.get("rpc_port"), self.host, "rpc_port") + self.protocol = str(config.get("protocol", "ascend")) + self.device_name = _resolve_device_name(config.get("device_name", "auto"), self.protocol) + self.pool_size = int(config.get("memory_pool_size", 1024**3)) + self.pool_device = _resolve_pool_device(config.get("memory_pool_device", "npu")) + self.sender_host = config.get("sender_host") + sender_zmq_port = config.get("sender_zmq_port") + self.sender_zmq_port = self._resolve_optional_port(sender_zmq_port, "sender_zmq_port") + + role = str(config.get("role", "sender")).lower() + if role not in {"sender", "receiver"}: + raise ValueError( + f"Invalid role={role!r} for YuanrongTransferEngineConnector. Expected 'sender' or 'receiver'." + ) + self.can_put = role == "sender" + + self.engine = TransferEngine() + local_endpoint = f"{self.host}:{self.rpc_port}" + self._ensure_result_ok( + self.engine.initialize(local_endpoint, self.protocol, self.device_name), + f"initialize transfer engine on {local_endpoint}", + ) + resolved_rpc_port = int(self.engine.get_rpc_port()) + if resolved_rpc_port > 0: + self.rpc_port = resolved_rpc_port + + logger.info( + "YuanrongTransferEngineConnector initialized at %s:%s (protocol=%s, device=%s)", + self.host, + self.rpc_port, + self.protocol, + self.device_name, + ) + + try: + self.pool = torch.empty(self.pool_size, dtype=torch.uint8, device=self.pool_device) + self.base_ptr = int(self.pool.data_ptr()) + self._ensure_result_ok( + self.engine.register_memory(self.base_ptr, self.pool_size), + f"register memory pool ({self.pool_size} bytes)", + ) + except Exception as exc: + logger.error("Failed to allocate/register transfer_engine pool: %s", exc) + raise + + self.allocator = BufferAllocator(self.pool_size, alignment=4096) + self.zmq_ctx = zmq.Context() + self._control_executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="yuanrong-transfer-control") + + if self.can_put: + self._listener_thread = threading.Thread(target=self._zmq_listener_loop, daemon=True) + self._listener_thread.start() + self._listener_ready.wait(timeout=1.0) + if self._bind_error is not None: + raise RuntimeError( + f"YuanrongTransferEngineConnector failed to bind ZMQ on " + f"{self.host}:{self.zmq_port}: {self._bind_error}" + ) from self._bind_error + else: + logger.info( + "YuanrongTransferEngineConnector started in receiver mode (sender_host=%r, sender_zmq_port=%r)", + self.sender_host, + self.sender_zmq_port, + ) + + @staticmethod + def _ensure_result_ok(result: Any, action: str) -> None: + if result.is_error(): + raise RuntimeError(f"{action} failed: {result.to_string()}") + + @staticmethod + def _resolve_optional_port(configured_port: Any, field_name: str) -> int | None: + configured_port = expand_env_value(configured_port, strict=True, field_name=field_name) + if configured_port is None or str(configured_port).strip().lower() in {"", "auto"}: + return None + return expand_env_int(configured_port, field_name) + + @classmethod + def _resolve_port(cls, configured_port: Any, host: str, field_name: str) -> int: + resolved_port = cls._resolve_optional_port(configured_port, field_name) + if resolved_port is not None: + return resolved_port + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind((host, 0)) + return int(sock.getsockname()[1]) + + @staticmethod + def _get_local_ip() -> str: + try: + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock: + sock.connect(("8.8.8.8", 80)) + return sock.getsockname()[0] + except Exception: + try: + return socket.gethostbyname(socket.gethostname()) + except Exception: + return "127.0.0.1" + + @staticmethod + def _sync_tensor_device(tensor: torch.Tensor | None) -> None: + if tensor is None: + return + device = getattr(tensor, "device", None) + if device is None or getattr(device, "type", "cpu") == "cpu": + return + + backend = getattr(torch, device.type, None) + if backend is None: + return + + current_stream = getattr(backend, "current_stream", None) + if callable(current_stream): + try: + current_stream(device).synchronize() + return + except TypeError: + try: + current_stream().synchronize() + return + except Exception: + pass + except Exception: + pass + + synchronize = getattr(backend, "synchronize", None) + if callable(synchronize): + try: + synchronize(device) + except TypeError: + synchronize() + + def get_connection_info(self) -> dict[str, Any]: + return { + "host": self.host, + "zmq_port": self.zmq_port, + "rpc_port": self.rpc_port, + "can_put": self.can_put, + } + + def update_sender_info( + self, + sender_host: str, + sender_zmq_port: int, + ) -> None: + self.sender_host = sender_host + self.sender_zmq_port = int(sender_zmq_port) + + def _get_req_socket(self, zmq_addr: str, timeout_ms: int) -> zmq.Socket: + cache: dict[str, zmq.Socket] = getattr(self._req_local, "cache", None) # type: ignore[assignment] + if cache is None: + cache = {} + self._req_local.cache = cache + + sock = cache.get(zmq_addr) + if sock is None: + sock = self.zmq_ctx.socket(zmq.REQ) + sock.connect(zmq_addr) + cache[zmq_addr] = sock + + sock.setsockopt(zmq.SNDTIMEO, timeout_ms) + sock.setsockopt(zmq.RCVTIMEO, timeout_ms) + return sock + + def _invalidate_req_socket(self, zmq_addr: str) -> None: + cache: dict[str, zmq.Socket] = getattr(self._req_local, "cache", None) # type: ignore[assignment] + if cache is None: + return + sock = cache.pop(zmq_addr, None) + if sock is not None: + try: + sock.close(linger=0) + except Exception: + pass + + def put(self, from_stage: str, to_stage: str, put_key: str, data: Any) -> tuple[bool, int, dict[str, Any] | None]: + if self._closed: + raise RuntimeError("Cannot put data: YuanrongTransferEngineConnector is closed") + + if not self.can_put: + self._metrics["errors"] += 1 + logger.warning("Rejecting put for %s: connector is in receiver mode", put_key) + return False, 0, None + + put_key = self._make_key(put_key, from_stage, to_stage) + + try: + src_addr = 0 + size = 0 + holder = None + is_fast_path = True + should_release_holder = False + + if isinstance(data, bytes) and len(data) == 0: + self._metrics["errors"] += 1 + logger.warning("Rejecting put for %s: empty bytes payload", put_key) + return False, 0, None + if isinstance(data, torch.Tensor) and data.nbytes == 0: + self._metrics["errors"] += 1 + logger.warning("Rejecting put for %s: zero-size tensor", put_key) + return False, 0, None + + if not isinstance(data, (ManagedBuffer, torch.Tensor, bytes)): + data = OmniSerializer.serialize(data) + is_fast_path = False + + if isinstance(data, ManagedBuffer): + if data.pool_tensor.data_ptr() != self.pool.data_ptr(): + logger.warning("ManagedBuffer from different pool detected. Falling back to copy path.") + data = data.tensor.contiguous() + else: + src_addr = self.base_ptr + data.offset + size = data.size + holder = data + should_release_holder = False + + if isinstance(data, (torch.Tensor, bytes)): + if isinstance(data, torch.Tensor): + size = data.nbytes + tensor_data = data + else: + size = len(data) + tensor_data = torch.frombuffer(data, dtype=torch.uint8) + + try: + offset = self.allocator.alloc(size) + holder = ManagedBuffer(self.allocator, offset, size, self.pool) + should_release_holder = True + except MemoryError: + self._metrics["errors"] += 1 + logger.error("Pool exhausted, cannot put data size %s", size) + return False, 0, None + + try: + dst_tensor = holder.tensor + if isinstance(data, torch.Tensor): + if not data.is_contiguous(): + data = data.contiguous() + src_view = data.view(torch.uint8).flatten() + dst_tensor.copy_(src_view, non_blocking=src_view.device != dst_tensor.device) + self._sync_tensor_device(src_view) + self._sync_tensor_device(dst_tensor) + else: + dst_tensor.copy_(tensor_data) + self._sync_tensor_device(dst_tensor) + except Exception as exc: + holder.release() + self._metrics["errors"] += 1 + logger.error("Failed to copy data to pool: %s", exc) + return False, 0, None + + src_addr = self.base_ptr + offset + + if size <= 0: + if should_release_holder and isinstance(holder, ManagedBuffer): + holder.release() + self._metrics["errors"] += 1 + logger.warning("Rejecting put for %s: final payload size is %s", put_key, size) + return False, 0, None + + with self._local_buffers_lock: + old_item = self._local_buffers.pop(put_key, None) + if old_item: + _, _, old_holder, old_should_release, _, _ = old_item + if old_should_release and isinstance(old_holder, ManagedBuffer): + old_holder.release() + self._local_buffers[put_key] = ( + [src_addr], + [size], + holder, + should_release_holder, + is_fast_path, + _time_mod.monotonic(), + ) + + metadata = { + "source_host": self.host, + "source_port": self.zmq_port, + "source_rpc_port": self.rpc_port, + "source_addrs": [src_addr], + "lengths": [size], + "data_size": size, + "is_fast_path": is_fast_path, + } + + self._metrics["puts"] += 1 + self._metrics["bytes_transferred"] += size + return True, size, metadata + except Exception as exc: + self._metrics["errors"] += 1 + logger.error("transfer_engine put failed for %s: %s", put_key, exc, exc_info=True) + return False, 0, None + + def _query_metadata_at(self, get_key: str, host: str, port: int) -> dict[str, Any] | None: + zmq_addr = f"tcp://{host}:{port}" + req_socket = self._get_req_socket(zmq_addr, timeout_ms=5000) + + try: + query = QueryRequest(request_id=get_key) + req_socket.send(QUERY_INFO + msgspec.msgpack.encode(query)) + resp = req_socket.recv() + if resp == INFO_NOT_FOUND: + return None + + query_resp = msgspec.msgpack.decode(resp, type=QueryResponse) + return { + "source_host": host, + "source_port": port, + "source_rpc_port": query_resp.rpc_port, + "source_addrs": query_resp.src_addrs, + "lengths": query_resp.lengths, + "data_size": query_resp.data_size, + "is_fast_path": query_resp.is_fast_path, + } + except Exception as exc: + self._invalidate_req_socket(zmq_addr) + logger.debug("Failed to query metadata at %s for %s: %s", zmq_addr, get_key, exc) + return None + + @classmethod + def _buffer_to_bytes(cls, buffer: ManagedBuffer) -> bytes: + tensor = buffer.tensor + cls._sync_tensor_device(tensor) + if getattr(tensor.device, "type", "cpu") != "cpu": + tensor = tensor.cpu() + return tensor.numpy().tobytes() + + def _request_sender_cleanup(self, request_id: str, source_host: str | None, source_port: int | None) -> bool: + if not source_host or not source_port: + return False + + zmq_addr = f"tcp://{source_host}:{source_port}" + req_socket = self._get_req_socket(zmq_addr, timeout_ms=3000) + + try: + payload = CleanupRequest(request_id=request_id) + req_socket.send(CLEANUP_KEY + msgspec.msgpack.encode(payload)) + return req_socket.recv() == ACK_OK + except Exception as exc: + self._invalidate_req_socket(zmq_addr) + logger.warning("Failed to cleanup sender buffer for %s: %s", request_id, exc) + return False + + def get( + self, + from_stage: str, + to_stage: str, + get_key: str, + metadata: dict[str, Any] | None = None, + ) -> tuple[Any, int] | None: + if self._closed: + raise RuntimeError("Cannot get data: YuanrongTransferEngineConnector is closed") + + get_key = self._make_key(get_key, from_stage, to_stage) + t0 = _time_mod.perf_counter() + + if not metadata: + if not self.sender_host or self.sender_zmq_port is None or str(self.sender_host).lower() == "auto": + self._metrics["errors"] += 1 + raise RuntimeError( + "get(metadata=None) requires sender_host and sender_zmq_port to be resolved. " + "Call update_sender_info(host, port) before polling." + ) + metadata = self._query_metadata_at(get_key, str(self.sender_host), int(self.sender_zmq_port)) + if not metadata: + return None + elif not (metadata.get("source_host") and metadata.get("source_rpc_port") and metadata.get("source_addrs")): + partial_host = metadata.get("source_host") + partial_port = metadata.get("source_port") + if not partial_host or not partial_port: + self._metrics["errors"] += 1 + logger.warning( + "get(%s): partial metadata missing source_host/source_port. metadata=%s", + get_key, + metadata, + ) + return None + queried = self._query_metadata_at(get_key, str(partial_host), int(partial_port)) + if not queried: + return None + metadata = queried + + t1 = _time_mod.perf_counter() + query_ms = (t1 - t0) * 1000 + + source_host = metadata.get("source_host") + source_port = metadata.get("source_port") + source_rpc_port = metadata.get("source_rpc_port") + source_addrs = list(metadata.get("source_addrs") or []) + lengths = list(metadata.get("lengths") or []) + data_size = int(metadata.get("data_size", 0)) + is_fast_path = bool(metadata.get("is_fast_path", False)) + + if not source_host or not source_rpc_port or str(source_host).lower() == "auto": + self._metrics["errors"] += 1 + logger.error( + "Invalid metadata for %s: source_host=%r, source_rpc_port=%r", + get_key, + source_host, + source_rpc_port, + ) + return None + if not source_addrs: + self._metrics["errors"] += 1 + logger.error("Invalid metadata for %s: missing source_addrs", get_key) + return None + if not lengths: + lengths = [data_size] + if len(source_addrs) != len(lengths): + self._metrics["errors"] += 1 + logger.error( + "Invalid metadata for %s: source_addrs/lengths mismatch (%s vs %s)", + get_key, + len(source_addrs), + len(lengths), + ) + return None + + total_size = sum(int(length) for length in lengths) + if data_size <= 0: + data_size = total_size + if data_size <= 0: + self._metrics["errors"] += 1 + logger.warning("Skipping get for %s: data_size is 0", get_key) + return None + + try: + offset = self.allocator.alloc(data_size) + recv_buffer = ManagedBuffer(self.allocator, offset, data_size, self.pool) + base_dst_ptr = self.base_ptr + offset + except MemoryError: + self._metrics["errors"] += 1 + logger.error("Failed to allocate %s bytes in receive pool", data_size) + return None + + dst_addrs: list[int] = [] + cursor = 0 + for length in lengths: + dst_addrs.append(base_dst_ptr + cursor) + cursor += int(length) + if cursor != data_size: + recv_buffer.release() + self._metrics["errors"] += 1 + logger.error("Length sum mismatch for %s: lengths_sum=%s, data_size=%s", get_key, cursor, data_size) + return None + + t2 = _time_mod.perf_counter() + alloc_ms = (t2 - t1) * 1000 + target_hostname = f"{source_host}:{int(source_rpc_port)}" + + try: + result = self.engine.batch_transfer_sync_read(target_hostname, dst_addrs, source_addrs, lengths) + self._ensure_result_ok(result, f"read from {target_hostname}") + self._sync_tensor_device(recv_buffer.tensor) + + t3 = _time_mod.perf_counter() + read_ms = (t3 - t2) * 1000 + + cleanup_ok = self._request_sender_cleanup(get_key, source_host, int(source_port) if source_port else None) + if not cleanup_ok: + logger.warning("Sender-side cleanup not confirmed for %s", get_key) + + if is_fast_path: + total_ms = (_time_mod.perf_counter() - t0) * 1000 + mbps = (data_size / 1024 / 1024) / (total_ms / 1000) if total_ms > 0 else 0 + logger.info( + "[YR GET] %s: query=%.1fms, alloc=%.1fms, read=%.1fms, total=%.1fms, %.1f MB/s " + "(fast_path, zero-copy)", + get_key, + query_ms, + alloc_ms, + read_ms, + total_ms, + mbps, + ) + self._metrics["gets"] += 1 + self._metrics["bytes_transferred"] += data_size + return recv_buffer, data_size + + try: + copy_start = _time_mod.perf_counter() + raw_bytes = self._buffer_to_bytes(recv_buffer) + copy_end = _time_mod.perf_counter() + value = OmniSerializer.deserialize(raw_bytes) + end = _time_mod.perf_counter() + total_ms = (end - t0) * 1000 + mbps = (data_size / 1024 / 1024) / (total_ms / 1000) if total_ms > 0 else 0 + logger.info( + "[YR GET] %s: query=%.1fms, alloc=%.1fms, read=%.1fms, copy=%.1fms, total=%.1fms, %.1f MB/s", + get_key, + query_ms, + alloc_ms, + read_ms, + (copy_end - copy_start) * 1000, + total_ms, + mbps, + ) + self._metrics["gets"] += 1 + self._metrics["bytes_transferred"] += data_size + return value, data_size + finally: + recv_buffer.release() + except Exception as exc: + self._metrics["errors"] += 1 + self._metrics["timeouts"] += 1 + logger.error("transfer_engine get failed for %s: %s", get_key, exc, exc_info=True) + recv_buffer.release() + return None + + def cleanup(self, request_id: str, from_stage: str | None = None, to_stage: str | None = None) -> None: + if (from_stage is None) != (to_stage is None): + raise ValueError( + "cleanup() requires both from_stage and to_stage, or neither. " + f"Got from_stage={from_stage!r}, to_stage={to_stage!r}" + ) + if from_stage is not None and to_stage is not None: + request_id = self._make_key(request_id, from_stage, to_stage) + + with self._local_buffers_lock: + item = self._local_buffers.pop(request_id, None) + if not item: + return + _, _, holder, should_release, _, _ = item + if should_release and isinstance(holder, ManagedBuffer): + holder.release() + + def health(self) -> dict[str, Any]: + if self._closed: + return {"status": "unhealthy", "error": "Connector is closed"} + + return { + "status": "healthy", + "host": self.host, + "zmq_port": self.zmq_port, + "rpc_port": self.rpc_port, + "protocol": self.protocol, + "pool_device": self.pool_device, + "pool_size": self.pool_size, + **self._metrics, + } + + def close(self) -> None: + if getattr(self, "_closed", True): + return + self._closed = True + + self._stop_event.set() + + if self._listener_thread is not None and self._listener_thread.is_alive(): + self._listener_thread.join(timeout=2.0) + + if self._control_executor is not None: + self._control_executor.shutdown(wait=True, cancel_futures=False) + + with self._local_buffers_lock: + for _, item in list(self._local_buffers.items()): + _, _, holder, should_release, _, _ = item + if should_release and isinstance(holder, ManagedBuffer): + holder.release() + self._local_buffers.clear() + + cache: dict[str, zmq.Socket] = getattr(self._req_local, "cache", None) # type: ignore[assignment] + if cache: + for sock in cache.values(): + try: + sock.close(linger=0) + except Exception: + pass + cache.clear() + + try: + if hasattr(self, "engine") and hasattr(self.engine, "unregister_memory"): + result = self.engine.unregister_memory(self.base_ptr) + if result.is_error(): + logger.warning("Failed to unregister memory: %s", result.to_string()) + except Exception as exc: + logger.warning("Failed to unregister memory: %s", exc) + + try: + if hasattr(self, "engine") and hasattr(self.engine, "finalize"): + result = self.engine.finalize() + if result.is_error(): + logger.warning("Failed to finalize transfer engine: %s", result.to_string()) + except Exception as exc: + logger.warning("Failed to finalize transfer engine: %s", exc) + + try: + if hasattr(self, "zmq_ctx"): + self.zmq_ctx.term() + except Exception as exc: + logger.warning("Failed to terminate ZMQ context: %s", exc) + + self.pool = None + + def _cleanup_stale_buffers(self) -> None: + now = _time_mod.monotonic() + with self._local_buffers_lock: + stale_keys = [k for k, v in self._local_buffers.items() if now - v[5] > _BUFFER_TTL_SECONDS] + for key in stale_keys: + item = self._local_buffers.pop(key) + _, _, holder, should_release, _, _ = item + if should_release and isinstance(holder, ManagedBuffer): + holder.release() + logger.warning("TTL expired (%ss): cleaned up stale buffer for %s", _BUFFER_TTL_SECONDS, key) + + def _zmq_listener_loop(self) -> None: + socket_obj = self.zmq_ctx.socket(zmq.ROUTER) + try: + socket_obj.bind(f"tcp://{self.host}:{self.zmq_port}") + except zmq.ZMQError as exc: + logger.error("ZMQ bind failed on %s:%s: %s", self.host, self.zmq_port, exc) + self.can_put = False + self._bind_error = exc + self._listener_ready.set() + return + + self._listener_ready.set() + + notify_addr = f"inproc://notify-{id(self)}" + notify_recv = self.zmq_ctx.socket(zmq.PULL) + notify_recv.bind(notify_addr) + + poller = zmq.Poller() + poller.register(socket_obj, zmq.POLLIN) + poller.register(notify_recv, zmq.POLLIN) + response_queue: queue.Queue = queue.Queue() + + try: + while not self._stop_event.is_set(): + try: + events = dict(poller.poll(1000)) + + if notify_recv in events: + while True: + try: + notify_recv.recv(zmq.NOBLOCK) + except zmq.Again: + break + + while True: + try: + identity, response = response_queue.get_nowait() + socket_obj.send_multipart([identity, b"", response]) + except queue.Empty: + break + + now = _time_mod.monotonic() + if now - self._last_ttl_check >= 10.0: + self._last_ttl_check = now + self._cleanup_stale_buffers() + + if socket_obj in events: + frames = socket_obj.recv_multipart() + if len(frames) >= 2: + self._control_executor.submit( + self._handle_control_request, + response_queue, + notify_addr, + frames[0], + frames[-1], + ) + except zmq.ContextTerminated: + break + except Exception: + logger.debug("Listener loop error", exc_info=True) + finally: + try: + notify_recv.close(linger=0) + socket_obj.close(linger=0) + except Exception: + pass + + def _handle_control_request(self, response_queue: queue.Queue, notify_addr: str, identity, payload: bytes) -> None: + try: + if payload.startswith(QUERY_INFO): + self._handle_query_request(response_queue, notify_addr, identity, payload[len(QUERY_INFO) :]) + return + if payload.startswith(CLEANUP_KEY): + self._handle_cleanup_request(response_queue, notify_addr, identity, payload[len(CLEANUP_KEY) :]) + return + + response_queue.put((identity, ACK_ERROR)) + except Exception as exc: + logger.error("Control request failed: %s", exc) + response_queue.put((identity, ACK_ERROR)) + + self._notify_listener(notify_addr) + + def _handle_query_request(self, response_queue: queue.Queue, notify_addr: str, identity, payload: bytes) -> None: + try: + query = msgspec.msgpack.decode(payload, type=QueryRequest) + with self._local_buffers_lock: + item = self._local_buffers.get(query.request_id) + + if not item: + response_queue.put((identity, INFO_NOT_FOUND)) + else: + src_addrs, src_lengths, _, _, is_fast_path, _ = item + resp = QueryResponse( + request_id=query.request_id, + rpc_port=self.rpc_port, + data_size=sum(int(length) for length in src_lengths), + is_fast_path=is_fast_path, + src_addrs=src_addrs, + lengths=src_lengths, + ) + response_queue.put((identity, msgspec.msgpack.encode(resp))) + except Exception as exc: + logger.error("Query request failed: %s", exc) + response_queue.put((identity, INFO_NOT_FOUND)) + + self._notify_listener(notify_addr) + + def _handle_cleanup_request(self, response_queue: queue.Queue, notify_addr: str, identity, payload: bytes) -> None: + try: + request = msgspec.msgpack.decode(payload, type=CleanupRequest) + self.cleanup(request.request_id) + response_queue.put((identity, ACK_OK)) + except Exception as exc: + logger.error("Cleanup request failed: %s", exc) + response_queue.put((identity, ACK_ERROR)) + + self._notify_listener(notify_addr) + + def _notify_listener(self, notify_addr: str) -> None: + try: + local = self._worker_local + sock = getattr(local, "notify_socket", None) + cached_addr = getattr(local, "notify_addr", None) + if sock is None or cached_addr != notify_addr: + if sock is not None: + sock.close(linger=0) + sock = self.zmq_ctx.socket(zmq.PUSH) + sock.connect(notify_addr) + local.notify_socket = sock + local.notify_addr = notify_addr + sock.send(b"", zmq.NOBLOCK) + except Exception: + local.notify_socket = None + local.notify_addr = None From fc8486cf20d9014bb189ffceb042c109f7e6d340 Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Wed, 20 May 2026 03:24:06 -0400 Subject: [PATCH 27/38] [Doc][TTS] CosyVoice3 online docs + residual TTS yaml cleanup + remove VoxCPM v1 (#3748) Signed-off-by: Yueqian Lin --- benchmarks/tts/README.md | 1 - benchmarks/tts/bench_voxcpm_offline.py | 922 ------------------ docs/api/README.md | 1 - docs/serving/speech_api.md | 7 + .../offline_inference/text_to_speech.md | 60 -- .../examples/online_serving/text_to_speech.md | 136 ++- .../text_to_speech/README.md | 55 +- .../text_to_speech/voxcpm/end2end.py | 220 ----- .../online_serving/text_to_speech/README.md | 71 +- .../text_to_speech/cosyvoice3/run_server.sh | 28 + .../cosyvoice3/speech_client.py | 142 +++ .../voxcpm/openai_speech_client.py | 155 --- .../text_to_speech/voxcpm/run_server.sh | 38 - pyproject.toml | 2 +- .../test_voxcpm_expansion.py | 99 -- tests/engine/test_arg_utils.py | 18 - .../openai_api/test_serving_speech_voxcpm.py | 275 ------ tests/entrypoints/test_utils.py | 33 - .../test_voxcpm_async_chunk.py | 85 -- vllm_omni/deploy/voxcpm.yaml | 69 -- vllm_omni/engine/arg_utils.py | 1 - .../entrypoints/openai/serving_speech.py | 69 +- vllm_omni/model_executor/models/registry.py | 6 - .../model_executor/models/voxcpm/__init__.py | 5 - .../model_executor/models/voxcpm/voxcpm.py | 886 ----------------- .../models/voxcpm/voxcpm_loader.py | 247 ----- .../models/voxcpm/voxcpm_runtime_utils.py | 44 - .../models/voxcpm/voxcpm_stage_wrappers.py | 185 ---- .../stage_configs/qwen3_tts_uniproc.yaml | 95 -- .../stage_configs/voxcpm_async_chunk.yaml | 98 -- .../stage_input_processors/voxcpm.py | 115 --- .../platforms/npu/stage_configs/voxcpm.yaml | 67 -- .../npu/stage_configs/voxcpm_async_chunk.yaml | 89 -- .../transformers_utils/configs/__init__.py | 3 - .../transformers_utils/configs/voxcpm.py | 68 -- 35 files changed, 248 insertions(+), 4147 deletions(-) delete mode 100644 benchmarks/tts/bench_voxcpm_offline.py delete mode 100644 examples/offline_inference/text_to_speech/voxcpm/end2end.py create mode 100755 examples/online_serving/text_to_speech/cosyvoice3/run_server.sh create mode 100644 examples/online_serving/text_to_speech/cosyvoice3/speech_client.py delete mode 100644 examples/online_serving/text_to_speech/voxcpm/openai_speech_client.py delete mode 100755 examples/online_serving/text_to_speech/voxcpm/run_server.sh delete mode 100644 tests/e2e/offline_inference/test_voxcpm_expansion.py delete mode 100644 tests/entrypoints/openai_api/test_serving_speech_voxcpm.py delete mode 100644 tests/model_executor/stage_input_processors/test_voxcpm_async_chunk.py delete mode 100644 vllm_omni/deploy/voxcpm.yaml delete mode 100644 vllm_omni/model_executor/models/voxcpm/__init__.py delete mode 100644 vllm_omni/model_executor/models/voxcpm/voxcpm.py delete mode 100644 vllm_omni/model_executor/models/voxcpm/voxcpm_loader.py delete mode 100644 vllm_omni/model_executor/models/voxcpm/voxcpm_runtime_utils.py delete mode 100644 vllm_omni/model_executor/models/voxcpm/voxcpm_stage_wrappers.py delete mode 100644 vllm_omni/model_executor/stage_configs/qwen3_tts_uniproc.yaml delete mode 100644 vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml delete mode 100644 vllm_omni/model_executor/stage_input_processors/voxcpm.py delete mode 100644 vllm_omni/platforms/npu/stage_configs/voxcpm.yaml delete mode 100644 vllm_omni/platforms/npu/stage_configs/voxcpm_async_chunk.yaml delete mode 100644 vllm_omni/transformers_utils/configs/voxcpm.py diff --git a/benchmarks/tts/README.md b/benchmarks/tts/README.md index d8841f382..f6930a866 100644 --- a/benchmarks/tts/README.md +++ b/benchmarks/tts/README.md @@ -233,7 +233,6 @@ sentinel for regressions in this area. benchmarks/tts/ ├── README.md (this file) ├── bench_tts.py CLI — serve-mode benchmark driver -├── bench_voxcpm_offline.py CLI — offline VoxCPM benchmark (sync + streaming) ├── plot_results.py Generate per-task / per-concurrency curves └── model_configs.yaml Model registry (supported tasks + extra body) ``` diff --git a/benchmarks/tts/bench_voxcpm_offline.py b/benchmarks/tts/bench_voxcpm_offline.py deleted file mode 100644 index 9ab0fa508..000000000 --- a/benchmarks/tts/bench_voxcpm_offline.py +++ /dev/null @@ -1,922 +0,0 @@ -"""Offline VoxCPM benchmark for vLLM Omni. - -Supports both: -- sync one-shot (Omni.generate) -- streaming (AsyncOmni.generate with async_chunk config) -- text-only synthesis -- voice cloning -- text/clone batch inputs from txt or jsonl - -Usage:: - - # Sync (default voice) - python benchmarks/tts/bench_voxcpm_offline.py \\ - --model /path/to/VoxCPM \\ - --text "Hello world" \\ - --output-dir results/audio/ - - # Streaming (async_chunk) - python benchmarks/tts/bench_voxcpm_offline.py \\ - --model /path/to/VoxCPM \\ - --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \\ - --txt-prompts prompts.txt \\ - --output-dir results/audio/ - - # Voice cloning batch via JSONL - python benchmarks/tts/bench_voxcpm_offline.py \\ - --model /path/to/VoxCPM \\ - --jsonl-prompts prompts.jsonl \\ - --output-dir results/audio/ -""" - -from __future__ import annotations - -import asyncio -import json -import logging -import os -import tempfile -import time -import uuid -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -import torch -from vllm.utils.argparse_utils import FlexibleArgumentParser - -from vllm_omni import AsyncOmni, Omni - - -def _find_repo_root(start: Path) -> Path: - """Walk up from ``start`` until a repo marker is found. - - Falls back to ``parents[2]`` for backwards compatibility if no marker hits - (which can only happen in unusual checkouts — the tree should always have - pyproject.toml + vllm_omni/ at the top level). - """ - for candidate in [start, *start.parents]: - if (candidate / "pyproject.toml").is_file() and (candidate / "vllm_omni").is_dir(): - return candidate - return start.parents[2] - - -REPO_ROOT = _find_repo_root(Path(__file__).resolve()) -DEFAULT_STAGE_ASYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm_async_chunk.yaml" -DEFAULT_STAGE_SYNC = REPO_ROOT / "vllm_omni" / "deploy" / "voxcpm.yaml" - -logger = logging.getLogger(__name__) - - -@dataclass(frozen=True, slots=True) -class PromptSpec: - text: str - label: str - ref_audio: str | None = None - ref_text: str | None = None - - -def _require_soundfile(): - try: - import soundfile as sf # type: ignore - except ModuleNotFoundError as exc: - raise RuntimeError( - "soundfile is required to write VoxCPM benchmark WAV outputs. Install it with: pip install soundfile" - ) from exc - return sf - - -def _build_prompt( - args, - *, - text: str, - ref_audio: str | None = None, - ref_text: str | None = None, - global_request_id: str | None = None, -) -> dict[str, Any]: - additional_information: dict[str, list[Any]] = { - "text": [text], - "cfg_value": [args.cfg_value], - "inference_timesteps": [args.inference_timesteps], - "min_len": [args.min_len], - "max_new_tokens": [args.max_new_tokens], - } - if args.streaming_prefix_len is not None: - additional_information["streaming_prefix_len"] = [args.streaming_prefix_len] - - if ref_audio: - additional_information["ref_audio"] = [ref_audio] - if ref_text: - additional_information["ref_text"] = [ref_text] - if global_request_id is not None: - additional_information["global_request_id"] = [global_request_id] - - return { - "prompt_token_ids": [1], - "additional_information": additional_information, - } - - -def _extract_audio_tensor(mm: dict[str, Any]) -> torch.Tensor: - audio = mm.get("audio", mm.get("model_outputs")) - if audio is None: - raise ValueError("No audio output found in multimodal output.") - if isinstance(audio, list): - parts = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio] - audio = torch.cat(parts, dim=-1) if parts else torch.zeros(0) - if not isinstance(audio, torch.Tensor): - audio = torch.as_tensor(audio) - return audio.float().cpu().reshape(-1) - - -def _extract_sample_rate(mm: dict[str, Any]) -> int: - sr_raw = mm.get("sr", 24000) - if isinstance(sr_raw, list) and sr_raw: - sr_raw = sr_raw[-1] - if hasattr(sr_raw, "item"): - return int(sr_raw.item()) - return int(sr_raw) - - -def _emit_offline_metrics( - *, - request_id: str, - elapsed_s: float, - first_audio_elapsed: float | None, - audio_duration_s: float, -) -> None: - metrics = { - "request_id": request_id, - "ttfp_ms": round(first_audio_elapsed * 1000.0, 3) if first_audio_elapsed is not None else None, - "audio_duration_s": round(audio_duration_s, 6), - "rtf": round(elapsed_s / audio_duration_s, 6) if audio_duration_s > 0 else None, - } - print(f"[OfflineMetrics] {metrics}") - - -def _write_audio_tensor(output_path: Path, audio_tensor: Any, sample_rate: int) -> None: - sf = _require_soundfile() - if isinstance(audio_tensor, torch.Tensor): - audio_np = audio_tensor.float().cpu().clamp(-1.0, 1.0).numpy() - else: - audio_np = torch.as_tensor(audio_tensor).float().cpu().clamp(-1.0, 1.0).numpy() - sf.write( - output_path, - audio_np, - sample_rate, - format="WAV", - subtype="PCM_16", - ) - - -def _save_wav(mm: dict[str, Any], output_dir: Path, request_id: str) -> Path: - output_dir.mkdir(parents=True, exist_ok=True) - output_path = output_dir / f"output_{request_id}.wav" - _write_audio_tensor(output_path, _extract_audio_tensor(mm), _extract_sample_rate(mm)) - return output_path - - -def _iter_request_multimodal_outputs(request_output: Any): - outputs = getattr(request_output, "outputs", None) - if outputs: - for output in outputs: - mm = getattr(output, "multimodal_output", None) - if isinstance(mm, dict): - yield mm - - mm = getattr(request_output, "multimodal_output", None) - if isinstance(mm, dict): - yield mm - - -def _read_non_empty_lines(path: str) -> list[str]: - with open(path, encoding="utf-8") as f: - return [line.strip() for line in f if line.strip()] - - -def _load_prompt_specs(args) -> list[PromptSpec]: - specs: list[PromptSpec] = [] - - if args.txt_prompts is not None: - texts = _read_non_empty_lines(args.txt_prompts) - if not texts: - raise ValueError(f"No prompts found in {args.txt_prompts}") - for idx, text in enumerate(texts, start=1): - specs.append( - PromptSpec( - text=text, - label=f"item{idx:03d}", - ref_audio=args.ref_audio, - ref_text=args.ref_text, - ) - ) - return specs - - if args.jsonl_prompts is not None: - with open(args.jsonl_prompts, encoding="utf-8") as f: - for line_no, raw_line in enumerate(f, start=1): - line = raw_line.strip() - if not line: - continue - try: - item = json.loads(line) - except json.JSONDecodeError as exc: - raise ValueError(f"{args.jsonl_prompts}:{line_no} is not valid JSON: {exc}") from exc - if not isinstance(item, dict): - raise ValueError(f"{args.jsonl_prompts}:{line_no} must be a JSON object") - - text = item.get("text") - if not isinstance(text, str) or not text.strip(): - raise ValueError(f"{args.jsonl_prompts}:{line_no} requires non-empty string field 'text'") - - ref_audio = item.get("ref_audio", args.ref_audio) - ref_text = item.get("ref_text", args.ref_text) - if (ref_audio is None) != (ref_text is None): - raise ValueError( - f"{args.jsonl_prompts}:{line_no} must provide both 'ref_audio' and 'ref_text' together" - ) - - specs.append( - PromptSpec( - text=text.strip(), - label=f"item{len(specs) + 1:03d}", - ref_audio=ref_audio, - ref_text=ref_text, - ) - ) - - if not specs: - raise ValueError(f"No prompts found in {args.jsonl_prompts}") - return specs - - specs.append( - PromptSpec( - text=args.text, - label="item001", - ref_audio=args.ref_audio, - ref_text=args.ref_text, - ) - ) - return specs - - -def _build_prompt_for_spec(args, spec: PromptSpec, *, global_request_id: str | None = None) -> dict[str, Any]: - return _build_prompt( - args, - text=spec.text, - ref_audio=spec.ref_audio, - ref_text=spec.ref_text, - global_request_id=global_request_id, - ) - - -def _count_voice_clone_prompts(prompt_specs: list[PromptSpec]) -> int: - return sum(1 for spec in prompt_specs if spec.ref_audio is not None) - - -def _get_warmup_specs(prompt_specs: list[PromptSpec]) -> list[PromptSpec]: - return prompt_specs[:1] - - -def _extract_stream_finished(stage_output: Any) -> bool: - request_output = getattr(stage_output, "request_output", None) - request_finished = getattr(request_output, "finished", None) - if request_finished is not None: - return bool(request_finished) - return bool(getattr(stage_output, "finished", False)) - - -def _build_profiled_stage_config( - stage_configs_path: str, - profiler_dir: str, -) -> str: - stage_config_path = Path(stage_configs_path) - yaml_text = stage_config_path.read_text(encoding="utf-8") - injected_lines: list[str] = [] - injected_count = 0 - - for line in yaml_text.splitlines(): - injected_lines.append(line) - if line.strip() != "engine_args:": - continue - indent = line[: len(line) - len(line.lstrip())] - child_indent = indent + " " - grandchild_indent = child_indent + " " - injected_lines.extend( - [ - f"{child_indent}profiler_config:", - f'{grandchild_indent}profiler: "torch"', - f'{grandchild_indent}torch_profiler_dir: "{profiler_dir}"', - f"{grandchild_indent}torch_profiler_with_stack: true", - ] - ) - injected_count += 1 - - if injected_count == 0: - raise ValueError(f"No engine_args block found in stage config: {stage_configs_path}") - - tmp = tempfile.NamedTemporaryFile( - mode="w", - encoding="utf-8", - delete=False, - suffix=".yaml", - prefix=f"{stage_config_path.stem}_profile_", - ) - tmp.write("\n".join(injected_lines) + "\n") - tmp.close() - return tmp.name - - -def parse_args(): - parser = FlexibleArgumentParser( - description="Offline split-stage VoxCPM inference with vLLM Omni (auto sync/streaming by stage config)" - ) - parser.add_argument( - "--model", - type=str, - default=os.environ.get("VOXCPM_MODEL"), - help="Local VoxCPM model directory. Defaults to $VOXCPM_MODEL.", - ) - parser.add_argument( - "--text", - type=str, - default="This is a split-stage VoxCPM synthesis example running on vLLM Omni.", - help="Text to synthesize. Ignored when --txt-prompts or --jsonl-prompts is used.", - ) - parser.add_argument( - "--txt-prompts", - type=str, - default=None, - help="Path to a .txt file with one synthesis text per line.", - ) - parser.add_argument( - "--jsonl-prompts", - type=str, - default=None, - help=( - "Path to a .jsonl file. Each line must contain at least {'text': ...}; " - "clone rows can also set ref_audio/ref_text, and ref_text must be the " - "real transcript of ref_audio." - ), - ) - parser.add_argument( - "--ref-audio", - type=str, - default=None, - help=( - "Optional reference audio path for voice cloning. With --txt-prompts, " - "the same reference is applied to every line." - ), - ) - parser.add_argument( - "--ref-text", - type=str, - default=None, - help=( - "Real transcript of the reference audio. Placeholder text or mismatched " - "text will usually produce noisy/electronic clone audio." - ), - ) - parser.add_argument( - "--stage-configs-path", - type=str, - default=str(DEFAULT_STAGE_SYNC), - help="Stage config YAML path. Routing is selected only from this path.", - ) - parser.add_argument( - "--cfg-value", - type=float, - default=2.0, - help="Classifier-free guidance value for VoxCPM.", - ) - parser.add_argument( - "--inference-timesteps", - type=int, - default=10, - help="Number of inference timesteps.", - ) - parser.add_argument( - "--min-len", - type=int, - default=2, - help="Minimum generated token length.", - ) - parser.add_argument( - "--max-new-tokens", - type=int, - default=4096, - help="Maximum generated token length.", - ) - parser.add_argument( - "--streaming-prefix-len", - type=int, - default=None, - help="VoxCPM streaming window (optional, streaming mode only).", - ) - parser.add_argument( - "--output-dir", - type=str, - default=None, - help="Directory for output WAV files.", - ) - parser.add_argument( - "--stage-init-timeout", - type=int, - default=600, - help="Stage initialization timeout in seconds.", - ) - parser.add_argument( - "--log-stats", - dest="log_stats", - action="store_true", - help="Enable vLLM Omni stats logging.", - ) - parser.add_argument( - "--no-log-stats", - dest="log_stats", - action="store_false", - help="Disable vLLM Omni stats logging.", - ) - parser.set_defaults(log_stats=True) - parser.add_argument( - "--num-runs", - type=int, - default=1, - help="Number of full inference runs (same prompt each time). Default 1.", - ) - parser.add_argument( - "--warmup-runs", - type=int, - default=0, - help=( - "Optional number of warmup passes before measured runs. Warmup uses only " - "the first prompt and does not save outputs." - ), - ) - parser.add_argument( - "--enable-profiler", - action="store_true", - help=( - "Enable torch profiler for the configured stages. A temporary profiled " - "stage config is generated automatically." - ), - ) - parser.add_argument( - "--profiler-dir", - type=str, - default=None, - help="Directory for profiler traces. Defaults to /profiler when profiling is enabled.", - ) - parser.add_argument( - "--profiler-stages", - type=int, - nargs="*", - default=None, - help="Optional stage ids to profile. Defaults to all stages that have profiler_config.", - ) - parser.add_argument( - "--profiler-wait-seconds", - type=float, - default=30.0, - help="Seconds to wait after stop_profile for trace files to flush.", - ) - args = parser.parse_args() - - if not args.model: - parser.error("--model is required unless $VOXCPM_MODEL is set") - if args.txt_prompts is not None and args.jsonl_prompts is not None: - parser.error("--txt-prompts and --jsonl-prompts are mutually exclusive") - if (args.ref_audio is None) != (args.ref_text is None): - parser.error("--ref-audio and --ref-text must be provided together") - if args.num_runs < 1: - parser.error("--num-runs must be >= 1") - if args.warmup_runs < 0: - parser.error("--warmup-runs must be >= 0") - if args.output_dir is None: - args.output_dir = ( - "output_audio_streaming" if _is_streaming_stage_config(args.stage_configs_path) else "output_audio" - ) - if args.enable_profiler and args.profiler_dir is None: - args.profiler_dir = str(Path(args.output_dir) / "profiler") - try: - args.prompt_specs = _load_prompt_specs(args) - except ValueError as exc: - parser.error(str(exc)) - - return args - - -def _is_streaming_stage_config(stage_configs_path: str) -> bool: - cfg_name = Path(stage_configs_path).name.lower() - return "async_chunk" in cfg_name - - -async def _collect_streaming_audio( - omni: AsyncOmni, - args: Any, - spec: PromptSpec, - request_id: str, - *, - phase_label: str, - prompt_index: int, - prompt_count: int, - print_prompt: bool = False, -) -> tuple[torch.Tensor, int, float, float | None]: - prompt = _build_prompt_for_spec(args, spec, global_request_id=request_id) - delta_chunks: list[torch.Tensor] = [] - sample_rate = 24000 - chunk_i = 0 - prev_total_samples = 0 - t_start = time.perf_counter() - first_audio_elapsed: float | None = None - - if print_prompt: - print(f"---prompt---:{prompt}") - - async for stage_output in omni.generate(prompt, request_id=request_id): - mm = getattr(stage_output, "multimodal_output", None) - if not isinstance(mm, dict): - ro = getattr(stage_output, "request_output", None) - if ro is None: - continue - mm = getattr(ro, "multimodal_output", None) - if not isinstance(mm, dict) and getattr(ro, "outputs", None): - seq = ro.outputs[0] - mm = getattr(seq, "multimodal_output", None) - if not isinstance(mm, dict): - continue - sample_rate = _extract_sample_rate(mm) - try: - w = _extract_audio_tensor(mm) - n = int(w.numel()) - if n == 0: - continue - finished = _extract_stream_finished(stage_output) - if n > prev_total_samples: - delta = w.reshape(-1)[prev_total_samples:] - prev_total_samples = n - elif finished and n == prev_total_samples: - delta = w.reshape(-1)[:0] - else: - delta = w.reshape(-1) - prev_total_samples += int(delta.numel()) - if int(delta.numel()) > 0: - delta_chunks.append(delta) - if first_audio_elapsed is None and int(delta.numel()) > 0: - first_audio_elapsed = time.perf_counter() - t_start - logger.info( - "%s prompt=%d/%d chunk=%d delta_samples=%d buf_len=%d finished=%s", - phase_label, - prompt_index + 1, - prompt_count, - chunk_i, - int(delta.numel()), - n, - finished, - ) - chunk_i += 1 - except ValueError: - if not _extract_stream_finished(stage_output): - logger.debug("skip non-audio partial output chunk=%d", chunk_i) - - if not delta_chunks: - raise RuntimeError("No audio chunks received; check stage config and logs.") - - audio_cat = torch.cat([c.reshape(-1) for c in delta_chunks], dim=0) - elapsed = time.perf_counter() - t_start - return audio_cat, sample_rate, elapsed, first_audio_elapsed - - -async def _abort_streaming_residual_work( - omni: AsyncOmni, - request_id: str, - *, - settle_seconds: float = 0.1, -) -> None: - """Stop any late stage-0 work once the final audio has been collected.""" - await omni.engine.abort_async([request_id]) - if settle_seconds > 0: - await asyncio.sleep(settle_seconds) - - -async def _run_streaming_single( - omni: AsyncOmni, - args: Any, - spec: PromptSpec, - output_dir: Path, - request_id: str, - *, - run_index: int, - num_runs: int, - prompt_index: int, - prompt_count: int, -) -> Path: - audio_cat, sample_rate, elapsed, first_audio_elapsed = await _collect_streaming_audio( - omni, - args, - spec, - request_id, - phase_label=f"run={run_index + 1}/{num_runs}", - prompt_index=prompt_index, - prompt_count=prompt_count, - print_prompt=(run_index == 0 and prompt_index == 0), - ) - await _abort_streaming_residual_work(omni, request_id) - output_path = output_dir / f"output_run{run_index + 1}_{spec.label}.wav" - _write_audio_tensor(output_path, audio_cat, sample_rate) - audio_duration_s = float(audio_cat.numel()) / float(sample_rate) if sample_rate > 0 else 0.0 - ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else "" - rtf_text = f", rtf={elapsed / audio_duration_s:.3f}" if audio_duration_s > 0 else "" - print( - f"Saved (streaming) run {run_index + 1}/{num_runs}, " - f"prompt {prompt_index + 1}/{prompt_count}: {output_path} ({elapsed:.2f}s{ttfp_text}{rtf_text})" - ) - _emit_offline_metrics( - request_id=request_id, - elapsed_s=elapsed, - first_audio_elapsed=first_audio_elapsed, - audio_duration_s=audio_duration_s, - ) - return output_path - - -async def _run_streaming_warmup(args, omni: AsyncOmni) -> None: - if args.warmup_runs == 0: - return - - warmup_specs = _get_warmup_specs(args.prompt_specs) - print( - f"Warmup: {args.warmup_runs} run(s) using the first prompt " - f"({len(warmup_specs)} prompt(s)); outputs will be discarded." - ) - for warmup_index in range(args.warmup_runs): - t_warmup = time.perf_counter() - tasks = [] - request_ids: list[str] = [] - for prompt_index, spec in enumerate(warmup_specs): - request_id = f"warmup_stream_{warmup_index + 1}_{spec.label}_{uuid.uuid4().hex[:8]}" - request_ids.append(request_id) - tasks.append( - _collect_streaming_audio( - omni, - args, - spec, - request_id, - phase_label=f"warmup={warmup_index + 1}/{args.warmup_runs}", - prompt_index=prompt_index, - prompt_count=len(warmup_specs), - ) - ) - results = await asyncio.gather(*tasks) - for request_id in request_ids: - await _abort_streaming_residual_work(omni, request_id) - total_samples = sum(int(audio.numel()) for audio, _, _, _ in results) - warmup_ttfps = [ttfp for _, _, _, ttfp in results if ttfp is not None] - ttfp_text = f", ttfp={min(warmup_ttfps):.2f}s" if warmup_ttfps else "" - print( - f"Warmup (streaming) {warmup_index + 1}/{args.warmup_runs} finished: " - f"{len(results)} prompt(s), {total_samples} sample(s) " - f"({time.perf_counter() - t_warmup:.2f}s{ttfp_text})" - ) - - -async def _run_streaming(args) -> list[Path]: - output_dir = Path(args.output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - - omni = AsyncOmni( - model=args.model, - stage_configs_path=args.stage_configs_path, - log_stats=args.log_stats, - stage_init_timeout=args.stage_init_timeout, - ) - - await _run_streaming_warmup(args, omni) - profiler_started = False - if args.enable_profiler: - profile_prefix = f"voxcpm_streaming_{int(time.time())}" - stages_text = args.profiler_stages if args.profiler_stages is not None else "all-configured" - print(f"Starting profiler (streaming): stages={stages_text}, dir={args.profiler_dir}") - await omni.start_profile(profile_prefix=profile_prefix, stages=args.profiler_stages) - profiler_started = True - t_total = time.perf_counter() - total_elapsed = 0.0 - paths: list[Path] = [] - prompt_specs: list[PromptSpec] = args.prompt_specs - try: - for run in range(args.num_runs): - for prompt_index, spec in enumerate(prompt_specs): - request_id = f"stream_{run + 1}_{spec.label}_{uuid.uuid4().hex[:8]}" - paths.append( - await _run_streaming_single( - omni, - args, - spec, - output_dir, - request_id, - run_index=run, - num_runs=args.num_runs, - prompt_index=prompt_index, - prompt_count=len(prompt_specs), - ) - ) - total_elapsed = time.perf_counter() - t_total - finally: - if profiler_started: - print("Stopping profiler (streaming)...") - await omni.stop_profile(stages=args.profiler_stages) - if args.profiler_wait_seconds > 0: - print(f"Waiting {args.profiler_wait_seconds:.1f}s for profiler traces to flush...") - await asyncio.sleep(args.profiler_wait_seconds) - - print( - f"All streaming runs finished: {args.num_runs} run(s), " - f"{len(prompt_specs)} prompt(s), {len(paths)} file(s) in {total_elapsed:.2f}s total" - ) - return paths - - -def _run_sync(args) -> list[Path]: - output_dir = Path(args.output_dir) - - omni = Omni( - model=args.model, - stage_configs_path=args.stage_configs_path, - log_stats=args.log_stats, - stage_init_timeout=args.stage_init_timeout, - ) - - def _run_sync_single( - spec: PromptSpec, - *, - request_prefix: str, - save_outputs: bool, - run_index: int | None = None, - ) -> tuple[list[Path], int, float | None, float, float, str]: - global_request_id = f"{request_prefix}_{spec.label}" - prompt = _build_prompt_for_spec(args, spec, global_request_id=global_request_id) - if save_outputs and run_index == 0 and spec.label == "item001": - print(f"---prompt---:{prompt}") - - saved_paths: list[Path] = [] - output_count = 0 - first_audio_elapsed: float | None = None - total_audio_duration_s = 0.0 - metrics_request_id = global_request_id - t_start = time.perf_counter() - for stage_outputs in omni.generate(prompt): - request_output = stage_outputs.request_output - if request_output is None: - continue - request_output_id = getattr(request_output, "request_id", None) - if isinstance(request_output_id, str) and request_output_id: - metrics_request_id = request_output_id - for j, mm in enumerate(_iter_request_multimodal_outputs(request_output)): - output_count += 1 - if first_audio_elapsed is None: - try: - audio_tensor = _extract_audio_tensor(mm) - if int(audio_tensor.numel()) > 0: - first_audio_elapsed = time.perf_counter() - t_start - total_audio_duration_s += float(audio_tensor.numel()) / float(_extract_sample_rate(mm)) - except ValueError: - pass - else: - try: - audio_tensor = _extract_audio_tensor(mm) - total_audio_duration_s += float(audio_tensor.numel()) / float(_extract_sample_rate(mm)) - except ValueError: - pass - if not save_outputs: - continue - save_stem = f"run{run_index + 1}_{spec.label}" if j == 0 else f"run{run_index + 1}_{spec.label}_{j}" - saved_paths.append(_save_wav(mm, output_dir, save_stem)) - - if output_count == 0: - raise RuntimeError("No output from Omni.generate") - elapsed_s = time.perf_counter() - t_start - return saved_paths, output_count, first_audio_elapsed, elapsed_s, total_audio_duration_s, metrics_request_id - - if args.warmup_runs: - warmup_specs = _get_warmup_specs(args.prompt_specs) - print( - f"Warmup: {args.warmup_runs} run(s) using the first prompt " - f"({len(warmup_specs)} prompt(s)); outputs will be discarded." - ) - for warmup_index in range(args.warmup_runs): - t_warmup = time.perf_counter() - _, output_count, first_audio_elapsed, elapsed_s, audio_duration_s, _ = _run_sync_single( - warmup_specs[0], - request_prefix=f"warmup_sync{warmup_index + 1}", - save_outputs=False, - ) - ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else "" - rtf_text = f", rtf={elapsed_s / audio_duration_s:.3f}" if audio_duration_s > 0 else "" - print( - f"Warmup (sync) {warmup_index + 1}/{args.warmup_runs} finished: " - f"{output_count} output(s) ({time.perf_counter() - t_warmup:.2f}s{ttfp_text}{rtf_text})" - ) - - profiler_started = False - if args.enable_profiler: - profile_prefix = f"voxcpm_sync_{int(time.time())}" - stages_text = args.profiler_stages if args.profiler_stages is not None else "all-configured" - print(f"Starting profiler (sync): stages={stages_text}, dir={args.profiler_dir}") - omni.start_profile(profile_prefix=profile_prefix, stages=args.profiler_stages) - profiler_started = True - - t_total = time.perf_counter() - total_elapsed = 0.0 - saved_paths: list[Path] = [] - prompt_specs: list[PromptSpec] = args.prompt_specs - try: - for run in range(args.num_runs): - t_run = time.perf_counter() - run_paths: list[Path] = [] - for prompt_index, spec in enumerate(prompt_specs): - prompt_paths, _, first_audio_elapsed, elapsed_s, audio_duration_s, metrics_request_id = ( - _run_sync_single( - spec, - request_prefix=f"sync_run{run + 1}_{prompt_index + 1:03d}", - save_outputs=True, - run_index=run, - ) - ) - run_paths.extend(prompt_paths) - ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else "" - rtf_text = f", rtf={elapsed_s / audio_duration_s:.3f}" if audio_duration_s > 0 else "" - print( - f"Saved (sync) run {run + 1}/{args.num_runs}, " - f"prompt {prompt_index + 1}/{len(prompt_specs)}: {len(prompt_paths)} file(s){ttfp_text}{rtf_text}" - ) - _emit_offline_metrics( - request_id=metrics_request_id, - elapsed_s=elapsed_s, - first_audio_elapsed=first_audio_elapsed, - audio_duration_s=audio_duration_s, - ) - - saved_paths.extend(run_paths) - print( - f"Run {run + 1}/{args.num_runs} finished: {len(run_paths)} file(s) ({time.perf_counter() - t_run:.2f}s)" - ) - for path in run_paths: - print(f" {path}") - - total_elapsed = time.perf_counter() - t_total - finally: - if profiler_started: - print("Stopping profiler (sync)...") - omni.stop_profile(stages=args.profiler_stages) - if args.profiler_wait_seconds > 0: - print(f"Waiting {args.profiler_wait_seconds:.1f}s for profiler traces to flush...") - time.sleep(args.profiler_wait_seconds) - - print( - f"All sync runs finished: {args.num_runs} run(s), " - f"{len(prompt_specs)} prompt(s), {len(saved_paths)} file(s) in {total_elapsed:.2f}s total" - ) - return saved_paths - - -def main(args) -> int: - logging.basicConfig(level=logging.INFO) - profiled_stage_config_path: str | None = None - original_stage_config_path = args.stage_configs_path - if args.enable_profiler: - Path(args.profiler_dir).mkdir(parents=True, exist_ok=True) - profiled_stage_config_path = _build_profiled_stage_config( - args.stage_configs_path, - str(Path(args.profiler_dir).resolve()), - ) - args.stage_configs_path = profiled_stage_config_path - - is_streaming = _is_streaming_stage_config(args.stage_configs_path) - voice_clone_count = _count_voice_clone_prompts(args.prompt_specs) - print(f"Model: {args.model}") - print(f"Stage config: {original_stage_config_path}") - print(f"Route: {'streaming' if is_streaming else 'sync'} (from stage-configs-path)") - print(f"Prompt count: {len(args.prompt_specs)}") - print("Batch mode: sequential (aligned with native VoxCPM)") - print(f"Warmup runs: {args.warmup_runs}") - print(f"Voice cloning prompts: {voice_clone_count}/{len(args.prompt_specs)}") - if args.enable_profiler: - print(f"Profiler: enabled (dir={args.profiler_dir}, stages={args.profiler_stages or 'all-configured'})") - print(f"Profiled stage config: {args.stage_configs_path}") - if voice_clone_count: - print("Voice cloning note: --ref-text/ref_text must match the spoken content of the reference audio.") - print(f"Num runs: {args.num_runs}") - try: - if is_streaming: - asyncio.run(_run_streaming(args)) - else: - _run_sync(args) - finally: - if profiled_stage_config_path is not None and os.path.exists(profiled_stage_config_path): - os.unlink(profiled_stage_config_path) - return 0 - - -if __name__ == "__main__": - os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - raise SystemExit(main(parse_args())) diff --git a/docs/api/README.md b/docs/api/README.md index 11401dc02..dce02e13d 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -121,7 +121,6 @@ Configuration classes. - [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLTextConfig][] - [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLVisionConfig][] - [vllm_omni.transformers_utils.configs.omnivoice.OmniVoiceConfig][] -- [vllm_omni.transformers_utils.configs.voxcpm.VoxCPMConfig][] - [vllm_omni.transformers_utils.configs.voxcpm2.VoxCPM2Config][] ## Workers diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index 60d6df718..23f89bad5 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -5,6 +5,9 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation. - **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Output: 24 kHz. - **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz. - **Voxtral TTS** (`mistralai/Voxtral-4B-TTS-2603`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz. +- **CosyVoice3** (`FunAudioLLM/Fun-CosyVoice3-0.5B-2512`) -- 2-stage talker + flow-matching code2wav. Voice cloning via `ref_audio` + `ref_text` (no presets). Output: 24 kHz. + +See the [Supported Models](#supported-models) section below for the full list, including OmniVoice, VoxCPM2, and MOSS-TTS-Nano. Each server instance runs a single model (specified at startup via `vllm serve --omni`). @@ -26,6 +29,10 @@ vllm serve fishaudio/s2-pro --omni --port 8091 # Voxtral TTS vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091 + +# CosyVoice3 (voice cloning only — supply ref_audio + ref_text per request) +vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 \ + --omni --port 8091 --trust-remote-code ``` ### Generate Speech diff --git a/docs/user_guide/examples/offline_inference/text_to_speech.md b/docs/user_guide/examples/offline_inference/text_to_speech.md index e88a68220..7ad3d2d40 100644 --- a/docs/user_guide/examples/offline_inference/text_to_speech.md +++ b/docs/user_guide/examples/offline_inference/text_to_speech.md @@ -17,7 +17,6 @@ list of supported architectures across all modalities, see | Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate | |---|---|---|---|---|---|---| | VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | — | continuation (`--ref-audio` + `--ref-text`) | 48 kHz | -| VoxCPM | local model dir | split | ✓ | ✓ (`voxcpm_async_chunk.yaml`) | — | 24 kHz | | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ (`async_chunk: true` default) | — | 22.05 kHz | | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ (`--streaming`) | — | 44.1 kHz | | OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design (`--instruct`), language (`--lang`) | 24 kHz | @@ -72,61 +71,6 @@ python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \ --- -## VoxCPM - -Split-stage TTS. The hub example covers single TTS, single voice cloning, and streaming. Use `benchmarks/voxcpm/` for warmup, batch JSONL prompts, profiler injection, and offline TTFP / RTF measurement. - -### Prerequisites -```bash -pip install voxcpm soundfile -# or use a local source tree: -export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src -``` - -If the native VoxCPM `config.json` does not contain HF metadata such as `model_type`, prepare a persistent HF-compatible config directory and point the stage configs to it via `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH`: - -```bash -export VOXCPM_MODEL=/path/to/voxcpm-model -export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config -mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH" -cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json" -cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true -python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)' -``` - -### Quick start -```bash -python examples/offline_inference/text_to_speech/voxcpm/end2end.py \ - --model "$VOXCPM_MODEL" \ - --text "This is a split-stage VoxCPM synthesis example running on vLLM Omni." -``` - -### Voice cloning -```bash -python examples/offline_inference/text_to_speech/voxcpm/end2end.py \ - --model "$VOXCPM_MODEL" \ - --text "This sentence is synthesized with a cloned voice." \ - --ref-audio /path/to/reference.wav \ - --ref-text "The exact transcript spoken in reference.wav." -``` - -### Streaming -Pass the async-chunk stage config: -```bash -python examples/offline_inference/text_to_speech/voxcpm/end2end.py \ - --model "$VOXCPM_MODEL" \ - --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \ - --text "This is a split-stage VoxCPM streaming example running on vLLM Omni." -``` - -### Notes -- `voxcpm.yaml` is the default non-streaming stage config; `voxcpm_async_chunk.yaml` enables streaming. -- Streaming is currently single-request oriented. -- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades quality. -- For online serving, see [`examples/online_serving/voxcpm`](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/voxcpm/README.md). For benchmark reporting, see [`benchmarks/voxcpm`](https://github.com/vllm-project/vllm-omni/tree/main/benchmarks/voxcpm/README.md). - ---- - ## CosyVoice3 2-stage TTS pipeline (`talker` + `code2wav`) at 22.05 kHz. @@ -381,10 +325,6 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T ``````py --8<-- "examples/offline_inference/text_to_speech/qwen3_tts/end2end.py" `````` -??? abstract "voxcpm/end2end.py" - ``````py - --8<-- "examples/offline_inference/text_to_speech/voxcpm/end2end.py" - `````` ??? abstract "voxcpm2/end2end.py" ``````py --8<-- "examples/offline_inference/text_to_speech/voxcpm2/end2end.py" diff --git a/docs/user_guide/examples/online_serving/text_to_speech.md b/docs/user_guide/examples/online_serving/text_to_speech.md index 703b7cf7a..ff22dd0fd 100644 --- a/docs/user_guide/examples/online_serving/text_to_speech.md +++ b/docs/user_guide/examples/online_serving/text_to_speech.md @@ -17,14 +17,14 @@ For the full list of supported architectures across all modalities, see | Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo | |---|---|---|---|---|---| +| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | — | | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ | | OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — | | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) | -| VoxCPM | local model dir | ✓ | ✓ (PCM stream) | — | — | | VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ | | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ | -CosyVoice3 is intentionally absent: no online example exists for it yet. See its [offline section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3) instead. +For offline inference of any of these models, see the [offline TTS hub](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md). ## Common Quick Start @@ -96,6 +96,54 @@ For full request-shape documentation (all parameters, response formats, error co --- +## CosyVoice3 + +2-stage TTS (`talker` + flow-matching `code2wav`) at 24 kHz. Voice cloning only — every request needs `ref_audio` + `ref_text`; there are no built-in voice presets. + +### Prerequisites +```bash +huggingface-cli download FunAudioLLM/Fun-CosyVoice3-0.5B-2512 +``` + +If your downloaded checkpoint lacks `config.json`, add one with `{"model_type": "cosyvoice3", "architectures": ["CosyVoice3Model"]}` (the loader reads `model_type` to select the class). + +### Launch +```bash +vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remote-code +# or: +./cosyvoice3/run_server.sh +``` + +Streaming is on by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path. + +### CLI client +The client defaults to the official upstream zero-shot prompt, so it runs without extra flags: +```bash +cd examples/online_serving/text_to_speech/cosyvoice3 +python speech_client.py --text "收到好友从远方寄来的生日礼物。" +``` + +Pass your own reference clip and transcript for a different voice: +```bash +python speech_client.py --text "Hello, this is a cloned voice." \ + --ref-audio /path/to/reference.wav \ + --ref-text "Transcript of the reference audio." +``` + +Stream PCM instead of WAV: +```bash +python speech_client.py --text "Hello world" --stream --output output.pcm +``` + +The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-text`, `--response-format`, `--stream`, `--output`. + +### Notes +- Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform. +- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config ` to override. +- For offline inference and the end-to-end script, see the [offline CosyVoice3 section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3). + +--- + ## Fish Speech S2 Pro 4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec. @@ -196,10 +244,10 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091 ./qwen3_tts/run_server.sh Base ``` -### Choosing an executor backend (uniproc vs mp) -Stage configs ship with the chunked-streaming default. To use the uniproc executor (lower IPC overhead for the Base cloning task), pass `--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_uniproc.yaml`. See [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) and [#2604](https://github.com/vllm-project/vllm-omni/pull/2604) for the full investigation. +### Executor backend +Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config ` to override. -To opt out of chunked streaming, pass `--no-async-chunk` instead — the pipeline auto-dispatches to the end-to-end codec processor. +To opt out of chunked streaming, pass `--no-async-chunk` — the pipeline auto-dispatches to the end-to-end codec processor. ### Sending requests ```bash @@ -281,71 +329,11 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000 `qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement. ### Notes -- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above. +- Base voice cloning has per-request reference-audio cost; the uniproc default keeps IPC overhead off the critical path. See the executor-backend section above for background. - `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N- `. --- -## VoxCPM - -Split-stage TTS at 24 kHz. - -### Prerequisites -```bash -pip install voxcpm -# or use a local source tree: -export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src -``` - -If the native VoxCPM `config.json` lacks HF `model_type`, set up an HF-compatible config dir: -```bash -export VOXCPM_MODEL=/path/to/voxcpm-model -export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config -mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH" -cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json" -cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true -python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)' -``` - -### Launch -```bash -export VOXCPM_MODEL=/path/to/voxcpm-model -./voxcpm/run_server.sh # async-chunk streaming (default) -./voxcpm/run_server.sh sync # non-streaming -``` -Or directly: -```bash -vllm serve "$VOXCPM_MODEL" \ - --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \ - --trust-remote-code --enforce-eager --omni --port 8091 -``` - -### Sending requests -```bash -# Basic TTS -python voxcpm/openai_speech_client.py \ - --model "$VOXCPM_MODEL" \ - --text "This is a VoxCPM online text-to-speech example." - -# Voice cloning -python voxcpm/openai_speech_client.py \ - --model "$VOXCPM_MODEL" \ - --text "This sentence is synthesized with a cloned voice." \ - --ref-audio /path/to/reference.wav \ - --ref-text "The exact transcript spoken in reference.wav." - -# Streaming PCM -python voxcpm/openai_speech_client.py \ - --model "$VOXCPM_MODEL" \ - --text "This is a streaming VoxCPM request." \ - --stream --output voxcpm_stream.pcm -``` - -### Notes -- `voxcpm.yaml` for one-shot decode; `voxcpm_async_chunk.yaml` for single-request streaming. Do not use the async-chunk config for concurrent requests or `/v1/audio/speech/batch`. -- Generic TTS fields not supported by VoxCPM: `voice`, `instructions`, `language`, `speaker_embedding`, `x_vector_only_mode`. -- For benchmark measurement, see [`benchmarks/voxcpm`](https://github.com/vllm-project/vllm-omni/tree/main/benchmarks/voxcpm/README.md). - --- ## VoxCPM2 @@ -407,6 +395,14 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts ## Example materials +??? abstract "cosyvoice3/run_server.sh" + + --8<-- "examples/online_serving/text_to_speech/cosyvoice3/run_server.sh" + +??? abstract "cosyvoice3/speech_client.py" + + --8<-- "examples/online_serving/text_to_speech/cosyvoice3/speech_client.py" + ??? abstract "fish_speech/gradio_demo.py" ``````py --8<-- "examples/online_serving/text_to_speech/fish_speech/gradio_demo.py" @@ -463,14 +459,6 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts ``````py --8<-- "examples/online_serving/text_to_speech/qwen3_tts/tts_common.py" `````` -??? abstract "voxcpm/openai_speech_client.py" - ``````py - --8<-- "examples/online_serving/text_to_speech/voxcpm/openai_speech_client.py" - `````` -??? abstract "voxcpm/run_server.sh" - ``````sh - --8<-- "examples/online_serving/text_to_speech/voxcpm/run_server.sh" - `````` ??? abstract "voxcpm2/gradio_demo.py" ``````py --8<-- "examples/online_serving/text_to_speech/voxcpm2/gradio_demo.py" diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index bfdee543c..722b9eeb3 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -14,13 +14,12 @@ list of supported architectures across all modalities, see | Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate | |---|---|---|---|---|---|---| -| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 22.05 kHz | +| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz | | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz | | Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz | | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | single (AR + codec) | ✓ (required) | ✓ | voice_clone, continuation | 48 kHz | | OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz | | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants | 24 kHz | -| VoxCPM | `openbmb/VoxCPM-0.5B` | split | ✓ | ✓ | — | 24 kHz | | VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | ✓ (online) | continuation | 48 kHz | | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz | @@ -41,7 +40,7 @@ python examples/offline_inference/text_to_speech//end2end.py \ ## CosyVoice3 -2-stage TTS pipeline (`talker` + `code2wav`) at 22.05 kHz. +2-stage TTS pipeline (`talker` + `code2wav`) at 24 kHz. ### Prerequisites ```bash @@ -308,56 +307,6 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \ --- -## VoxCPM - -Split-stage TTS. The hub example covers single TTS, single voice cloning, and streaming. Use `benchmarks/voxcpm/` for warmup, batch JSONL prompts, profiler injection, and offline TTFP / RTF measurement. - -### Prerequisites -```bash -pip install voxcpm soundfile -# or use a local source tree: -export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src -``` - -### Quick start -```bash -python examples/offline_inference/text_to_speech/voxcpm/end2end.py \ - --model openbmb/VoxCPM-0.5B \ - --text "This is a split-stage VoxCPM synthesis example running on vLLM Omni." -``` - -### Voice cloning -```bash -python examples/offline_inference/text_to_speech/voxcpm/end2end.py \ - --model openbmb/VoxCPM-0.5B \ - --text "This sentence is synthesized with a cloned voice." \ - --ref-audio /path/to/reference.wav \ - --ref-text "The exact transcript spoken in reference.wav." -``` - -### Streaming -Pass `--streaming` together with the legacy async-chunk stage config: -```bash -python examples/offline_inference/text_to_speech/voxcpm/end2end.py \ - --model openbmb/VoxCPM-0.5B \ - --streaming \ - --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \ - --text "This is a split-stage VoxCPM streaming example running on vLLM Omni." -``` - -### Persistent HF config (optional) -The engine auto-patches a missing `model_type` in the checkpoint's `config.json` in-memory at startup (see `vllm_omni/engine/arg_utils.py`), so no manual setup is required. - -If you'd rather maintain a persistent patched config dir (e.g. to share across runs or to add fields the auto-patcher doesn't set), export `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH` to its location — both the bundled deploy and async-chunk yamls read it via `${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,}`. - -### Notes -- Non-streaming auto-loads `vllm_omni/deploy/voxcpm.yaml` from the model directory name / HF `model_type`. Streaming uses the legacy `voxcpm_async_chunk.yaml` passed via `--stage-configs-path`. -- Streaming is currently single-request oriented. -- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades quality. -- For online serving, see the [VoxCPM section in the online hub](../../online_serving/text_to_speech/README.md#voxcpm). For benchmark reporting, see [`benchmarks/voxcpm`](../../../benchmarks/voxcpm/README.md). - ---- - ## VoxCPM2 Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → FSQ → residual_lm → LocDiT → AudioVAE`. diff --git a/examples/offline_inference/text_to_speech/voxcpm/end2end.py b/examples/offline_inference/text_to_speech/voxcpm/end2end.py deleted file mode 100644 index 79efa3123..000000000 --- a/examples/offline_inference/text_to_speech/voxcpm/end2end.py +++ /dev/null @@ -1,220 +0,0 @@ -"""Minimal offline VoxCPM example for vLLM Omni.""" - -from __future__ import annotations - -import asyncio -import time -from pathlib import Path -from typing import Any - -import soundfile as sf -import torch -from vllm.utils.argparse_utils import FlexibleArgumentParser - -from vllm_omni import AsyncOmni, Omni -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults - - -def _build_prompt(args) -> dict[str, Any]: - additional_information: dict[str, list[Any]] = { - "text": [args.text], - "cfg_value": [args.cfg_value], - "inference_timesteps": [args.inference_timesteps], - "min_len": [args.min_len], - "max_new_tokens": [args.max_new_tokens], - } - if args.streaming_prefix_len is not None: - additional_information["streaming_prefix_len"] = [args.streaming_prefix_len] - if args.ref_audio is not None: - additional_information["ref_audio"] = [args.ref_audio] - if args.ref_text is not None: - additional_information["ref_text"] = [args.ref_text] - return { - "prompt_token_ids": [1], - "additional_information": additional_information, - } - - -def _extract_audio_tensor(mm: dict[str, Any]) -> torch.Tensor: - audio = mm.get("audio", mm.get("model_outputs")) - if audio is None: - raise ValueError("No audio output found in multimodal output.") - if isinstance(audio, list): - parts = [torch.as_tensor(item).float().cpu().reshape(-1) for item in audio] - audio = torch.cat(parts, dim=-1) if parts else torch.zeros(0) - if not isinstance(audio, torch.Tensor): - audio = torch.as_tensor(audio) - return audio.float().cpu().reshape(-1) - - -def _extract_sample_rate(mm: dict[str, Any]) -> int: - sr_raw = mm.get("sr", 24000) - if isinstance(sr_raw, list) and sr_raw: - sr_raw = sr_raw[-1] - if hasattr(sr_raw, "item"): - return int(sr_raw.item()) - return int(sr_raw) - - -def _save_audio(audio: torch.Tensor, sample_rate: int, output_dir: Path, request_id: str) -> Path: - output_dir.mkdir(parents=True, exist_ok=True) - output_path = output_dir / f"output_{request_id}.wav" - sf.write( - output_path, - audio.float().cpu().clamp(-1.0, 1.0).numpy(), - sample_rate, - format="WAV", - subtype="PCM_16", - ) - return output_path - - -async def _run_streaming(args) -> Path: - prompt = _build_prompt(args) - output_dir = Path(args.output_dir) if args.output_dir is not None else Path("output_audio_streaming") - request_id = "streaming_example" - sample_rate = 24000 - buffered_samples = 0 - chunks: list[torch.Tensor] = [] - started = time.perf_counter() - omni = AsyncOmni( - model=args.model, - stage_configs_path=args.stage_configs_path, - deploy_config=args.deploy_config, - log_stats=args.log_stats, - stage_init_timeout=args.stage_init_timeout, - ) - try: - async for stage_output in omni.generate(prompt, request_id=request_id): - mm = getattr(stage_output, "multimodal_output", None) - if not isinstance(mm, dict): - request_output = getattr(stage_output, "request_output", None) - if request_output is None: - continue - mm = getattr(request_output, "multimodal_output", None) - if not isinstance(mm, dict) and getattr(request_output, "outputs", None): - mm = getattr(request_output.outputs[0], "multimodal_output", None) - if not isinstance(mm, dict): - continue - audio = _extract_audio_tensor(mm) - if audio.numel() == 0: - continue - sample_rate = _extract_sample_rate(mm) - if audio.numel() > buffered_samples: - delta = audio[buffered_samples:] - buffered_samples = int(audio.numel()) - else: - delta = audio - buffered_samples += int(delta.numel()) - if delta.numel() > 0: - chunks.append(delta) - if not chunks: - raise RuntimeError("No streaming audio chunks received from VoxCPM.") - output_audio = torch.cat(chunks, dim=0) - output_path = _save_audio(output_audio, sample_rate, output_dir, request_id) - print(f"Saved streaming audio to: {output_path} ({time.perf_counter() - started:.2f}s)") - return output_path - finally: - omni.shutdown() - - -def _run_sync(args) -> Path: - prompt = _build_prompt(args) - output_dir = Path(args.output_dir) if args.output_dir is not None else Path("output_audio") - request_id = "sync_example" - started = time.perf_counter() - last_mm: dict[str, Any] | None = None - omni = Omni( - model=args.model, - stage_configs_path=args.stage_configs_path, - deploy_config=args.deploy_config, - log_stats=args.log_stats, - stage_init_timeout=args.stage_init_timeout, - ) - for stage_outputs in omni.generate(prompt): - request_output = getattr(stage_outputs, "request_output", None) - if request_output is None: - continue - outputs = getattr(request_output, "outputs", None) - if outputs: - for output in outputs: - mm = getattr(output, "multimodal_output", None) - if isinstance(mm, dict): - last_mm = mm - mm = getattr(request_output, "multimodal_output", None) - if isinstance(mm, dict): - last_mm = mm - if last_mm is None: - raise RuntimeError("No audio output received from VoxCPM.") - output_path = _save_audio( - _extract_audio_tensor(last_mm), - _extract_sample_rate(last_mm), - output_dir, - request_id, - ) - print(f"Saved audio to: {output_path} ({time.perf_counter() - started:.2f}s)") - return output_path - - -def parse_args(): - parser = FlexibleArgumentParser(description="Minimal offline VoxCPM example for vLLM Omni.") - parser.add_argument("--model", type=str, required=True, help="Local VoxCPM model directory.") - parser.add_argument( - "--deploy-config", - type=str, - default=None, - help=("Override the deploy config path."), - ) - parser.add_argument( - "--stage-configs-path", - type=str, - default=None, - help=( - "Legacy stage_args yaml path. Required for streaming " - "(vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml); " - "leave unset for non-streaming to use the auto-resolved deploy config." - ), - ) - parser.add_argument( - "--streaming", - action="store_true", - default=False, - help="Stream audio chunks as they arrive via AsyncOmni.", - ) - parser.add_argument("--text", type=str, required=True, help="Input text for synthesis.") - parser.add_argument("--ref-audio", type=str, default=None, help="Reference audio path for voice cloning.") - parser.add_argument("--ref-text", type=str, default=None, help="Transcript of the reference audio.") - parser.add_argument("--output-dir", type=str, default=None, help="Output directory for generated wav files.") - parser.add_argument("--cfg-value", type=float, default=2.0, help="Guidance value passed to VoxCPM.") - parser.add_argument("--inference-timesteps", type=int, default=10, help="Number of diffusion timesteps.") - parser.add_argument("--min-len", type=int, default=2, help="Minimum latent length.") - parser.add_argument("--max-new-tokens", type=int, default=4096, help="Maximum latent length.") - parser.add_argument( - "--streaming-prefix-len", - type=int, - default=3, - help="Streaming prefix length used by voxcpm_async_chunk.yaml.", - ) - parser.add_argument("--stage-init-timeout", type=int, default=600, help="Stage initialization timeout in seconds.") - parser.add_argument("--log-stats", action="store_true", help="Enable vLLM Omni stats logging.") - nullify_stage_engine_defaults(parser) - args = parser.parse_args() - if (args.ref_audio is None) != (args.ref_text is None): - raise ValueError("Voice cloning requires --ref-audio and --ref-text together.") - return args - - -def main(args) -> None: - route = "streaming" if args.streaming else "sync" - print(f"Model: {args.model}") - print(f"Deploy config: {args.deploy_config or ''}") - print(f"Stage configs path: {args.stage_configs_path or ''}") - print(f"Route: {route}") - if route == "streaming": - asyncio.run(_run_streaming(args)) - else: - _run_sync(args) - - -if __name__ == "__main__": - main(parse_args()) diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index 8481e723e..75b85580d 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -19,7 +19,6 @@ For the full list of supported architectures across all modalities, see | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✓ (`ref_audio` required) | ✓ (PCM stream) | — | ✓ | | OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — | | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) | -| VoxCPM | `openbmb/VoxCPM-0.5B` | ✓ | ✓ (PCM stream) | — | — | | VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ | | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ | @@ -294,10 +293,10 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091 ./qwen3_tts/run_server.sh Base ``` -### Choosing an executor backend (uniproc vs mp) -Stage configs ship with the chunked-streaming default. To use the uniproc executor (lower IPC overhead for the Base cloning task), pass `--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_uniproc.yaml`. See [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) and [#2604](https://github.com/vllm-project/vllm-omni/pull/2604) for the full investigation. +### Executor backend +Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config ` to override. -To opt out of chunked streaming, pass `--no-async-chunk` instead — the pipeline auto-dispatches to the end-to-end codec processor. +To opt out of chunked streaming, pass `--no-async-chunk` — the pipeline auto-dispatches to the end-to-end codec processor. ### Sending requests ```bash @@ -375,73 +374,11 @@ python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-de `qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement. ### Notes -- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above. +- Base voice cloning has per-request reference-audio cost; the uniproc default keeps IPC overhead off the critical path. See the executor-backend section above for background. - `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N- `. --- -## VoxCPM - -Split-stage TTS at 24 kHz. - -### Prerequisites -```bash -pip install voxcpm -# or use a local source tree: -export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src -``` - -If the native VoxCPM `config.json` lacks HF `model_type`, set up an HF-compatible config dir: -```bash -export VOXCPM_MODEL=/path/to/voxcpm-model -export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config -mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH" -cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json" -cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true -python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)' -``` - -### Launch -```bash -export VOXCPM_MODEL=/path/to/voxcpm-model -./voxcpm/run_server.sh # async-chunk streaming (default) -./voxcpm/run_server.sh sync # non-streaming -``` -Or directly: -```bash -vllm serve "$VOXCPM_MODEL" \ - --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \ - --trust-remote-code --enforce-eager --omni --port 8091 -``` - -### Sending requests -```bash -# Basic TTS -python voxcpm/openai_speech_client.py \ - --model "$VOXCPM_MODEL" \ - --text "This is a VoxCPM online text-to-speech example." - -# Voice cloning -python voxcpm/openai_speech_client.py \ - --model "$VOXCPM_MODEL" \ - --text "This sentence is synthesized with a cloned voice." \ - --ref-audio /path/to/reference.wav \ - --ref-text "The exact transcript spoken in reference.wav." - -# Streaming PCM -python voxcpm/openai_speech_client.py \ - --model "$VOXCPM_MODEL" \ - --text "This is a streaming VoxCPM request." \ - --stream --output voxcpm_stream.pcm -``` - -### Notes -- `voxcpm.yaml` for one-shot decode; `voxcpm_async_chunk.yaml` for single-request streaming. Do not use the async-chunk config for concurrent requests or `/v1/audio/speech/batch`. -- Generic TTS fields not supported by VoxCPM: `voice`, `instructions`, `language`, `speaker_embedding`, `x_vector_only_mode`. -- For benchmark measurement, see [`benchmarks/voxcpm`](../../../benchmarks/voxcpm/README.md). - ---- - ## VoxCPM2 Single-stage native AR TTS at 48 kHz. diff --git a/examples/online_serving/text_to_speech/cosyvoice3/run_server.sh b/examples/online_serving/text_to_speech/cosyvoice3/run_server.sh new file mode 100755 index 000000000..e40b1bb13 --- /dev/null +++ b/examples/online_serving/text_to_speech/cosyvoice3/run_server.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# Launch vLLM-Omni server for CosyVoice3 TTS +# +# Usage: +# ./run_server.sh +# CUDA_VISIBLE_DEVICES=0 ./run_server.sh +# +# Streaming (async-chunk) is on by default via vllm_omni/deploy/cosyvoice3.yaml. +# Set NO_ASYNC_CHUNK=1 to use the legacy synchronous path. + +set -e + +MODEL="${MODEL:-FunAudioLLM/Fun-CosyVoice3-0.5B-2512}" +PORT="${PORT:-8091}" + +EXTRA_ARGS=() +if [[ -n "${NO_ASYNC_CHUNK:-}" ]]; then + EXTRA_ARGS+=(--no-async-chunk) +fi + +echo "Starting CosyVoice3 server with model: $MODEL" + +vllm serve "$MODEL" \ + --host 0.0.0.0 \ + --port "$PORT" \ + --trust-remote-code \ + --omni \ + "${EXTRA_ARGS[@]}" diff --git a/examples/online_serving/text_to_speech/cosyvoice3/speech_client.py b/examples/online_serving/text_to_speech/cosyvoice3/speech_client.py new file mode 100644 index 000000000..a837d9243 --- /dev/null +++ b/examples/online_serving/text_to_speech/cosyvoice3/speech_client.py @@ -0,0 +1,142 @@ +"""Client for CosyVoice3 TTS via /v1/audio/speech endpoint. + +CosyVoice3 has no built-in voice presets: every request is voice cloning +driven by ``ref_audio`` + ``ref_text``. The defaults below point at the +official upstream zero-shot prompt so the script runs out of the box. + +Examples: + # Voice cloning with the default upstream prompt + python speech_client.py --text "收到好友从远方寄来的生日礼物。" + + # Custom reference clip + transcript + python speech_client.py --text "Hello, this is a cloned voice." \ + --ref-audio /path/to/reference.wav \ + --ref-text "Transcript of the reference audio." + + # Streaming PCM output + python speech_client.py --text "Hello world" --stream --output output.pcm +""" + +import argparse +import base64 +import os + +import httpx + +DEFAULT_API_BASE = "http://localhost:8091" +DEFAULT_API_KEY = "EMPTY" +DEFAULT_MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512" + +# Official CosyVoice zero-shot prompt and its transcript. +DEFAULT_REF_AUDIO = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav" +DEFAULT_REF_TEXT = "希望你以后能够做的比我还好呦。" + + +def encode_audio_to_base64(audio_path: str) -> str: + """Encode a local audio file to a base64 data URL.""" + if not os.path.exists(audio_path): + raise FileNotFoundError(f"Audio file not found: {audio_path}") + ext = audio_path.lower().rsplit(".", 1)[-1] + mime_map = {"wav": "audio/wav", "mp3": "audio/mpeg", "flac": "audio/flac", "ogg": "audio/ogg"} + mime_type = mime_map.get(ext, "audio/wav") + with open(audio_path, "rb") as f: + audio_b64 = base64.b64encode(f.read()).decode("utf-8") + return f"data:{mime_type};base64,{audio_b64}" + + +def run_tts(args) -> None: + """Generate speech via the /v1/audio/speech API.""" + payload = { + "model": args.model, + "input": args.text, + "response_format": args.response_format, + } + + if args.ref_audio.startswith(("http://", "https://")): + payload["ref_audio"] = args.ref_audio + else: + payload["ref_audio"] = encode_audio_to_base64(args.ref_audio) + payload["ref_text"] = args.ref_text + + if args.stream: + payload["stream"] = True + payload["response_format"] = "pcm" + + print(f"Model: {args.model}") + print(f"Text: {args.text}") + print(f"Voice cloning: ref_audio={args.ref_audio}, ref_text={args.ref_text}") + print("Generating audio...") + + api_url = f"{args.api_base}/v1/audio/speech" + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {args.api_key}", + } + + if args.stream: + output_path = args.output or "output.pcm" + with httpx.Client(timeout=300.0) as client: + with client.stream("POST", api_url, json=payload, headers=headers) as resp: + if resp.status_code != 200: + print(f"Error: {resp.status_code}") + print(resp.read().decode()) + return + total_bytes = 0 + with open(output_path, "wb") as f: + for chunk in resp.iter_bytes(): + f.write(chunk) + total_bytes += len(chunk) + print(f"Streamed {total_bytes} bytes to: {output_path}") + else: + with httpx.Client(timeout=300.0) as client: + response = client.post(api_url, json=payload, headers=headers) + + if response.status_code != 200: + print(f"Error: {response.status_code}") + print(response.text) + return + + try: + text = response.content.decode("utf-8") + if text.startswith('{"error"'): + print(f"Error: {text}") + return + except UnicodeDecodeError: + pass + + output_path = args.output or "output.wav" + with open(output_path, "wb") as f: + f.write(response.content) + print(f"Audio saved to: {output_path}") + + +def main(): + parser = argparse.ArgumentParser(description="CosyVoice3 TTS client") + parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL") + parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key") + parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name") + parser.add_argument("--text", required=True, help="Text to synthesize") + parser.add_argument( + "--ref-audio", + default=DEFAULT_REF_AUDIO, + help="Reference audio for voice cloning (path or URL)", + ) + parser.add_argument( + "--ref-text", + default=DEFAULT_REF_TEXT, + help="Transcript of the reference audio", + ) + parser.add_argument("--stream", action="store_true", help="Enable streaming (PCM output)") + parser.add_argument( + "--response-format", + default="wav", + choices=["wav", "mp3", "flac", "pcm", "aac", "opus"], + help="Audio format (default: wav)", + ) + parser.add_argument("--output", "-o", default=None, help="Output file path") + args = parser.parse_args() + run_tts(args) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/text_to_speech/voxcpm/openai_speech_client.py b/examples/online_serving/text_to_speech/voxcpm/openai_speech_client.py deleted file mode 100644 index c400114e8..000000000 --- a/examples/online_serving/text_to_speech/voxcpm/openai_speech_client.py +++ /dev/null @@ -1,155 +0,0 @@ -"""OpenAI-compatible client for VoxCPM via /v1/audio/speech. - -Examples: - # Basic text-to-speech - python openai_speech_client.py --text "Hello from VoxCPM" - - # Voice cloning - python openai_speech_client.py \ - --text "This sentence uses the cloned voice." \ - --ref-audio /path/to/reference.wav \ - --ref-text "The exact transcript spoken in the reference audio." - - # Streaming PCM output - python openai_speech_client.py \ - --text "This is a streaming VoxCPM request." \ - --stream \ - --output output.pcm -""" - -import argparse -import base64 -import os - -import httpx - -DEFAULT_API_BASE = "http://localhost:8091" -DEFAULT_API_KEY = "EMPTY" -DEFAULT_MODEL = "OpenBMB/VoxCPM1.5" - - -def encode_audio_to_base64(audio_path: str) -> str: - """Encode a local audio file to base64 data URL.""" - if not os.path.exists(audio_path): - raise FileNotFoundError(f"Audio file not found: {audio_path}") - - ext = audio_path.lower().rsplit(".", 1)[-1] - mime_map = { - "wav": "audio/wav", - "mp3": "audio/mpeg", - "flac": "audio/flac", - "ogg": "audio/ogg", - } - mime_type = mime_map.get(ext, "audio/wav") - - with open(audio_path, "rb") as f: - audio_b64 = base64.b64encode(f.read()).decode("utf-8") - return f"data:{mime_type};base64,{audio_b64}" - - -def build_payload(args) -> dict[str, object]: - payload: dict[str, object] = { - "model": args.model, - "input": args.text, - "response_format": "pcm" if args.stream else args.response_format, - } - - if args.ref_audio: - if args.ref_audio.startswith(("http://", "https://", "data:")): - payload["ref_audio"] = args.ref_audio - else: - payload["ref_audio"] = encode_audio_to_base64(args.ref_audio) - if args.ref_text: - payload["ref_text"] = args.ref_text - if args.max_new_tokens is not None: - payload["max_new_tokens"] = args.max_new_tokens - if args.stream: - payload["stream"] = True - - return payload - - -def run_tts(args) -> None: - payload = build_payload(args) - api_url = f"{args.api_base}/v1/audio/speech" - headers = { - "Content-Type": "application/json", - "Authorization": f"Bearer {args.api_key}", - } - - print(f"Model: {args.model}") - print(f"Text: {args.text}") - if args.ref_audio: - print("Mode: voice cloning") - print(f"Reference audio: {args.ref_audio}") - else: - print("Mode: text-to-speech") - - if args.stream: - output_path = args.output or "voxcpm_output.pcm" - with httpx.Client(timeout=300.0) as client: - with client.stream("POST", api_url, json=payload, headers=headers) as response: - if response.status_code != 200: - print(f"Error: {response.status_code}") - print(response.read().decode("utf-8", errors="ignore")) - return - - total_bytes = 0 - with open(output_path, "wb") as f: - for chunk in response.iter_bytes(): - if not chunk: - continue - f.write(chunk) - total_bytes += len(chunk) - print(f"Streamed {total_bytes} bytes to: {output_path}") - return - - with httpx.Client(timeout=300.0) as client: - response = client.post(api_url, json=payload, headers=headers) - - if response.status_code != 200: - print(f"Error: {response.status_code}") - print(response.text) - return - - try: - text = response.content.decode("utf-8") - if text.startswith('{"error"'): - print(f"Error: {text}") - return - except UnicodeDecodeError: - pass - - output_path = args.output or "voxcpm_output.wav" - with open(output_path, "wb") as f: - f.write(response.content) - print(f"Audio saved to: {output_path}") - - -def main(): - parser = argparse.ArgumentParser(description="VoxCPM OpenAI-compatible speech client") - parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL") - parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key") - parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name or path") - parser.add_argument("--text", required=True, help="Text to synthesize") - parser.add_argument("--ref-audio", default=None, help="Reference audio path, URL, or data URL") - parser.add_argument( - "--ref-text", - default=None, - help="The exact transcript spoken in the reference audio", - ) - parser.add_argument("--stream", action="store_true", help="Enable streaming PCM output") - parser.add_argument( - "--response-format", - default="wav", - choices=["wav", "pcm", "flac", "mp3", "aac", "opus"], - help="Audio format for non-streaming mode (default: wav)", - ) - parser.add_argument("--max-new-tokens", type=int, default=None, help="Maximum tokens to generate") - parser.add_argument("--output", "-o", default=None, help="Output file path") - args = parser.parse_args() - run_tts(args) - - -if __name__ == "__main__": - main() diff --git a/examples/online_serving/text_to_speech/voxcpm/run_server.sh b/examples/online_serving/text_to_speech/voxcpm/run_server.sh deleted file mode 100755 index c5ac08b9e..000000000 --- a/examples/online_serving/text_to_speech/voxcpm/run_server.sh +++ /dev/null @@ -1,38 +0,0 @@ -#!/bin/bash -# Launch vLLM-Omni server for VoxCPM online speech serving. -# -# Usage: -# ./run_server.sh # default: async_chunk stage config -# ./run_server.sh async # async_chunk stage config -# ./run_server.sh sync # no-async-chunk stage config -# VOXCPM_MODEL=/path/to/model ./run_server.sh - -set -e - -MODE="${1:-async}" -MODEL="${VOXCPM_MODEL:-OpenBMB/VoxCPM1.5}" - -case "$MODE" in - async) - STAGE_CONFIG="vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml" - ;; - sync) - STAGE_CONFIG="vllm_omni/deploy/voxcpm.yaml" - ;; - *) - echo "Unknown mode: $MODE" - echo "Supported: async, sync" - exit 1 - ;; -esac - -echo "Starting VoxCPM server with model: $MODEL" -echo "Stage config: $STAGE_CONFIG" - -vllm serve "$MODEL" \ - --stage-configs-path "$STAGE_CONFIG" \ - --host 0.0.0.0 \ - --port 8091 \ - --trust-remote-code \ - --enforce-eager \ - --omni diff --git a/pyproject.toml b/pyproject.toml index 95e95fed8..35d333fd1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,7 +63,7 @@ dev = [ "funasr>=1.0.0", "lpips>=0.1.4", # for quantization fp8 "orjson>=3.11.7", # to make v1/embeddings API fast - "voxcpm>=2.0.2", # for voxcpm model Test + "voxcpm>=2.0.2", # for VoxCPM2 path-locating in tests "torchdiffeq>=0.2.5" # for covo_audio Test ] diff --git a/tests/e2e/offline_inference/test_voxcpm_expansion.py b/tests/e2e/offline_inference/test_voxcpm_expansion.py deleted file mode 100644 index 66639a17f..000000000 --- a/tests/e2e/offline_inference/test_voxcpm_expansion.py +++ /dev/null @@ -1,99 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""E2E test for VoxCPM offline inference.""" - -from typing import Any - -import numpy as np -import pytest -import torch - -from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniRunner -from tests.helpers.stage_config import get_deploy_config_path - -VOXCPM_MODEL = "OpenBMB/VoxCPM1.5" -STAGE_CONFIG = get_deploy_config_path("voxcpm.yaml") -SAMPLE_RATE = 24000 - -# (model, stage_config_path) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)`` -_OMNI_RUNNER_PARAM = (VOXCPM_MODEL, STAGE_CONFIG) - -pytestmark = [ - pytest.mark.full_model, - pytest.mark.tts, - pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), -] - - -def _build_prompt(text: str) -> dict[str, Any]: - return { - "prompt_token_ids": [1], - "additional_information": { - "text": [text], - "cfg_value": [2.0], - "inference_timesteps": [10], - "min_len": [2], - "max_new_tokens": [1024], - }, - } - - -def _extract_audio_tensor(multimodal_output: dict[str, Any]) -> torch.Tensor: - audio = multimodal_output.get("audio", multimodal_output.get("model_outputs")) - assert audio is not None, f"No audio output found, keys={list(multimodal_output.keys())}" - - if isinstance(audio, list): - parts: list[torch.Tensor] = [] - for item in audio: - if item is None: - continue - tensor = torch.as_tensor(item) - if tensor.numel() == 0: - continue - parts.append(tensor.float().cpu().reshape(-1)) - return torch.cat(parts, dim=-1) if parts else torch.zeros((0,), dtype=torch.float32) - - return torch.as_tensor(audio).float().cpu().reshape(-1) - - -def _extract_final_multimodal_output(outputs) -> dict[str, Any]: - for item in reversed(outputs): - request_output = getattr(item, "request_output", None) - if request_output is not None: - multimodal_output = getattr(request_output, "multimodal_output", None) - if isinstance(multimodal_output, dict): - return multimodal_output - completions = getattr(request_output, "outputs", None) or [] - for completion in completions: - multimodal_output = getattr(completion, "multimodal_output", None) - if isinstance(multimodal_output, dict): - return multimodal_output - - multimodal_output = getattr(item, "multimodal_output", None) - if isinstance(multimodal_output, dict): - return multimodal_output - - raise AssertionError("No multimodal audio output found in VoxCPM generate results") - - -@hardware_test(res={"cuda": "L4"}, num_cards=1) -def test_voxcpm_zero_shot_001(omni_runner: OmniRunner) -> None: - outputs = list(omni_runner.omni.generate(_build_prompt("Hello, this is a VoxCPM offline inference test."))) - - assert outputs, "No outputs returned" - - multimodal_output = _extract_final_multimodal_output(outputs) - audio = _extract_audio_tensor(multimodal_output) - assert audio.numel() > SAMPLE_RATE // 2, f"Audio too short: {audio.numel()} samples" - - duration_s = audio.shape[0] / SAMPLE_RATE - assert 0.5 < duration_s < 30.0, f"Audio duration out of range: {duration_s:.2f}s" - - peak = float(torch.max(torch.abs(audio)).item()) if audio.numel() > 0 else 0.0 - assert peak > 0.01, "Generated audio appears to be silence" - - audio_np = audio.numpy() - rms = float(np.sqrt(np.mean(np.square(audio_np)))) if audio_np.size else 0.0 - assert rms > 1e-4, "Generated audio RMS too low" diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 8dec9c5a9..cf77e2256 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -232,24 +232,6 @@ def test_stage_configs_path_field(): assert args.stage_configs_path == "/some/path.yaml" -def test_voxcpm_model_arch_injects_model_type_override(mocker): - """Ensure VoxCPM model_arch injects hf_overrides for config resolution.""" - mocker.patch.object(OmniEngineArgs, "_ensure_omni_models_registered", return_value=True) - mocker.patch.object(OmniEngineArgs, "_patch_empty_hf_config") - mocker.patch.object(EngineArgs, "create_model_config", return_value=Mock()) - mocker.patch.object(OmniModelConfig, "from_vllm_model_config", return_value=Mock()) - - args = OmniEngineArgs( - model="OpenBMB/VoxCPM1.5", - model_arch="VoxCPMForConditionalGeneration", - ) - args.create_model_config() - - assert args.hf_overrides["architectures"] == ["VoxCPMForConditionalGeneration"] - assert args.hf_overrides["model_type"] == "voxcpm" - args._patch_empty_hf_config.assert_called_once_with("voxcpm") - - def test_strip_single_engine_args(): """_strip_single_engine_args should remove EngineArgs fields but keep omni fields.""" kwargs = { diff --git a/tests/entrypoints/openai_api/test_serving_speech_voxcpm.py b/tests/entrypoints/openai_api/test_serving_speech_voxcpm.py deleted file mode 100644 index f49c1d9ce..000000000 --- a/tests/entrypoints/openai_api/test_serving_speech_voxcpm.py +++ /dev/null @@ -1,275 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""UTs for VoxCPM OpenAI speech serving behavior.""" - -import asyncio -from types import SimpleNamespace -from unittest.mock import AsyncMock - -import pytest -from pytest_mock import MockerFixture - -from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest -from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -@pytest.fixture -def voxcpm_server(mocker: MockerFixture): - mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set()) - mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None) - - mock_engine_client = mocker.MagicMock() - mock_engine_client.errored = False - mock_engine_client.model_config = mocker.MagicMock(model="OpenBMB/VoxCPM1.5") - mock_engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)] - mock_engine_client.tts_batch_max_items = 32 - mock_engine_client.generate = mocker.MagicMock(return_value="generator") - mock_engine_client.stage_configs = [ - SimpleNamespace( - engine_args=SimpleNamespace( - model_stage="latent_generator", - model_arch="VoxCPMForConditionalGeneration", - ), - tts_args={}, - ), - SimpleNamespace( - engine_args=SimpleNamespace(model_stage="vae"), - tts_args={}, - ), - ] - - mock_models = mocker.MagicMock() - mock_models.is_base_model.return_value = True - - return OmniOpenAIServingSpeech( - engine_client=mock_engine_client, - models=mock_models, - request_logger=mocker.MagicMock(), - ) - - -@pytest.fixture -def voxcpm2_server(mocker: MockerFixture): - mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None) - - mock_engine_client = mocker.MagicMock() - mock_engine_client.errored = False - mock_engine_client.model_config = mocker.MagicMock(model="OpenBMB/VoxCPM2") - mock_engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)] - mock_engine_client.tts_batch_max_items = 32 - mock_engine_client.generate = mocker.MagicMock(return_value="generator") - mock_engine_client.stage_configs = [ - SimpleNamespace( - engine_args=SimpleNamespace( - model_stage="latent_generator", - model_arch="VoxCPM2TalkerForConditionalGeneration", - ), - tts_args={}, - ), - SimpleNamespace( - engine_args=SimpleNamespace(model_stage="vae"), - tts_args={}, - ), - ] - - mock_models = mocker.MagicMock() - mock_models.is_base_model.return_value = True - - return OmniOpenAIServingSpeech( - engine_client=mock_engine_client, - models=mock_models, - request_logger=mocker.MagicMock(), - ) - - -class TestVoxCPMServing: - def test_voxcpm_model_type_detection(self, voxcpm_server): - assert voxcpm_server._tts_model_type == "voxcpm" - assert voxcpm_server._is_tts is True - assert voxcpm_server.supported_speakers == set() - - @pytest.mark.parametrize( - ("request_kwargs", "expected_substring"), - [ - ({"voice": "alice"}, "voice"), - ({"instructions": "whisper"}, "instructions"), - ({"language": "en"}, "language"), - ({"task_type": "CustomVoice"}, "plain tts"), - ({"x_vector_only_mode": True}, "x_vector_only_mode"), - ({"speaker_embedding": [0.1, 0.2]}, "speaker_embedding"), - ({"initial_codec_chunk_frames": 4}, "initial_codec_chunk_frames"), - ({"ref_text": "reference"}, "ref_audio"), - ], - ) - def test_validate_voxcpm_rejects_unsupported_fields(self, voxcpm_server, request_kwargs, expected_substring): - request = OpenAICreateSpeechRequest(input="hello voxcpm", **request_kwargs) - error = voxcpm_server._validate_voxcpm_request(request) - assert error is not None - assert expected_substring in error.lower() - - def test_validate_voxcpm_accepts_plain_tts_request(self, voxcpm_server): - request = OpenAICreateSpeechRequest(input="hello voxcpm", max_new_tokens=256) - assert voxcpm_server._validate_voxcpm_request(request) is None - - def test_validate_voxcpm_accepts_voice_clone_request(self, voxcpm_server): - request = OpenAICreateSpeechRequest( - input="clone this voice", - ref_audio="data:audio/wav;base64,QUJD", - ref_text="reference transcript", - max_new_tokens=256, - ) - assert voxcpm_server._validate_voxcpm_request(request) is None - - def test_prepare_speech_generation_voxcpm_text_only(self, voxcpm_server): - request = OpenAICreateSpeechRequest(input="hello voxcpm", max_new_tokens=321) - - request_id, generator, tts_params = asyncio.run(voxcpm_server._prepare_speech_generation(request)) - - assert request_id.startswith("speech-") - assert generator == "generator" - assert tts_params == { - "text": ["hello voxcpm"], - "cfg_value": [2.0], - "inference_timesteps": [10], - "min_len": [2], - "max_new_tokens": [321], - } - - voxcpm_server.engine_client.generate.assert_called_once() - call = voxcpm_server.engine_client.generate.call_args - assert call.kwargs["prompt"] == { - "prompt_token_ids": [1], - "additional_information": tts_params, - "type": "token", - } - assert call.kwargs["output_modalities"] == ["audio"] - - def test_prepare_speech_generation_voxcpm_voice_clone_resolves_ref_audio(self, voxcpm_server): - voxcpm_server._resolve_ref_audio = AsyncMock(return_value=([0.1, -0.1, 0.2], 16000)) - request = OpenAICreateSpeechRequest( - input="clone this voice", - ref_audio="data:audio/wav;base64,QUJD", - ref_text="reference transcript", - max_new_tokens=512, - ) - - request_id, generator, tts_params = asyncio.run(voxcpm_server._prepare_speech_generation(request)) - - assert request_id.startswith("speech-") - assert generator == "generator" - assert tts_params == { - "text": ["clone this voice"], - "cfg_value": [2.0], - "inference_timesteps": [10], - "min_len": [2], - "max_new_tokens": [512], - "ref_text": ["reference transcript"], - "ref_audio": [[[0.1, -0.1, 0.2], 16000]], - } - - voxcpm_server._resolve_ref_audio.assert_awaited_once_with("data:audio/wav;base64,QUJD") - call = voxcpm_server.engine_client.generate.call_args - assert call.kwargs["prompt"] == { - "prompt_token_ids": [1], - "additional_information": tts_params, - "type": "token", - } - - -class TestVoxCPM2Serving: - """Regression tests for VoxCPM2 serving behavior.""" - - def test_voxcpm2_model_type_detection(self, voxcpm2_server): - assert voxcpm2_server._tts_model_type == "voxcpm2" - assert voxcpm2_server._is_tts is True - - def test_voxcpm2_default_voice_is_supported(self, voxcpm2_server): - """VoxCPM2 should accept voice='default' (zero-shot mode, no voice cloning). - - Regression: previously returned 400 Invalid voice 'default'. - """ - assert "default" in voxcpm2_server.supported_speakers - - def test_voxcpm2_inline_ref_audio_does_not_set_voice_name(self, voxcpm2_server): - """Inline ref_audio must not populate voice_name in additional_information. - - Regression: the speaker cache keys on voice_name, so two requests with - the same voice but different inline ref_audio would collide — the second - request reused stale cached features from the first, causing a prefill - length mismatch and engine crash. - """ - mock_prompt = {"prompt_token_ids": [1, 1, 1], "additional_information": {}, "type": "token"} - voxcpm2_server._build_voxcpm2_prompt = AsyncMock(return_value=mock_prompt) - - request = OpenAICreateSpeechRequest( - input="test", - voice="default", - ref_audio="data:audio/wav;base64,QUJD", - ) - - asyncio.run(voxcpm2_server._prepare_speech_generation(request)) - - additional = mock_prompt["additional_information"] - assert "voice_name" not in additional, ( - "voice_name must not be set for inline ref_audio — " - "it causes speaker cache collisions across different audio clips" - ) - - def test_voxcpm2_uploaded_voice_sets_voice_name(self, voxcpm2_server): - """Uploaded voice without inline ref_audio must populate voice_name. - - This ensures the speaker cache still works for the intended use case: - a voice uploaded once and referenced by name in subsequent requests. - """ - mock_prompt = {"prompt_token_ids": [1, 1, 1], "additional_information": {}, "type": "token"} - voxcpm2_server._build_voxcpm2_prompt = AsyncMock(return_value=mock_prompt) - voxcpm2_server.uploaded_speakers["my_voice"] = { - "name": "my_voice", - "file_path": "/tmp/fake.wav", - "created_at": 12345, - } - voxcpm2_server.supported_speakers.add("my_voice") - - request = OpenAICreateSpeechRequest( - input="test", - voice="my_voice", - ) - - asyncio.run(voxcpm2_server._prepare_speech_generation(request)) - - additional = mock_prompt["additional_information"] - assert additional.get("voice_name") == "my_voice" - assert "voice_created_at" in additional - - def test_voxcpm2_uploaded_voice_with_inline_override_does_not_set_voice_name(self, voxcpm2_server): - """Uploaded voice with inline ref_audio override must not populate voice_name. - - Prevents cache poisoning: an attacker could send an uploaded voice name - with crafted ref_audio, caching malicious features under the uploaded - voice's name for subsequent legitimate requests. - """ - mock_prompt = {"prompt_token_ids": [1, 1, 1], "additional_information": {}, "type": "token"} - voxcpm2_server._build_voxcpm2_prompt = AsyncMock(return_value=mock_prompt) - voxcpm2_server.uploaded_speakers["my_voice"] = { - "name": "my_voice", - "file_path": "/tmp/fake.wav", - "created_at": 12345, - } - voxcpm2_server.supported_speakers.add("my_voice") - - request = OpenAICreateSpeechRequest( - input="test", - voice="my_voice", - ref_audio="data:audio/wav;base64,QUJD", - ) - - asyncio.run(voxcpm2_server._prepare_speech_generation(request)) - - additional = mock_prompt["additional_information"] - assert "voice_name" not in additional, ( - "voice_name must not be set when inline ref_audio overrides an " - "uploaded voice — it would poison the cache for future requests" - ) diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index b52b49d68..479ada963 100644 --- a/tests/entrypoints/test_utils.py +++ b/tests/entrypoints/test_utils.py @@ -313,39 +313,6 @@ class TestResolveModelConfigPath: assert result is not None assert "glm_image.yaml" in result - def test_voxcpm_transformers_format_resolution(self, mocker: MockerFixture): - """Test VoxCPM transformers config resolves to the voxcpm stage config.""" - mocker.patch( - "vllm_omni.entrypoints.utils.get_config", - side_effect=ValueError("missing transformers config"), - ) - mocker.patch( - "vllm_omni.entrypoints.utils.file_or_path_exists", - side_effect=lambda _model, filename, revision=None: filename == "config.json", - ) - mocker.patch( - "vllm_omni.entrypoints.utils.get_hf_file_to_dict", - return_value={"model_type": "voxcpm"}, - ) - mocker.patch( - "vllm_omni.entrypoints.utils.current_omni_platform.get_default_stage_config_path", - return_value="vllm_omni/deploy", - ) - - original_exists = os.path.exists - - def mock_exists(path): - if "voxcpm.yaml" in str(path): - return True - return original_exists(path) - - mocker.patch("os.path.exists", side_effect=mock_exists) - - result = resolve_model_config_path("OpenBMB/VoxCPM1.5") - - assert result is not None - assert "voxcpm.yaml" in result - class TestLoadAndResolveStageConfigs: def test_load_and_resolve_with_kwargs(self): diff --git a/tests/model_executor/stage_input_processors/test_voxcpm_async_chunk.py b/tests/model_executor/stage_input_processors/test_voxcpm_async_chunk.py deleted file mode 100644 index 6738c763a..000000000 --- a/tests/model_executor/stage_input_processors/test_voxcpm_async_chunk.py +++ /dev/null @@ -1,85 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""UTs for VoxCPM async-chunk stage input processing.""" - -from types import SimpleNamespace - -import pytest -import torch - -from vllm_omni.model_executor.stage_input_processors.voxcpm import ( - _VOXCPM_LATENT_MAGIC, - _coerce_finished_flag, - latent2vae_async_chunk, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def _request(*, finished): - return SimpleNamespace(is_finished=lambda: finished) - - -def _decode_serialized_latent(codes: list[int]) -> torch.Tensor: - assert codes[0] == _VOXCPM_LATENT_MAGIC - latent_dim = codes[1] - time_dim = codes[2] - payload = torch.tensor(codes[3:], dtype=torch.int32).to(torch.uint16) - return payload.view(torch.bfloat16).to(torch.float32).reshape(1, latent_dim, time_dim) - - -@pytest.mark.parametrize( - ("value", "expected"), - [ - (None, False), - (False, False), - (True, True), - (torch.tensor(False), False), - (torch.tensor(True), True), - ([torch.tensor(True)], True), - (([True],), True), - ([], False), - ], -) -def test_coerce_finished_flag(value, expected): - assert _coerce_finished_flag(value) is expected - - -def test_latent2vae_async_chunk_serializes_latent_payload(): - latent = torch.arange(6, dtype=torch.float32).reshape(2, 3) - - payload = latent2vae_async_chunk( - transfer_manager=None, - pooling_output={"latent_audio_feat": latent}, - request=_request(finished=False), - is_finished=torch.tensor(False), - ) - - assert payload is not None - assert torch.equal(payload.meta.finished, torch.tensor(False, dtype=torch.bool)) - recovered = _decode_serialized_latent(payload.codes.audio.tolist()) - torch.testing.assert_close(recovered, latent.to(torch.bfloat16).to(torch.float32).unsqueeze(0)) - - -def test_latent2vae_async_chunk_returns_terminal_marker_without_latent(): - payload = latent2vae_async_chunk( - transfer_manager=None, - pooling_output=None, - request=_request(finished=[torch.tensor(True)]), - is_finished=False, - ) - - assert payload.codes.audio.tolist() == [] - assert torch.equal(payload.meta.finished, torch.tensor(True, dtype=torch.bool)) - - -def test_latent2vae_async_chunk_returns_none_for_nonterminal_empty_chunk(): - payload = latent2vae_async_chunk( - transfer_manager=None, - pooling_output={"latent_audio_feat": torch.zeros((0,), dtype=torch.float32)}, - request=_request(finished=False), - is_finished=False, - ) - - assert payload is None diff --git a/vllm_omni/deploy/voxcpm.yaml b/vllm_omni/deploy/voxcpm.yaml deleted file mode 100644 index a5f324f66..000000000 --- a/vllm_omni/deploy/voxcpm.yaml +++ /dev/null @@ -1,69 +0,0 @@ -# VoxCPM two-stage (latent → VAE) without async_chunk: one-shot latent then decode. -stage_args: - - stage_id: 0 - stage_type: llm - is_comprehension: true - runtime: - devices: "0" - max_batch_size: 1 - engine_args: - dtype: bfloat16 - model_stage: latent_generator - model_arch: VoxCPMForConditionalGeneration - # Optional persistent HF-compatible config dir for native VoxCPM models. - hf_config_path: ${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,} - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - enforce_eager: true - trust_remote_code: true - async_scheduling: false - enable_prefix_caching: false - engine_output_type: latent - gpu_memory_utilization: 0.7 - distributed_executor_backend: "mp" - max_num_batched_tokens: 4096 - max_model_len: 4096 - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 4096 - stop_token_ids: [2] - seed: 42 - detokenize: false - repetition_penalty: 1.0 - final_output: false - - - stage_id: 1 - stage_type: llm - runtime: - devices: "0" - max_batch_size: 1 - engine_args: - dtype: float32 - model_stage: vae - model_arch: VoxCPMForConditionalGeneration - hf_config_path: ${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,} - worker_type: generation - scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler - enforce_eager: true - trust_remote_code: true - async_scheduling: false - enable_prefix_caching: false - engine_output_type: audio - gpu_memory_utilization: 0.15 - distributed_executor_backend: "mp" - max_num_batched_tokens: 8192 - max_model_len: 4096 - engine_input_source: [0] - custom_process_input_func: vllm_omni.model_executor.stage_input_processors.voxcpm.latent2vae - final_output: true - final_output_type: audio - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 1 - seed: 42 - detokenize: true - repetition_penalty: 1.0 diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 14d0d43e8..12a5a450a 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -22,7 +22,6 @@ _ARCH_TO_MODEL_TYPE: dict[str, str] = { "CosyVoice3Model": "cosyvoice3", "OmniVoiceModel": "omnivoice", "VoxCPM2TalkerForConditionalGeneration": "voxcpm2", - "VoxCPMForConditionalGeneration": "voxcpm", } # Maps model architecture names to tokenizer subfolder paths within HF repos. diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 70a491a94..3b7a2e747 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -65,7 +65,6 @@ _FISH_TTS_MODEL_STAGES = {"fish_speech_slow_ar"} _COSYVOICE3_TTS_MODEL_STAGES = {"cosyvoice3_talker"} _OMNIVOICE_TTS_MODEL_STAGES = {"omnivoice_generator"} _COVO_AUDIO_MODEL_STAGES = {"fused_thinker_talker"} -_VOXCPM_TTS_MODEL_STAGES = {"latent_generator", "vae"} _VOXCPM2_TTS_MODEL_STAGES = {"latent_generator"} _MING_TTS_MODEL_STAGES = {"ming_tts"} _MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"} @@ -76,7 +75,6 @@ _TTS_MODEL_STAGES: set[str] = ( | _COSYVOICE3_TTS_MODEL_STAGES | _OMNIVOICE_TTS_MODEL_STAGES | _COVO_AUDIO_MODEL_STAGES - | _VOXCPM_TTS_MODEL_STAGES | _VOXCPM2_TTS_MODEL_STAGES | _MING_TTS_MODEL_STAGES | _MOSS_TTS_MODEL_STAGES @@ -481,8 +479,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): model_arch = getattr(self._tts_stage.engine_args, "model_arch", None) if model_arch == "VoxCPM2TalkerForConditionalGeneration": return "voxcpm2" - if model_arch == "VoxCPMForConditionalGeneration": - return "voxcpm" if model_stage in _QWEN3_TTS_MODEL_STAGES: return "qwen3_tts" if model_stage in _VOXTRAL_TTS_MODEL_STAGES: @@ -497,12 +493,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): model_arch = getattr(self._tts_stage.engine_args, "model_arch", None) if model_arch and "CovoAudio" in model_arch: return "covo_audio" - if model_stage in (_VOXCPM_TTS_MODEL_STAGES | _VOXCPM2_TTS_MODEL_STAGES): - has_vae_stage = any( - getattr(getattr(stage, "engine_args", None), "model_stage", None) == "vae" - for stage in self.engine_client.stage_configs - ) - return "voxcpm" if has_vae_stage or model_stage == "vae" else "voxcpm2" + if model_stage in _VOXCPM2_TTS_MODEL_STAGES: + return "voxcpm2" if model_stage in _MING_TTS_MODEL_STAGES: return "ming_flash_omni_tts" if model_stage in _MOSS_TTS_MODEL_STAGES: @@ -536,8 +528,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # is no static speaker list to surface here. return set() try: - if self._tts_model_type == "voxcpm": - return set() if self._tts_model_type == "voxcpm2": return {"default"} if self._tts_model_type == "voxtral_tts": @@ -599,8 +589,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _estimate_prompt_len(self, tts_params: dict[str, Any]) -> int: """Estimate prompt length so the placeholder matches model-side embeddings.""" try: - if self._tts_model_type == "voxcpm": - return 1 from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( Qwen3TTSTalkerForConditionalGeneration, ) @@ -1136,8 +1124,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return self._validate_fish_tts_request(request) if self._tts_model_type == "cosyvoice3": return self._validate_cosyvoice3_request(request) - if self._tts_model_type == "voxcpm": - return self._validate_voxcpm_request(request) if self._tts_model_type == "voxcpm2": if request.max_new_tokens is not None: if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN: @@ -1242,43 +1228,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return None - def _validate_voxcpm_request(self, request: OpenAICreateSpeechRequest) -> str | None: - """Validate VoxCPM request parameters. Returns error message or None.""" - if not request.input or not request.input.strip(): - return "Input text cannot be empty" - - if request.voice is not None: - return "'voice' is not supported for VoxCPM" - if request.instructions is not None: - return "'instructions' is not supported for VoxCPM" - if request.language is not None: - return "'language' is not supported for VoxCPM" - if request.task_type not in (None, "Base"): - return "VoxCPM only supports plain TTS or voice cloning with ref_audio/ref_text" - if request.x_vector_only_mode is not None: - return "'x_vector_only_mode' is not supported for VoxCPM" - if request.speaker_embedding is not None: - return "'speaker_embedding' is not supported for VoxCPM" - if request.initial_codec_chunk_frames is not None: - return "'initial_codec_chunk_frames' is not supported for VoxCPM" - - if request.ref_audio is not None: - fmt_err = self._validate_ref_audio_format(request.ref_audio) - if fmt_err: - return fmt_err - if not request.ref_text or not request.ref_text.strip(): - return "Voice cloning requires 'ref_text' (transcript of the reference audio)" - elif request.ref_text is not None: - return "'ref_text' requires 'ref_audio' for VoxCPM voice cloning" - - if request.max_new_tokens is not None: - if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN: - return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}" - if request.max_new_tokens > _TTS_MAX_NEW_TOKENS_MAX: - return f"max_new_tokens cannot exceed {_TTS_MAX_NEW_TOKENS_MAX}" - - return None - def _validate_qwen_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: """Validate Qwen TTS request parameters. Returns error message or None.""" # Infer Base task when ref_audio or ref_text is provided without explicit task_type. @@ -1643,18 +1592,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): Processes each parameter if present, skips if not. Values are wrapped in lists as required by the model. """ - if self._tts_model_type == "voxcpm": - params: dict[str, Any] = { - "text": [request.input], - "cfg_value": [2.0], - "inference_timesteps": [10], - "min_len": [2], - "max_new_tokens": [request.max_new_tokens or 4096], - } - if request.ref_text is not None: - params["ref_text"] = [request.ref_text] - return params - params: dict[str, Any] = {} # Text content (always required) @@ -2169,8 +2106,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): model_type = "voxtral_tts" elif self._tts_model_type == "cosyvoice3": model_type = "cosyvoice3" - elif self._tts_model_type == "voxcpm": - model_type = "voxcpm" elif self._tts_model_type == "voxcpm2": model_type = "voxcpm2" elif self._tts_model_type == "ming_flash_omni_tts": diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index cc900d471..4f29c0ae3 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -150,12 +150,6 @@ _OMNI_MODELS = { "fish_speech_dac_decoder", "FishSpeechDACDecoder", ), - ## VoxCPM - "VoxCPMForConditionalGeneration": ( - "voxcpm", - "voxcpm", - "VoxCPMForConditionalGeneration", - ), ## VoxCPM2 "VoxCPM2TalkerForConditionalGeneration": ( "voxcpm2", diff --git a/vllm_omni/model_executor/models/voxcpm/__init__.py b/vllm_omni/model_executor/models/voxcpm/__init__.py deleted file mode 100644 index ffde6dee6..000000000 --- a/vllm_omni/model_executor/models/voxcpm/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from .voxcpm import VoxCPMForConditionalGeneration - -__all__ = [ - "VoxCPMForConditionalGeneration", -] diff --git a/vllm_omni/model_executor/models/voxcpm/voxcpm.py b/vllm_omni/model_executor/models/voxcpm/voxcpm.py deleted file mode 100644 index 6fa36fc42..000000000 --- a/vllm_omni/model_executor/models/voxcpm/voxcpm.py +++ /dev/null @@ -1,886 +0,0 @@ -from __future__ import annotations - -import json -import os -import sys -import tempfile -import warnings -import wave -from collections.abc import Callable, Generator, Iterable -from pathlib import Path -from typing import Any - -import numpy as np -import torch -import torch.nn as nn -from einops import rearrange -from tqdm import tqdm -from vllm.config import VllmConfig -from vllm.logger import init_logger -from vllm.sequence import IntermediateTensors - -from vllm_omni.model_executor.models.output_templates import OmniOutput - -from .voxcpm_loader import ( - _build_prompt_cache_with_soundfile, - _device_to_string, - _force_cuda_available_for_npu, - _import_voxcpm_audio_vae_classes, - _import_voxcpm_base_model_class, - _is_torchcodec_load_error, - _normalize_dtype_name, - _prepare_runtime_model_dir, - _resolve_runtime_device, -) -from .voxcpm_runtime_utils import resolve_voxcpm_model_dir -from .voxcpm_stage_wrappers import _DirectVoxCPMAudioVAE, _DirectVoxCPMLatentGenerator - -logger = init_logger(__name__) -_VOXCPM_LATENT_MAGIC = 131071 - - -def _make_voxcpm_model_for_omni(base: type[Any]) -> type[Any]: - """Subclass upstream VoxCPMModel: local ``_inference`` + ``latents_only`` prompt-cache generation.""" - - from voxcpm.model.utils import get_dtype - - class VoxCPMModelForOmni(base): - @torch.inference_mode() - def build_prompt_cache(self, *args: Any, **kwargs: Any): - try: - return super().build_prompt_cache(*args, **kwargs) - except (ImportError, ModuleNotFoundError, RuntimeError) as exc: - if not _is_torchcodec_load_error(exc): - raise - return _build_prompt_cache_with_soundfile(self, *args, **kwargs) - - @torch.inference_mode() - def _inference( - self, - text: torch.Tensor, - text_mask: torch.Tensor, - feat: torch.Tensor, - feat_mask: torch.Tensor, - min_len: int = 2, - max_len: int = 2000, - inference_timesteps: int = 10, - cfg_value: float = 2.0, - streaming: bool = False, - streaming_prefix_len: int = 3, - ) -> Generator[tuple[torch.Tensor, torch.Tensor | list[torch.Tensor]], None, None]: - B, _, _, _ = feat.shape - - feat_embed = self.feat_encoder(feat) - feat_embed = self.enc_to_lm_proj(feat_embed) - - scale_emb = self.config.lm_config.scale_emb if self.config.lm_config.use_mup else 1.0 - text_embed = self.base_lm.embed_tokens(text) * scale_emb - combined_embed = text_mask.unsqueeze(-1) * text_embed + feat_mask.unsqueeze(-1) * feat_embed - - prefix_feat_cond = feat[:, -1, ...] - pred_feat_seq: list[torch.Tensor] = [] - - audio_patch_count = int(feat_mask.sum().item()) - if audio_patch_count > 0: - context_len = min(streaming_prefix_len - 1, audio_patch_count) - prompt_context_patches = list(feat[:, -context_len:, :, :].split(1, dim=1)) - pred_feat_seq = prompt_context_patches + pred_feat_seq - - enc_outputs, kv_cache_tuple = self.base_lm( - inputs_embeds=combined_embed, - is_causal=True, - ) - self.base_lm.kv_cache.fill_caches(kv_cache_tuple) - - enc_outputs = self.fsq_layer(enc_outputs) * feat_mask.unsqueeze(-1) + enc_outputs * text_mask.unsqueeze(-1) - lm_hidden = enc_outputs[:, -1, :] - - residual_enc_outputs, residual_kv_cache_tuple = self.residual_lm( - inputs_embeds=enc_outputs + feat_mask.unsqueeze(-1) * feat_embed, - is_causal=True, - ) - self.residual_lm.kv_cache.fill_caches(residual_kv_cache_tuple) - residual_hidden = residual_enc_outputs[:, -1, :] - - for step_idx in tqdm(range(max_len)): - dit_hidden = self.lm_to_dit_proj(lm_hidden) + self.res_to_dit_proj(residual_hidden) - pred_feat = self.feat_decoder( - mu=dit_hidden, - patch_size=self.patch_size, - cond=prefix_feat_cond.transpose(1, 2).contiguous(), - n_timesteps=inference_timesteps, - cfg_value=cfg_value, - ).transpose(1, 2) - - curr_embed = self.enc_to_lm_proj(self.feat_encoder(pred_feat.unsqueeze(1))) - pred_feat_seq.append(pred_feat.unsqueeze(1)) - prefix_feat_cond = pred_feat - - if streaming: - pred_feat_chunk = torch.cat(pred_feat_seq[-streaming_prefix_len:], dim=1) - feat_pred = rearrange(pred_feat_chunk, "b t p d -> b d (t p)", b=B, p=self.patch_size) - yield feat_pred, pred_feat_seq - - stop_flag = self.stop_head(self.stop_actn(self.stop_proj(lm_hidden))).argmax(dim=-1)[0].cpu().item() - if step_idx > min_len and stop_flag == 1: - break - - lm_hidden = self.base_lm.forward_step( - curr_embed[:, 0, :], - torch.tensor([self.base_lm.kv_cache.step()], device=curr_embed.device), - ).clone() - lm_hidden = self.fsq_layer(lm_hidden) - residual_hidden = self.residual_lm.forward_step( - lm_hidden + curr_embed[:, 0, :], - torch.tensor([self.residual_lm.kv_cache.step()], device=curr_embed.device), - ).clone() - - if not streaming: - pred_feat_seq_cat = torch.cat(pred_feat_seq, dim=1) - feat_pred = rearrange(pred_feat_seq_cat, "b t p d -> b d (t p)", b=B, p=self.patch_size) - yield feat_pred, pred_feat_seq_cat.squeeze(0).cpu() - - @torch.inference_mode() - def generate_latents_with_prompt_cache( - self, - target_text: str, - prompt_cache: dict, - min_len: int = 2, - max_len: int = 2000, - inference_timesteps: int = 10, - cfg_value: float = 2.0, - retry_badcase: bool = False, - retry_badcase_max_times: int = 3, - retry_badcase_ratio_threshold: float = 6.0, - streaming_prefix_len: int = 3, - ) -> tuple[None, torch.Tensor, torch.Tensor]: - return next( - self._generate_with_prompt_cache( - target_text=target_text, - prompt_cache=prompt_cache, - min_len=min_len, - max_len=max_len, - inference_timesteps=inference_timesteps, - cfg_value=cfg_value, - retry_badcase=retry_badcase, - retry_badcase_max_times=retry_badcase_max_times, - retry_badcase_ratio_threshold=retry_badcase_ratio_threshold, - streaming=False, - streaming_prefix_len=streaming_prefix_len, - latents_only=True, - ) - ) - - @torch.inference_mode() - def generate_latents_with_prompt_cache_streaming( - self, - target_text: str, - prompt_cache: dict, - min_len: int = 2, - max_len: int = 2000, - inference_timesteps: int = 10, - cfg_value: float = 2.0, - retry_badcase: bool = False, - retry_badcase_max_times: int = 3, - retry_badcase_ratio_threshold: float = 6.0, - streaming_prefix_len: int = 3, - ) -> Generator[tuple[None, torch.Tensor, torch.Tensor], None, None]: - return self._generate_with_prompt_cache( - target_text=target_text, - prompt_cache=prompt_cache, - min_len=min_len, - max_len=max_len, - inference_timesteps=inference_timesteps, - cfg_value=cfg_value, - retry_badcase=retry_badcase, - retry_badcase_max_times=retry_badcase_max_times, - retry_badcase_ratio_threshold=retry_badcase_ratio_threshold, - streaming=True, - streaming_prefix_len=streaming_prefix_len, - latents_only=True, - ) - - @torch.inference_mode() - def _generate_with_prompt_cache( - self, - target_text: str, - prompt_cache: dict, - min_len: int = 2, - max_len: int = 2000, - inference_timesteps: int = 10, - cfg_value: float = 2.0, - retry_badcase: bool = False, - retry_badcase_max_times: int = 3, - retry_badcase_ratio_threshold: float = 6.0, - streaming: bool = False, - streaming_prefix_len: int = 3, - latents_only: bool = False, - ) -> Generator[tuple[torch.Tensor | None, torch.Tensor, torch.Tensor | list[torch.Tensor]], None, None]: - if retry_badcase and streaming: - warnings.warn("Retry on bad cases is not supported in streaming mode, setting retry_badcase=False.") - retry_badcase = False - if prompt_cache is None: - prompt_audio_feat = torch.empty((0, self.patch_size, self.audio_vae.latent_dim), dtype=torch.float32) - text = target_text - else: - prompt_audio_feat = prompt_cache["audio_feat"] - prompt_text = prompt_cache["prompt_text"] - text = prompt_text + target_text - - text_token = torch.LongTensor(self.text_tokenizer(text)) - text_token = torch.cat( - [ - text_token, - torch.tensor([self.audio_start_token], dtype=torch.int32, device=text_token.device), - ], - dim=-1, - ) - target_text_token = torch.LongTensor(self.text_tokenizer(target_text)) - - audio_length = prompt_audio_feat.size(0) - text_length = text_token.shape[0] - text_pad_token = torch.zeros(audio_length, dtype=torch.int32, device=text_token.device) - audio_pad_feat = torch.zeros( - (text_token.shape[0], self.patch_size, self.audio_vae.latent_dim), - dtype=torch.float32, - device=text_token.device, - ) - text_token = torch.cat([text_token, text_pad_token]) - audio_feat = torch.cat([audio_pad_feat, prompt_audio_feat], dim=0) - text_mask = ( - torch.cat([torch.ones(text_length), torch.zeros(audio_length)]).type(torch.int32).to(text_token.device) - ) - audio_mask = ( - torch.cat([torch.zeros(text_length), torch.ones(audio_length)]).type(torch.int32).to(text_token.device) - ) - - text_token = text_token.unsqueeze(0).to(self.device) - text_mask = text_mask.unsqueeze(0).to(self.device) - audio_feat = audio_feat.unsqueeze(0).to(self.device).to(get_dtype(self.config.dtype)) - audio_mask = audio_mask.unsqueeze(0).to(self.device) - - target_text_length = len(self.text_tokenizer(target_text)) - retry_badcase_times = 0 - while retry_badcase_times < retry_badcase_max_times: - inference_result = self._inference( - text_token, - text_mask, - audio_feat, - audio_mask, - min_len=min_len, - max_len=min(int(target_text_length * retry_badcase_ratio_threshold + 10), max_len), - inference_timesteps=inference_timesteps, - cfg_value=cfg_value, - streaming=streaming, - streaming_prefix_len=streaming_prefix_len, - ) - if streaming: - patch_len = self.patch_size * self.chunk_size - for latent_pred, pred_audio_feat in inference_result: - if latents_only: - decode_audio = None - yield (decode_audio, target_text_token, latent_pred) - else: - decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32)) - decode_audio = decode_audio[..., -patch_len:].squeeze(1).cpu() - yield (decode_audio, target_text_token, pred_audio_feat) - break - - latent_pred, pred_audio_feat = next(inference_result) - if retry_badcase and pred_audio_feat.shape[0] >= target_text_length * retry_badcase_ratio_threshold: - ratio = pred_audio_feat.shape[0] / target_text_length - print(f" Badcase detected, audio_text_ratio={ratio}, retrying...", file=sys.stderr) - retry_badcase_times += 1 - continue - break - - if not streaming: - if latents_only: - decode_audio = None - else: - decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32)) - patch_len = self.patch_size * self.chunk_size - if audio_mask.sum().item() > 0: - decode_audio = decode_audio[..., patch_len * (streaming_prefix_len - 1) :].squeeze(1).cpu() - else: - decode_audio = decode_audio[..., :].squeeze(1).cpu() - yield (decode_audio, target_text_token, pred_audio_feat) - - VoxCPMModelForOmni.__name__ = "VoxCPMModelForOmni" - VoxCPMModelForOmni.__qualname__ = "VoxCPMModelForOmni" - return VoxCPMModelForOmni - - -def _import_voxcpm_model_class() -> type[Any]: - base = _import_voxcpm_base_model_class() - return _make_voxcpm_model_for_omni(base) - - -def _load_native_voxcpm_model( - model_path: str, - *, - device: torch.device, - dtype: str | None, -): - VoxCPMModel = _import_voxcpm_model_class() - model_dir = resolve_voxcpm_model_dir(model_path) - runtime_model_path = _prepare_runtime_model_dir(model_dir, target_device=device, target_dtype=dtype) - - if device.type == "npu" and hasattr(torch, "npu"): - torch.npu.set_device(device) - - with _force_cuda_available_for_npu(device): - return VoxCPMModel.from_local( - runtime_model_path, - optimize=device.type == "cuda", - ) - - -def _load_native_voxcpm_latent_generator( - model_path: str, - *, - device: torch.device, - dtype: str | None, -) -> _DirectVoxCPMLatentGenerator: - return _DirectVoxCPMLatentGenerator(_load_native_voxcpm_model(model_path, device=device, dtype=dtype)) - - -def _load_native_voxcpm_audio_vae( - model_path: str, - *, - device: torch.device, -) -> _DirectVoxCPMAudioVAE: - AudioVAE, AudioVAEConfig = _import_voxcpm_audio_vae_classes() - model_dir = resolve_voxcpm_model_dir(model_path) - runtime_model_path = _prepare_runtime_model_dir(model_dir, target_device=device, target_dtype="float32") - config_dict = json.loads((Path(runtime_model_path) / "config.json").read_text()) - audio_vae_config = config_dict.get("audio_vae_config") - audio_vae = AudioVAE(config=AudioVAEConfig(**audio_vae_config)) if audio_vae_config is not None else AudioVAE() - - state_dict = torch.load( - Path(runtime_model_path) / "audiovae.pth", - map_location="cpu", - weights_only=True, - )["state_dict"] - audio_vae.load_state_dict(state_dict, strict=True) - audio_vae = audio_vae.to(device=device, dtype=torch.float32).eval() - if device.type == "npu" and hasattr(torch, "npu"): - torch.npu.set_device(device) - patch_size = int(config_dict.get("patch_size", 2)) - return _DirectVoxCPMAudioVAE(audio_vae, patch_size=patch_size) - - -class VoxCPMForConditionalGeneration(nn.Module): - input_modalities = "audio" - _LATENT_STAGES = {"latent_generator", "latent", "ar_dit"} - _VAE_STAGES = {"vae", "audio_vae"} - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - del prefix - self.vllm_config = vllm_config - self.model_path = vllm_config.model_config.model - self.model_stage = getattr(vllm_config.model_config, "model_stage", "latent_generator") - self.have_multimodal_outputs = True - self.has_preprocess = False - self.has_postprocess = False - self.enable_update_additional_information = True - self.requires_raw_input_tokens = True - self.inject_omni_request_id_into_runtime_info = True - self._pipeline = None - self._latent_stream_gens: dict[str, Any] = {} - self._latent_stream_terminal_pending: dict[str, int] = {} - self._latent_stream_completed: set[str] = set() - self._next_local_stream_key = 0 - self._ar_emit_stop_token = True - - def _runner_hidden_device_dtype(self) -> tuple[torch.device, torch.dtype]: - device = _resolve_runtime_device(self.vllm_config) - model_config = getattr(self.vllm_config, "model_config", None) - dtype = getattr(model_config, "dtype", torch.float32) if model_config is not None else torch.float32 - return device, dtype - - def _ensure_model_loaded(self): - if self._pipeline is not None: - return - - target_device = _resolve_runtime_device(self.vllm_config) - model_dtype = getattr(self.vllm_config.model_config, "dtype", None) - normalized_dtype = _normalize_dtype_name(model_dtype) - if self.model_stage in self._LATENT_STAGES: - self._pipeline = _load_native_voxcpm_latent_generator( - self.model_path, - device=target_device, - dtype=normalized_dtype, - ) - elif self.model_stage in self._VAE_STAGES: - self._pipeline = _load_native_voxcpm_audio_vae( - self.model_path, - device=target_device, - ) - else: - raise ValueError( - f"Unsupported VoxCPM model_stage: {self.model_stage}. " - "pure_voxcpm only supports split-stage latent_generator/vae inference." - ) - - logger.info("Loaded VoxCPM stage '%s' on %s", self.model_stage, _device_to_string(target_device)) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - del weights - self._ensure_model_loaded() - return set() - - @staticmethod - def _extract_val(info: dict[str, Any], key: str, default: Any) -> Any: - value = info.get(key, default) - if isinstance(value, list): - return value[0] if value else default - return value - - def _resolve_stream_request_key(self, info: dict[str, Any]) -> str: - request_key = info.get("__voxcpm_stream_key") - if request_key is not None: - return str(request_key) - - request_key = info.get("_omni_req_id") - if request_key is not None: - request_key = str(request_key) - info["__voxcpm_stream_key"] = request_key - return request_key - - request_key = f"voxcpm-local-{self._next_local_stream_key}" - self._next_local_stream_key += 1 - info["__voxcpm_stream_key"] = request_key - return str(request_key) - - def _recover_latent_from_input_ids(self, input_ids: torch.Tensor | None) -> torch.Tensor | None: - if input_ids is None or input_ids.numel() == 0: - return None - flat_ids = input_ids.detach().reshape(-1).to("cpu") - if flat_ids.numel() < 4 or int(flat_ids[0].item()) != _VOXCPM_LATENT_MAGIC: - return None - latent_dim = int(flat_ids[1].item()) - time_dim = int(flat_ids[2].item()) - payload = flat_ids[3:] - expected = latent_dim * time_dim - if latent_dim <= 0 or time_dim <= 0: - raise ValueError(f"Invalid VoxCPM latent header: latent_dim={latent_dim}, time_dim={time_dim}") - if int(payload.numel()) != expected: - raise ValueError( - "Invalid VoxCPM latent payload size: " - f"expected={expected}, actual={int(payload.numel())}, " - f"latent_dim={latent_dim}, time_dim={time_dim}" - ) - packed = payload.to(dtype=torch.int32).to(torch.uint16) - return packed.view(torch.bfloat16).to(torch.float32).reshape(1, latent_dim, time_dim) - - def _maybe_recover_vae_infos( - self, - infos: list[dict[str, Any]], - input_ids: torch.Tensor | None, - *, - async_chunk: bool, - ) -> list[dict[str, Any]]: - if not async_chunk: - return infos - if any(self._extract_val(info, "latent_audio_feat", None) is not None for info in infos): - return infos - recovered = self._recover_latent_from_input_ids(input_ids) - if recovered is None: - return infos - return [{"latent_audio_feat": recovered}] - - @staticmethod - def _normalize_audio_samples(samples: Any) -> np.ndarray: - if isinstance(samples, torch.Tensor): - return samples.detach().cpu().float().reshape(-1).numpy() - return np.asarray(samples, dtype=np.float32).reshape(-1) - - @classmethod - def _normalize_ref_audio(cls, ref_audio: Any) -> tuple[np.ndarray, int]: - if isinstance(ref_audio, str): - raise TypeError("String ref_audio should be handled as a path before waveform normalization.") - - if isinstance(ref_audio, dict): - sample_rate = ref_audio.get("sample_rate") or ref_audio.get("sampling_rate") or ref_audio.get("sr") - samples = None - for key in ("audio", "wav", "samples", "array", "waveform"): - if key in ref_audio and ref_audio[key] is not None: - samples = ref_audio[key] - break - if sample_rate is None or samples is None: - raise ValueError("ref_audio dict must contain waveform data and sample rate.") - return cls._normalize_audio_samples(samples), int(sample_rate) - - if isinstance(ref_audio, (list, tuple)): - if len(ref_audio) == 1: - return cls._normalize_ref_audio(ref_audio[0]) - if len(ref_audio) == 2 and np.isscalar(ref_audio[1]): - return cls._normalize_audio_samples(ref_audio[0]), int(ref_audio[1]) - - raise TypeError(f"Unsupported ref_audio format: {type(ref_audio)!r}") - - @staticmethod - def _write_temp_prompt_wav(waveform: np.ndarray, sample_rate: int) -> str: - prompt_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav") - prompt_file.close() - - wav = np.asarray(waveform, dtype=np.float32).reshape(-1) - wav = np.clip(wav, -1.0, 1.0) - pcm16 = (wav * 32767.0).astype(np.int16) - with wave.open(prompt_file.name, "wb") as wav_file: - wav_file.setnchannels(1) - wav_file.setsampwidth(2) - wav_file.setframerate(int(sample_rate)) - wav_file.writeframes(pcm16.tobytes()) - - return prompt_file.name - - @classmethod - def _resolve_prompt_inputs(cls, info: dict[str, Any]) -> tuple[str | None, str | None, str | None]: - prompt_text = cls._extract_val(info, "prompt_text", None) - prompt_wav_path = cls._extract_val(info, "prompt_wav_path", None) - if prompt_wav_path: - if prompt_text is None: - prompt_text = cls._extract_val(info, "ref_text", None) - return prompt_wav_path, prompt_text, None - - ref_audio = cls._extract_val(info, "ref_audio", None) - ref_text = cls._extract_val(info, "ref_text", None) - if ref_audio is None or ref_text is None: - return None, None, None - if isinstance(ref_audio, str): - return ref_audio, ref_text, None - - waveform, sample_rate = cls._normalize_ref_audio(ref_audio) - temp_prompt_wav = cls._write_temp_prompt_wav(waveform, sample_rate) - return temp_prompt_wav, ref_text, temp_prompt_wav - - def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: - if input_ids.numel() == 0: - return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32) - return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32) - - def _get_vocab_size(self) -> int: - model_config = getattr(self.vllm_config, "model_config", None) - if model_config is not None: - getter = getattr(model_config, "get_vocab_size", None) - if callable(getter): - try: - return int(getter()) - except Exception: - pass - hf_config = getattr(model_config, "hf_text_config", None) - if hf_config is not None and hasattr(hf_config, "vocab_size"): - return int(hf_config.vocab_size) - return 32000 - - def _make_empty_output( - self, - *, - output_key: str, - payload_factory: Callable[[], torch.Tensor], - infos: list[dict[str, Any]], - sample_rate: int, - out_device: torch.device, - out_dtype: torch.dtype, - hidden_rows: int | None = None, - ) -> OmniOutput: - if hidden_rows is None: - hidden_rows = len(infos) - return OmniOutput( - text_hidden_states=torch.zeros((hidden_rows, 1), device=out_device, dtype=out_dtype), - multimodal_outputs={ - output_key: [payload_factory() for _ in infos], - "sr": [torch.tensor(sample_rate, dtype=torch.int32) for _ in infos], - }, - ) - - def _finalize_stage_output( - self, - *, - output_key: str, - outputs: list[torch.Tensor], - sample_rates: list[torch.Tensor], - out_device: torch.device, - out_dtype: torch.dtype, - hidden_rows: int | None = None, - ) -> OmniOutput: - multimodal_outputs: dict[str, Any] = {output_key: outputs, "sr": sample_rates} - if hidden_rows is not None: - text_hidden_states = torch.zeros((hidden_rows, 1), device=out_device, dtype=out_dtype) - elif outputs: - outputs_tensor = torch.stack(outputs) - text_hidden_states = ( - outputs_tensor.unsqueeze(-1) - if outputs_tensor.ndim == 1 - else outputs_tensor.reshape(-1, outputs_tensor.shape[-1]) - ) - else: - text_hidden_states = torch.zeros((0, 1), device=out_device, dtype=out_dtype) - text_hidden_states = text_hidden_states.to(device=out_device, dtype=out_dtype) - return OmniOutput( - text_hidden_states=text_hidden_states, - multimodal_outputs=multimodal_outputs, - ) - - def _forward_vae_stage( - self, - infos: list[dict[str, Any]], - *, - sample_rate: int, - async_chunk: bool, - out_device: torch.device, - out_dtype: torch.dtype, - ) -> OmniOutput: - if all(self._extract_val(info, "latent_audio_feat", None) is None for info in infos): - self._ar_emit_stop_token = True - return self._make_empty_output( - output_key="model_outputs", - payload_factory=lambda: torch.zeros((0,), dtype=torch.float32), - infos=infos, - sample_rate=sample_rate, - out_device=out_device, - out_dtype=out_dtype, - ) - - outputs: list[torch.Tensor] = [] - sample_rates: list[torch.Tensor] = [] - for info in infos: - latent_audio_feat = self._extract_val(info, "latent_audio_feat", None) - audio_tensor = self._pipeline.decode(latent_audio_feat, trim_streaming_patch=async_chunk) - outputs.append(audio_tensor.float().cpu()) - sample_rates.append(torch.tensor(sample_rate, dtype=torch.int32)) - - self._ar_emit_stop_token = True - return self._finalize_stage_output( - output_key="model_outputs", - outputs=outputs, - sample_rates=sample_rates, - out_device=out_device, - out_dtype=out_dtype, - ) - - def _forward_latent_stage( - self, - infos: list[dict[str, Any]], - *, - sample_rate: int, - async_chunk: bool, - out_device: torch.device, - out_dtype: torch.dtype, - hidden_rows: int, - ) -> OmniOutput: - texts = [self._extract_val(info, "text", "") for info in infos] - if all(not text for text in texts): - self._ar_emit_stop_token = True - return self._make_empty_output( - output_key="latent_audio_feat", - payload_factory=lambda: torch.zeros((0,), dtype=torch.float32), - infos=infos, - sample_rate=sample_rate, - out_device=out_device, - out_dtype=out_dtype, - hidden_rows=hidden_rows, - ) - - outputs: list[torch.Tensor] = [] - sample_rates: list[torch.Tensor] = [] - last_chunk_flags: list[bool] | None = [] if async_chunk else None - payload_finished_flags: list[bool] | None = [] if async_chunk else None - for info in infos: - text = self._extract_val(info, "text", "") - cfg_value = float(self._extract_val(info, "cfg_value", 2.0)) - inference_timesteps = int(self._extract_val(info, "inference_timesteps", 10)) - min_len = int(self._extract_val(info, "min_len", 2)) - max_len = int(self._extract_val(info, "max_len", self._extract_val(info, "max_new_tokens", 4096))) - retry_badcase = bool(self._extract_val(info, "retry_badcase", True)) - retry_badcase_max_times = int(self._extract_val(info, "retry_badcase_max_times", 3)) - retry_badcase_ratio_threshold = float(self._extract_val(info, "retry_badcase_ratio_threshold", 6.0)) - streaming_prefix_len = int(self._extract_val(info, "streaming_prefix_len", 3)) - - request_key = self._resolve_stream_request_key(info) - created_temp: str | None = None - - if async_chunk: - terminal_pending = self._latent_stream_terminal_pending.get(request_key, 0) - if terminal_pending > 0: - outputs.append(torch.zeros((0,), dtype=torch.float32)) - assert last_chunk_flags is not None - last_chunk_flags.append(True) - assert payload_finished_flags is not None - payload_finished_flags.append(terminal_pending == 1) - if terminal_pending == 1: - self._latent_stream_terminal_pending.pop(request_key, None) - else: - self._latent_stream_terminal_pending[request_key] = terminal_pending - 1 - sample_rates.append(torch.tensor(sample_rate, dtype=torch.int32)) - continue - - if request_key in self._latent_stream_completed: - outputs.append(torch.zeros((0,), dtype=torch.float32)) - assert last_chunk_flags is not None - last_chunk_flags.append(True) - assert payload_finished_flags is not None - payload_finished_flags.append(False) - sample_rates.append(torch.tensor(sample_rate, dtype=torch.int32)) - continue - - if request_key not in self._latent_stream_gens: - prompt_wav_path, prompt_text, temp_prompt_wav = self._resolve_prompt_inputs(info) - created_temp = temp_prompt_wav - self._latent_stream_gens[request_key] = self._pipeline.iter_latent_chunks_streaming( - text=text, - prompt_wav_path=prompt_wav_path, - prompt_text=prompt_text, - cfg_value=cfg_value, - inference_timesteps=inference_timesteps, - min_len=min_len, - max_len=max_len, - streaming_prefix_len=streaming_prefix_len, - retry_badcase=False, - retry_badcase_max_times=retry_badcase_max_times, - retry_badcase_ratio_threshold=retry_badcase_ratio_threshold, - ) - generator = self._latent_stream_gens[request_key] - try: - chunk_latent, is_last = next(generator) - except StopIteration: - self._latent_stream_gens.pop(request_key, None) - self._latent_stream_terminal_pending[request_key] = 1 - self._latent_stream_completed.add(request_key) - outputs.append(torch.zeros((0,), dtype=torch.float32)) - assert last_chunk_flags is not None - last_chunk_flags.append(True) - assert payload_finished_flags is not None - payload_finished_flags.append(True) - else: - if is_last: - self._latent_stream_gens.pop(request_key, None) - self._latent_stream_terminal_pending[request_key] = 1 - self._latent_stream_completed.add(request_key) - outputs.append(chunk_latent.detach().float().cpu()) - assert last_chunk_flags is not None - last_chunk_flags.append(bool(is_last)) - assert payload_finished_flags is not None - payload_finished_flags.append(False) - finally: - if created_temp is not None and os.path.exists(created_temp): - os.unlink(created_temp) - sample_rates.append(torch.tensor(sample_rate, dtype=torch.int32)) - continue - - prompt_wav_path, prompt_text, temp_prompt_wav = self._resolve_prompt_inputs(info) - try: - latent_audio_feat = self._pipeline.generate_latents( - text=text, - prompt_wav_path=prompt_wav_path, - prompt_text=prompt_text, - cfg_value=cfg_value, - inference_timesteps=inference_timesteps, - min_len=min_len, - max_len=max_len, - retry_badcase=retry_badcase, - retry_badcase_max_times=retry_badcase_max_times, - retry_badcase_ratio_threshold=retry_badcase_ratio_threshold, - ) - outputs.append(latent_audio_feat.float().cpu()) - finally: - if temp_prompt_wav is not None and os.path.exists(temp_prompt_wav): - os.unlink(temp_prompt_wav) - - sample_rates.append(torch.tensor(sample_rate, dtype=torch.int32)) - - self._ar_emit_stop_token = all(last_chunk_flags) if async_chunk and last_chunk_flags else True - output = self._finalize_stage_output( - output_key="latent_audio_feat", - outputs=outputs, - sample_rates=sample_rates, - out_device=out_device, - out_dtype=out_dtype, - hidden_rows=hidden_rows, - ) - if async_chunk and payload_finished_flags is not None: - output.multimodal_outputs["finished"] = [ - torch.tensor(flag, dtype=torch.bool) for flag in payload_finished_flags - ] - return output - - def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> torch.Tensor: - del sampling_metadata - if isinstance(hidden_states, OmniOutput): - hidden_states = hidden_states.text_hidden_states - if hidden_states is None: - device, dtype = self._runner_hidden_device_dtype() - hidden_states = torch.zeros((0, 1), device=device, dtype=dtype) - if hidden_states.ndim == 1: - hidden_states = hidden_states.unsqueeze(-1) - elif hidden_states.ndim > 2: - hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) - - vocab_size = self._get_vocab_size() - num_rows = int(hidden_states.shape[0]) - logits = torch.zeros((num_rows, vocab_size), dtype=torch.float32, device=hidden_states.device) - eos_id = 2 if vocab_size > 2 else 0 - safe_id = 1 if vocab_size > 1 and 1 != eos_id else 0 - emit_stop = getattr(self, "_ar_emit_stop_token", True) - if num_rows > 0: - if emit_stop: - logits[:, eos_id] = 1.0e6 - else: - logits[:, eos_id] = -1.0e9 - logits[:, safe_id] = 1.0e6 - return logits - - @torch.no_grad() - def forward( - self, - input_ids: torch.Tensor | None = None, - positions: torch.Tensor | None = None, - intermediate_tensors: Any = None, - inputs_embeds: torch.Tensor | None = None, - runtime_additional_information: list[dict[str, Any]] | None = None, - model_intermediate_buffer: list[dict[str, Any]] | None = None, - **kwargs: Any, - ) -> OmniOutput: - del positions, intermediate_tensors, inputs_embeds, kwargs - self._ensure_model_loaded() - out_device, out_dtype = self._runner_hidden_device_dtype() - if input_ids is not None and input_ids.device.type == out_device.type: - out_device = input_ids.device - - infos = model_intermediate_buffer or runtime_additional_information or [{}] - hidden_rows = len(infos) - if input_ids is not None and len(input_ids.shape) > 0: - hidden_rows = max(hidden_rows, int(input_ids.shape[0])) - sample_rate = int(getattr(self._pipeline, "sample_rate", 24000)) - async_chunk = bool(getattr(self.vllm_config.model_config, "async_chunk", False)) - if self.model_stage in self._VAE_STAGES: - infos = self._maybe_recover_vae_infos(infos, input_ids, async_chunk=async_chunk) - return self._forward_vae_stage( - infos, - sample_rate=sample_rate, - async_chunk=async_chunk, - out_device=out_device, - out_dtype=out_dtype, - ) - if self.model_stage in self._LATENT_STAGES: - return self._forward_latent_stage( - infos, - sample_rate=sample_rate, - async_chunk=async_chunk, - out_device=out_device, - out_dtype=out_dtype, - hidden_rows=hidden_rows, - ) - raise ValueError(f"Unsupported VoxCPM model_stage at runtime: {self.model_stage}") - - def make_empty_intermediate_tensors( - self, batch_size: int, dtype: torch.dtype, device: torch.device - ) -> IntermediateTensors: - del batch_size, dtype, device - return {} - - -__all__ = ["VoxCPMForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/voxcpm/voxcpm_loader.py b/vllm_omni/model_executor/models/voxcpm/voxcpm_loader.py deleted file mode 100644 index dac7117ca..000000000 --- a/vllm_omni/model_executor/models/voxcpm/voxcpm_loader.py +++ /dev/null @@ -1,247 +0,0 @@ -from __future__ import annotations - -import importlib -import json -import os -import shutil -import sys -import tempfile -from contextlib import contextmanager -from hashlib import sha256 -from pathlib import Path -from typing import Any -from unittest.mock import patch - -import numpy as np -import torch -from vllm.config import VllmConfig -from vllm.logger import init_logger - -logger = init_logger(__name__) - - -def _iter_voxcpm_src_candidates() -> list[Path]: - candidates: list[Path] = [] - env_path = os.environ.get("VLLM_OMNI_VOXCPM_CODE_PATH") - if env_path: - candidates.append(Path(env_path).expanduser()) - - repo_root = Path(__file__).resolve().parents[4] - candidates.append(repo_root.parent / "VoxCPM" / "src") - - unique_candidates: list[Path] = [] - seen: set[str] = set() - for candidate in candidates: - candidate_key = str(candidate) - if candidate_key in seen: - continue - seen.add(candidate_key) - unique_candidates.append(candidate) - return unique_candidates - - -def _prepend_voxcpm_src(candidate: Path) -> None: - candidate_str = str(candidate) - if candidate_str not in sys.path: - sys.path.insert(0, candidate_str) - - -def _import_voxcpm_attrs(module_name: str, *attr_names: str) -> tuple[Any, ...]: - last_exc: ImportError | None = None - for candidate in _iter_voxcpm_src_candidates(): - if not candidate.exists(): - continue - _prepend_voxcpm_src(candidate) - try: - module = importlib.import_module(module_name) - return tuple(getattr(module, attr_name) for attr_name in attr_names) - except ImportError as exc: - last_exc = exc - - try: - module = importlib.import_module(module_name) - return tuple(getattr(module, attr_name) for attr_name in attr_names) - except ImportError as exc: - last_exc = exc - - raise ImportError(f"Failed to import {module_name}.") from last_exc - - -def _import_voxcpm_base_model_class(): - """Import upstream ``VoxCPMModel`` from ``VoxCPM/src/voxcpm`` (env, sibling tree, or pip).""" - try: - (VoxCPMModel,) = _import_voxcpm_attrs("voxcpm.model.voxcpm", "VoxCPMModel") - return VoxCPMModel - except ImportError as exc: - raise ImportError( - "Failed to import VoxCPMModel. Install the `voxcpm` package or set " - "`VLLM_OMNI_VOXCPM_CODE_PATH` to the VoxCPM repository `src` directory " - "(the parent of the `voxcpm` package that contains `model/` and `modules/`)." - ) from exc - - -def _import_voxcpm_audio_vae_classes(): - try: - return _import_voxcpm_attrs("voxcpm.modules.audiovae", "AudioVAE", "AudioVAEConfig") - except ImportError as exc: - raise ImportError( - "Failed to import VoxCPM AudioVAE. Install the `voxcpm` package or set " - "`VLLM_OMNI_VOXCPM_CODE_PATH` to the VoxCPM repository `src` directory." - ) from exc - - -def _device_to_string(device: torch.device) -> str: - if device.index is None: - return device.type - return f"{device.type}:{device.index}" - - -def _normalize_dtype_name(dtype: Any) -> str | None: - if dtype is None: - return None - if isinstance(dtype, torch.dtype): - mapping = { - torch.bfloat16: "bfloat16", - torch.float16: "float16", - torch.float32: "float32", - } - return mapping.get(dtype, str(dtype).removeprefix("torch.")) - dtype_str = str(dtype) - return dtype_str.removeprefix("torch.") - - -def _resolve_runtime_device(vllm_config: VllmConfig) -> torch.device: - try: - from vllm_omni.platforms import current_omni_platform - - return current_omni_platform.get_torch_device() - except Exception: - pass - - device = getattr(getattr(vllm_config, "device_config", None), "device", None) - if isinstance(device, torch.device): - return device - if device: - return torch.device(device) - return torch.device("cpu") - - -def _prepare_runtime_model_dir( - model_path: str | Path, - *, - target_device: torch.device, - target_dtype: str | None, -) -> str: - source_dir = Path(model_path) - config_path = source_dir / "config.json" - if not config_path.exists(): - return str(source_dir) - - config_text = config_path.read_text() - config_dict = json.loads(config_text) - desired_device = target_device.type - desired_dtype = target_dtype or config_dict.get("dtype") - - if config_dict.get("device") == desired_device and config_dict.get("dtype") == desired_dtype: - return str(source_dir) - - digest = sha256(f"{source_dir.resolve()}:{config_text}:{desired_device}:{desired_dtype}".encode()).hexdigest()[:16] - runtime_dir = Path(tempfile.gettempdir()) / "vllm_omni_voxcpm_runtime" / digest - runtime_dir.mkdir(parents=True, exist_ok=True) - - for entry in source_dir.iterdir(): - target = runtime_dir / entry.name - if entry.name == "config.json" or target.exists(): - continue - try: - target.symlink_to(entry, target_is_directory=entry.is_dir()) - except OSError as exc: - logger.warning( - "Falling back to copying VoxCPM runtime artifact %s into %s because symlink creation failed: %s", - entry, - runtime_dir, - exc, - ) - if entry.is_dir(): - shutil.copytree(entry, target, dirs_exist_ok=True) - else: - shutil.copy2(entry, target) - - patched_config = dict(config_dict) - patched_config["device"] = desired_device - if desired_dtype is not None: - patched_config["dtype"] = desired_dtype - (runtime_dir / "config.json").write_text(json.dumps(patched_config, indent=2, sort_keys=True)) - return str(runtime_dir) - - -@contextmanager -def _force_cuda_available_for_npu(device: torch.device): - if device.type != "npu": - yield - return - - with patch("torch.cuda.is_available", return_value=True): - yield - - -def _is_torchcodec_load_error(exc: BaseException) -> bool: - message = str(exc).lower() - return "torchcodec" in message or "load_with_torchcodec" in message - - -def _load_audio_with_soundfile( - prompt_wav_path: str, - *, - sample_rate: int, -) -> torch.Tensor: - try: - import soundfile as sf - except ImportError: - raise - - audio_np, source_sr = sf.read(prompt_wav_path, dtype="float32", always_2d=True) - audio = torch.from_numpy(np.ascontiguousarray(audio_np.T)) - - if audio.size(0) > 1: - audio = audio.mean(dim=0, keepdim=True) - - if int(source_sr) != int(sample_rate): - try: - import torchaudio - except ImportError as exc: - raise ImportError("torchaudio is required for resampling prompt audio.") from exc - audio = torchaudio.functional.resample(audio, int(source_sr), int(sample_rate)) - - return audio - - -def _build_prompt_cache_with_soundfile(model: Any, *args: Any, **kwargs: Any) -> dict[str, Any]: - if args: - prompt_text = args[0] - prompt_wav_path = args[1] if len(args) > 1 else kwargs.get("prompt_wav_path") - else: - prompt_text = kwargs.get("prompt_text") - prompt_wav_path = kwargs.get("prompt_wav_path") - - if not prompt_text or not prompt_wav_path: - raise ValueError("prompt_text and prompt_wav_path are required") - - audio = _load_audio_with_soundfile(prompt_wav_path, sample_rate=int(model.sample_rate)) - - patch_len = model.patch_size * model.chunk_size - if audio.size(1) % patch_len != 0: - padding_size = patch_len - audio.size(1) % patch_len - audio = torch.nn.functional.pad(audio, (padding_size, 0)) - - audio_feat = model.audio_vae.encode(audio.to(model.device), model.sample_rate).cpu() - audio_feat = audio_feat.view( - model.audio_vae.latent_dim, - -1, - model.patch_size, - ).permute(1, 2, 0) - - return { - "prompt_text": prompt_text, - "audio_feat": audio_feat, - } diff --git a/vllm_omni/model_executor/models/voxcpm/voxcpm_runtime_utils.py b/vllm_omni/model_executor/models/voxcpm/voxcpm_runtime_utils.py deleted file mode 100644 index 36b4282c2..000000000 --- a/vllm_omni/model_executor/models/voxcpm/voxcpm_runtime_utils.py +++ /dev/null @@ -1,44 +0,0 @@ -from __future__ import annotations - -import json -import shutil -from pathlib import Path - - -def resolve_voxcpm_model_dir(model: str) -> Path: - model_path = Path(model).expanduser() - if model_path.exists(): - return model_path - - from huggingface_hub import snapshot_download - - return Path(snapshot_download(repo_id=model)) - - -def prepare_voxcpm_hf_config_dir(model_dir: str | Path, hf_config_dir: str | Path) -> Path: - model_dir = Path(model_dir).expanduser() - hf_config_dir = Path(hf_config_dir).expanduser() - hf_config_dir.mkdir(parents=True, exist_ok=True) - - source_config_path = model_dir / "config.json" - if not source_config_path.exists(): - raise FileNotFoundError(f"VoxCPM config.json not found under {model_dir}") - - config_path = hf_config_dir / "config.json" - shutil.copy2(source_config_path, config_path) - - source_generation_config_path = model_dir / "generation_config.json" - if source_generation_config_path.exists(): - shutil.copy2(source_generation_config_path, hf_config_dir / "generation_config.json") - - config_dict = json.loads(config_path.read_text(encoding="utf-8")) - config_dict["model_type"] = "voxcpm" - config_dict.setdefault("architectures", ["VoxCPMForConditionalGeneration"]) - config_path.write_text(json.dumps(config_dict, indent=2, ensure_ascii=False), encoding="utf-8") - return hf_config_dir - - -__all__ = [ - "prepare_voxcpm_hf_config_dir", - "resolve_voxcpm_model_dir", -] diff --git a/vllm_omni/model_executor/models/voxcpm/voxcpm_stage_wrappers.py b/vllm_omni/model_executor/models/voxcpm/voxcpm_stage_wrappers.py deleted file mode 100644 index f4446c796..000000000 --- a/vllm_omni/model_executor/models/voxcpm/voxcpm_stage_wrappers.py +++ /dev/null @@ -1,185 +0,0 @@ -from __future__ import annotations - -import os -from collections.abc import Generator -from typing import Any - -import torch -import torch.nn as nn -from einops import rearrange - - -class _DirectVoxCPMLatentGenerator: - def __init__(self, tts_model: Any): - self.tts_model = tts_model - self.sample_rate = int(getattr(tts_model, "sample_rate", 24000)) - - def generate_latents( - self, - *, - text: str, - prompt_wav_path: str | None = None, - prompt_text: str | None = None, - cfg_value: float = 2.0, - inference_timesteps: int = 10, - min_len: int = 2, - max_len: int = 4096, - retry_badcase: bool = True, - retry_badcase_max_times: int = 3, - retry_badcase_ratio_threshold: float = 6.0, - ) -> torch.Tensor: - if not isinstance(text, str) or not text.strip(): - raise ValueError("target text must be a non-empty string") - if (prompt_wav_path is None) != (prompt_text is None): - raise ValueError("prompt_wav_path and prompt_text must both be provided or both be None") - if prompt_wav_path is not None and not os.path.exists(prompt_wav_path): - raise FileNotFoundError(f"prompt_wav_path does not exist: {prompt_wav_path}") - - prompt_cache = None - if prompt_wav_path is not None and prompt_text is not None: - prompt_cache = self.tts_model.build_prompt_cache( - prompt_text=prompt_text, - prompt_wav_path=prompt_wav_path, - ) - - gen_kw = dict( - target_text=" ".join(text.split()), - prompt_cache=prompt_cache, - min_len=min_len, - max_len=max_len, - inference_timesteps=inference_timesteps, - cfg_value=cfg_value, - retry_badcase=retry_badcase, - retry_badcase_max_times=retry_badcase_max_times, - retry_badcase_ratio_threshold=retry_badcase_ratio_threshold, - ) - latent_entry = getattr(self.tts_model, "generate_latents_with_prompt_cache", None) - if latent_entry is not None: - _, _, pred_audio_feat = latent_entry(**gen_kw) - else: - try: - _, _, pred_audio_feat = self.tts_model.generate_with_prompt_cache( - **gen_kw, - latents_only=True, - ) - except TypeError: - _, _, pred_audio_feat = self.tts_model.generate_with_prompt_cache(**gen_kw) - return pred_audio_feat.detach().cpu().to(torch.float32) - - def iter_latent_chunks_streaming( - self, - *, - text: str, - prompt_wav_path: str | None = None, - prompt_text: str | None = None, - cfg_value: float = 2.0, - inference_timesteps: int = 10, - min_len: int = 2, - max_len: int = 4096, - streaming_prefix_len: int = 3, - retry_badcase: bool = False, - retry_badcase_max_times: int = 3, - retry_badcase_ratio_threshold: float = 6.0, - ) -> Generator[tuple[torch.Tensor, bool], None, None]: - """Yield ``(latent_window, is_last_chunk)`` for Omni async_chunk latent to VAE.""" - if not isinstance(text, str) or not text.strip(): - raise ValueError("target text must be a non-empty string") - if (prompt_wav_path is None) != (prompt_text is None): - raise ValueError("prompt_wav_path and prompt_text must both be provided or both be None") - if prompt_wav_path is not None and not os.path.exists(prompt_wav_path): - raise FileNotFoundError(f"prompt_wav_path does not exist: {prompt_wav_path}") - - prompt_cache = None - if prompt_wav_path is not None and prompt_text is not None: - prompt_cache = self.tts_model.build_prompt_cache( - prompt_text=prompt_text, - prompt_wav_path=prompt_wav_path, - ) - - gen_kw = dict( - target_text=" ".join(text.split()), - prompt_cache=prompt_cache, - min_len=min_len, - max_len=max_len, - inference_timesteps=inference_timesteps, - cfg_value=cfg_value, - retry_badcase=retry_badcase, - retry_badcase_max_times=retry_badcase_max_times, - retry_badcase_ratio_threshold=retry_badcase_ratio_threshold, - streaming_prefix_len=streaming_prefix_len, - ) - stream_entry = getattr(self.tts_model, "generate_latents_with_prompt_cache_streaming", None) - if stream_entry is not None: - gen = stream_entry(**gen_kw) - else: - fallback_stream_entry = getattr(self.tts_model, "generate_with_prompt_cache_streaming", None) - if fallback_stream_entry is not None: - gen = fallback_stream_entry(**gen_kw, latents_only=True) - else: - gen = self.tts_model._generate_with_prompt_cache(streaming=True, latents_only=True, **gen_kw) - - iterator = iter(gen) - previous = next(iterator, None) - while previous is not None: - current = next(iterator, None) - _, _target_tok, chunk_latent = previous - if not isinstance(chunk_latent, torch.Tensor): - chunk_latent = torch.as_tensor(chunk_latent) - yield chunk_latent, current is None - previous = current - - -class _DirectVoxCPMAudioVAE: - def __init__(self, audio_vae: nn.Module, *, patch_size: int = 2): - self.audio_vae = audio_vae - self.sample_rate = int(getattr(audio_vae, "sample_rate", 24000)) - self.latent_dim = int(getattr(audio_vae, "latent_dim", 64)) - self.patch_size = int(patch_size) - self._chunk_size = int(getattr(audio_vae, "chunk_size", 1)) - self._stream_audio_patch_samples = max(1, self.patch_size * self._chunk_size) - - def _prepare_latents_for_decode(self, latent_audio_feat: Any) -> torch.Tensor: - latents = latent_audio_feat - if not isinstance(latents, torch.Tensor): - latents = torch.tensor(latents, dtype=torch.float32) - latents = latents.detach().to(torch.float32) - - if latents.ndim == 3: - if latents.shape[-1] == self.latent_dim: - latents = rearrange(latents, "t p d -> 1 d (t p)") - elif latents.shape[1] == self.latent_dim: - latents = latents.contiguous() - else: - raise ValueError(f"Unsupported latent_audio_feat shape: {tuple(latents.shape)}") - elif latents.ndim == 2: - if latents.shape[0] == self.latent_dim: - latents = latents.unsqueeze(0) - elif latents.shape[1] == self.latent_dim: - latents = rearrange(latents, "t d -> 1 d t") - else: - raise ValueError(f"Unsupported latent_audio_feat shape: {tuple(latents.shape)}") - else: - raise ValueError(f"Unsupported latent_audio_feat ndim: {latents.ndim}") - - return latents - - @torch.no_grad() - def decode(self, latent_audio_feat: Any, *, trim_streaming_patch: bool = False) -> torch.Tensor: - latents = self._prepare_latents_for_decode(latent_audio_feat) - device = next(self.audio_vae.parameters()).device - raw = self.audio_vae.decode(latents.to(device=device, dtype=torch.float32)) - if isinstance(raw, dict): - audio = raw.get("audio") - if audio is None: - audio = next(v for v in raw.values() if isinstance(v, torch.Tensor)) - else: - audio = raw - if audio.dim() == 3: - stream = audio.squeeze(1) - elif audio.dim() == 2: - stream = audio - else: - stream = audio.reshape(audio.shape[0], -1) - if trim_streaming_patch: - stream = stream[..., -self._stream_audio_patch_samples :] - return stream.reshape(-1).detach().cpu().to(torch.float32) diff --git a/vllm_omni/model_executor/stage_configs/qwen3_tts_uniproc.yaml b/vllm_omni/model_executor/stage_configs/qwen3_tts_uniproc.yaml deleted file mode 100644 index c370646a1..000000000 --- a/vllm_omni/model_executor/stage_configs/qwen3_tts_uniproc.yaml +++ /dev/null @@ -1,95 +0,0 @@ -async_chunk: true -stage_args: - - stage_id: 0 - stage_type: llm - is_comprehension: true - runtime: - devices: "0" - engine_args: - model_stage: qwen3_tts - max_num_seqs: 10 - model_arch: Qwen3TTSTalkerForConditionalGeneration - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - enforce_eager: false - trust_remote_code: true - async_scheduling: true - enable_prefix_caching: false - engine_output_type: latent - gpu_memory_utilization: 0.3 - max_num_batched_tokens: 512 - max_model_len: 4096 - custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk - # Use named connector to apply runtime.connectors.extra. - output_connectors: - to_stage_1: connector_of_shared_memory - default_sampling_params: - temperature: 0.9 - top_k: 50 - max_tokens: 4096 - seed: 42 - detokenize: false - repetition_penalty: 1.05 - stop_token_ids: [2150] - - - stage_id: 1 - stage_type: llm - runtime: - devices: "0" - engine_args: - model_stage: code2wav - max_num_seqs: 10 - model_arch: Qwen3TTSCode2Wav - worker_type: generation - scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler - enforce_eager: true - trust_remote_code: true - async_scheduling: true - enable_prefix_caching: false - engine_output_type: audio - gpu_memory_utilization: 0.3 - # Must be divisible by num_code_groups and cover (left_context + chunk). - # Prefill length is Q * num_frames (e.g. 16 * 2148 = 34368); keep headroom past 32k. - max_num_batched_tokens: 65536 - # async_chunk appends windows per step; max_model_len must cover accumulated flat codec stream. - max_model_len: 65536 - engine_input_source: [0] - final_output: true - final_output_type: audio - # Distributed connector configuration - input_connectors: - from_stage_0: connector_of_shared_memory - tts_args: - max_instructions_length: 500 - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 65536 - seed: 42 - detokenize: true - repetition_penalty: 1.0 - -runtime: - enabled: true - - connectors: - connector_of_shared_memory: - name: SharedMemoryConnector - extra: - shm_threshold_bytes: 65536 - # Frame-aligned codec streaming transport. - codec_streaming: true - # Connector polling / timeout (unit: loop count, sleep interval in seconds). - connector_get_sleep_s: 0.01 - connector_get_max_wait_first_chunk: 3000 - connector_get_max_wait: 300 - # Match the decoder sliding attention window to avoid chunk-boundary noise. - codec_chunk_frames: 25 - codec_left_context_frames: 72 - # Emit only the first audio chunk early, then return to codec_chunk_frames. - initial_codec_chunk_frames: 1 - - edges: - - from: 0 - to: 1 diff --git a/vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml b/vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml deleted file mode 100644 index c6fd177a3..000000000 --- a/vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml +++ /dev/null @@ -1,98 +0,0 @@ -# VoxCPM two-stage streaming (align with qwen3_tts.yaml async_chunk pattern). -# Stage0 (latent_generator) emits latent in time chunks; Stage1 (VAE) decodes as chunks arrive. -async_chunk: true -stage_args: - - stage_id: 0 - stage_type: llm - is_comprehension: true - runtime: - devices: "0" - max_batch_size: 1 - engine_args: - dtype: bfloat16 - model_stage: latent_generator - model_arch: VoxCPMForConditionalGeneration - # Optional persistent HF-compatible config dir for native VoxCPM models. - hf_config_path: ${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,} - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - enforce_eager: true - trust_remote_code: true - async_scheduling: true - enable_prefix_caching: false - engine_output_type: latent - gpu_memory_utilization: 0.7 - distributed_executor_backend: "mp" - max_num_batched_tokens: 4096 - max_model_len: 4096 - custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.voxcpm.latent2vae_async_chunk - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 4096 - stop_token_ids: [2] - seed: 42 - detokenize: false - repetition_penalty: 1.0 - final_output: false - output_connectors: - to_stage_1: voxcpm_shm - - - stage_id: 1 - stage_type: llm - runtime: - devices: "0" - max_batch_size: 1 - engine_args: - dtype: float32 - model_stage: vae - model_arch: VoxCPMForConditionalGeneration - hf_config_path: ${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,} - worker_type: generation - scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler - enforce_eager: true - trust_remote_code: true - async_scheduling: false - enable_prefix_caching: false - engine_output_type: audio - gpu_memory_utilization: 0.15 - distributed_executor_backend: "mp" - max_num_batched_tokens: 8192 - max_model_len: 4096 - engine_input_source: [0] - final_output: true - final_output_type: audio - input_connectors: - from_stage_0: voxcpm_shm - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 128 - seed: 42 - detokenize: true - repetition_penalty: 1.0 - - -runtime: - enabled: true - - connectors: - voxcpm_shm: - name: SharedMemoryConnector - extra: - shm_threshold_bytes: 65536 - # Frame-aligned codec streaming transport. - codec_streaming: true - # Connector polling / timeout (unit: loop count, sleep interval in seconds). - connector_get_sleep_s: 0.01 - connector_get_max_wait_first_chunk: 3000 - connector_get_max_wait: 300 - # Align with Omni: small chunks with sufficient context overlap. - codec_chunk_frames: 1 - codec_left_context_frames: 1 - - edges: - - from: 0 - to: 1 diff --git a/vllm_omni/model_executor/stage_input_processors/voxcpm.py b/vllm_omni/model_executor/stage_input_processors/voxcpm.py deleted file mode 100644 index 125335cd2..000000000 --- a/vllm_omni/model_executor/stage_input_processors/voxcpm.py +++ /dev/null @@ -1,115 +0,0 @@ -from __future__ import annotations - -from typing import Any - -import torch -from vllm.inputs import TextPrompt - -from vllm_omni.data_entry_keys import ( - CodesStruct, - MetaStruct, - OmniPayloadStruct, -) -from vllm_omni.inputs.data import OmniTokensPrompt - -_VOXCPM_LATENT_MAGIC = 131071 - - -def _serialize_latent_to_codes(latent: Any) -> list[int]: - latent_tensor = latent if isinstance(latent, torch.Tensor) else torch.as_tensor(latent) - latent_tensor = latent_tensor.detach().cpu().contiguous() - if latent_tensor.ndim == 3: - if latent_tensor.shape[0] != 1: - raise ValueError(f"Expected batch=1 latent tensor, got shape={tuple(latent_tensor.shape)}") - latent_tensor = latent_tensor.squeeze(0) - if latent_tensor.ndim != 2: - raise ValueError(f"Unsupported latent_audio_feat shape for async chunk: {tuple(latent_tensor.shape)}") - latent_dim, time_dim = int(latent_tensor.shape[0]), int(latent_tensor.shape[1]) - packed = latent_tensor.to(torch.bfloat16).contiguous().view(torch.uint16).reshape(-1).to(torch.int32) - return [_VOXCPM_LATENT_MAGIC, latent_dim, time_dim, *packed.tolist()] - - -def _coerce_finished_flag(value: Any) -> bool: - """Normalize VoxCPM async-chunk finished markers to a Python bool.""" - if value is None: - return False - if isinstance(value, torch.Tensor): - if value.numel() != 1: - raise ValueError(f"finished tensor must be scalar, got shape={tuple(value.shape)}") - return bool(value.detach().cpu().item()) - if isinstance(value, (list, tuple)): - if not value: - return False - if len(value) != 1: - raise ValueError(f"finished container must have one element, got len={len(value)}") - return _coerce_finished_flag(value[0]) - return bool(value) - - -def latent2vae( - source_outputs: list[Any], - _prompt: OmniTokensPrompt | TextPrompt | None = None, - _requires_multimodal_data: bool = False, -) -> list[OmniTokensPrompt]: - vae_inputs: list[OmniTokensPrompt] = [] - for source_output in source_outputs: - output = source_output.outputs[0] - multimodal_output = getattr(output, "multimodal_output", None) - if not isinstance(multimodal_output, dict) or "latent_audio_feat" not in multimodal_output: - raise ValueError( - "VoxCPM latent stage output missing 'latent_audio_feat'. " - f"request_id={getattr(source_output, 'request_id', None)}" - ) - - additional_information = { - "latent_audio_feat": multimodal_output["latent_audio_feat"], - } - if "sr" in multimodal_output: - additional_information["sample_rate"] = [int(multimodal_output["sr"])] - - vae_inputs.append( - OmniTokensPrompt( - prompt_token_ids=[0], - additional_information=additional_information, - multi_modal_data=None, - mm_processor_kwargs=None, - ) - ) - - return vae_inputs - - -def _eof_payload() -> OmniPayloadStruct: - return OmniPayloadStruct( - codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), - meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)), - ) - - -def latent2vae_async_chunk( - transfer_manager: Any, - pooling_output: dict[str, Any] | None, - request: Any, - is_finished: bool = False, -) -> OmniPayloadStruct | None: - """Stage-0 latent → stage-1 VAE under ``async_chunk`` (connector payload).""" - # Kept for callback signature compatibility with OmniChunkTransferAdapter. - _ = transfer_manager - finished_request = _coerce_finished_flag(is_finished) - if callable(getattr(request, "is_finished", None)): - finished_request = finished_request or _coerce_finished_flag(request.is_finished()) - if not isinstance(pooling_output, dict): - return _eof_payload() if finished_request else None - - latent = pooling_output.get("latent_audio_feat") - if isinstance(latent, torch.Tensor) and latent.numel() == 0: - latent = None - - if latent is None: - return _eof_payload() if finished_request else None - - serialized_codes = _serialize_latent_to_codes(latent) - return OmniPayloadStruct( - codes=CodesStruct(audio=torch.tensor(serialized_codes, dtype=torch.long)), - meta=MetaStruct(finished=torch.tensor(finished_request, dtype=torch.bool)), - ) diff --git a/vllm_omni/platforms/npu/stage_configs/voxcpm.yaml b/vllm_omni/platforms/npu/stage_configs/voxcpm.yaml deleted file mode 100644 index dcd1f4051..000000000 --- a/vllm_omni/platforms/npu/stage_configs/voxcpm.yaml +++ /dev/null @@ -1,67 +0,0 @@ -stage_args: - - stage_id: 0 - stage_type: llm - is_comprehension: true - runtime: - devices: "0" - max_batch_size: 1 - engine_args: - dtype: bfloat16 - model_stage: latent_generator - model_arch: VoxCPMForConditionalGeneration - # Optional persistent HF-compatible config dir for native VoxCPM models. - hf_config_path: ${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,} - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - enforce_eager: true - trust_remote_code: true - async_scheduling: false - enable_prefix_caching: false - engine_output_type: latent - gpu_memory_utilization: 0.75 - distributed_executor_backend: "mp" - max_num_batched_tokens: 4096 - max_model_len: 4096 - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 4096 - seed: 42 - detokenize: false - repetition_penalty: 1.0 - final_output: false - - - stage_id: 1 - stage_type: llm - runtime: - devices: "0" - max_batch_size: 1 - engine_args: - dtype: float32 - model_stage: vae - model_arch: VoxCPMForConditionalGeneration - hf_config_path: ${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,} - worker_type: generation - scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler - enforce_eager: true - trust_remote_code: true - async_scheduling: false - enable_prefix_caching: false - engine_output_type: audio - gpu_memory_utilization: 0.1 - distributed_executor_backend: "mp" - max_num_batched_tokens: 8192 - max_model_len: 4096 - engine_input_source: [0] - custom_process_input_func: vllm_omni.model_executor.stage_input_processors.voxcpm.latent2vae - final_output: true - final_output_type: audio - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 1 - seed: 42 - detokenize: true - repetition_penalty: 1.0 diff --git a/vllm_omni/platforms/npu/stage_configs/voxcpm_async_chunk.yaml b/vllm_omni/platforms/npu/stage_configs/voxcpm_async_chunk.yaml deleted file mode 100644 index 87843634c..000000000 --- a/vllm_omni/platforms/npu/stage_configs/voxcpm_async_chunk.yaml +++ /dev/null @@ -1,89 +0,0 @@ -async_chunk: true -stage_args: - - stage_id: 0 - stage_type: llm - is_comprehension: true - runtime: - devices: "0" - max_batch_size: 1 - engine_args: - dtype: bfloat16 - model_stage: latent_generator - model_arch: VoxCPMForConditionalGeneration - # Optional persistent HF-compatible config dir for native VoxCPM models. - hf_config_path: ${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,} - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - enforce_eager: true - trust_remote_code: true - async_scheduling: false - enable_prefix_caching: false - engine_output_type: latent - gpu_memory_utilization: 0.75 - distributed_executor_backend: "mp" - max_num_batched_tokens: 4096 - max_model_len: 4096 - custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.voxcpm.latent2vae_async_chunk - output_connectors: - to_stage_1: connector_of_shared_memory - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 4096 - seed: 42 - detokenize: false - repetition_penalty: 1.0 - final_output: false - - - stage_id: 1 - stage_type: llm - runtime: - devices: "0" - max_batch_size: 1 - engine_args: - dtype: float32 - model_stage: vae - model_arch: VoxCPMForConditionalGeneration - hf_config_path: ${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,} - worker_type: generation - scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler - enforce_eager: true - trust_remote_code: true - async_scheduling: false - enable_prefix_caching: false - engine_output_type: audio - gpu_memory_utilization: 0.1 - distributed_executor_backend: "mp" - max_num_batched_tokens: 8192 - max_model_len: 4096 - engine_input_source: [0] - input_connectors: - from_stage_0: connector_of_shared_memory - final_output: true - final_output_type: audio - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 1 - seed: 42 - detokenize: true - repetition_penalty: 1.0 - -runtime: - enabled: true - - connectors: - connector_of_shared_memory: - name: SharedMemoryConnector - extra: - shm_threshold_bytes: 65536 - codec_streaming: false - connector_get_sleep_s: 0.01 - connector_get_max_wait_first_chunk: 3000 - connector_get_max_wait: 300 - - edges: - - from: 0 - to: 1 diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index 4f258bbc0..708d9c2b6 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -17,7 +17,6 @@ _CLASS_TO_MODULE: dict[str, str] = { "FishSpeechConfig": "vllm_omni.transformers_utils.configs.fish_speech", "FishSpeechSlowARConfig": "vllm_omni.transformers_utils.configs.fish_speech", "FishSpeechFastARConfig": "vllm_omni.transformers_utils.configs.fish_speech", - "VoxCPMConfig": "vllm_omni.transformers_utils.configs.voxcpm", "VoxCPM2Config": "vllm_omni.transformers_utils.configs.voxcpm2", "VoxtralTTSConfig": "vllm_omni.transformers_utils.configs.voxtral_tts", "CosyVoice3Config": "vllm_omni.transformers_utils.configs.cosyvoice3", @@ -37,7 +36,6 @@ __all__ = [ "FishSpeechConfig", "FishSpeechSlowARConfig", "FishSpeechFastARConfig", - "VoxCPMConfig", "VoxCPM2Config", "VoxtralTTSConfig", "CosyVoice3Config", @@ -70,6 +68,5 @@ from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech # from vllm_omni.transformers_utils.configs import mammoth_moda2 as _mammoth_moda2 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import ming_flash_omni as _ming_flash_omni # noqa: F401, E402 from vllm_omni.transformers_utils.configs import omnivoice as _omnivoice # noqa: F401, E402 -from vllm_omni.transformers_utils.configs import voxcpm as _voxcpm # noqa: F401, E402 from vllm_omni.transformers_utils.configs import voxcpm2 as _voxcpm2 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import voxtral_tts as _voxtral_tts # noqa: F401, E402 diff --git a/vllm_omni/transformers_utils/configs/voxcpm.py b/vllm_omni/transformers_utils/configs/voxcpm.py deleted file mode 100644 index 026783891..000000000 --- a/vllm_omni/transformers_utils/configs/voxcpm.py +++ /dev/null @@ -1,68 +0,0 @@ -from transformers import AutoConfig -from transformers.configuration_utils import PretrainedConfig - - -class VoxCPMConfig(PretrainedConfig): - model_type = "voxcpm" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - bos_token_id: int = 1, - eos_token_id: int = 2, - vocab_size: int = 32000, - hidden_size: int = 1024, - intermediate_size: int = 4096, - max_position_embeddings: int = 4096, - num_attention_heads: int = 16, - num_hidden_layers: int = 24, - num_key_value_heads: int = 16, - rms_norm_eps: float = 1e-6, - rope_theta: float = 10000.0, - rope_scaling: dict | None = None, - lm_config: dict | None = None, - encoder_config: dict | None = None, - dit_config: dict | None = None, - audio_vae_config: dict | None = None, - patch_size: int = 2, - feat_dim: int = 64, - residual_lm_num_layers: int = 6, - scalar_quantization_latent_dim: int = 256, - scalar_quantization_scale: int = 9, - max_length: int = 4096, - device: str = "cuda", - dtype: str = "bfloat16", - dit_mean_mode: bool = False, - **kwargs, - ): - super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs) - self.vocab_size = vocab_size - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.max_position_embeddings = max_position_embeddings - self.num_attention_heads = num_attention_heads - self.num_hidden_layers = num_hidden_layers - self.num_key_value_heads = num_key_value_heads - self.rms_norm_eps = rms_norm_eps - self.rope_theta = rope_theta - self.rope_scaling = rope_scaling - - self.lm_config = lm_config or {} - self.encoder_config = encoder_config or {} - self.dit_config = dit_config or {} - self.audio_vae_config = audio_vae_config - - self.patch_size = patch_size - self.feat_dim = feat_dim - self.residual_lm_num_layers = residual_lm_num_layers - self.scalar_quantization_latent_dim = scalar_quantization_latent_dim - self.scalar_quantization_scale = scalar_quantization_scale - self.max_length = max_length - self.device = device - self.dtype = dtype - self.dit_mean_mode = dit_mean_mode - - -AutoConfig.register("voxcpm", VoxCPMConfig) - -__all__ = ["VoxCPMConfig"] From e1c6bc7597eeffbc84df1d9162a55c860e168833 Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Wed, 20 May 2026 15:26:59 +0800 Subject: [PATCH 28/38] [Test] add run_nightly_jobs.sh for local nightly pytest parity (#3670) Signed-off-by: wangyu <410167048@qq.com> --- pyproject.toml | 6 + tests/e2e/accuracy/test_hunyuan_image3.py | 35 +- tools/nightly/run_nightly_jobs.sh | 754 ++++++++++++++++++++++ 3 files changed, 786 insertions(+), 9 deletions(-) create mode 100644 tools/nightly/run_nightly_jobs.sh diff --git a/pyproject.toml b/pyproject.toml index 35d333fd1..bb9cdc688 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,6 +71,11 @@ demo = [ "gradio>=6.7.0", ] +# Local-test-only extras +local = [ + "FlagEmbedding", +] + docs = [ "mkdocs>=1.5.0", "mkdocs-api-autonav", @@ -184,6 +189,7 @@ markers = [ "core_model: L1&L2 tests (run in each PR)", "advanced_model: L3 level tests (run on each merge)", "full_model: L4 level tests (run nightly)", + "local_model: tests (run locally)", # function module markers "diffusion: Diffusion model tests", "omni: Omni model tests", diff --git a/tests/e2e/accuracy/test_hunyuan_image3.py b/tests/e2e/accuracy/test_hunyuan_image3.py index bf230c07f..6541866d4 100644 --- a/tests/e2e/accuracy/test_hunyuan_image3.py +++ b/tests/e2e/accuracy/test_hunyuan_image3.py @@ -23,18 +23,31 @@ from tests.e2e.accuracy.helpers import ( download_images, model_output_dir, ) -from tests.helpers.runtime import OmniRunner +from tests.helpers.mark import hardware_marks from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens, resolve_stop_token_ids os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA" -pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] +pytestmark = [pytest.mark.local_model, pytest.mark.diffusion, *hardware_marks(res={"cuda": "H100"}, num_cards=8)] # ============================================================================ # Configurable Parameters # ============================================================================ -AR_DEVICES = "0,1" -DIT_DEVICES = "2,3" +# Comma-separated logical CUDA device ids per stage: split visible GPUs (0..n-1), first half -> AR, second -> DiT. + + +def _default_ar_dit_devices() -> tuple[str, str]: + """First floor(n/2) logical devices -> AR, rest -> DiT. ``device_count`` respects ``CUDA_VISIBLE_DEVICES``.""" + n = torch.accelerator.device_count() + if n < 2: + return "0,1", "2,3" + split = n // 2 + ar = ",".join(str(i) for i in range(split)) + dit = ",".join(str(i) for i in range(split, n)) + return ar, dit + + +AR_DEVICES, DIT_DEVICES = _default_ar_dit_devices() MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct" NUM_INFERENCE_STEPS = 50 GUIDANCE_SCALE = 2.5 @@ -75,7 +88,7 @@ _BASE_CONFIG = { "scheduler_cls": "vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler", "gpu_memory_utilization": 0.95, "enforce_eager": True, "trust_remote_code": True, "engine_output_type": "latent", "enable_prefix_caching": False, - "max_num_batched_tokens": 32768, "tensor_parallel_size": 2, "pipeline_parallel_size": 1, + "max_num_batched_tokens": 32768, "tensor_parallel_size": AR_TP_SIZE, "pipeline_parallel_size": 1, "hf_overrides": {"rope_parameters": {"mrope_section": [0, 32, 32], "rope_type": "default"}}, }, "is_comprehension": False, "final_output": True, "final_output_type": "text", @@ -91,7 +104,7 @@ _BASE_CONFIG = { "engine_args": { "model_stage": "dit", "model_arch": "HunyuanImage3ForCausalMM", "enforce_eager": True, "trust_remote_code": True, "distributed_executor_backend": "mp", - "parallel_config": {"tensor_parallel_size": 2, "enable_expert_parallel": True}, + "parallel_config": {"tensor_parallel_size": DIT_TP_SIZE, "enable_expert_parallel": True}, }, "engine_input_source": [0], "custom_process_input_func": "vllm_omni.model_executor.stage_input_processors.hunyuan_image3.ar2diffusion", @@ -134,12 +147,13 @@ def _make_config(enable_kv_reuse: bool, path: Path) -> None: def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Image, str, float]: from transformers import AutoTokenizer + from tests.helpers.runtime import OmniRunner from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType from vllm_omni.platforms import current_omni_platform build_kwargs: dict = {"task": "it2i", "bot_task": "think_recaption", "sys_type": "en_unified", "num_images": 2} - tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True) result = build_prompt_tokens( PROMPT, tokenizer, @@ -149,7 +163,7 @@ def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Imag system_prompt_type = result.system_prompt_type ar_stop_token_ids = resolve_stop_token_ids(task="it2i", bot_task="think_recaption", tokenizer=tokenizer) - with OmniRunner(MODEL_NAME, stage_configs_path=stage_configs_path) as runner: + with OmniRunner(MODEL_PATH, stage_configs_path=stage_configs_path) as runner: params_list = list(runner.omni.default_sampling_params_list) for sp in params_list: if isinstance(sp, OmniDiffusionSamplingParams): @@ -206,7 +220,10 @@ def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Imag return image, cot_text, elapsed -@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Needs 4+ GPUs (2 AR + 2 DiT)") +@pytest.mark.skipif( + torch.accelerator.device_count() < AR_TP_SIZE + DIT_TP_SIZE, + reason=f"Needs {AR_TP_SIZE + DIT_TP_SIZE}+ GPUs ({AR_TP_SIZE} AR + {DIT_TP_SIZE} DiT)", +) def test_image_to_image_alignment(accuracy_artifact_root: Path, accuracy_assets_root: Path) -> None: if importlib.util.find_spec("FlagEmbedding") is None: raise ImportError("Missing dependency: FlagEmbedding\nInstall with: pip install FlagEmbedding") diff --git a/tools/nightly/run_nightly_jobs.sh b/tools/nightly/run_nightly_jobs.sh new file mode 100644 index 000000000..b115a9ee2 --- /dev/null +++ b/tools/nightly/run_nightly_jobs.sh @@ -0,0 +1,754 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Extract steps from .buildkite/test-nightly.yml that contain pytest, synthesize +# small bash wrappers (exports + pytest), run them, and tee output to logs named +# after each step's Buildkite "key" when present (otherwise a slug of the label). +# YAML steps whose label contains "Perf Test" run first, then +# python tools/nightly/generate_nightly_perf_excel.py runs (even if some perf jobs failed; +# the script only aggregates JSON already present under tests/dfx/perf/results). +# Excel and generate_nightly_perf_excel.log are written under $REPO_ROOT/logs/; +# per-job tee logs use LOG_DIR (see --log-dir). Then remaining jobs run. +# +# Test kind (--test-type / TEST_TYPE), multiple allowed (OR semantics): +# Repeat the flag and/or use comma-separated values, e.g. +# --test-type perf,acc OR --test-type perf --test-type function +# If any selected value is all, YAML steps are unconstrained except that perf/acc/function tokens +# are ignored; all,stability / all,local / all,stability,local still run those extra bundles. +# perf — label contains "Perf Test" (throughput / benchmark jobs) +# acc — label contains "Accuracy Test" (incl. GEBench / GEdit-Bench style) +# function — label has neither "Perf Test" nor "Accuracy Test" (incl. Doc, Multi-Replica, etc.) +# stability— fixed dfx stability scripts under tests/dfx/stability/scripts/ (see below) +# local — pytest -sv -m " and local_model" from repo root (no YAML step extract) +# all — no test-kind filter for YAML steps (any leaf step with pytest in test-nightly.yml) +# +# Model area (--model-type / MODEL_TYPE), multiple allowed (OR semantics): +# e.g. --model-type omni,tts +# omni — label contains "Omni ·" +# tts — label contains "TTS ·" +# diffusion— label contains "Diffusion" +# all — no model filter (e.g. Quantization-only steps match only when test-type allows) +# +# Combining stability and/or local with perf/acc/function/all runs each enabled bundle; each enabled +# mode must match at least one job or the script exits with status 2. +# +# local (when included in TEST_TYPE): +# From repo root: pytest -sv -m " and local_model" (markers from MODEL_TYPE: omni, tts, +# diffusion; all → "(omni or tts or diffusion) and local_model"). Not filtered by nightly YAML. +# +# stability (when included in TEST_TYPE): +# From repo root: pytest -s -v tests/dfx/stability/scripts/test_stability_*.py +# model_type: omni → qwen3_omni; tts → qwen3_tts; diffusion → qwen_image + wan22; all → all four +# LABEL_SUBSTR: if set, script path / job key / filename must contain it +# +# Requirements: bash, python3, PyYAML (pip install pyyaml) +# +# Usage: +# bash path/to/run_nightly_jobs.sh +# REPO_ROOT=/path/to/vllm-omni bash path/to/run_nightly_jobs.sh --test-type perf,acc --dry-run +# YML=/path/to/vllm-omni/.buildkite/test-nightly.yml bash path/to/run_nightly_jobs.sh +# +# Repository / YAML (no dependency on where this script lives): +# • Set REPO_ROOT (or pass --repo-root) — default YAML is $REPO_ROOT/.buildkite/test-nightly.yml +# • Or set YML (or --yaml) — repo root is inferred as parent of the .buildkite directory +# • Or run from inside the clone: git rev-parse --show-toplevel, else walk up from $PWD, +# then from the script's directory, until .buildkite/test-nightly.yml exists +# +# Optional environment: +# REPO_ROOT - vllm-omni root (working directory for pytest); see above +# YML - path to test-nightly.yml (default: $REPO_ROOT/.buildkite/test-nightly.yml) +# LOG_DIR - logs + generated job scripts (default: $REPO_ROOT/logs/nightly_jobs) +# TEST_TYPE - comma-separated and/or repeated flags (default: all); see above +# MODEL_TYPE - comma-separated and/or repeated flags (default: all); see above +# LABEL_SUBSTR - YAML mode: substring of Buildkite step label; stability: substring of path/key/filename +# DRY_RUN=1 - print extracted commands only; do not write scripts or run pytest +# +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +LABEL_SUBSTR="${LABEL_SUBSTR:-}" +TEST_TYPE_ENV="${TEST_TYPE:-all}" +MODEL_TYPE_ENV="${MODEL_TYPE:-all}" +TEST_TYPE_CLI_PARTS=() +MODEL_TYPE_CLI_PARTS=() +TEST_TYPE_FROM_CLI=0 +MODEL_TYPE_FROM_CLI=0 +DRY_RUN="${DRY_RUN:-0}" +# REPO_ROOT / YML / LOG_DIR resolved after CLI (do not assume script path vs repo root) + +usage() { + sed -n '2,66p' "$0" | sed 's/^# \{0,1\}//' +} + +# Append comma-separated tokens from $2 into array named $1 (lowercased, trimmed). +_split_append_csv_array() { + local -n _arr="${1}" + local _raw="${2}" + local _ifs="${IFS}" + local _part + IFS=',' + for _part in ${_raw}; do + IFS="${_ifs}" + _part="$(printf '%s' "${_part}" | tr '[:upper:]' '[:lower:]' | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')" + [[ -n "${_part}" ]] || continue + _arr+=("${_part}") + done +} + +# Build comma-separated TEST_TYPE from CLI parts or env; collapse to a single "all" if present. +_finalize_test_type_csv() { + if [[ "${TEST_TYPE_FROM_CLI}" -eq 1 ]]; then + if ((${#TEST_TYPE_CLI_PARTS[@]} == 0)); then + echo "--test-type requires a non-empty value" >&2 + exit 2 + fi + local _has_all=0 _has_stability=0 _has_local=0 + local _x + for _x in "${TEST_TYPE_CLI_PARTS[@]}"; do + if [[ "${_x}" == all ]]; then + _has_all=1 + fi + if [[ "${_x}" == stability ]]; then + _has_stability=1 + fi + if [[ "${_x}" == local ]]; then + _has_local=1 + fi + done + if [[ "${_has_all}" -eq 1 ]]; then + local _out="all" + if [[ "${_has_stability}" -eq 1 ]]; then + _out="${_out},stability" + fi + if [[ "${_has_local}" -eq 1 ]]; then + _out="${_out},local" + fi + printf '%s' "${_out}" + return 0 + fi + local -A _seen=() + local -a _out=() + for _x in "${TEST_TYPE_CLI_PARTS[@]}"; do + [[ ${_seen["${_x}"]+isset} ]] && continue + _seen["${_x}"]=1 + _out+=("${_x}") + done + (IFS=','; printf '%s' "${_out[*]}") + else + printf '%s' "${TEST_TYPE_ENV}" + fi +} + +_finalize_model_type_csv() { + if [[ "${MODEL_TYPE_FROM_CLI}" -eq 1 ]]; then + if ((${#MODEL_TYPE_CLI_PARTS[@]} == 0)); then + echo "--model-type requires a non-empty value" >&2 + exit 2 + fi + local -A _seen=() + local -a _out=() + local _x + for _x in "${MODEL_TYPE_CLI_PARTS[@]}"; do + if [[ "${_x}" == all ]]; then + printf '%s' "all" + return 0 + fi + [[ ${_seen["${_x}"]+isset} ]] && continue + _seen["${_x}"]=1 + _out+=("${_x}") + done + (IFS=','; printf '%s' "${_out[*]}") + else + printf '%s' "${MODEL_TYPE_ENV}" + fi +} + +while [[ $# -gt 0 ]]; do + case "$1" in + -h | --help) + usage + exit 0 + ;; + --dry-run) + DRY_RUN=1 + shift + ;; + --log-dir) + LOG_DIR="$2" + shift 2 + ;; + --yaml) + YML="$2" + shift 2 + ;; + --repo-root) + REPO_ROOT="$2" + shift 2 + ;; + --label-substr) + LABEL_SUBSTR="$2" + shift 2 + ;; + --model-type) + MODEL_TYPE_FROM_CLI=1 + _split_append_csv_array MODEL_TYPE_CLI_PARTS "$2" + shift 2 + ;; + --test-type) + TEST_TYPE_FROM_CLI=1 + _split_append_csv_array TEST_TYPE_CLI_PARTS "$2" + shift 2 + ;; + *) + echo "Unknown option: $1" >&2 + usage >&2 + exit 2 + ;; + esac +done + +TEST_TYPE="$(_finalize_test_type_csv)" +MODEL_TYPE="$(_finalize_model_type_csv)" + +BUILDKITE_REL=".buildkite/test-nightly.yml" + +_find_repo_containing_nightly() { + local dir="${1:-}" + [[ -n "$dir" ]] || return 1 + dir="$(cd "$dir" && pwd)" || return 1 + while true; do + if [[ -f "${dir}/${BUILDKITE_REL}" ]]; then + printf '%s\n' "$dir" + return 0 + fi + [[ "${dir}" == "/" ]] && return 1 + dir="$(dirname "${dir}")" + done +} + +_derive_repo_root_from_yml() { + local yml="$1" + local d + d="$(cd "$(dirname "${yml}")" && pwd)" + [[ "$(basename "${d}")" == ".buildkite" ]] || return 1 + printf '%s\n' "$(dirname "${d}")" +} + +# Resolve REPO_ROOT, YML (no relative path between script and repo) +if [[ -n "${YML:-}" && -n "${REPO_ROOT:-}" ]]; then + REPO_ROOT="$(cd "${REPO_ROOT}" && pwd)" + YML="$(cd "$(dirname "${YML}")" && pwd)/$(basename "${YML}")" +elif [[ -n "${YML:-}" ]]; then + YML="$(cd "$(dirname "${YML}")" && pwd)/$(basename "${YML}")" + if ! REPO_ROOT="$(_derive_repo_root_from_yml "${YML}")"; then + echo "Could not derive REPO_ROOT from YML=${YML} (expected file at /.buildkite/test-nightly.yml)." >&2 + echo "Set REPO_ROOT explicitly (or pass --repo-root) for pytest working directory." >&2 + exit 2 + fi +elif [[ -n "${REPO_ROOT:-}" ]]; then + REPO_ROOT="$(cd "${REPO_ROOT}" && pwd)" + YML="${REPO_ROOT}/${BUILDKITE_REL}" +else + REPO_ROOT="" + if command -v git >/dev/null 2>&1; then + REPO_ROOT="$(git -C "${PWD}" rev-parse --show-toplevel 2>/dev/null || true)" + fi + if [[ -z "${REPO_ROOT}" ]]; then + REPO_ROOT="$(_find_repo_containing_nightly "${PWD}" || true)" + fi + if [[ -z "${REPO_ROOT}" ]]; then + REPO_ROOT="$(_find_repo_containing_nightly "${SCRIPT_DIR}" || true)" + fi + if [[ -z "${REPO_ROOT}" ]]; then + echo "Could not locate ${BUILDKITE_REL}. Set REPO_ROOT or YML, run from inside the vllm-omni clone," >&2 + echo "or place this script (or run from a cwd) under the repository tree." >&2 + exit 2 + fi + YML="${REPO_ROOT}/${BUILDKITE_REL}" +fi + +LOG_DIR="${LOG_DIR:-${REPO_ROOT}/logs/nightly_jobs}" + +# Require nightly YAML only when at least one non-stability test kind needs it (validated in Python too). +_needs_nightly_yml=0 +IFS=',' read -ra _TTOK <<< "${TEST_TYPE}" +for _t in "${_TTOK[@]}"; do + _t="$(printf '%s' "${_t}" | tr '[:upper:]' '[:lower:]' | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')" + case "${_t}" in + perf | acc | function | all) + _needs_nightly_yml=1 + ;; + esac +done +if [[ "${_needs_nightly_yml}" -eq 1 ]] && [[ ! -f "${YML}" ]]; then + echo "YAML not found: ${YML}" >&2 + exit 2 +fi + +if ! command -v python3 >/dev/null 2>&1; then + echo "python3 is required." >&2 + exit 1 +fi + +mkdir -p "${LOG_DIR}/jobs" +LOG_DIR="$(cd "${LOG_DIR}" && pwd)" +export REPO_ROOT LOG_DIR YML LABEL_SUBSTR TEST_TYPE MODEL_TYPE DRY_RUN + +# Drop stale wrapper scripts, old tee logs, and perf manifest so a filtered run does not +# mix with the previous outputs. +if [[ "${DRY_RUN}" != "1" ]]; then + shopt -s nullglob + _stale=( "${LOG_DIR}/jobs"/*.sh ) + _stale_logs=( "${LOG_DIR}"/*.log ) + shopt -u nullglob + if ((${#_stale[@]})); then + rm -f "${_stale[@]}" + fi + if ((${#_stale_logs[@]})); then + rm -f "${_stale_logs[@]}" + fi + rm -f "${LOG_DIR}/jobs/.perf_job_keys" + rm -f "${REPO_ROOT}/logs/nightly_perf_manual.xlsx" + rm -f "${LOG_DIR}/nightly_perf_manual.xlsx" +fi + +# shellcheck disable=SC2016,SC1078,SC1079 +python3 - <<'PY' +from __future__ import annotations + +import os +import re +import stat +import sys +from pathlib import Path + +try: + import yaml +except ImportError: + print("Missing PyYAML. Install with: pip install pyyaml", file=sys.stderr) + sys.exit(1) + +REPO_ROOT = Path(os.environ["REPO_ROOT"]).resolve() +LOG_DIR = Path(os.environ["LOG_DIR"]).resolve() +YML = Path(os.environ["YML"]).resolve() +LABEL_SUBSTR = (os.environ.get("LABEL_SUBSTR") or "").strip() +DRY_RUN = os.environ.get("DRY_RUN", "0") == "1" + +ALLOWED_TEST_TYPES = frozenset({"all", "perf", "acc", "function", "stability", "local"}) +ALLOWED_MODEL_TYPES = frozenset({"all", "omni", "tts", "diffusion"}) + + +def parse_model_types(raw: str) -> list[str]: + parts = [p.strip().lower() for p in (raw or "").split(",") if p.strip()] + if not parts: + parts = ["all"] + bad = [p for p in parts if p not in ALLOWED_MODEL_TYPES] + if bad: + print( + f"Invalid MODEL_TYPE / --model-type value(s): {bad!r} " + f"(allowed: {', '.join(sorted(ALLOWED_MODEL_TYPES))})", + file=sys.stderr, + ) + sys.exit(2) + if "all" in parts: + return ["all"] + out: list[str] = [] + seen: set[str] = set() + for p in parts: + if p not in seen: + seen.add(p) + out.append(p) + return out + + +def parse_test_types(raw: str) -> list[str]: + """If 'all' appears, YAML dimension is unconstrained; extras like stability/local are kept.""" + parts = [p.strip().lower() for p in (raw or "").split(",") if p.strip()] + if not parts: + parts = ["all"] + bad = [p for p in parts if p not in ALLOWED_TEST_TYPES] + if bad: + print( + f"Invalid TEST_TYPE / --test-type value(s): {bad!r} " + f"(allowed: {', '.join(sorted(ALLOWED_TEST_TYPES))})", + file=sys.stderr, + ) + sys.exit(2) + if "all" in parts: + out: list[str] = ["all"] + if "stability" in parts: + out.append("stability") + if "local" in parts: + out.append("local") + return out + out: list[str] = [] + seen: set[str] = set() + for p in parts: + if p not in seen: + seen.add(p) + out.append(p) + return out + + +PYTEST_CMD_RE = re.compile( + r"(?:timeout\s+\S+\s+)?(?:python3? -m\s+)?pytest\s+[^\n&|;]*" +) + + +def iter_leaf_steps(steps, group: str | None = None): + """Flatten Buildkite steps; yield (leaf_step_dict, group_title).""" + for raw in steps or []: + if not isinstance(raw, dict): + continue + nested = raw.get("steps") + if isinstance(nested, list) and nested: + g = raw.get("group") + next_group: str | None + if isinstance(g, str): + next_group = g + elif g is not None: + next_group = str(g) + else: + next_group = group + yield from iter_leaf_steps(nested, next_group) + continue + if raw.get("commands"): + yield raw, group + + +def raw_command_text(step: dict) -> str: + raw = step.get("commands") or [] + if isinstance(raw, str): + raw = [raw] + text = "\n".join((c.strip() if isinstance(c, str) else "") for c in raw if c) + return text.replace("$$", "$") + + +def export_lines(text: str) -> list[str]: + out: list[str] = [] + for line in text.splitlines(): + s = line.strip() + if not s or s.startswith("#"): + continue + if s.startswith("export "): + out.append(s) + return out + + +def pytest_lines(text: str) -> list[str]: + out: list[str] = [] + for m in PYTEST_CMD_RE.finditer(text): + line = m.group(0).strip() + line_start = text.rfind("\n", 0, m.start()) + 1 + before = text[line_start : m.start()] + if before.lstrip().startswith("#"): + continue + if line: + out.append(line) + return out + + +def label_matches_test_type(label: str, test_types: list[str]) -> bool: + """perf / acc / function / all — aligned with nightly step labels (OR across test_types).""" + if "all" in test_types: + return True + has_perf = "Perf Test" in label + has_acc = "Accuracy Test" in label + is_function = not has_perf and not has_acc + for t in test_types: + if t in ("stability", "local"): + continue + if t == "perf" and has_perf: + return True + if t == "acc" and has_acc: + return True + if t == "function" and is_function: + return True + return False + + +def label_matches_model_type(label: str, model_types: list[str]) -> bool: + """Omni · / TTS · / Diffusion — OR across model_types.""" + if "all" in model_types: + return True + if "omni" in model_types and "Omni ·" in label: + return True + if "tts" in model_types and "TTS ·" in label: + return True + if "diffusion" in model_types and "Diffusion" in label: + return True + return False + + +def job_key_from_step(step: dict, label: str) -> str: + k = step.get("key") + if isinstance(k, str) and k.strip(): + return k.strip() + slug = re.sub(r"[^\w\-.]+", "_", label or "job", flags=re.UNICODE) + slug = re.sub(r"_+", "_", slug).strip("_") or "job" + return slug + + +# Fixed stability scripts (when stability is selected); MODEL_TYPES narrows which run. +STABILITY_CASES: list[tuple[str, str, tuple[str, ...]]] = [ + ("stability_qwen3_omni", "tests/dfx/stability/scripts/test_stability_qwen3_omni.py", ("omni",)), + ("stability_qwen3_tts", "tests/dfx/stability/scripts/test_stability_qwen3_tts.py", ("tts",)), + ("stability_qwen_image", "tests/dfx/stability/scripts/test_stability_qwen_image.py", ("diffusion",)), + ("stability_wan22", "tests/dfx/stability/scripts/test_stability_wan22.py", ("diffusion",)), +] + + +def _stability_model_matches(model_types: list[str], families: tuple[str, ...]) -> bool: + if "all" in model_types: + return True + return bool(set(model_types) & set(families)) + + +def _write_job_script(key: str, script_lines: list[str], jobs_dir: Path) -> None: + body = "\n".join(script_lines) + "\n" + if DRY_RUN: + print(f"=== {key} ===") + print(body) + return + job_path = jobs_dir / f"{key}.sh" + job_path.write_text(body, encoding="utf-8") + try: + mode = job_path.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH + job_path.chmod(mode) + except OSError: + pass + print(f"generated {job_path}", file=sys.stderr) + + +def local_pytest_marker_expr(model_types: list[str]) -> str: + """Build -m expression: (omni|tts|diffusion) markers combined with local_model.""" + if "all" in model_types: + families: tuple[str, ...] = ("omni", "tts", "diffusion") + else: + families = tuple(model_types) + if len(families) == 1: + return f"{families[0]} and local_model" + return f"({' or '.join(families)}) and local_model" + + +def run_local_mode(jobs_dir: Path, model_types: list[str]) -> int: + """Single job: pytest -sv -m ' and local_model' (testpaths from pytest config).""" + marker_expr = local_pytest_marker_expr(model_types) + key = "local_pytest" + pytest_line = f'pytest -sv -m "{marker_expr}"' + script_lines = [ + "#!/usr/bin/env bash", + f"# Local marker tests — MODEL_TYPE={','.join(model_types)}", + "set -euo pipefail", + f'cd "{REPO_ROOT}"', + pytest_line, + ] + _write_job_script(key, script_lines, jobs_dir) + return 1 + + +def run_stability_mode(jobs_dir: Path, model_types: list[str]) -> int: + matched = 0 + for key, rel_posix, families in STABILITY_CASES: + if not _stability_model_matches(model_types, families): + continue + if LABEL_SUBSTR: + fn = Path(rel_posix).name + if ( + LABEL_SUBSTR not in rel_posix + and LABEL_SUBSTR not in key + and LABEL_SUBSTR not in fn + ): + continue + rel_path = REPO_ROOT / rel_posix + if not rel_path.is_file(): + print(f"# skip (missing file): {rel_path}", file=sys.stderr) + continue + pytest_line = f"pytest -s -v {rel_posix}" + matched += 1 + script_lines = [ + "#!/usr/bin/env bash", + f"# Stability: {key} — {rel_posix}", + "set -euo pipefail", + f'cd "{REPO_ROOT}"', + pytest_line, + ] + _write_job_script(key, script_lines, jobs_dir) + return matched + + +def main() -> None: + test_types = parse_test_types(os.environ.get("TEST_TYPE", "all")) + model_types = parse_model_types(os.environ.get("MODEL_TYPE", "all")) + + want_stability = "stability" in test_types + want_local = "local" in test_types + yaml_test_types = [t for t in test_types if t not in ("stability", "local")] + needs_yaml = bool(yaml_test_types) + + jobs_dir = LOG_DIR / "jobs" + if not DRY_RUN: + jobs_dir.mkdir(parents=True, exist_ok=True) + + matched_stability = 0 + matched_yaml = 0 + perf_job_keys: list[str] = [] + + if want_stability: + matched_stability = run_stability_mode(jobs_dir, model_types) + + if want_local: + run_local_mode(jobs_dir, model_types) + + if needs_yaml: + if not YML.is_file(): + print(f"YAML not found: {YML}", file=sys.stderr) + sys.exit(1) + + data = yaml.safe_load(YML.read_text(encoding="utf-8")) + top_steps = (data or {}).get("steps") or [] + + for step, _grp in iter_leaf_steps(top_steps): + label = step.get("label") or "" + if LABEL_SUBSTR and LABEL_SUBSTR not in label: + continue + if not label_matches_test_type(label, yaml_test_types): + continue + if not label_matches_model_type(label, model_types): + continue + text = raw_command_text(step) + exports = export_lines(text) + pys = pytest_lines(text) + if not pys: + print(f"# skip (no pytest line): {label!r}", file=sys.stderr) + continue + + key = job_key_from_step(step, label) + matched_yaml += 1 + + script_lines = [ + "#!/usr/bin/env bash", + f'# From Buildkite label: {label.replace(chr(10), " ")}', + "set -euo pipefail", + f'cd "{REPO_ROOT}"', + ] + script_lines.extend(exports) + script_lines.extend(pys) + _write_job_script(key, script_lines, jobs_dir) + if "Perf Test" in label and key not in perf_job_keys: + perf_job_keys.append(key) + + if not DRY_RUN: + manifest_path = LOG_DIR / "jobs" / ".perf_job_keys" + if perf_job_keys: + manifest_path.write_text("\n".join(perf_job_keys) + "\n", encoding="utf-8") + elif manifest_path.is_file(): + manifest_path.unlink() + + errs: list[str] = [] + if want_stability and matched_stability == 0: + errs.append( + f"No stability jobs matched TEST_TYPE={test_types!r} MODEL_TYPE={model_types!r} " + f"LABEL_SUBSTR={LABEL_SUBSTR!r} under {REPO_ROOT}" + ) + if needs_yaml and matched_yaml == 0: + errs.append( + f"No YAML steps matched TEST_TYPE={test_types!r} MODEL_TYPE={model_types!r} " + f"LABEL_SUBSTR={LABEL_SUBSTR!r} in {YML}" + ) + if errs: + for e in errs: + print(e, file=sys.stderr) + sys.exit(2) + + +main() +PY + +if [[ "${DRY_RUN}" == "1" ]]; then + exit 0 +fi + +# Run jobs: YAML perf steps first (if manifest exists), then generate Excel (always when +# manifest non-empty, even if some perf jobs failed; Excel uses whatever JSON exists on disk), +# then remaining jobs. Paths align with nightly Buildkite (tests/dfx/perf/results). +run_generated_jobs_with_tee() { + set -o pipefail + local any_fail=0 + local perf_list="${LOG_DIR}/jobs/.perf_job_keys" + local -a perf_keys=() + if [[ -f "${perf_list}" ]] && [[ -s "${perf_list}" ]]; then + mapfile -t perf_keys < "${perf_list}" + fi + + _run_one_job() { + local _job="$1" + local base out job_status + base="$(basename "${_job}" .sh)" + out="${LOG_DIR}/${base}.log" + echo "==> ${_job} (tee ${out})" >&2 + (cd "${REPO_ROOT}" && bash "${_job}") 2>&1 | tee "${out}" + job_status="${PIPESTATUS[0]}" + return "${job_status}" + } + + local k + for k in "${perf_keys[@]}"; do + [[ -n "${k}" ]] || continue + if [[ ! -f "${LOG_DIR}/jobs/${k}.sh" ]]; then + echo "warning: perf job script missing, skip: ${LOG_DIR}/jobs/${k}.sh" >&2 + any_fail=1 + continue + fi + _run_one_job "${LOG_DIR}/jobs/${k}.sh" || any_fail=1 + done + + if ((${#perf_keys[@]})); then + local excel_out excel_log excel_status excel_repo_logs + excel_repo_logs="${REPO_ROOT}/logs" + mkdir -p "${excel_repo_logs}" + excel_out="${excel_repo_logs}/nightly_perf_$(date -u +%Y%m%d-%H%M%S).xlsx" + excel_log="${excel_repo_logs}/generate_nightly_perf_excel.log" + if [[ "${any_fail}" -ne 0 ]]; then + echo "Note: one or more perf jobs failed; still running generate_nightly_perf_excel.py " \ + "(report reflects JSON already under tests/dfx/perf/results)." >&2 + fi + echo "==> python3 tools/nightly/generate_nightly_perf_excel.py -> ${excel_out} (tee ${excel_log})" >&2 + ( + cd "${REPO_ROOT}" && python3 tools/nightly/generate_nightly_perf_excel.py \ + --input-dir "${REPO_ROOT}/tests/dfx/perf/results" \ + --diffusion-input-dir "${REPO_ROOT}/tests/dfx/perf/results" \ + --output-file "${excel_out}" + ) 2>&1 | tee "${excel_log}" + excel_status="${PIPESTATUS[0]}" + if [[ "${excel_status}" -ne 0 ]]; then + any_fail=1 + echo "generate_nightly_perf_excel.py failed (exit ${excel_status}). See ${excel_log}" >&2 + fi + fi + + local -A _perf_done=() + for k in "${perf_keys[@]}"; do + [[ -n "${k}" ]] && _perf_done["${k}"]=1 + done + + local _job base + shopt -s nullglob + for _job in "${LOG_DIR}/jobs"/*.sh; do + base="$(basename "${_job}" .sh)" + [[ ${_perf_done["${base}"]+isset} ]] && continue + _run_one_job "${_job}" || any_fail=1 + done + shopt -u nullglob + + if [[ "${any_fail}" -ne 0 ]]; then + echo "One or more selected jobs failed. See logs under ${LOG_DIR}." >&2 + return 1 + fi + echo "All selected jobs finished OK. Logs: ${LOG_DIR}/*.log" >&2 + return 0 +} + +run_generated_jobs_with_tee || exit 1 From 5cf4605bf918ffa2f0f9b3b83042557d9ad46c05 Mon Sep 17 00:00:00 2001 From: bjf-frz Date: Wed, 20 May 2026 15:31:41 +0800 Subject: [PATCH 29/38] [Bugfix]Fix distributed stage0 multimodal cache routing (#3740) Signed-off-by: bjf-frz <240110042+bjf-frz@users.noreply.github.com> Co-authored-by: bjf-frz <240110042+bjf-frz@users.noreply.github.com> --- tests/engine/test_async_omni_engine_input.py | 126 +++++++++++++++++++ vllm_omni/engine/async_omni_engine.py | 31 +++-- vllm_omni/engine/stage_pool.py | 45 +++++++ 3 files changed, 191 insertions(+), 11 deletions(-) diff --git a/tests/engine/test_async_omni_engine_input.py b/tests/engine/test_async_omni_engine_input.py index 9dd577c93..a1ec472f6 100644 --- a/tests/engine/test_async_omni_engine_input.py +++ b/tests/engine/test_async_omni_engine_input.py @@ -3,6 +3,7 @@ from pytest_mock import MockerFixture from vllm.sampling_params import SamplingParams from vllm.v1.engine import EngineCoreRequest +from vllm_omni.distributed.omni_coordinator import ReplicaInfo, ReplicaStatus from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.async_omni_engine import AsyncOmniEngine, StageRuntimeInfo from vllm_omni.engine.stage_pool import StagePool @@ -95,6 +96,44 @@ class _FakeStageClient: stage_type = "llm" final_output = False + def __init__(self, input_address: str | None = None): + if input_address is not None: + self.client_addresses = {"input_address": input_address} + + +class _FakeHub: + def __init__(self, replicas: list[ReplicaInfo]): + self._replicas = replicas + + def get_replicas_for_stage(self, stage_id: int): + return type( + "ReplicaList", + (), + {"replicas": [rep for rep in self._replicas if rep.stage_id == stage_id]}, + )() + + +class _RoundRobinLB: + def __init__(self): + self._next = 0 + + def select(self, task, replicas): # noqa: ARG002 + idx = self._next % len(replicas) + self._next += 1 + return idx + + +def _replica(input_addr: str) -> ReplicaInfo: + return ReplicaInfo( + input_addr=input_addr, + output_addr=input_addr.replace("input", "output"), + stage_id=0, + status=ReplicaStatus.UP, + queue_length=0, + last_heartbeat=0.0, + registered_at=0.0, + ) + def test_build_add_request_message_scopes_mm_uuids_to_selected_stage0_replica(mocker: MockerFixture): engine = object.__new__(AsyncOmniEngine) @@ -132,6 +171,93 @@ def test_build_add_request_message_scopes_mm_uuids_to_selected_stage0_replica(mo assert seen_uuids[0].removeprefix("stage0:rep0:") == seen_uuids[1].removeprefix("stage0:rep1:") +@pytest.mark.asyncio +async def test_build_add_request_message_scopes_mm_uuids_to_distributed_stage0_replica(mocker: MockerFixture): + engine = object.__new__(AsyncOmniEngine) + params = SamplingParams(max_tokens=8) + engine.model = "test-model" + engine.default_sampling_params_list = [params] + engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")] + engine.supported_tasks = ("generate",) + + addr0 = "tcp://host-a:1000/input" + addr1 = "tcp://host-b:1000/input" + stage_pool = StagePool(0, [_FakeStageClient(addr0), _FakeStageClient(addr1)]) + stage_pool.attach_hub(_FakeHub([_replica(addr0), _replica(addr1)])) + stage_pool.attach_load_balancer(_RoundRobinLB()) + engine.stage_pools = [stage_pool] + + seen_uuids: list[str] = [] + + def process_inputs(**kwargs): + prompt = kwargs["prompt"] + seen_uuids.append(prompt["multi_modal_uuids"]["image"][0]) + return _make_engine_core_request(kwargs["request_id"]) + + input_processor = mocker.Mock() + input_processor.process_inputs.side_effect = process_inputs + engine.input_processor = input_processor + + for request_id in ("req-1", "req-2"): + engine._build_add_request_message( + request_id=request_id, + prompt={ + "prompt": "describe", + "multi_modal_data": {"image": "same-image"}, + }, + sampling_params_list=[params], + final_stage_id=0, + ) + + assert seen_uuids[0].startswith("stage0:rep0:") + assert seen_uuids[1].startswith("stage0:rep1:") + assert stage_pool.get_bound_replica_id("req-1") == 0 + assert stage_pool.get_bound_replica_id("req-2") == 1 + assert await stage_pool.pick("req-1") == 0 + assert await stage_pool.pick("req-2") == 1 + + +def test_build_add_request_message_skips_distributed_mm_scope_when_no_replica(mocker: MockerFixture): + engine = object.__new__(AsyncOmniEngine) + params = SamplingParams(max_tokens=8) + engine.model = "test-model" + engine.default_sampling_params_list = [params] + engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")] + engine.supported_tasks = ("generate",) + + addr0 = "tcp://host-a:1000/input" + addr1 = "tcp://host-b:1000/input" + stage_pool = StagePool(0, [_FakeStageClient(addr0), _FakeStageClient(addr1)]) + stage_pool.attach_hub(_FakeHub([])) + stage_pool.attach_load_balancer(_RoundRobinLB()) + engine.stage_pools = [stage_pool] + + seen_prompt: dict | None = None + + def process_inputs(**kwargs): + nonlocal seen_prompt + seen_prompt = kwargs["prompt"] + return _make_engine_core_request(kwargs["request_id"]) + + input_processor = mocker.Mock() + input_processor.process_inputs.side_effect = process_inputs + engine.input_processor = input_processor + + engine._build_add_request_message( + request_id="req-no-replica", + prompt={ + "prompt": "describe", + "multi_modal_data": {"image": "same-image"}, + }, + sampling_params_list=[params], + final_stage_id=0, + ) + + assert seen_prompt is not None + assert "multi_modal_uuids" not in seen_prompt + assert stage_pool.get_bound_replica_id("req-no-replica") is None + + def test_stage_pool_replica_count_falls_back_to_clients(): class PoolWithoutLiveNumReplicas: clients = [object(), None, object()] diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 420f808dc..ad8dbff6f 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1557,21 +1557,30 @@ class AsyncOmniEngine: if stage0_pool.stage_type == "diffusion" or self._stage_pool_replica_count(stage0_pool) <= 1: return None - # This synchronous request path can safely pre-bind local in-process - # replicas. Distributed head mode still uses the async picker inside - # StagePool.submit_initial(). - if self._stage_pool_is_distributed(stage0_pool): - logger.debug( - "[AsyncOmniEngine] Skipping stage-0 multimodal cache scoping for distributed routing req=%s", - request_id, - ) - return None - prompts = prompt if isinstance(prompt, list) else [prompt] if not any(isinstance(p, dict) and p.get("multi_modal_data") for p in prompts): return None - replica_id = stage0_pool.select_replica_id(request_id) + if self._stage_pool_is_distributed(stage0_pool): + preselect_replica_id = getattr(stage0_pool, "preselect_replica_id", None) + if not callable(preselect_replica_id): + logger.debug( + "[AsyncOmniEngine] Skipping stage-0 multimodal cache scoping for distributed routing " + "without preselect support req=%s", + request_id, + ) + return None + replica_id = preselect_replica_id(request_id) + if replica_id is None: + logger.debug( + "[AsyncOmniEngine] Skipping stage-0 multimodal cache scoping for distributed routing " + "because no serviceable replica is available yet req=%s", + request_id, + ) + return None + else: + replica_id = stage0_pool.select_replica_id(request_id) + for p in prompts: self._ensure_stage_replica_mm_uuids( p, diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 68de78788..effe64705 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -301,6 +301,51 @@ class StagePool: raise RuntimeError(f"no UP replica for stage {self.stage_id} after {self.DISPATCH_WAIT_TIMEOUT_S:.1f}s") await asyncio.sleep(min(self.DISPATCH_RETRY_INTERVAL_S, deadline - now)) + def preselect_replica_id( + self, + request_id: str, + task: Task | None = None, + *, + affinity_request_id: str | None = None, + ) -> int | None: + """Synchronously pick and bind a replica before request preprocessing. + + The main-thread input preprocessing path cannot await :meth:`pick`, but + multimodal cache UUID scoping needs to know the same replica that + :meth:`submit_initial` will later use. In distributed mode this checks + the hub's cached replica snapshot once and records the selected input + address in ``_affinity`` so the async submit path reuses the route. If + no replica is currently serviceable, return ``None`` and let the async + submit-time router wait without blocking the caller. + """ + if self._hub is None or self._lb is None: + return self.select_replica_id(request_id, affinity_request_id=affinity_request_id) + + bound_addr = self._affinity.get(request_id) + if bound_addr is not None: + replica_id = self._serviceable_replica_id_for_addr(bound_addr) + if replica_id is not None: + return replica_id + self._affinity.pop(request_id, None) + + if affinity_request_id is not None: + parent_addr = self._affinity.get(affinity_request_id) + if parent_addr is not None: + replica_id = self._serviceable_replica_id_for_addr(parent_addr) + if replica_id is not None: + self._affinity[request_id] = parent_addr + return replica_id + + task = task or Task(request_id=request_id) + candidates = self._collect_serviceable_replicas() + if not candidates: + return None + + lb_idx = self._lb.select(task, [rep for rep, _ in candidates]) + replica_info, replica_id = candidates[lb_idx] + self._affinity[request_id] = replica_info.input_addr + return replica_id + def _collect_serviceable_replicas(self) -> list[tuple[ReplicaInfo, int]]: """Return list of ``(ReplicaInfo, replica_id)`` for UP, attached replicas.""" if self._hub is None: From 2aafe0bf94989e1946c7e901c4a71391e7492ed2 Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Wed, 20 May 2026 15:40:34 +0800 Subject: [PATCH 30/38] [Perf] Optimize sampler D2H sync for HY-Image (#3617) Signed-off-by: gcanlin --- .../models/hunyuan_image3/hunyuan_image3.py | 44 +++++++++++++------ 1 file changed, 31 insertions(+), 13 deletions(-) diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py index b68a6fc3a..3ed3e8f89 100644 --- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py +++ b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py @@ -1590,6 +1590,10 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo self._sampler: Sampler | None = None self._eos_token_id: int = tokenizer.eos_token_id + # Lazily built on first sample() call so we can pick up logits.device + # without guessing during init. See `sample()` for the comprehension + # fast path that uses this. + self._blocked_token_ids_tensor: torch.Tensor | None = None self._replace_rotary_embeddings() self._patch_moe_blocks() @@ -1992,25 +1996,39 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo min_score = torch.finfo(logits.dtype).min + if self._is_comprehension: + # Comprehension path is stateless: we only need to mask a fixed + # set of blocked token ids on every step. Do it in one batched + # index_fill_ instead of a per-(req, id) Python loop of scalar + # GPU writes, and skip reading `output_token_ids` from CPU + # (unused here) so this branch never forces a D2H sync. + if self._blocked_token_ids_tensor is None and self._blocked_token_ids: + self._blocked_token_ids_tensor = torch.tensor( + sorted(self._blocked_token_ids), + dtype=torch.long, + device=logits.device, + ) + if self._blocked_token_ids_tensor is not None: + logits.index_fill_(-1, self._blocked_token_ids_tensor, min_score) + return self._sampler(logits=logits, sampling_metadata=sampling_metadata) + + # Generation path retains the per-request stateful logic for forced + # stage-transition tokens and ratio-restriction. for req_idx in range(logits.shape[0]): decoded_tokens: list[int] = ( sampling_metadata.output_token_ids[req_idx] if req_idx < len(sampling_metadata.output_token_ids) else [] ) last_token = decoded_tokens[-1] if decoded_tokens else -1 - if self._is_comprehension: - for tid in self._blocked_token_ids: - logits[req_idx, tid] = min_score - else: - forced = self._get_forced_token(decoded_tokens) - if forced is not None: - logits[req_idx].fill_(min_score) - logits[req_idx, forced] = 0 - elif last_token == self._size_token_id: - self._apply_ratio_restriction(logits, req_idx, min_score) - elif last_token in self._all_ratio_ids: - logits[req_idx].fill_(min_score) - logits[req_idx, self._eos_token_id] = 0 + forced = self._get_forced_token(decoded_tokens) + if forced is not None: + logits[req_idx].fill_(min_score) + logits[req_idx, forced] = 0 + elif last_token == self._size_token_id: + self._apply_ratio_restriction(logits, req_idx, min_score) + elif last_token in self._all_ratio_ids: + logits[req_idx].fill_(min_score) + logits[req_idx, self._eos_token_id] = 0 return self._sampler(logits=logits, sampling_metadata=sampling_metadata) From f74cba84ac31657751130d7f1c6f2d4d07e7d241 Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Wed, 20 May 2026 15:52:40 +0800 Subject: [PATCH 31/38] [Docs] Complete quantization nav and online guide (#3764) Signed-off-by: David Chen <530634352@qq.com> --- docs/.nav.yml | 2 ++ docs/user_guide/quantization/mxfp4.md | 4 +-- docs/user_guide/quantization/online.md | 37 ++++++++++++++++++-------- 3 files changed, 30 insertions(+), 13 deletions(-) diff --git a/docs/.nav.yml b/docs/.nav.yml index a60a4368d..562a49e84 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -80,6 +80,8 @@ nav: - Quantized KV Cache: user_guide/quantization/quantized_kvcache.md - FP8 W8A8: user_guide/quantization/fp8.md - Int8 W8A8: user_guide/quantization/int8.md + - MXFP8 W8A8: user_guide/quantization/mxfp8.md + - MXFP4 W4A4: user_guide/quantization/mxfp4.md - ModelOpt: user_guide/quantization/modelopt.md - GGUF: user_guide/quantization/gguf.md - AutoRound: user_guide/quantization/autoround.md diff --git a/docs/user_guide/quantization/mxfp4.md b/docs/user_guide/quantization/mxfp4.md index 2fd4a8395..7463ada23 100644 --- a/docs/user_guide/quantization/mxfp4.md +++ b/docs/user_guide/quantization/mxfp4.md @@ -134,7 +134,7 @@ can be pinned to BF16 via the Python API: ```python omni = Omni( model="", - quantization={ + quantization_config={ "method": "mxfp4_dualscale", "ignored_layers": ["blocks.10.attn1.to_q"], # explicit per-layer override }, @@ -397,7 +397,7 @@ names** discovered in Step 1. No code changes to the model are required. ```python omni = Omni( model="/path/to/your-model", - quantization={ + quantization_config={ "method": "mxfp4_dualscale", "ignored_layers": [ "blocks.0.attn1.to_qkv", # runtime name, not diffusers name diff --git a/docs/user_guide/quantization/online.md b/docs/user_guide/quantization/online.md index 3dcb8297e..b9a6861d1 100644 --- a/docs/user_guide/quantization/online.md +++ b/docs/user_guide/quantization/online.md @@ -9,20 +9,24 @@ separate quantized checkpoint. This mode is different from pre-quantized checkpoint formats such as GGUF, AutoRound, msModelSlim, or serialized Int8 checkpoints. Those formats are prepared before serving and are documented in their method-specific guides. +For MXFP8 and MXFP4, use this page for load-time quantization from BF16 +checkpoints, and use the method-specific pages for offline checkpoints produced +by msModelSlim and the merge tools. ## Hardware Support -| Device | FP8 W8A8 | Int8 W8A8 | -|--------|----------|-----------| -| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | -| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | -| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | -| AMD ROCm | ⭕ | ⭕ | -| Intel XPU | ⭕ | ⭕ | -| Ascend NPU | ❌ | ✅ | +| Device | FP8 W8A8 | Int8 W8A8 | MXFP8 W8A8 | MXFP4 W4A4 | +|--------|----------|-----------|------------|------------| +| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ⭕ | ⭕ | +| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ⭕ | ⭕ | +| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | +| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | +| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | +| Ascend NPU | ❌ | ✅ | ✅ | ✅ | Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this -guide. FP8 on Ampere may use a weight-only path where available. +guide. FP8 on Ampere may use a weight-only path where available. MXFP8 and +MXFP4 are documented for the Ascend NPU path. ## Model Type Support @@ -32,6 +36,8 @@ guide. FP8 on Ampere may use a weight-only path where available. |--------|-------|----------------|--------| | FP8 W8A8 | [FP8](fp8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models | | Int8 W8A8 | [Int8](int8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image | +| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B, Wan2.2-TI2V-5B | Ascend NPU only | +| MXFP4 W4A4 | [MXFP4](mxfp4.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B | Ascend NPU only; TI2V-5B is not supported | ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS) @@ -53,6 +59,9 @@ from vllm_omni import Omni omni_fp8 = Omni(model="", quantization="fp8") omni_int8 = Omni(model="", quantization="int8") +omni_mxfp8 = Omni(model="", quantization="mxfp8") +omni_mxfp4 = Omni(model="", quantization="mxfp4") +omni_mxfp4_dualscale = Omni(model="", quantization="mxfp4_dualscale") ``` CLI: @@ -60,6 +69,9 @@ CLI: ```bash vllm serve --omni --quantization fp8 vllm serve --omni --quantization int8 +vllm serve --omni --quantization mxfp8 +vllm serve --omni --quantization mxfp4 +vllm serve --omni --quantization mxfp4_dualscale ``` Per-component routing: @@ -77,10 +89,11 @@ config = build_quant_config({ | Parameter | Methods | Description | |-----------|---------|-------------| -| `method` | FP8, Int8 | Quantization method, usually `"fp8"` or `"int8"` | -| `ignored_layers` | FP8, Int8 | Layer name patterns to keep in BF16/FP16 | +| `method` | FP8, Int8, MXFP8, MXFP4 | Quantization method: `"fp8"`, `"int8"`, `"mxfp8"`, `"mxfp4"`, or `"mxfp4_dualscale"` | +| `ignored_layers` | FP8, Int8, MXFP8, MXFP4 | Layer name patterns to keep in BF16/FP16 | | `activation_scheme` | FP8, Int8 | The runtime value `"dynamic"` selects online activation scaling | | `weight_block_size` | FP8 | Optional block-wise FP8 weight quantization size | +| `num_bf16_fallback_layers` | MXFP4 DualScale | Leading transformer blocks to keep in BF16 for online `mxfp4_dualscale`; defaults to `5` | ## Validation and Notes @@ -91,3 +104,5 @@ config = build_quant_config({ model as supported. 4. If a model already ships quantized weights, use the matching pre-quantized method guide instead of online quantization. +5. For Ascend MXFP4 deployments, prefer offline `mxfp4_dualscale` checkpoints + when production quality is more important than avoiding preprocessing. From 083b5e363fbd32203e30cfaa3de2199370badec3 Mon Sep 17 00:00:00 2001 From: Samit <285365963@qq.com> Date: Wed, 20 May 2026 22:26:18 +0800 Subject: [PATCH 32/38] [Diffusion] Support LoRA in step-wise execution (#3639) Signed-off-by: samithuang <285365963@qq.com> --- .../feature/diffusion_continuous_batching.md | 5 +- docs/user_guide/diffusion/step_execution.md | 3 + docs/user_guide/diffusion_features.md | 4 +- tests/diffusion/test_diffusion_scheduler.py | 156 ++++++++++++++++ .../diffusion/test_diffusion_step_pipeline.py | 166 +++++++++++++----- vllm_omni/diffusion/sched/base_scheduler.py | 10 +- vllm_omni/diffusion/sched/interface.py | 5 + .../diffusion/worker/diffusion_worker.py | 49 +++++- 8 files changed, 338 insertions(+), 60 deletions(-) diff --git a/docs/design/feature/diffusion_continuous_batching.md b/docs/design/feature/diffusion_continuous_batching.md index ab8ff69fe..c443dea78 100644 --- a/docs/design/feature/diffusion_continuous_batching.md +++ b/docs/design/feature/diffusion_continuous_batching.md @@ -83,7 +83,10 @@ There are two important details: Today that compatibility rule is still shape-sensitive. `height`, `width`, `num_frames`, and CFG-related fields remain part of the key, so different -resolutions or incompatible guidance settings do **not** co-batch yet. +resolutions or incompatible guidance settings do **not** co-batch yet. The +key also covers LoRA identity (`lora_int_id`, `lora_scale`), so requests +targeting different adapters or scales run in separate batches and the +worker can activate exactly one adapter per step. The current batching unit is one `OmniDiffusionRequest`. Requests with multiple prompts do not participate in batching today. diff --git a/docs/user_guide/diffusion/step_execution.md b/docs/user_guide/diffusion/step_execution.md index 189510916..07184a787 100644 --- a/docs/user_guide/diffusion/step_execution.md +++ b/docs/user_guide/diffusion/step_execution.md @@ -60,6 +60,9 @@ batching only becomes relevant when `--max-num-seqs > 1`. - `cache_backend` is not supported together with step execution. - Unsupported pipelines fail early during model loading. - Request-mode extras such as KV transfer are not wired into step mode yet. +- LoRA is supported in step mode, but each batch must use a single adapter: + requests with different `lora_request` or `lora_scale` are scheduled into + separate batches. ## When To Use It diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 26b4d9bf7..61694ac9d 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -184,7 +184,7 @@ The following tables show which models support each feature: | **💾VAE Patch Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | | | | | | **💾FP8 Quant** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ✅ | ✅ | | | | | **🔧LoRA Inference** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | | -| **🔄Step Execution** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ❓ | ✅ | ✅ | ❌ | | +| **🔄Step Execution** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ❓ | ✅ | ✅ | ✅ | | !!! info @@ -193,7 +193,7 @@ The following tables show which models support each feature: 3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible. 4. CPU Offloading (Layerwise) supports single-card for now. 5. Using FP8-Quant as an example of qunatization methods. - 6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT) or LoRA. + 6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT). LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches). ## Multi-Thread Weight Loading diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index f0bc1da35..ca28b2629 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/tests/diffusion/test_diffusion_scheduler.py @@ -153,6 +153,54 @@ class _StubScheduler(SchedulerInterface): return None +class TestGetSamplingParamsKey: + """Pure-function tests for the batch-compatibility key builder.""" + + @staticmethod + def _make(lora_int_id: int | None = None, lora_scale: float = 1.0) -> OmniDiffusionRequest: + from vllm_omni.lora.request import LoRARequest + + sp = OmniDiffusionSamplingParams(num_inference_steps=2) + if lora_int_id is not None: + sp.lora_request = LoRARequest( + lora_name=f"adapter-{lora_int_id}", + lora_int_id=lora_int_id, + lora_path=f"/tmp/lora-{lora_int_id}", + ) + sp.lora_scale = lora_scale + return OmniDiffusionRequest( + prompts=["prompt"], + sampling_params=sp, + request_ids=[f"req-{lora_int_id}-{lora_scale}"], + ) + + def test_distinguishes_lora_id(self) -> None: + from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key + + assert get_sampling_params_key(self._make(lora_int_id=1)) != get_sampling_params_key(self._make(lora_int_id=2)) + + def test_distinguishes_lora_scale(self) -> None: + from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key + + assert get_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5)) != get_sampling_params_key( + self._make(lora_int_id=1, lora_scale=1.0) + ) + + def test_treats_no_lora_as_distinct_bucket(self) -> None: + from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key + + assert get_sampling_params_key(self._make(lora_int_id=None)) != get_sampling_params_key( + self._make(lora_int_id=1) + ) + + def test_equal_for_same_lora_identity(self) -> None: + from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key + + a = get_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5)) + b = get_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5)) + assert a == b + + class TestRequestScheduler: def setup_method(self) -> None: self.scheduler: RequestScheduler = RequestScheduler() @@ -728,6 +776,114 @@ class TestStepScheduler: assert second.num_running_reqs == 1 assert second.num_waiting_reqs == 1 + def test_step_batch_co_schedules_requests_sharing_lora(self) -> None: + """Multiple requests with the same LoRA (id + scale) co-batch.""" + from vllm_omni.lora.request import LoRARequest + + scheduler = StepScheduler() + scheduler.initialize(SimpleNamespace(max_num_seqs=3)) + + lora = LoRARequest(lora_name="adapter", lora_int_id=42, lora_path="/tmp/lora") + + def _with_lora(req_id: str) -> OmniDiffusionRequest: + sp = OmniDiffusionSamplingParams(num_inference_steps=4) + sp.lora_request = lora + sp.lora_scale = 0.5 + return _make_step_request(req_id, sampling_params=sp) + + req_a = scheduler.add_request(_with_lora("a")) + req_b = scheduler.add_request(_with_lora("b")) + req_c = scheduler.add_request(_with_lora("c")) + + sched_output = scheduler.schedule() + + assert _new_ids(sched_output) == [req_a, req_b, req_c] + assert sched_output.num_running_reqs == 3 + assert sched_output.num_waiting_reqs == 0 + + def test_step_batch_separates_requests_with_different_lora_ids(self) -> None: + """Different LoRA adapters → distinct batches admitted in FIFO order.""" + from vllm_omni.lora.request import LoRARequest + + scheduler = StepScheduler() + scheduler.initialize(SimpleNamespace(max_num_seqs=4)) + + lora_a = LoRARequest(lora_name="adapter-A", lora_int_id=1, lora_path="/tmp/lora-a") + lora_b = LoRARequest(lora_name="adapter-B", lora_int_id=2, lora_path="/tmp/lora-b") + + def _build(req_id: str, lora: LoRARequest) -> OmniDiffusionRequest: + sp = OmniDiffusionSamplingParams(num_inference_steps=2) + sp.lora_request = lora + return _make_step_request(req_id, sampling_params=sp) + + req_a1 = scheduler.add_request(_build("a1", lora_a)) + req_b1 = scheduler.add_request(_build("b1", lora_b)) + req_a2 = scheduler.add_request(_build("a2", lora_a)) + + # Strict FIFO admission: a1 starts; b1 (different LoRA) blocks the + # queue head, so a2 (compatible with a1) is *not* skipped ahead. + first = scheduler.schedule() + assert _new_ids(first) == [req_a1] + assert first.num_waiting_reqs == 2 + + # Drain a1 → b1 becomes head-of-line and is admitted with its LoRA. + scheduler.update_from_output(first, _make_step_output(req_a1, step_index=2, finished=True)) + second = scheduler.schedule() + assert _new_ids(second) == [req_b1] + assert second.num_waiting_reqs == 1 + + # Drain b1 → a2 is admitted next; LoRA-A is re-activated for it. + scheduler.update_from_output(second, _make_step_output(req_b1, step_index=2, finished=True)) + third = scheduler.schedule() + assert _new_ids(third) == [req_a2] + assert third.num_waiting_reqs == 0 + + def test_step_batch_separates_requests_with_different_lora_scale(self) -> None: + """Same adapter id but different scales → still separate batches.""" + from vllm_omni.lora.request import LoRARequest + + scheduler = StepScheduler() + scheduler.initialize(SimpleNamespace(max_num_seqs=4)) + + lora = LoRARequest(lora_name="adapter", lora_int_id=7, lora_path="/tmp/lora") + + def _build(req_id: str, scale: float) -> OmniDiffusionRequest: + sp = OmniDiffusionSamplingParams(num_inference_steps=2) + sp.lora_request = lora + sp.lora_scale = scale + return _make_step_request(req_id, sampling_params=sp) + + req_full = scheduler.add_request(_build("full", 1.0)) + req_half = scheduler.add_request(_build("half", 0.5)) + + sched_output = scheduler.schedule() + + admitted = _new_ids(sched_output) + assert admitted == [req_full] + assert req_half not in admitted + assert sched_output.num_waiting_reqs == 1 + + def test_step_batch_separates_lora_from_no_lora(self) -> None: + """A LoRA request and a no-LoRA request do not share a batch.""" + from vllm_omni.lora.request import LoRARequest + + scheduler = StepScheduler() + scheduler.initialize(SimpleNamespace(max_num_seqs=4)) + + lora = LoRARequest(lora_name="adapter", lora_int_id=3, lora_path="/tmp/lora") + + sp_with = OmniDiffusionSamplingParams(num_inference_steps=2) + sp_with.lora_request = lora + req_with = scheduler.add_request(_make_step_request("with", sampling_params=sp_with)) + req_without = scheduler.add_request(_make_step_request("without", num_inference_steps=2)) + + sched_output = scheduler.schedule() + + admitted = _new_ids(sched_output) + assert admitted == [req_with] + assert req_without not in admitted + assert sched_output.num_waiting_reqs == 1 + def test_preempt_request_preserves_step_index(self) -> None: request = _make_step_request("preempt", num_inference_steps=3) req_id = self.scheduler.add_request(request) diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py index 1954f4a05..2722de7e3 100644 --- a/tests/diffusion/test_diffusion_step_pipeline.py +++ b/tests/diffusion/test_diffusion_step_pipeline.py @@ -285,6 +285,19 @@ def _make_scheduler_output(req, sched_req_id="req-1", step_id=0, finished_req_id ) +def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None): + """Scheduler output for a homogeneous batch (one NewRequestData per req).""" + new_reqs = [NewRequestData(sched_req_id=r.request_ids[0], req=r) for r in reqs] + return DiffusionSchedulerOutput( + step_id=step_id, + scheduled_new_reqs=new_reqs, + scheduled_cached_reqs=CachedRequestData.make_empty(), + finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids), + num_running_reqs=len(new_reqs), + num_waiting_reqs=0, + ) + + def _make_cached_scheduler_output(sched_req_id="req-1", step_id=1, finished_req_ids=None): return DiffusionSchedulerOutput( step_id=step_id, @@ -489,71 +502,128 @@ class TestRunner: DiffusionModelRunner.load_model(runner) +class _RecordingLoRAManager: + def __init__(self) -> None: + self.calls: list[tuple[object | None, float]] = [] + + def set_active_adapter(self, adapter, scale: float = 1.0) -> None: + self.calls.append((adapter, scale)) + + +def _make_step_worker(lora_manager=None, *, expected_output=None): + """Build a bare DiffusionWorker primed for execute_stepwise tests.""" + worker = object.__new__(DiffusionWorker) + worker.lora_manager = lora_manager + worker._step_lora_state = {} + output = expected_output if expected_output is not None else RunnerOutput(req_id="req-1") + worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: output) + return worker + + @pytest.mark.cpu class TestWorker: """DiffusionWorker.execute_stepwise""" def test_delegates_to_model_runner(self): - worker = object.__new__(DiffusionWorker) expected = RunnerOutput(req_id="req-1", step_index=1, finished=False, result=None) - scheduler_output = SimpleNamespace( - scheduled_new_reqs=[ - SimpleNamespace( - req=SimpleNamespace( - sampling_params=SimpleNamespace(lora_request=None), - ) - ) - ] - ) - worker.lora_manager = None - worker.model_runner = SimpleNamespace( - execute_stepwise=lambda arg: expected if arg is scheduler_output else None - ) + worker = _make_step_worker(expected_output=expected) + scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), sched_req_id="req-1") output = DiffusionWorker.execute_stepwise(worker, scheduler_output) assert output is expected - def test_clears_active_lora_before_stepwise_execution(self): - worker = object.__new__(DiffusionWorker) - scheduler_output = SimpleNamespace( - scheduled_new_reqs=[ - SimpleNamespace( - req=SimpleNamespace( - sampling_params=SimpleNamespace(lora_request=None), - ) - ) - ] - ) - calls: list[object | None] = [] - - class _FakeLoRAManager: - def set_active_adapter(self, adapter): - calls.append(adapter) - - worker.lora_manager = _FakeLoRAManager() - worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: RunnerOutput(req_id="req-1")) + def test_deactivates_lora_when_request_has_no_adapter(self): + manager = _RecordingLoRAManager() + worker = _make_step_worker(lora_manager=manager) + scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), sched_req_id="req-1") DiffusionWorker.execute_stepwise(worker, scheduler_output) - assert calls == [None] + assert manager.calls == [(None, 1.0)] - def test_rejects_lora_requests_in_step_mode(self): - worker = object.__new__(DiffusionWorker) - scheduler_output = SimpleNamespace( - scheduled_new_reqs=[ - SimpleNamespace( - req=SimpleNamespace( - sampling_params=SimpleNamespace(lora_request=object()), - ) - ) - ] + def test_activates_lora_for_step_requests(self): + from vllm_omni.lora.request import LoRARequest + + lora_request = LoRARequest(lora_name="adapter", lora_int_id=7, lora_path="/tmp/lora") + request = _make_engine_request("req-1") + request.sampling_params.lora_request = lora_request + request.sampling_params.lora_scale = 0.75 + + manager = _RecordingLoRAManager() + worker = _make_step_worker(lora_manager=manager) + scheduler_output = _make_scheduler_output(request, sched_req_id="req-1") + + DiffusionWorker.execute_stepwise(worker, scheduler_output) + + assert manager.calls == [(lora_request, 0.75)] + + def test_recovers_lora_for_cached_step_requests(self): + from vllm_omni.lora.request import LoRARequest + + lora_request = LoRARequest(lora_name="adapter", lora_int_id=11, lora_path="/tmp/lora") + request = _make_engine_request("req-1") + request.sampling_params.lora_request = lora_request + request.sampling_params.lora_scale = 0.5 + + manager = _RecordingLoRAManager() + worker = _make_step_worker(lora_manager=manager) + first = _make_scheduler_output(request, sched_req_id="req-1") + second = _make_cached_scheduler_output(sched_req_id="req-1", step_id=1) + + DiffusionWorker.execute_stepwise(worker, first) + DiffusionWorker.execute_stepwise(worker, second) + + assert manager.calls == [(lora_request, 0.5), (lora_request, 0.5)] + + def test_activates_single_lora_for_homogeneous_batch(self): + """Multiple requests sharing the same LoRA → exactly one activation, + and every request id is registered in ``_step_lora_state``.""" + from vllm_omni.lora.request import LoRARequest + + lora_request = LoRARequest(lora_name="adapter", lora_int_id=9, lora_path="/tmp/lora") + reqs = [] + for rid in ("req-1", "req-2", "req-3"): + r = _make_engine_request(rid) + r.sampling_params.lora_request = lora_request + r.sampling_params.lora_scale = 0.6 + reqs.append(r) + + manager = _RecordingLoRAManager() + worker = _make_step_worker(lora_manager=manager) + scheduler_output = _make_batch_scheduler_output(reqs) + + DiffusionWorker.execute_stepwise(worker, scheduler_output) + + assert manager.calls == [(lora_request, 0.6)] + assert set(worker._step_lora_state) == {"req-1", "req-2", "req-3"} + for entry in worker._step_lora_state.values(): + assert entry == (lora_request, 0.6) + + def test_evicts_step_lora_state_for_finished_requests(self): + from vllm_omni.lora.request import LoRARequest + + lora_request = LoRARequest(lora_name="adapter", lora_int_id=3, lora_path="/tmp/lora") + finishing = _make_engine_request("req-1") + finishing.sampling_params.lora_request = lora_request + next_request = _make_engine_request("req-2") + next_request.sampling_params.lora_request = lora_request + + worker = _make_step_worker(lora_manager=_RecordingLoRAManager()) + first = _make_scheduler_output(finishing, sched_req_id="req-1") + next_batch = _make_scheduler_output( + next_request, + sched_req_id="req-2", + step_id=1, + finished_req_ids={"req-1"}, ) - worker.lora_manager = None - worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: RunnerOutput(req_id="req-1")) - with pytest.raises(ValueError, match="does not support LoRA"): - DiffusionWorker.execute_stepwise(worker, scheduler_output) + DiffusionWorker.execute_stepwise(worker, first) + assert "req-1" in worker._step_lora_state + + DiffusionWorker.execute_stepwise(worker, next_batch) + assert "req-1" not in worker._step_lora_state + assert worker._step_lora_state == {"req-2": (lora_request, 1.0)} @pytest.mark.cpu diff --git a/vllm_omni/diffusion/sched/base_scheduler.py b/vllm_omni/diffusion/sched/base_scheduler.py index f21c03e98..0bd16597f 100644 --- a/vllm_omni/diffusion/sched/base_scheduler.py +++ b/vllm_omni/diffusion/sched/base_scheduler.py @@ -22,7 +22,9 @@ from vllm_omni.diffusion.sched.interface import ( logger = init_logger(__name__) -_KEY_FIELD_NAMES = frozenset(f.name for f in fields(SamplingParamsKey)) +# LoRA identity is derived from `sampling.lora_request`, not a same-named field +# on sampling params, so it must be resolved separately from the bulk lookup. +_KEY_FIELD_NAMES = frozenset(f.name for f in fields(SamplingParamsKey)) - {"lora_int_id"} def get_sampling_params_key(request: OmniDiffusionRequest) -> SamplingParamsKey | None: @@ -31,7 +33,11 @@ def get_sampling_params_key(request: OmniDiffusionRequest) -> SamplingParamsKey return None sampling = request.sampling_params - return SamplingParamsKey(**{name: getattr(sampling, name) for name in _KEY_FIELD_NAMES}) + lora_request = getattr(sampling, "lora_request", None) + return SamplingParamsKey( + lora_int_id=lora_request.lora_int_id if lora_request is not None else None, + **{name: getattr(sampling, name) for name in _KEY_FIELD_NAMES}, + ) class _BaseScheduler(SchedulerInterface): diff --git a/vllm_omni/diffusion/sched/interface.py b/vllm_omni/diffusion/sched/interface.py index 8eef211ff..1bfb69458 100644 --- a/vllm_omni/diffusion/sched/interface.py +++ b/vllm_omni/diffusion/sched/interface.py @@ -65,6 +65,11 @@ class SamplingParamsKey: true_cfg_scale: float | None = None cfg_normalize: bool = False + # LoRA identity. Requests with different adapters or scales must run in + # separate batches so the worker can activate exactly one adapter per step. + lora_int_id: int | None = None + lora_scale: float = 1.0 + @dataclass class DiffusionRequestState: diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index e2c4b97c1..cd5fe246d 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -185,6 +185,10 @@ class DiffusionWorker: self.model_runner: DiffusionModelRunner | None = None self._sleep_saved_buffers: dict[str, torch.Tensor] = {} self.lora_manager: DiffusionLoRAManager | None = None + # Worker-side cache of (lora_request, lora_scale) per scheduled + # request id. Used by step mode to recover LoRA identity for cached + # requests, which only carry their sched_req_id in subsequent ticks. + self._step_lora_state: dict[str, tuple[LoRARequest | None, float]] = {} self.stage_id = getattr(od_config, "stage_id", 0) self.init_device() # Create model runner using the platform-specified class @@ -367,13 +371,7 @@ class DiffusionWorker: def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput: """Execute one diffusion step by delegating to the model runner.""" assert self.model_runner is not None, "Model runner not initialized" - if self.lora_manager is not None: - # Step mode does not support LoRA yet. Clear any previously active - # adapter first so worker-local LoRA state cannot leak in. - self.lora_manager.set_active_adapter(None) - - if any(new_req.req.sampling_params.lora_request is not None for new_req in scheduler_output.scheduled_new_reqs): - raise ValueError("Step mode does not support LoRA yet.") + self._activate_step_lora(scheduler_output) profiler = self._get_profiler() ctx = profiler.annotate_context_manager("diffusion_step") if profiler else nullcontext() with ctx: @@ -382,6 +380,43 @@ class DiffusionWorker: profiler.step() return output + def _activate_step_lora(self, scheduler_output: DiffusionSchedulerOutput) -> None: + """Activate the LoRA adapter for the scheduled step batch. + + Newly scheduled requests register their (lora_request, lora_scale) + in ``_step_lora_state`` so cached requests can resolve to the same + adapter on later ticks. Finished requests are evicted. Batch + homogeneity is enforced by the scheduler via ``SamplingParamsKey``, + so any scheduled request id resolves to the active LoRA identity. + """ + for sched_req_id in scheduler_output.finished_req_ids: + self._step_lora_state.pop(sched_req_id, None) + + for new_req in scheduler_output.scheduled_new_reqs: + sampling = new_req.req.sampling_params + self._step_lora_state[new_req.sched_req_id] = ( + sampling.lora_request, + sampling.lora_scale, + ) + + if self.lora_manager is None: + return + + lora_request: LoRARequest | None = None + lora_scale = 1.0 + for sched_req_id in scheduler_output.scheduled_req_ids: + entry = self._step_lora_state.get(sched_req_id) + if entry is not None: + lora_request, lora_scale = entry + break + + try: + self.lora_manager.set_active_adapter(lora_request, lora_scale) + except Exception as exc: + if lora_request is not None: + raise + logger.warning("LoRA activation skipped: %s", exc) + def load_weights(self, weights) -> set[str]: """Load weights by delegating to the model runner.""" assert self.model_runner is not None, "Model runner not initialized" From 7d9d8aee79d02d4616018ff277e9d32971c6da17 Mon Sep 17 00:00:00 2001 From: Siyuan Kong <101997501+ksiyuan@users.noreply.github.com> Date: Thu, 21 May 2026 16:43:33 +0800 Subject: [PATCH 33/38] [Bugfix] Fix qwen2_5_omni weight loading (#3598) Signed-off-by: ksiyuan --- .../qwen2_5_omni/test_weight_loading.py | 49 +++++++++++++++++++ .../models/qwen2_5_omni/qwen2_5_omni.py | 1 - .../qwen2_5_omni/qwen2_5_omni_thinker.py | 10 ++++ 3 files changed, 59 insertions(+), 1 deletion(-) create mode 100644 tests/model_executor/models/qwen2_5_omni/test_weight_loading.py diff --git a/tests/model_executor/models/qwen2_5_omni/test_weight_loading.py b/tests/model_executor/models/qwen2_5_omni/test_weight_loading.py new file mode 100644 index 000000000..1205dedb6 --- /dev/null +++ b/tests/model_executor/models/qwen2_5_omni/test_weight_loading.py @@ -0,0 +1,49 @@ +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestQwen2_5OmniWeightLoading: + def test_qkv_weight_mapping(self): + mock_quant_description = { + "visual.blocks.0.attn.q.weight": {"quant_type": "FLOAT"}, + "visual.blocks.0.attn.k.weight": {"quant_type": "FLOAT"}, + "visual.blocks.0.attn.v.weight": {"quant_type": "FLOAT"}, + } + + prefix = "visual.blocks.0.attn.qkv" + packed_modules_mapping = { + "attn_akv_proj": [ + "attn_q_proj", + "attn_k_proj", + "attn_v_proj", + ], + "qkv": [ + "q", + "k", + "v", + ], + } + + try: + proj_name = prefix.split(".")[-1] + if proj_name in packed_modules_mapping: + quant_type = None + shard_prefixes = [ + prefix.replace(proj_name, shard_proj_name) for shard_proj_name in packed_modules_mapping[proj_name] + ] + for shard_prefix in shard_prefixes: + shard_quant_type = mock_quant_description[shard_prefix + ".weight"] + + if quant_type is None: + quant_type = shard_quant_type + elif shard_quant_type != quant_type: + raise ValueError( + f"Not all shards of {prefix} are quantized with same quant type." + f"Shard {proj_name} uses {shard_quant_type}, but another shard" + f"use {quant_type}. Please check quantization config." + ) + else: + quant_type = mock_quant_description[prefix + ".weight"] + except KeyError as e: + pytest.fail(f"KeyError was raised: {e}\n") diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py index 0b5d78560..712515f27 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py @@ -81,7 +81,6 @@ class Qwen2_5OmniForConditionalGeneration( # Initialize thinker model (multimodal processing) self.thinker = init_vllm_registered_model( vllm_config=vllm_config, - prefix=maybe_prefix(prefix, "thinker"), hf_config=thinker_config, # Use registry architecture key architectures=["Qwen2_5OmniThinkerModel"], diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py index c4be88a88..931049d8d 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py @@ -324,6 +324,16 @@ class Qwen2_5OmniThinkerForConditionalGeneration( "gate_proj", "up_proj", ], + "attn_qkv_proj": [ + "attn_q_proj", + "attn_k_proj", + "attn_v_proj", + ], + "qkv": [ + "q", + "k", + "v", + ], } @classmethod From 474710dda90beb628f94255b8a3d5de54f033656 Mon Sep 17 00:00:00 2001 From: NumberWan Date: Thu, 21 May 2026 17:18:02 +0800 Subject: [PATCH 34/38] [Benchmark] Route i2i/ti2i to POST /v1/images/edits in diffusion_benchmark_serving (#3693) Signed-off-by: NumberWan --- benchmarks/diffusion/backends.py | 93 +++++++++++++++++++ .../diffusion/diffusion_benchmark_serving.py | 40 ++++++-- 2 files changed, 124 insertions(+), 9 deletions(-) diff --git a/benchmarks/diffusion/backends.py b/benchmarks/diffusion/backends.py index b7e533b2d..368dc5dcb 100644 --- a/benchmarks/diffusion/backends.py +++ b/benchmarks/diffusion/backends.py @@ -10,6 +10,8 @@ from typing import Any import aiohttp from tqdm import tqdm +DEFAULT_EDITS_BOT_TASK = "think" + @dataclass class RequestFuncInput: @@ -27,6 +29,7 @@ class RequestFuncInput: extra_body: dict[str, Any] = field(default_factory=dict) image_paths: list[str] | None = None request_id: str = field(default_factory=lambda: str(uuid.uuid4())) + default_bot_task: str | None = DEFAULT_EDITS_BOT_TASK @dataclass @@ -53,6 +56,95 @@ def _encode_image_as_data_url(path: str) -> str: return f"data:{mime};base64,{encoded}" +async def async_request_image_edits( + input: RequestFuncInput, + session: aiohttp.ClientSession, + pbar: tqdm | None = None, + enable_diffusion_pipeline_profiler: bool = False, +) -> RequestFuncOutput: + """POST /v1/images/edits (multipart).""" + del enable_diffusion_pipeline_profiler + output = RequestFuncOutput() + output.start_time = time.perf_counter() + + extra_body = dict(input.extra_body) + width = input.width or extra_body.get("width") or 1024 + height = input.height or extra_body.get("height") or 1024 + edits_url = input.api_url + + form = aiohttp.FormData() + form.add_field("model", input.model) + form.add_field("prompt", input.prompt) + form.add_field("size", f"{width}x{height}") + form.add_field("response_format", "b64_json") + + if input.num_inference_steps is not None: + form.add_field("num_inference_steps", str(input.num_inference_steps)) + elif extra_body.get("num_inference_steps") is not None: + form.add_field("num_inference_steps", str(extra_body["num_inference_steps"])) + + if input.seed is not None: + form.add_field("seed", str(input.seed)) + elif extra_body.get("seed") is not None: + form.add_field("seed", str(extra_body["seed"])) + + if extra_body.get("guidance_scale") is not None: + form.add_field("guidance_scale", str(extra_body["guidance_scale"])) + if extra_body.get("negative_prompt") is not None: + form.add_field("negative_prompt", str(extra_body["negative_prompt"])) + if extra_body.get("true_cfg_scale") is not None: + form.add_field("true_cfg_scale", str(extra_body["true_cfg_scale"])) + if extra_body.get("sys_type") is not None: + form.add_field("sys_type", str(extra_body["sys_type"])) + if extra_body.get("system_prompt") is not None: + form.add_field("system_prompt", str(extra_body["system_prompt"])) + + bot_task = extra_body.get("bot_task") + if bot_task is None and input.default_bot_task is not None: + bot_task = input.default_bot_task + if bot_task is not None: + form.add_field("bot_task", str(bot_task)) + + assert input.image_paths is not None + for img_path in input.image_paths: + if not os.path.exists(img_path): + output.error = f"Image file not found: {img_path}" + output.success = False + if pbar: + pbar.update(1) + return output + with open(img_path, "rb") as img_f: + image_bytes = img_f.read() + form.add_field( + "image", + image_bytes, + filename=os.path.basename(img_path), + content_type=_guess_mime_type(img_path), + ) + + try: + async with session.post(edits_url, data=form) as response: + if response.status == 200: + resp_json = await response.json() + output.response_body = resp_json + output.success = True + else: + output.error = f"HTTP {response.status}: {await response.text()}" + output.success = False + except Exception as e: + output.error = str(e) + output.success = False + + output.latency = time.perf_counter() - output.start_time + + if output.success and input.slo_ms is not None: + output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms) + + if pbar: + pbar.update(1) + return output + + async def async_request_chat_completions( input: RequestFuncInput, session: aiohttp.ClientSession, @@ -378,6 +470,7 @@ backends_function_mapping = { "2i": { "/v1/chat/completions": (async_request_chat_completions, "/v1/chat/completions"), "/v1/images/generations": (async_request_openai_image_generations, "/v1/images/generations"), + "/v1/images/edits": (async_request_image_edits, "/v1/images/edits"), }, "2v": { "/v1/videos": (async_request_v1_videos, "/v1/videos"), diff --git a/benchmarks/diffusion/diffusion_benchmark_serving.py b/benchmarks/diffusion/diffusion_benchmark_serving.py index c59844f65..7e8ea97e1 100644 --- a/benchmarks/diffusion/diffusion_benchmark_serving.py +++ b/benchmarks/diffusion/diffusion_benchmark_serving.py @@ -7,10 +7,13 @@ Benchmark online serving for diffusion models (Image/Video Generation). If you want to use i2v, i2i dataset, you should `uv pip install gdown` first Supports multiple endpoints: - - /v1/chat/completions: OpenAI chat-compatible image requests + - /v1/chat/completions: OpenAI chat-compatible image requests (e.g. t2i, Qwen i2i) + - /v1/images/edits: OpenAI image edit / IT2I (multipart; e.g. Hunyuan --bot-task think) - /v1/images/generations: OpenAI image generation requests - /v1/videos: Async video jobs +Legacy --backend vllm-omni and openai are aliases for chat/completions and images/generations. + Usage: # Video (/v1/videos endpoint) t2v: @@ -39,7 +42,12 @@ Usage: {"width":1536,"height":1536,"num_inference_steps":35,"weight":0.15} ]' - i2i: + ti2i (Hunyuan / OpenAI image edit API): + python3 benchmarks/diffusion/diffusion_benchmark_serving.py \ + --endpoint /v1/images/edits --dataset random --task ti2i --num-prompts 10 \ + --bot-task think + + i2i (chat-based models such as Qwen-Image-Edit): python3 benchmarks/diffusion/diffusion_benchmark_serving.py \ --endpoint /v1/chat/completions --dataset vbench --task i2i --num-prompts 10 @@ -132,7 +140,7 @@ class VBenchDataset(BaseDataset): def _load_data(self) -> list[dict[str, Any]]: if self.args.task == "t2v": return self._load_t2v_prompts() - elif self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]: + elif self.args.task in ["i2v", "ti2v", "ti2i", "i2i", "it2i"]: return self._load_i2v_data() else: return self._load_t2v_prompts() @@ -524,7 +532,7 @@ class TraceDataset(BaseDataset): single = row.get("image_path") image_paths = [single] if single else None - if not image_paths and self.args.task in ["i2v", "i2i", "ti2v", "ti2i"]: + if not image_paths and self.args.task in ["i2v", "i2i", "ti2v", "ti2i", "it2i"]: raise ValueError( f"Task {self.args.task} requires image input, but no image_path or image_paths found in trace row." ) @@ -585,7 +593,7 @@ class RandomDataset(BaseDataset): self._sampled_requests = None # Random image generate - if self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]: + if self.args.task in ["i2v", "ti2v", "ti2i", "i2i", "it2i"]: self._random_image_path = self._generate_random_image_paths() else: self._random_image_path = None @@ -891,7 +899,9 @@ def wait_for_service(base_url: str, timeout: int = 120) -> None: def _default_endpoint_for_task(task: str) -> str: if task in {"t2v", "i2v", "ti2v"}: return "/v1/videos" - if task in {"t2i", "i2i", "ti2i"}: + if task in {"i2i", "ti2i", "it2i"}: + return "/v1/images/edits" + if task == "t2i": return "/v1/chat/completions" raise ValueError(f"Unsupported task for endpoint resolution: {task}") @@ -902,7 +912,7 @@ async def benchmark(args): args.base_url = f"http://{args.host}:{args.port}" VIDEO_TASKS = {"t2v", "i2v", "ti2v"} - IMAGE_TASKS = {"t2i", "i2i", "ti2i"} + IMAGE_TASKS = {"t2i", "i2i", "ti2i", "it2i"} if args.task in VIDEO_TASKS: task_type = "2v" @@ -947,6 +957,10 @@ async def benchmark(args): requests_list = dataset.get_requests() print(f"Prepared {len(requests_list)} requests from {args.dataset} dataset.") + if args.endpoint == "/v1/images/edits": + for req in requests_list: + req.default_bot_task = args.bot_task + # Limit concurrency if args.max_concurrency is not None: semaphore = asyncio.Semaphore(args.max_concurrency) @@ -998,6 +1012,8 @@ async def benchmark(args): metrics["model"] = args.model metrics["dataset"] = args.dataset metrics["task"] = args.task + if args.endpoint == "/v1/images/edits": + metrics["bot_task"] = args.bot_task print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="=")) @@ -1088,7 +1104,7 @@ if __name__ == "__main__": "--task", type=str, default="t2v", - choices=["t2v", "i2v", "ti2v", "ti2i", "i2i", "t2i"], + choices=["t2v", "i2v", "ti2v", "ti2i", "i2i", "it2i", "t2i"], help="Task type.", ) parser.add_argument( @@ -1195,9 +1211,15 @@ if __name__ == "__main__": default=1, help=( "Number of synthetic input images to attach for image-conditioned tasks " - "(i2v, ti2v, ti2i, i2i) when using random dataset." + "(i2v, ti2v, ti2i, i2i, it2i) when using random dataset." ), ) + parser.add_argument( + "--bot-task", + type=str, + default="think", + help=("bot_task form field for --endpoint /v1/images/edits (think, recaption, think_recaption, vanilla)."), + ) args = parser.parse_args() asyncio.run(benchmark(args)) From 8297570c2b8deca12d7dcf4fb12ee417128984ba Mon Sep 17 00:00:00 2001 From: Liang Lv Date: Thu, 21 May 2026 17:54:41 +0800 Subject: [PATCH 35/38] [AutoRound] Support WAN2.2 W4A16 quantization model (#3353) Signed-off-by: lvliang-intel Signed-off-by: hyh_hh Co-authored-by: hxhhhlalala Co-authored-by: hyh_hh --- .buildkite/test-nightly.yml | 2 +- docs/user_guide/quantization/autoround.md | 4 +- docs/user_guide/quantization/mxfp4.md | 2 +- .../test_wan22_quant_config_propagation.py | 301 +++++++++++++++++ .../test_wan22_autoround_w4a16_expansion.py | 311 ++++++++++++++++++ .../models/wan2_2/pipeline_wan2_2.py | 5 +- .../models/wan2_2/pipeline_wan2_2_i2v.py | 19 +- .../models/wan2_2/pipeline_wan2_2_vace.py | 3 + vllm_omni/quantization/factory.py | 41 ++- 9 files changed, 680 insertions(+), 8 deletions(-) create mode 100644 tests/diffusion/models/wan2_2/test_wan22_quant_config_propagation.py create mode 100644 tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index 1b61044af..ae93d2353 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -959,7 +959,7 @@ steps: - label: ":full_moon: Diffusion X2V · Function Test" timeout_in_minutes: 90 commands: - - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m "full_model and cuda" --run-level "full_model" + - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model" agents: queue: "mithril-h100-pool" plugins: diff --git a/docs/user_guide/quantization/autoround.md b/docs/user_guide/quantization/autoround.md index 2261d79a5..88fed3b62 100644 --- a/docs/user_guide/quantization/autoround.md +++ b/docs/user_guide/quantization/autoround.md @@ -32,7 +32,9 @@ guide. AutoRound is Intel-supported. |-------|------------|-------|--------|---------| | FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend | | Qwen-Image | Not listed | Diffusion transformer | W4A16 | Not validated | -| Wan2.2 | Not listed | Diffusion transformer | W4A16 | Not validated | +| Wan2.2-I2V | `Intel/Wan2.2-I2V-A14B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend | +| Wan2.2-T2V | `Intel/Wan2.2-T2V-A14B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend | +| Wan2.2-TI2V | `Intel/Wan2.2-TI2V-5B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend | ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS) diff --git a/docs/user_guide/quantization/mxfp4.md b/docs/user_guide/quantization/mxfp4.md index 7463ada23..401a55ad4 100644 --- a/docs/user_guide/quantization/mxfp4.md +++ b/docs/user_guide/quantization/mxfp4.md @@ -397,7 +397,7 @@ names** discovered in Step 1. No code changes to the model are required. ```python omni = Omni( model="/path/to/your-model", - quantization_config={ + quantization={ "method": "mxfp4_dualscale", "ignored_layers": [ "blocks.0.attn1.to_qkv", # runtime name, not diffusers name diff --git a/tests/diffusion/models/wan2_2/test_wan22_quant_config_propagation.py b/tests/diffusion/models/wan2_2/test_wan22_quant_config_propagation.py new file mode 100644 index 000000000..b14059933 --- /dev/null +++ b/tests/diffusion/models/wan2_2/test_wan22_quant_config_propagation.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for Wan2.2 quant_config propagation through transformer creation. + +Tests cover: +- create_transformer_from_config passes quant_config and prefix +- create_vace_transformer_from_config passes quant_config and prefix +- set_tf_model_config propagates quant_config to OmniDiffusionConfig +- patch_wan_rms_norm safely iterates sys.modules with concurrent modifications +- I2V transformer_2 quant_config is built from config dict +""" + +import sys +from types import SimpleNamespace + +import pytest +from pytest_mock import MockerFixture + +import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module +import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace as wan22_vace_module +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( + create_transformer_from_config, +) +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import ( + create_vace_transformer_from_config, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +# --------------------------------------------------------------------------- +# create_transformer_from_config: quant_config / prefix forwarding +# --------------------------------------------------------------------------- + + +class TestCreateTransformerQuant: + """Verify quant_config and prefix are forwarded to WanTransformer3DModel.""" + + def test_quant_config_passed_through(self, mocker: MockerFixture): + captured = {} + + class FakeTransformer: + def __init__(self, **kwargs): + captured.update(kwargs) + + mocker.patch.object(wan22_module, "WanTransformer3DModel", FakeTransformer) + + fake_qc = mocker.MagicMock() + create_transformer_from_config( + {"patch_size": [1, 2, 2], "num_layers": 2}, + quant_config=fake_qc, + ) + assert captured.get("quant_config") is fake_qc + + def test_prefix_passed_through(self, mocker: MockerFixture): + captured = {} + + class FakeTransformer: + def __init__(self, **kwargs): + captured.update(kwargs) + + mocker.patch.object(wan22_module, "WanTransformer3DModel", FakeTransformer) + + create_transformer_from_config( + {"patch_size": [1, 2, 2]}, + prefix="model.transformer.", + ) + assert captured.get("prefix") == "model.transformer." + + def test_quant_config_none_by_default(self, mocker: MockerFixture): + captured = {} + + class FakeTransformer: + def __init__(self, **kwargs): + captured.update(kwargs) + + mocker.patch.object(wan22_module, "WanTransformer3DModel", FakeTransformer) + + create_transformer_from_config({"patch_size": [1, 2, 2]}) + # When quant_config is None and prefix is "", they are not added + assert "quant_config" not in captured or captured["quant_config"] is None + + def test_quant_config_and_prefix_together(self, mocker: MockerFixture): + captured = {} + + class FakeTransformer: + def __init__(self, **kwargs): + captured.update(kwargs) + + mocker.patch.object(wan22_module, "WanTransformer3DModel", FakeTransformer) + + fake_qc = mocker.MagicMock() + create_transformer_from_config( + {"patch_size": [1, 2, 2], "num_attention_heads": 4}, + quant_config=fake_qc, + prefix="blocks.", + ) + assert captured["quant_config"] is fake_qc + assert captured["prefix"] == "blocks." + + +# --------------------------------------------------------------------------- +# create_vace_transformer_from_config: quant_config / prefix forwarding +# --------------------------------------------------------------------------- + + +class TestCreateVaceTransformerQuant: + """Verify quant_config and prefix are forwarded to WanVACETransformer3DModel.""" + + def test_quant_config_passed_through(self, mocker: MockerFixture): + captured = {} + + class FakeVACETransformer: + def __init__(self, **kwargs): + captured.update(kwargs) + + mocker.patch.object(wan22_vace_module, "WanVACETransformer3DModel", FakeVACETransformer) + + fake_qc = mocker.MagicMock() + create_vace_transformer_from_config( + {"patch_size": [1, 2, 2], "num_layers": 2}, + quant_config=fake_qc, + ) + assert captured.get("quant_config") is fake_qc + + def test_prefix_passed_through(self, mocker: MockerFixture): + captured = {} + + class FakeVACETransformer: + def __init__(self, **kwargs): + captured.update(kwargs) + + mocker.patch.object(wan22_vace_module, "WanVACETransformer3DModel", FakeVACETransformer) + + create_vace_transformer_from_config( + {"patch_size": [1, 2, 2]}, + prefix="vace.", + ) + assert captured.get("prefix") == "vace." + + +# --------------------------------------------------------------------------- +# set_tf_model_config: propagation of quant_config +# --------------------------------------------------------------------------- + + +class TestSetTfModelConfig: + """Test that set_tf_model_config propagates quant_config correctly.""" + + def _make_od_config(self): + """Create a minimal OmniDiffusionConfig-like object for testing.""" + from vllm_omni.diffusion.data import OmniDiffusionConfig + + cfg = object.__new__(OmniDiffusionConfig) + cfg.quantization_config = None + cfg.tf_model_config = None + return cfg + + def test_propagates_quant_config_when_none(self, mocker: MockerFixture): + cfg = self._make_od_config() + fake_qc = mocker.MagicMock() + tf_config = SimpleNamespace(quant_config=fake_qc, quant_method="auto-round") + + cfg.set_tf_model_config(tf_config) + + assert cfg.tf_model_config is tf_config + assert cfg.quantization_config is fake_qc + + def test_does_not_overwrite_existing_quantization_config(self, mocker: MockerFixture): + cfg = self._make_od_config() + existing_qc = mocker.MagicMock() + cfg.quantization_config = existing_qc + tf_config = SimpleNamespace(quant_config=mocker.MagicMock()) + + cfg.set_tf_model_config(tf_config) + + assert cfg.tf_model_config is tf_config + assert cfg.quantization_config is existing_qc # not overwritten + + def test_no_propagation_when_tf_quant_config_is_none(self, mocker: MockerFixture): + cfg = self._make_od_config() + tf_config = SimpleNamespace(quant_config=None) + + cfg.set_tf_model_config(tf_config) + + assert cfg.tf_model_config is tf_config + assert cfg.quantization_config is None + + +# --------------------------------------------------------------------------- +# patch_wan_rms_norm: sys.modules snapshot safety +# --------------------------------------------------------------------------- + + +class TestPatchWanRmsNorm: + """Test that patch_wan_rms_norm doesn't raise on concurrent module registration.""" + + def test_patches_modules_with_wan_rms_norm(self): + from vllm_omni.diffusion.layers.norm import RMSNormVAE + from vllm_omni.diffusion.models.wan2_2.patch_diffusers import patch_wan_rms_norm + + # Create a fake module that has WanRMS_norm + fake_module = SimpleNamespace(WanRMS_norm=lambda x: x) + sys.modules["_test_fake_wan_module"] = fake_module + + try: + patch_wan_rms_norm() + assert fake_module.WanRMS_norm is RMSNormVAE + finally: + del sys.modules["_test_fake_wan_module"] + + def test_no_error_when_modules_change_during_iteration(self): + """Regression test: list() snapshot prevents RuntimeError.""" + from vllm_omni.diffusion.models.wan2_2.patch_diffusers import patch_wan_rms_norm + + # Simulate a module being added during iteration by a side effect + original_items = sys.modules.items + + def items_with_side_effect(): + # This would cause RuntimeError without list() snapshot + result = list(original_items()) + # Add a new module to simulate concurrent modification + sys.modules["_test_dynamic_module"] = SimpleNamespace() + return result + + try: + # The function uses list(sys.modules.items()) so it takes a snapshot + # Just verify it doesn't raise + patch_wan_rms_norm() + finally: + sys.modules.pop("_test_dynamic_module", None) + + +# --------------------------------------------------------------------------- +# I2V transformer_2 quant_config extraction +# --------------------------------------------------------------------------- + + +class TestI2VTransformer2QuantConfig: + """Test the transformer_2 quant_config build logic from pipeline_wan2_2_i2v.""" + + def test_transformer_2_quant_config_built_from_dict(self): + """When transformer_2 config has quantization_config dict, build_quant_config is called.""" + from vllm_omni.quantization.factory import build_quant_config + + t2_config = { + "patch_size": [1, 2, 2], + "num_layers": 2, + "quantization_config": { + "quant_method": "auto-round", + "bits": 4, + "group_size": 128, + "sym": True, + "packing_format": "auto_round:auto_gptq", + }, + } + + # Replicate the logic from pipeline_wan2_2_i2v.py + t2_quant = t2_config.get("quantization_config") + if isinstance(t2_quant, dict) and "quant_method" in t2_quant: + method = t2_quant["quant_method"] + kwargs = {k: v for k, v in t2_quant.items() if k != "quant_method"} + t2_quant = build_quant_config(method, **kwargs) + else: + t2_quant = None + + from vllm.model_executor.layers.quantization.inc import INCConfig + + assert isinstance(t2_quant, INCConfig) + assert t2_quant.weight_bits == 4 + assert t2_quant.group_size == 128 + + def test_transformer_2_quant_config_none_when_missing(self): + """When transformer_2 config has no quantization_config, result is None.""" + t2_config = { + "patch_size": [1, 2, 2], + "num_layers": 2, + } + + t2_quant = t2_config.get("quantization_config") + if isinstance(t2_quant, dict) and "quant_method" in t2_quant: + pass # won't enter + else: + t2_quant = None + + assert t2_quant is None + + def test_transformer_2_quant_config_none_when_dict_lacks_method(self): + """When quantization_config is a dict but missing quant_method, result is None.""" + t2_config = { + "patch_size": [1, 2, 2], + "quantization_config": {"bits": 4}, # no quant_method key + } + + t2_quant = t2_config.get("quantization_config") + if isinstance(t2_quant, dict) and "quant_method" in t2_quant: + pass + else: + t2_quant = None + + assert t2_quant is None diff --git a/tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py new file mode 100644 index 000000000..044435cb5 --- /dev/null +++ b/tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py @@ -0,0 +1,311 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E tests for Wan2.2 AutoRound W4A16 quantized inference. + +These tests cover I2V (image-to-video) and T2V (text-to-video) generation +with quantized weights. + +Requirements: + - CUDA GPU (H100 or equivalent, ~36 GiB for quantized model) + - The quantized model checkpoint (Intel/Wan2.2-I2V-A14B-Diffusers-int4-AutoRound, + Intel/Wan2.2-T2V-A14B-Diffusers-int4-AutoRound) +""" + +import gc +import os as _os + +import numpy as np +import pytest +import torch +from PIL import Image + +from tests.helpers.env import DeviceMemoryMonitor +from tests.helpers.mark import hardware_test +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.platforms import current_omni_platform + +_os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +QUANTIZED_MODEL_I2V = "Intel/Wan2.2-I2V-A14B-Diffusers-int4-AutoRound" +BASELINE_MODEL_I2V = "Wan-AI/Wan2.2-I2V-A14B-Diffusers" +QUANTIZED_MODEL_T2V = "Intel/Wan2.2-T2V-A14B-Diffusers-int4-AutoRound" +BASELINE_MODEL_T2V = "Wan-AI/Wan2.2-T2V-A14B-Diffusers" + +QUANTIZED_MODEL_I2V = _os.environ.get("WAN22_I2V_AUTOROUND_MODEL", QUANTIZED_MODEL_I2V) +BASELINE_MODEL_I2V = _os.environ.get("WAN22_I2V_BASELINE_MODEL", BASELINE_MODEL_I2V) +QUANTIZED_MODEL_T2V = _os.environ.get("WAN22_T2V_AUTOROUND_MODEL", QUANTIZED_MODEL_T2V) +BASELINE_MODEL_T2V = _os.environ.get("WAN22_T2V_BASELINE_MODEL", BASELINE_MODEL_T2V) + +pytestmark = [ + pytest.mark.full_model, + pytest.mark.diffusion, +] + +# Small resolution to keep GPU memory & time manageable +HEIGHT = 480 +WIDTH = 640 +NUM_FRAMES = 5 # must satisfy num_frames % 4 == 1 for Wan2.2 +NUM_STEPS = 2 # minimal for smoke-test + +# Parametrise: (model, stage_config_path=None, extra_omni_kwargs) +# When stage_config_path is None, the engine auto-resolves from the model's own config. +quant_i2v_params = [(QUANTIZED_MODEL_I2V, None, {"enforce_eager": True})] +baseline_i2v_params = [(BASELINE_MODEL_I2V, None, {"enforce_eager": True})] +quant_t2v_params = [(QUANTIZED_MODEL_T2V, None, {"enforce_eager": True})] +baseline_t2v_params = [(BASELINE_MODEL_T2V, None, {"enforce_eager": True})] + +# Module-level storage for peak memory results across tests +_memory_results: dict[str, float] = {} + + +def _sampling_params_i2v() -> OmniDiffusionSamplingParams: + """Create sampling params for I2V generation.""" + return OmniDiffusionSamplingParams( + height=HEIGHT, + width=WIDTH, + num_frames=NUM_FRAMES, + num_inference_steps=NUM_STEPS, + guidance_scale=5.0, + guidance_scale_2=6.0, + boundary_ratio=0.875, + generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42), + ) + + +def _sampling_params_t2v() -> OmniDiffusionSamplingParams: + """Create sampling params for T2V generation.""" + return OmniDiffusionSamplingParams( + height=HEIGHT, + width=WIDTH, + num_frames=NUM_FRAMES, + num_inference_steps=NUM_STEPS, + guidance_scale=4.0, + generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42), + ) + + +def _create_test_image(width: int = WIDTH, height: int = HEIGHT) -> Image.Image: + """Create a deterministic test image for I2V tests.""" + rng = np.random.RandomState(42) + arr = rng.randint(0, 256, (height, width, 3), dtype=np.uint8) + return Image.fromarray(arr) + + +def _generate_i2v_video(omni_runner_handler, prompt: str = "A cat sitting on a table, smooth motion") -> tuple: + """Generate one I2V video, return (frames, peak_memory_mb).""" + gc.collect() + current_omni_platform.empty_cache() + device_index = current_omni_platform.current_device() + current_omni_platform.reset_peak_memory_stats() + monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02) + monitor.start() + + image = _create_test_image() + response = omni_runner_handler.send_diffusion_request( + { + "prompt": prompt, + "images": image, + "sampling_params": _sampling_params_i2v(), + }, + ) + + peak = monitor.peak_used_mb + monitor.stop() + + assert response.success, f"Request failed: {response.error_message}" + assert response.images is not None and len(response.images) > 0, "Expected image output" + frames = response.images[0] + + gc.collect() + current_omni_platform.empty_cache() + + return frames, peak + + +def _generate_t2v_video(omni_runner_handler, prompt: str = "A cat sitting on a table") -> tuple: + """Generate one T2V video, return (frames, peak_memory_mb).""" + gc.collect() + current_omni_platform.empty_cache() + device_index = current_omni_platform.current_device() + current_omni_platform.reset_peak_memory_stats() + monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02) + monitor.start() + + response = omni_runner_handler.send_diffusion_request( + { + "prompt": prompt, + "sampling_params": _sampling_params_t2v(), + }, + ) + + peak = monitor.peak_used_mb + monitor.stop() + + assert response.success, f"Request failed: {response.error_message}" + assert response.images is not None and len(response.images) > 0, "Expected image output" + frames = response.images[0] + + gc.collect() + current_omni_platform.empty_cache() + + return frames, peak + + +# ------------------------------------------------------------------ +# Test: I2V quantized model generates valid video +# ------------------------------------------------------------------ + + +@hardware_test(res={"cuda": "H100"}) +@pytest.mark.parametrize("omni_runner", quant_i2v_params, indirect=True) +def test_wan22_i2v_autoround_w4a16_generates_video(omni_runner, omni_runner_handler): + """Load the W4A16 quantized Wan2.2 I2V model and verify it produces a valid video.""" + frames, _ = _generate_i2v_video(omni_runner_handler) + + assert frames is not None, "Expected video frames output" + assert hasattr(frames, "shape"), "Expected frames to have a shape attribute" + + # frames shape: (batch, num_frames, height, width, channels) + assert frames.shape[1] == NUM_FRAMES, f"Expected {NUM_FRAMES} frames, got {frames.shape[1]}" + assert frames.shape[2] == HEIGHT, f"Expected height {HEIGHT}, got {frames.shape[2]}" + assert frames.shape[3] == WIDTH, f"Expected width {WIDTH}, got {frames.shape[3]}" + + # Sanity: video should not be blank (frames are [0, 1] floats) + arr = np.asarray(frames) + assert arr.std() > 0.01, "Generated video appears blank (std ≈ 0)" + + +# ------------------------------------------------------------------ +# Test: T2V quantized model generates valid video +# ------------------------------------------------------------------ + + +@hardware_test(res={"cuda": "H100"}) +@pytest.mark.parametrize("omni_runner", quant_t2v_params, indirect=True) +def test_wan22_t2v_autoround_w4a16_generates_video(omni_runner, omni_runner_handler): + """Load the W4A16 quantized Wan2.2 T2V model and verify it produces a valid video.""" + frames, _ = _generate_t2v_video(omni_runner_handler) + + assert frames is not None, "Expected video frames output" + assert hasattr(frames, "shape"), "Expected frames to have a shape attribute" + + assert frames.shape[1] == NUM_FRAMES, f"Expected {NUM_FRAMES} frames, got {frames.shape[1]}" + assert frames.shape[2] == HEIGHT, f"Expected height {HEIGHT}, got {frames.shape[2]}" + assert frames.shape[3] == WIDTH, f"Expected width {WIDTH}, got {frames.shape[3]}" + + arr = np.asarray(frames) + assert arr.std() > 0.01, "Generated video appears blank (std ≈ 0)" + + +# ------------------------------------------------------------------ +# Test: I2V quantized peak memory +# ------------------------------------------------------------------ + + +@hardware_test(res={"cuda": "H100"}) +@pytest.mark.parametrize("omni_runner", quant_i2v_params, indirect=True) +def test_wan22_i2v_autoround_w4a16_quant_peak(omni_runner, omni_runner_handler): + """Measure peak GPU memory of W4A16 quantized I2V model.""" + frames, peak = _generate_i2v_video(omni_runner_handler) + + assert frames is not None, "Expected video frames output" + _memory_results["quant_i2v"] = peak + print(f"\nQuantized I2V (W4A16) peak memory: {peak:.0f} MB") + + +# ------------------------------------------------------------------ +# Test: I2V baseline peak memory +# ------------------------------------------------------------------ + + +@hardware_test(res={"cuda": "H100"}) +@pytest.mark.parametrize("omni_runner", baseline_i2v_params, indirect=True) +def test_wan22_i2v_autoround_w4a16_baseline_peak(omni_runner, omni_runner_handler): + """Measure peak GPU memory of BF16 baseline I2V model.""" + frames, peak = _generate_i2v_video(omni_runner_handler) + + assert frames is not None, "Expected video frames output" + _memory_results["baseline_i2v"] = peak + print(f"\nBaseline I2V (BF16) peak memory: {peak:.0f} MB") + + +# ------------------------------------------------------------------ +# Test: I2V memory savings +# ------------------------------------------------------------------ + + +@hardware_test(res={"cuda": "H100"}) +def test_wan22_i2v_autoround_w4a16_memory_savings(): + """Assert quantized I2V model uses meaningfully less memory than BF16 baseline.""" + quant_peak = _memory_results["quant_i2v"] + baseline_peak = _memory_results["baseline_i2v"] + + savings = baseline_peak - quant_peak + print(f"\nQuantized I2V (W4A16) peak memory: {quant_peak:.0f} MB") + print(f"Baseline I2V (BF16) peak memory: {baseline_peak:.0f} MB") + print(f"Savings: {savings:.0f} MB") + + # Wan2.2 I2V A14B transformer is ~28 GB in BF16; W4A16 should save ~20 GB. + # Use a conservative threshold to account for activations and overhead. + min_savings_mb = 5000 + assert quant_peak + min_savings_mb < baseline_peak, ( + f"Quantized model ({quant_peak:.0f} MB) should use at least " + f"{min_savings_mb} MB less than baseline ({baseline_peak:.0f} MB)" + ) + + +# ------------------------------------------------------------------ +# Test: T2V quantized peak memory +# ------------------------------------------------------------------ + + +@hardware_test(res={"cuda": "H100"}) +@pytest.mark.parametrize("omni_runner", quant_t2v_params, indirect=True) +def test_wan22_t2v_autoround_w4a16_quant_peak(omni_runner, omni_runner_handler): + """Measure peak GPU memory of W4A16 quantized T2V model.""" + frames, peak = _generate_t2v_video(omni_runner_handler) + + assert frames is not None, "Expected video frames output" + _memory_results["quant_t2v"] = peak + print(f"\nQuantized T2V (W4A16) peak memory: {peak:.0f} MB") + + +# ------------------------------------------------------------------ +# Test: T2V baseline peak memory +# ------------------------------------------------------------------ + + +@hardware_test(res={"cuda": "H100"}) +@pytest.mark.parametrize("omni_runner", baseline_t2v_params, indirect=True) +def test_wan22_t2v_autoround_w4a16_baseline_peak(omni_runner, omni_runner_handler): + """Measure peak GPU memory of BF16 baseline T2V model.""" + frames, peak = _generate_t2v_video(omni_runner_handler) + + assert frames is not None, "Expected video frames output" + _memory_results["baseline_t2v"] = peak + print(f"\nBaseline T2V (BF16) peak memory: {peak:.0f} MB") + + +# ------------------------------------------------------------------ +# Test: T2V memory savings +# ------------------------------------------------------------------ + + +@hardware_test(res={"cuda": "H100"}) +def test_wan22_t2v_autoround_w4a16_memory_savings(): + """Assert quantized T2V model uses meaningfully less memory than BF16 baseline.""" + quant_peak = _memory_results["quant_t2v"] + baseline_peak = _memory_results["baseline_t2v"] + + savings = baseline_peak - quant_peak + print(f"\nQuantized T2V (W4A16) peak memory: {quant_peak:.0f} MB") + print(f"Baseline T2V (BF16) peak memory: {baseline_peak:.0f} MB") + print(f"Savings: {savings:.0f} MB") + + # Wan2.2 T2V A14B transformer is ~28 GB in BF16; W4A16 should save ~20 GB. + # Use a conservative threshold to account for activations and overhead. + min_savings_mb = 5000 + assert quant_peak + min_savings_mb < baseline_peak, ( + f"Quantized model ({quant_peak:.0f} MB) should use at least " + f"{min_savings_mb} MB less than baseline ({baseline_peak:.0f} MB)" + ) diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index 3ee46ffb0..2d8c752a4 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -122,8 +122,7 @@ def load_transformer_config(model_path: str, subfolder: str = "transformer", loc def create_transformer_from_config( - config: dict, - quant_config: QuantizationConfig | None = None, + config: dict, quant_config: QuantizationConfig | None = None, prefix: str = "" ) -> WanTransformer3DModel: """Create WanTransformer3DModel from config dict.""" kwargs: dict = {} @@ -166,6 +165,8 @@ def create_transformer_from_config( if quant_config is not None: kwargs["quant_config"] = quant_config + if prefix: + kwargs["prefix"] = prefix return WanTransformer3DModel(**kwargs) diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index 42c4eff6a..5ff374205 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -231,10 +231,25 @@ class Wan22I2VPipeline( # Transformers (weights loaded via load_weights) # Load config from model directory or HF Hub to get correct in_channels for I2V models transformer_config = load_transformer_config(model, "transformer", local_files_only) - self.transformer = self._create_transformer(transformer_config) + self.transformer = create_transformer_from_config( + transformer_config, + quant_config=od_config.quantization_config, + ) if self.has_transformer_2: transformer_2_config = load_transformer_config(model, "transformer_2", local_files_only) - self.transformer_2 = self._create_transformer(transformer_2_config) + t2_quant = transformer_2_config.get("quantization_config") + if isinstance(t2_quant, dict) and "quant_method" in t2_quant: + from vllm_omni.quantization.factory import build_quant_config + + method = t2_quant["quant_method"] + kwargs = {k: v for k, v in t2_quant.items() if k != "quant_method"} + t2_quant = build_quant_config(method, **kwargs) + else: + t2_quant = None + self.transformer_2 = create_transformer_from_config( + transformer_2_config, + quant_config=t2_quant, + ) else: self.transformer_2 = None diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py index 75bdac27f..5ba2c6c69 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py @@ -45,6 +45,7 @@ logger = init_logger(__name__) def create_vace_transformer_from_config( config: dict, quant_config: QuantizationConfig | None = None, + prefix: str = "", ) -> WanVACETransformer3DModel: """Create WanVACETransformer3DModel from config dict.""" kwargs = {} @@ -84,6 +85,8 @@ def create_vace_transformer_from_config( kwargs["vace_in_channels"] = config["vace_in_channels"] if quant_config is not None: kwargs["quant_config"] = quant_config + if prefix: + kwargs["prefix"] = prefix return WanVACETransformer3DModel(**kwargs) diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py index 955f97cef..3766e4596 100644 --- a/vllm_omni/quantization/factory.py +++ b/vllm_omni/quantization/factory.py @@ -99,6 +99,43 @@ _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = { SUPPORTED_QUANTIZATION_METHODS: list[str] = list(dict.fromkeys(QUANTIZATION_METHODS + list(_OVERRIDES.keys()))) +def _build_reverse_alias_map() -> dict[str, str]: + """Build a mapping from normalized method aliases to canonical names. + + All keys in _OVERRIDES that share the same builder function are considered + aliases of each other. The canonical name is the first key (in definition + order) that maps to a given builder — i.e. the one returned by + builder().get_name(). + """ + builder_to_first_key: dict[Callable[..., QuantizationConfig], str] = {} + for key in _OVERRIDES: + builder = _OVERRIDES[key] + if builder not in builder_to_first_key: + builder_to_first_key[builder] = key + + result: dict[str, str] = {} + for key, builder in _OVERRIDES.items(): + canonical = builder_to_first_key[builder] + result[key.lower().replace("-", "_")] = canonical + return result + + +_CACHED_ALIAS_MAP: dict[str, str] | None = None + + +def _normalize_quant_method_alias(method: str | None) -> str | None: + """Map a method name (or any of its aliases) to its canonical internal name. + Returns the input unchanged if it is not a known alias. + """ + if method is None: + return None + global _CACHED_ALIAS_MAP + if _CACHED_ALIAS_MAP is None: + _CACHED_ALIAS_MAP = _build_reverse_alias_map() + normalized = method.lower().replace("-", "_") + return _CACHED_ALIAS_MAP.get(normalized, normalized) + + _MODEL_OPT_METHODS = { "modelopt", } @@ -334,7 +371,9 @@ def resolve_quant_config_from_disk( ) return build_quant_config(qc_method, **qc_kwargs) - if quant_config.get_name() != qc_method: + active_method = _normalize_quant_method_alias(quant_config.get_name()) + disk_method = _normalize_quant_method_alias(qc_method) + if active_method != disk_method: raise ValueError( f"Checkpoint config.json declares quant_method={qc_method!r} but the " f"active quantization config is {quant_config.get_name()!r}. " From 2bc9cb7f13e2db487c35f05a5973db620043ec9d Mon Sep 17 00:00:00 2001 From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> Date: Thu, 21 May 2026 19:52:21 +0800 Subject: [PATCH 36/38] [Feat] Support online quantization (fp8/int8) for LTX-2 (#3700) Signed-off-by: Yuanheng Zhao Signed-off-by: yuanheng --- benchmarks/diffusion/quantization_quality.py | 37 +++++-- .../text_to_video/text_to_video.py | 25 +++++ .../quantization/test_quantization_fp8.py | 98 +++++++++++++++++++ .../quantization/test_quantization_quality.py | 40 ++++++-- .../diffusion/models/ltx2/ltx2_transformer.py | 87 ++++++++++++++-- .../diffusion/models/ltx2/pipeline_ltx2.py | 18 +++- .../diffusion/models/ltx2/pipeline_ltx2_3.py | 3 +- 7 files changed, 285 insertions(+), 23 deletions(-) diff --git a/benchmarks/diffusion/quantization_quality.py b/benchmarks/diffusion/quantization_quality.py index 09a9fa74e..b47fae41f 100644 --- a/benchmarks/diffusion/quantization_quality.py +++ b/benchmarks/diffusion/quantization_quality.py @@ -34,6 +34,16 @@ Video example (text-to-video): --height 720 --width 1280 \ --num-frames 81 --num-inference-steps 40 --seed 42 +LTX-2 example (text-to-video; audio output is generated but not scored by LPIPS): + python benchmarks/diffusion/quantization_quality.py \ + --model Lightricks/LTX-Video-2 \ + --task t2v \ + --quantization fp8 int8 \ + --prompts \ + "A serene lakeside sunrise with mist over the water" \ + --height 704 --width 1216 \ + --num-frames 121 --num-inference-steps 40 --seed 42 + Multiple quantization methods: python benchmarks/diffusion/quantization_quality.py \ --model Tongyi-MAI/Z-Image-Turbo \ @@ -204,10 +214,19 @@ def _generate_video(omni, args, prompt, seed): else: frames = inner elif hasattr(first, "images") and first.images: - frames = first.images + frames = first.images[0] else: raise ValueError("Could not extract video frames from output.") + # LTX-2 (and similar audio+video models) may surface a dict or (video, audio) tuple + if isinstance(frames, dict): + frames = frames.get("video") or frames.get("frames") + elif isinstance(frames, tuple) and len(frames) == 2: + frames = frames[0] + + if frames is None: + raise ValueError("Could not extract video frames from output.") + if isinstance(frames, torch.Tensor): video = frames.detach().cpu() if video.dim() == 5: @@ -225,9 +244,11 @@ def _generate_video(omni, args, prompt, seed): return frames_array, elapsed, peak_mem -def _unload_omni(omni): - """Delete Omni instance and free GPU memory.""" - del omni +def _free_gpu_memory(): + """Force GC and release cached GPU memory. + + Must be called AFTER the caller has dropped (i.e., via `del`) + every reference to the Omni instance""" gc.collect() if torch.cuda.is_available(): torch.accelerator.empty_cache() @@ -267,7 +288,9 @@ def run_benchmark(args): bl_avg_time = np.mean([v[1] for v in baseline_outputs.values()]) bl_mem = baseline_outputs[prompts[0]][2] # use first prompt's memory - _unload_omni(omni_bl) + omni_bl.shutdown() + del omni_bl + _free_gpu_memory() # Save baseline outputs bl_dir = output_dir / "baseline" @@ -307,7 +330,9 @@ def run_benchmark(args): qt_avg_time = np.mean([v[1] for v in qt_outputs.values()]) qt_mem = qt_outputs[prompts[0]][2] - _unload_omni(omni_qt) + omni_qt.shutdown() + del omni_qt + _free_gpu_memory() # Save quantized outputs qt_dir = output_dir / config_label.replace(" ", "_") diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 8bbc871ea..23b9afe82 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -237,6 +237,27 @@ def parse_args() -> argparse.Namespace: return parser.parse_args() +def _extract_peak_memory_mb(result: Any) -> float: + """Pull worker-reported peak VRAM (MiB) generation result. + + Mirrors vllm_omni/entrypoints/openai/serving_video.py:_extract_peak_memory_mb. + """ + if isinstance(result, list): + result = result[0] if result else None + if result is None: + return 0.0 + val = getattr(result, "peak_memory_mb", 0.0) + if not val: + inner = getattr(result, "request_output", None) + if isinstance(inner, list): + inner = inner[0] if inner else None + val = getattr(inner, "peak_memory_mb", 0.0) + try: + return float(val or 0.0) + except (TypeError, ValueError): + return 0.0 + + def main(): args = parse_args() model_class_name = args.model_class_name @@ -348,6 +369,10 @@ def main(): # Print profiling results print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)") + peak_mb = _extract_peak_memory_mb(frames) + if peak_mb: + print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)") + audio = None if isinstance(frames, list): frames = frames[0] if frames else None diff --git a/tests/diffusion/quantization/test_quantization_fp8.py b/tests/diffusion/quantization/test_quantization_fp8.py index 85466f68e..58f91742d 100644 --- a/tests/diffusion/quantization/test_quantization_fp8.py +++ b/tests/diffusion/quantization/test_quantization_fp8.py @@ -109,6 +109,85 @@ def _generate_single_stage_image( return images, peak_mem +def _generate_single_stage_video( + model: str, + quantization: str | None = None, + height: int = 256, + width: int = 256, + num_frames: int = 25, + num_inference_steps: int = 8, + guidance_scale: float = 4.0, + seed: int = 42, + prompt: str = "A serene lakeside sunrise with mist over the water", + **extra_omni_kwargs: Any, +) -> tuple[int, float]: + """Generate a t2v output with a single-stage diffusion model. + + Returns (num_frames_produced, peak_memory_gib) + """ + omni_kwargs: dict[str, Any] = dict(extra_omni_kwargs) + if quantization: + omni_kwargs["quantization"] = quantization + + with OmniRunner(model, **omni_kwargs) as runner: + torch.accelerator.reset_peak_memory_stats() + + generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed) + outputs = runner.omni.generate( + {"prompt": prompt, "negative_prompt": ""}, + OmniDiffusionSamplingParams( + height=height, + width=width, + num_frames=num_frames, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + generator=generator, + ), + ) + + first = outputs[0] + + # Unwrap pipeline-style outputs (multi-stage / OmniRequestOutput.request_output). + frames: Any = None + if hasattr(first, "request_output") and isinstance(first.request_output, list): + inner = first.request_output[0] + if isinstance(inner, OmniRequestOutput) and inner.images: + frames = inner.images[0] + if frames is None and hasattr(first, "images") and first.images: + frames = first.images[0] + assert frames is not None, "No video frames returned from generate()" + + # LTX-2 (audio+video) may surface (video, audio) tuples or {"video": ...} dicts. + if isinstance(frames, dict): + frames = frames.get("video") or frames.get("frames") + elif isinstance(frames, tuple) and len(frames) == 2: + frames = frames[0] + assert frames is not None, "Could not extract video frames from output" + + if isinstance(frames, torch.Tensor): + video = frames.detach().cpu() + if video.dim() == 5: + video = video[0] + if video.dim() == 4 and video.shape[0] in (3, 4): + video = video.permute(1, 2, 3, 0) + num_frames_produced = int(video.shape[0]) + else: + import numpy as np + + arr = np.asarray(frames) + if arr.ndim == 5: + arr = arr[0] + num_frames_produced = int(arr.shape[0]) + + peak_mem_mb = getattr(first, "peak_memory_mb", None) + if peak_mem_mb: + peak_mem = float(peak_mem_mb) / 1024.0 + else: + peak_mem = torch.accelerator.max_memory_allocated() / (1024**3) + + return num_frames_produced, peak_mem + + def _generate_bagel_image( quantization_config: str | None = None, num_inference_steps: int = 15, @@ -277,6 +356,25 @@ def test_single_stage_flux_fp8_uses_less_memory(): assert mem_fp8 < mem_bf16, f"FP8 ({mem_fp8:.2f} GiB) should use less memory than BF16 ({mem_bf16:.2f} GiB)" +@hardware_test(res={"cuda": "H100"}) +def test_single_stage_ltx2_fp8_uses_less_memory(): + """FP8 should use less peak memory than BF16 for LTX-2.""" + _, mem_bf16 = _generate_single_stage_video( + model="Lightricks/LTX-2", + quantization=None, + ) + torch.accelerator.empty_cache() + + _, mem_fp8 = _generate_single_stage_video( + model="Lightricks/LTX-2", + quantization="fp8", + ) + + print(f"LTX-2 BF16 peak memory: {mem_bf16:.2f} GiB") + print(f"LTX-2 FP8 peak memory: {mem_fp8:.2f} GiB") + assert mem_fp8 < mem_bf16, f"FP8 ({mem_fp8:.2f} GiB) should use less memory than BF16 ({mem_bf16:.2f} GiB)" + + # ─── Multi-stage model tests (BAGEL) ───────────────────────────────────────── diff --git a/tests/diffusion/quantization/test_quantization_quality.py b/tests/diffusion/quantization/test_quantization_quality.py index 02ef0400f..692d83c12 100644 --- a/tests/diffusion/quantization/test_quantization_quality.py +++ b/tests/diffusion/quantization/test_quantization_quality.py @@ -143,6 +143,18 @@ QUALITY_CONFIGS = [ seed=142, num_inference_steps=20, ), + QualityTestConfig( + id="fp8_ltx2", + model="Lightricks/LTX-2", + quantization="fp8", + task="t2v", + prompt="A serene lakeside sunrise with mist over the water", + max_lpips=0.10, + height=256, + width=256, + num_frames=25, + num_inference_steps=8, + ), ] @@ -261,10 +273,19 @@ def _generate_video(omni, config: QualityTestConfig): else: frames = inner elif hasattr(first, "images") and first.images: - frames = first.images + frames = first.images[0] else: raise ValueError("Could not extract video frames from output.") + # LTX-2 (audio+video) may surface (video, audio) tuples or {"video": ...} dicts + if isinstance(frames, dict): + frames = frames.get("video") or frames.get("frames") + elif isinstance(frames, tuple) and len(frames) == 2: + frames = frames[0] + + if frames is None: + raise ValueError("Could not extract video frames from output.") + if isinstance(frames, torch.Tensor): video = frames.detach().cpu() if video.dim() == 5: @@ -275,7 +296,11 @@ def _generate_video(omni, config: QualityTestConfig): video = video.clamp(-1, 1) * 0.5 + 0.5 return video.float().numpy(), peak_mem - return np.asarray(frames), peak_mem + frames_array = np.asarray(frames) + if frames_array.ndim == 5: + # strip the leading batch dim + frames_array = frames_array[0] + return frames_array, peak_mem def _compute_lpips(baseline, quantized, task: str) -> float: @@ -318,8 +343,7 @@ def _compute_psnr_and_mae(baseline, quantized, task: str) -> tuple[float, float] return psnr, mae -def _unload(omni): - del omni +def _free_gpu_memory(): gc.collect() if torch.cuda.is_available(): torch.accelerator.empty_cache() @@ -395,7 +419,9 @@ def test_quantization_quality(config: QualityTestConfig): # --- BF16 baseline --- omni_bl = Omni(model=config.baseline_ref()) baseline_out, bl_mem = generate_fn(omni_bl, config) - _unload(omni_bl) + omni_bl.shutdown() + del omni_bl + _free_gpu_memory() _maybe_save_output(_OUTPUT_DIR, config, "baseline", baseline_out) # --- Quantized --- @@ -405,7 +431,9 @@ def test_quantization_quality(config: QualityTestConfig): else: omni_qt = Omni(model=config.quantized_ref(), quantization_config=quantization) quant_out, qt_mem = generate_fn(omni_qt, config) - _unload(omni_qt) + omni_qt.shutdown() + del omni_qt + _free_gpu_memory() _maybe_save_output(_OUTPUT_DIR, config, "quantized", quant_out) # --- Similarity metrics --- diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py index cd8eee9e0..2a65ad6b0 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py @@ -17,7 +17,7 @@ import inspect from collections.abc import Iterable from dataclasses import dataclass from types import SimpleNamespace -from typing import Any +from typing import TYPE_CHECKING, Any import torch import torch.distributed @@ -39,6 +39,9 @@ from vllm.model_executor.layers.linear import ( ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module @@ -173,7 +176,16 @@ class LTX2AdaLayerNormSingle(nn.Module): class ColumnParallelApproxGELU(nn.Module): - def __init__(self, dim_in: int, dim_out: int, *, approximate: str, bias: bool = True): + def __init__( + self, + dim_in: int, + dim_out: int, + *, + approximate: str, + bias: bool = True, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", + ): super().__init__() self.proj = ColumnParallelLinear( dim_in, @@ -181,6 +193,8 @@ class ColumnParallelApproxGELU(nn.Module): bias=bias, gather_output=False, return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.proj" if prefix else "proj", ) self.approximate = approximate @@ -200,6 +214,8 @@ class LTX2FeedForward(nn.Module): bias: bool = True, dropout: float = 0.0, final_dropout: bool = False, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ) -> None: super().__init__() @@ -211,13 +227,22 @@ class LTX2FeedForward(nn.Module): dropout_layer: nn.Module = nn.Dropout(dropout) if dropout > 0 else nn.Identity() layers: list[nn.Module] = [ - ColumnParallelApproxGELU(dim, inner_dim, approximate="tanh", bias=bias), + ColumnParallelApproxGELU( + dim, + inner_dim, + approximate="tanh", + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.net.0" if prefix else "net.0", + ), dropout_layer, RowParallelLinear( inner_dim, dim_out, input_is_parallel=True, return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.net.2" if prefix else "net.2", ), ] if final_dropout: @@ -515,6 +540,9 @@ class LTX2Attention(torch.nn.Module): rope_type: str = "interleaved", apply_gated_attention: bool = False, processor=None, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", + disable_kv_quant: bool = False, ): super().__init__() # LTX-2 uses "rms_norm_across_heads", LTX-2.3 uses "rms_norm" -- both @@ -548,6 +576,8 @@ class LTX2Attention(torch.nn.Module): head_size=self.head_dim, total_num_heads=heads, bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.to_qkv" if prefix else "to_qkv", ) self.query_num_heads = self.to_qkv.num_heads self.kv_num_heads = self.to_qkv.num_kv_heads @@ -562,6 +592,8 @@ class LTX2Attention(torch.nn.Module): bias=bias, gather_output=False, return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_q" if prefix else "to_q", ) self.to_k = ColumnParallelLinear( self.cross_attention_dim, @@ -569,6 +601,8 @@ class LTX2Attention(torch.nn.Module): bias=bias, gather_output=False, return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_k" if prefix else "to_k", ) self.to_v = ColumnParallelLinear( self.cross_attention_dim, @@ -576,6 +610,8 @@ class LTX2Attention(torch.nn.Module): bias=bias, gather_output=False, return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_v" if prefix else "to_v", ) self.heads = self.query_num_heads @@ -617,6 +653,8 @@ class LTX2Attention(torch.nn.Module): bias=out_bias, input_is_parallel=True, return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_out.0" if prefix else "to_out.0", ), torch.nn.Dropout(dropout) if dropout > 0 else torch.nn.Identity(), ] @@ -627,9 +665,12 @@ class LTX2Attention(torch.nn.Module): num_kv_heads=self.kv_num_heads, softmax_scale=1.0 / (dim_head**0.5), causal=False, + prefix=prefix, + disable_kv_quant=disable_kv_quant, ) # LTX-2.3: per-head gated attention + # leave unquantized for this linear if apply_gated_attention: self.to_gate_logits = nn.Linear(query_dim, heads, bias=True) else: @@ -748,6 +789,8 @@ class LTX2VideoTransformerBlock(nn.Module): elementwise_affine: bool = False, rope_type: str = "interleaved", perturbed_attn: bool = False, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ): super().__init__() self.video_cross_attn_adaln = video_cross_attn_adaln @@ -767,6 +810,8 @@ class LTX2VideoTransformerBlock(nn.Module): qk_norm=qk_norm, rope_type=rope_type, apply_gated_attention=video_gated_attn, + quant_config=quant_config, + prefix=f"{prefix}.attn1" if prefix else "attn1", ) self.audio_norm1 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine) @@ -781,6 +826,8 @@ class LTX2VideoTransformerBlock(nn.Module): qk_norm=qk_norm, rope_type=rope_type, apply_gated_attention=audio_gated_attn, + quant_config=quant_config, + prefix=f"{prefix}.audio_attn1" if prefix else "audio_attn1", ) # 2. Prompt Cross-Attention @@ -796,6 +843,9 @@ class LTX2VideoTransformerBlock(nn.Module): qk_norm=qk_norm, rope_type=rope_type, apply_gated_attention=video_gated_attn, + quant_config=quant_config, + prefix=f"{prefix}.attn2" if prefix else "attn2", + disable_kv_quant=True, ) self.audio_norm2 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine) @@ -810,9 +860,14 @@ class LTX2VideoTransformerBlock(nn.Module): qk_norm=qk_norm, rope_type=rope_type, apply_gated_attention=audio_gated_attn, + quant_config=quant_config, + prefix=f"{prefix}.audio_attn2" if prefix else "audio_attn2", + disable_kv_quant=True, ) # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention + # K/V here come from the other modality's latents (long sequences), + # so KV-cache quant remains enabled. self.audio_to_video_norm = _make_rms_norm(dim, eps=eps, elementwise_affine=elementwise_affine) self.audio_to_video_attn = LTX2Attention( query_dim=dim, @@ -825,6 +880,8 @@ class LTX2VideoTransformerBlock(nn.Module): qk_norm=qk_norm, rope_type=rope_type, apply_gated_attention=video_gated_attn, + quant_config=quant_config, + prefix=f"{prefix}.audio_to_video_attn" if prefix else "audio_to_video_attn", ) self.video_to_audio_norm = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine) @@ -839,14 +896,26 @@ class LTX2VideoTransformerBlock(nn.Module): qk_norm=qk_norm, rope_type=rope_type, apply_gated_attention=audio_gated_attn, + quant_config=quant_config, + prefix=f"{prefix}.video_to_audio_attn" if prefix else "video_to_audio_attn", ) # 4. Feedforward layers self.norm3 = _make_rms_norm(dim, eps=eps, elementwise_affine=elementwise_affine) - self.ff = LTX2FeedForward(dim, activation_fn=activation_fn) + self.ff = LTX2FeedForward( + dim, + activation_fn=activation_fn, + quant_config=quant_config, + prefix=f"{prefix}.ff" if prefix else "ff", + ) self.audio_norm3 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine) - self.audio_ff = LTX2FeedForward(audio_dim, activation_fn=activation_fn) + self.audio_ff = LTX2FeedForward( + audio_dim, + activation_fn=activation_fn, + quant_config=quant_config, + prefix=f"{prefix}.audio_ff" if prefix else "audio_ff", + ) # 5. Per-Layer Modulation Parameters # LTX-2.3 with cross_attn_adaln uses 9 params (extra 3 for cross-attn modulation); @@ -1365,6 +1434,9 @@ class LTX2VideoTransformer3DModel(nn.Module): _layerwise_offload_blocks_attrs = ["transformer_blocks"] _hsdp_shard_conditions = [is_transformer_block_module] _sp_plan: dict[str, Any] | None = None + packed_modules_mapping = { + "to_qkv": ["to_q", "to_k", "to_v"], + } @staticmethod def _build_sp_plan(rope_type: str) -> dict[str, Any]: @@ -1453,6 +1525,7 @@ class LTX2VideoTransformer3DModel(nn.Module): cross_attn_mod: bool = False, audio_gated_attn: bool = False, audio_cross_attn_mod: bool = False, + quant_config: "QuantizationConfig | None" = None, ) -> None: super().__init__() self.perturbed_attn = perturbed_attn @@ -1647,8 +1720,10 @@ class LTX2VideoTransformer3DModel(nn.Module): elementwise_affine=norm_elementwise_affine, rope_type=rope_type, perturbed_attn=perturbed_attn, + quant_config=quant_config, + prefix=f"transformer_blocks.{layer_idx}", ) - for _ in range(num_layers) + for layer_idx in range(num_layers) ] ) diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py index 3c2322940..c92fd687f 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py @@ -9,10 +9,13 @@ import json import os from collections.abc import Iterable from contextlib import nullcontext -from typing import Any, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar import numpy as np import torch + +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler from diffusers.pipelines.ltx2 import LTX2TextConnectors from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES @@ -67,14 +70,20 @@ def load_transformer_config(model_path: str, subfolder: str = "transformer", loc return {} -def create_transformer_from_config(config: dict) -> LTX2VideoTransformer3DModel: +def create_transformer_from_config( + config: dict, + quant_config: QuantizationConfig | None = None, +) -> LTX2VideoTransformer3DModel: """Create LTX2VideoTransformer3DModel from config dict.""" - if not config: + if not config and quant_config is None: return LTX2VideoTransformer3DModel() signature = inspect.signature(LTX2VideoTransformer3DModel.__init__) allowed_keys = set(signature.parameters.keys()) kwargs = {k: v for k, v in config.items() if k in allowed_keys} + if quant_config is not None: + kwargs["quant_config"] = quant_config + return LTX2VideoTransformer3DModel(**kwargs) @@ -214,7 +223,8 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): ).to(self.device) transformer_config = load_transformer_config(model, "transformer", local_files_only) - self.transformer = create_transformer_from_config(transformer_config) + quant_config = getattr(self.od_config, "quantization_config", None) + self.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config) self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( model, subfolder="scheduler", local_files_only=local_files_only diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py index 84ef29305..ec3b495f5 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py @@ -174,7 +174,8 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin): # --- Transformer: created empty, weights loaded via AutoWeightsLoader --- transformer_config = load_transformer_config(model, "transformer", local_files_only) - self.transformer = create_transformer_from_config(transformer_config) + quant_config = getattr(self.od_config, "quantization_config", None) + self.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config) # --- Scheduler --- self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( From bd1617ce2b96d65d9a808ad2e327410c7028de23 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Thu, 21 May 2026 22:06:01 +0800 Subject: [PATCH 37/38] Add new committers to governance page (#3749) Signed-off-by: hsliuustc0106 Co-authored-by: Cursor --- docs/community/governance.md | 29 ++++++++++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/docs/community/governance.md b/docs/community/governance.md index 6af578e2d..83374a155 100644 --- a/docs/community/governance.md +++ b/docs/community/governance.md @@ -39,11 +39,34 @@ Committers have write access and merge rights. They typically have deep expertis - [@linyueqian](https://github.com/linyueqian): TTS and Omni Support - [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration - [@princepride](https://github.com/princepride): Diffusion and Omni Support -- [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion +- [@RuixiangMa](https://github.com/RuixiangMa): Diffusion models, parallel, cache, and docs +- [@SamitHuang](https://github.com/SamitHuang): RL, Diffusion, and cache - [@tzhouam](https://github.com/tzhouam): Engine and New Model Support -- [@wtomin](https://github.com/wtomin): Diffusion and Parallelism +- [@wtomin](https://github.com/wtomin): Diffusion models, parallel, and docs - [@ZeldaHuang](https://github.com/ZeldaHuang): Omni Support -- [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp +- [@ZJY0516](https://github.com/ZJY0516): Diffusion attention backend, kernel fusion, and CustomOp +- [@yuanheng-zhao](https://github.com/yuanheng-zhao): Diffusion cache, offload, and Omni Support + +## Diffusion Workload Division + +This section breaks down diffusion-related responsibilities. If you have PRs touching these areas, please ping the listed owners for review. + +### Runtime-related + +- Diffusion models: @RuixiangMa, @wtomin + +### Optimization + +- Parallel: @RuixiangMa, @wtomin +- Attention backend: @ZJY0516 +- Cache: @yuanheng-zhao, @RuixiangMa, @SamitHuang +- Offload: @yuanheng-zhao +- Quantization: @lishunyang12, @david6666666 +- Kernel fusion / communication-computation: @ZJY0516 + +### Docs & Test + +- Diffusion docs: @RuixiangMa, @wtomin ## Meetings From e949ccf087b4fd0f3b45293d71e603511fa3a9e8 Mon Sep 17 00:00:00 2001 From: Jialong Liu <88185941+Galleons2029@users.noreply.github.com> Date: Thu, 21 May 2026 22:18:00 +0800 Subject: [PATCH 38/38] [Bugfix] Fix MiMo-Audio voice instability: stochastic local_sampler + codec streaming context (#3686) Signed-off-by: Galleons2029 Signed-off-by: Jialong Liu <88185941+Galleons2029@users.noreply.github.com> Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> --- vllm_omni/deploy/mimo_audio.yaml | 5 +- .../models/mimo_audio/mimo_audio_code2wav.py | 42 +++++++++++-- .../models/mimo_audio/mimo_audio_llm.py | 18 +++++- .../stage_configs/mimo_audio.yaml | 61 +++++++++++++++++++ .../stage_input_processors/mimo_audio.py | 41 ++++++++++--- 5 files changed, 150 insertions(+), 17 deletions(-) create mode 100644 vllm_omni/model_executor/stage_configs/mimo_audio.yaml diff --git a/vllm_omni/deploy/mimo_audio.yaml b/vllm_omni/deploy/mimo_audio.yaml index 1e4fd7505..6ebe59209 100644 --- a/vllm_omni/deploy/mimo_audio.yaml +++ b/vllm_omni/deploy/mimo_audio.yaml @@ -1,3 +1,4 @@ + # MiMo Audio deploy: 2-stage thinker+talker → code2wav. # # Default mode is async-chunk streaming on a single GPU (both stages on @@ -17,8 +18,8 @@ connectors: connector_get_sleep_s: 0.001 connector_get_max_wait_first_chunk: 3000 connector_get_max_wait: 300 - codec_chunk_frames: 3 - codec_left_context_frames: 3 + codec_chunk_frames: 30 + codec_left_context_frames: 40 stages: - stage_id: 0 diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py index 449425a2f..bd178afe6 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py @@ -26,6 +26,13 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput logger = logging.getLogger(__name__) +# Minimum safe values for codec streaming parameters. Mirrors the constants +# in stage_input_processors/mimo_audio.py — keep in sync. +_MIN_CODEC_CHUNK_FRAMES = 3 +_MIN_CODEC_LEFT_CONTEXT_FRAMES = 40 # must cover vocoder_attn_window_size[0] +_DEFAULT_CODEC_CHUNK_FRAMES = 10 +_DEFAULT_CODEC_LEFT_CONTEXT_FRAMES = 40 + def flat_codec_group_element_count(group_size: int, audio_channels: int) -> int: """Flat token count for one MiMo talker codec group on the code2wav wire. @@ -485,12 +492,37 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): if connector_cfg else None ) - self._codec_chunk_frames = int(extra_cfg.get("codec_chunk_frames", 3)) if isinstance(extra_cfg, dict) else 3 - if self._codec_chunk_frames <= 0: - raise ValueError(f"codec_chunk_frames must be positive, got {self._codec_chunk_frames}") - self._codec_left_context_frames = ( - int(extra_cfg.get("codec_left_context_frames", 3)) if isinstance(extra_cfg, dict) else 3 + raw_chunk = ( + int(extra_cfg.get("codec_chunk_frames", _DEFAULT_CODEC_CHUNK_FRAMES)) + if isinstance(extra_cfg, dict) + else _DEFAULT_CODEC_CHUNK_FRAMES ) + if raw_chunk < _MIN_CODEC_CHUNK_FRAMES: + logger.warning( + "codec_chunk_frames=%d is below minimum %d; falling back to %d.", + raw_chunk, + _MIN_CODEC_CHUNK_FRAMES, + _DEFAULT_CODEC_CHUNK_FRAMES, + ) + raw_chunk = _DEFAULT_CODEC_CHUNK_FRAMES + self._codec_chunk_frames = raw_chunk + + raw_left = ( + int(extra_cfg.get("codec_left_context_frames", _DEFAULT_CODEC_LEFT_CONTEXT_FRAMES)) + if isinstance(extra_cfg, dict) + else _DEFAULT_CODEC_LEFT_CONTEXT_FRAMES + ) + if raw_left < _MIN_CODEC_LEFT_CONTEXT_FRAMES: + logger.warning( + "codec_left_context_frames=%d is below minimum %d (must cover vocoder attention " + "window %s); falling back to %d to prevent voice instability.", + raw_left, + _MIN_CODEC_LEFT_CONTEXT_FRAMES, + getattr(self.config, "vocoder_attn_window_size", [40, 10]), + _DEFAULT_CODEC_LEFT_CONTEXT_FRAMES, + ) + raw_left = _DEFAULT_CODEC_LEFT_CONTEXT_FRAMES + self._codec_left_context_frames = raw_left def load_weights( self, diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py index 0a886bd5f..fed27c2c5 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py @@ -542,8 +542,20 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor ) self.device = current_omni_platform.get_torch_device() + # global_sampler MUST stay greedy (do_sample=False) so its token decision + # matches vLLM's external sampler (SamplingParams temperature=0.0). Both + # run argmax on the same logits, so they always agree on whether the next + # token is <|empty|> (audio step) or a real text token. Enabling + # do_sample=True here without also routing vLLM's sampled token back into + # this gate check would cause the two to diverge and corrupt KV-cache state. self.global_sampler = MiMoSampler(do_sample=False, temperature=0.6, top_p=0.95) - self.local_sampler = MiMoSampler(do_sample=False, temperature=0.9, top_p=0.95) + # local_sampler drives audio-code generation inside local_forward. It is + # entirely internal and NOT subject to vLLM's SamplingParams, so stochastic + # sampling is safe here and is required to produce natural, varied speech. + # Setting do_sample=True also disables the CUDA-graph path (use_cg gate in + # local_forward checks `do_sample is False`), preventing MiMoLocalSamplerTensor + # from silently forcing argmax even when temperature > 0. + self.local_sampler = MiMoSampler(do_sample=True, temperature=0.9, top_p=0.95) self.removed_tokens = None self.speech_vocab_sizes = config.parsed_speech_vocab_sizes() @@ -806,7 +818,7 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor device=tokens_device, ) if local_sampler is None: - local_sampler = MiMoSampler(do_sample=False, temperature=0.6, top_p=0.9) + local_sampler = MiMoSampler(do_sample=True, temperature=0.9, top_p=0.95) past_key_values = DynamicCache() for t in range(delay_iters): @@ -852,7 +864,7 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor local_sampler: MiMoSampler | None = None, ): if local_sampler is None: - local_sampler = MiMoSampler(do_sample=False, temperature=0.6, top_p=0.9) + local_sampler = MiMoSampler(do_sample=True, temperature=0.9, top_p=0.95) b = int(local_embeds.shape[0]) use_cg = (local_sampler.do_sample is None or local_sampler.do_sample is False) and bool( diff --git a/vllm_omni/model_executor/stage_configs/mimo_audio.yaml b/vllm_omni/model_executor/stage_configs/mimo_audio.yaml new file mode 100644 index 000000000..493c05f00 --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/mimo_audio.yaml @@ -0,0 +1,61 @@ + +# MiMo Audio deploy: 2-stage thinker+talker → code2wav. +# +# Default mode is async-chunk streaming on a single GPU (both stages on +# device 0) through SharedMemoryConnector. For the legacy 2-GPU sync +# pipeline (stage 1 on a second card), pass ``--no-async-chunk +# --stage-1-devices 1 --stage-1-max-model-len 18192 +# --stage-1-max-num-batched-tokens 18192`` to the serve CLI. +async_chunk: true +dtype: bfloat16 + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 365536 + codec_streaming: true + connector_get_sleep_s: 0.001 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + codec_chunk_frames: 30 + codec_left_context_frames: 40 + +stages: + - stage_id: 0 + max_num_seqs: 1 + gpu_memory_utilization: 0.5 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + max_num_batched_tokens: 8192 + max_model_len: 8192 + devices: "0" + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.6 + top_p: 0.95 + top_k: 50 + max_tokens: 18192 + seed: 42 + repetition_penalty: 1.1 + + - stage_id: 1 + max_num_seqs: 1 + gpu_memory_utilization: 0.4 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + async_scheduling: false + max_num_batched_tokens: 8192 + max_model_len: 8192 + devices: "0" + input_connectors: + from_stage_0: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 18192 + seed: 42 diff --git a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py b/vllm_omni/model_executor/stage_input_processors/mimo_audio.py index b1df30cb1..55aa003f6 100644 --- a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py +++ b/vllm_omni/model_executor/stage_input_processors/mimo_audio.py @@ -23,6 +23,16 @@ logger = init_logger(__name__) # example ``examples/offline_inference/mimo_audio/end2end.py``. MAX_CODE2WAV_TOKENS = 18192 +# Minimum safe values for codec streaming parameters. +# codec_left_context_frames must cover the vocoder attention window +# (vocoder_attn_window_size defaults to [40, 10]). Values below the minimum +# cause acoustic-state resets at chunk boundaries, producing voice instability +# (multiple speakers / timbre shifts in the output audio). +_MIN_CODEC_CHUNK_FRAMES = 3 +_MIN_CODEC_LEFT_CONTEXT_FRAMES = 40 +_DEFAULT_CODEC_CHUNK_FRAMES = 10 +_DEFAULT_CODEC_LEFT_CONTEXT_FRAMES = 40 + def prepend_and_flatten_colmajor(x: torch.Tensor, pad_vec: torch.Tensor) -> torch.Tensor: """ @@ -80,6 +90,10 @@ def _flush_remaining_codes( length = len(accumulated) chunk_length = length % chunk_size + # When the accumulated length aligns with chunk_size boundary (remainder == 0), + # we still need to flush the final chunk with full context to give the vocoder + # enough attention window — otherwise the tail audio cuts off and produces + # voice instability. Fall back to chunk_size as the context length. context_length = chunk_length if chunk_length != 0 else chunk_size end_index = min(length, left_context_size + context_length) @@ -136,8 +150,26 @@ def llm2code2wav_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", 3)) - left_context_size = int(cfg.get("codec_left_context_frames", 3)) + chunk_size = int(cfg.get("codec_chunk_frames", _DEFAULT_CODEC_CHUNK_FRAMES)) + if chunk_size < _MIN_CODEC_CHUNK_FRAMES: + logger.warning( + "codec_chunk_frames=%d is below minimum %d; falling back to %d.", + chunk_size, + _MIN_CODEC_CHUNK_FRAMES, + _DEFAULT_CODEC_CHUNK_FRAMES, + ) + chunk_size = _DEFAULT_CODEC_CHUNK_FRAMES + + left_context_size = int(cfg.get("codec_left_context_frames", _DEFAULT_CODEC_LEFT_CONTEXT_FRAMES)) + if left_context_size < _MIN_CODEC_LEFT_CONTEXT_FRAMES: + logger.warning( + "codec_left_context_frames=%d is below minimum %d (must cover vocoder attention window); " + "falling back to %d to prevent voice instability.", + left_context_size, + _MIN_CODEC_LEFT_CONTEXT_FRAMES, + _DEFAULT_CODEC_LEFT_CONTEXT_FRAMES, + ) + left_context_size = _DEFAULT_CODEC_LEFT_CONTEXT_FRAMES request_id = getattr(request, "external_req_id", None) @@ -157,11 +189,6 @@ def llm2code2wav_async_chunk( pad_vec = torch.tensor([TALKER_CODEC_PAD_TOKEN_ID] * 4, device=code_tensor.device, dtype=code_tensor.dtype) code_list = prepend_and_flatten_colmajor(code_tensor, pad_vec).tolist() - if sum(code_list) == 0: - if is_finished: - return _flush_remaining_codes(transfer_manager, request_id, chunk_size, left_context_size) - return None - if request_id is None: return None
PipelineTest SuiteCasesTest Files
PipelineGroupTest SuiteCasesTest Files