Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7db3a0cd5e |
No files matched your search
@@ -192,7 +192,7 @@ steps:
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 45m bash -c '
|
||||
timeout 20m bash -c '
|
||||
pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "slow" --run-level "core_model"
|
||||
'
|
||||
|
||||
|
||||
@@ -169,6 +169,7 @@ steps:
|
||||
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
|
||||
- vllm_omni/deploy/minicpmo_4_5.yaml
|
||||
- tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
|
||||
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
|
||||
- vllm_omni/experimental/fullduplex/
|
||||
- vllm_omni/model_executor/models/minicpmo_4_5/
|
||||
- vllm_omni/model_executor/models/step_audio2/
|
||||
|
||||
@@ -208,6 +208,7 @@ steps:
|
||||
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
|
||||
- vllm_omni/deploy/minicpmo_4_5.yaml
|
||||
- tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
|
||||
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
|
||||
- vllm_omni/experimental/fullduplex/
|
||||
- vllm_omni/model_executor/models/minicpmo_4_5/
|
||||
- vllm_omni/model_executor/models/step_audio2/
|
||||
@@ -306,24 +307,6 @@ steps:
|
||||
"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: "Diffusion · MiniMax-H3 DLO DP2 T2VA Test"
|
||||
timeout_in_minutes: 90
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py
|
||||
- tests/helpers/assertions.py
|
||||
- tests/helpers/runtime.py
|
||||
- vllm_omni/diffusion/models/minimax_h3/
|
||||
- vllm_omni/diffusion/offloader/
|
||||
- vllm_omni/diffusion/distributed/
|
||||
- vllm_omni/entrypoints/openai/
|
||||
commands:
|
||||
- |
|
||||
timeout 80m bash -c "
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py -m 'core_model and cuda and diffusion' --run-level 'core_model'
|
||||
"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: "Diffusion · Qwen Image Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image.py
|
||||
|
||||
@@ -105,7 +105,7 @@ For custom loop structures, follow `vllm_omni/diffusion/models/progress_bar.py`
|
||||
|
||||
## Path B: Custom/Private Repo Model
|
||||
|
||||
For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: BAGEL, HunyuanImage3. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns.
|
||||
For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns.
|
||||
|
||||
### B1. Understand the reference repo
|
||||
|
||||
@@ -126,9 +126,9 @@ This is the key design decision for custom models. Follow these placement rules:
|
||||
| **Pipeline orchestration** (init, forward, denoise loop) | `vllm_omni/diffusion/models/<name>/pipeline_<name>.py` | Always required |
|
||||
| **Custom transformer/backbone** (ported and adapted to vllm-omni) | `vllm_omni/diffusion/models/<name>/<name>_transformer.py` or similar | `wan2_2.py`, `fusion.py`, `bagel_transformer.py` |
|
||||
| **Custom sub-models** (VAE, fusion, autoencoder) | `vllm_omni/diffusion/models/<name>/` as separate files | `autoencoder.py`, `fusion.py` |
|
||||
| **External dependency code** (original repo utilities) | **External repo**, installed via download script or pip | original-repo package via git clone |
|
||||
| **External dependency code** (original repo utilities) | **External repo**, installed via download script or pip | `dreamid_omni` package via git clone |
|
||||
| **Hardcoded model configs** | Module-level dicts in pipeline file | `VIDEO_CONFIG`, `AUDIO_CONFIG` dicts |
|
||||
| **Download/setup script** | `examples/offline_inference/<name>/download_<name>.py` | `download_<name>.py` |
|
||||
| **Download/setup script** | `examples/offline_inference/<name>/download_<name>.py` | `download_dreamid_omni.py` |
|
||||
| **Custom `model_index.json`** | Generated by download script, placed at model root | Minimal: `{"_class_name": "YourPipeline", ...}` |
|
||||
|
||||
### B3. Handle external dependencies
|
||||
@@ -157,7 +157,7 @@ Copy the essential model files into `vllm_omni/diffusion/models/<name>/` and ada
|
||||
|
||||
Custom models have two common patterns for weight loading:
|
||||
|
||||
**Pattern 1: Bypass standard loader** (eager custom init)
|
||||
**Pattern 1: Bypass standard loader** (DreamID-Omni style)
|
||||
|
||||
When the original model has complex custom init functions that load weights in `__init__`:
|
||||
|
||||
|
||||
@@ -16,19 +16,20 @@ vllm_omni/diffusion/models/wan2_2/
|
||||
|
||||
The transformer is loaded separately via `weights_sources` + `load_weights()`. Non-transformer components (VAE, text encoder) are loaded in `__init__` via `from_pretrained()`.
|
||||
|
||||
### Custom model with external deps
|
||||
### Custom model with external deps (e.g., DreamID-Omni)
|
||||
|
||||
```
|
||||
vllm_omni/diffusion/models/<name>/
|
||||
vllm_omni/diffusion/models/dreamid_omni/
|
||||
├── __init__.py # Exports pipeline only
|
||||
├── pipeline_<name>.py # Pipeline: loads ALL weights in __init__ via custom helpers
|
||||
└── <backbone>.py # Custom transformer / fusion architecture
|
||||
├── pipeline_dreamid_omni.py # Pipeline: loads ALL weights in __init__ via custom helpers
|
||||
├── fusion.py # Custom fusion architecture (video + audio cross-attention)
|
||||
└── wan2_2.py # Re-implemented Wan backbone with split API
|
||||
|
||||
examples/offline_inference/<task>/
|
||||
└── download_<name>.py # Downloads weights and/or clones code repo
|
||||
examples/offline_inference/x_to_video_audio/
|
||||
└── download_dreamid_omni.py # Downloads weights from 3 HF repos + clones code repo
|
||||
```
|
||||
|
||||
All weights loaded eagerly in `__init__`. `load_weights()` is a no-op. External dependency imported with try/except.
|
||||
All weights loaded eagerly in `__init__`. `load_weights()` is a no-op. External dependency (`dreamid_omni` package) imported with try/except.
|
||||
|
||||
### Custom model with ported code (e.g., BAGEL)
|
||||
|
||||
@@ -75,7 +76,7 @@ self.weights_sources = [
|
||||
]
|
||||
```
|
||||
|
||||
### Pattern 3: Fully custom loading
|
||||
### Pattern 3: Fully custom loading (DreamID-Omni)
|
||||
|
||||
```
|
||||
init → load ALL weights eagerly via custom helpers → load_weights() = no-op
|
||||
@@ -107,8 +108,8 @@ Standard diffusers `model_index.json`:
|
||||
Custom model `model_index.json` (minimal):
|
||||
```json
|
||||
{
|
||||
"_class_name": "YourCustomPipeline",
|
||||
"fusion": "weights/model.safetensors"
|
||||
"_class_name": "DreamIDOmniPipeline",
|
||||
"fusion": "DreamID-Omni/dreamid_omni.safetensors"
|
||||
}
|
||||
```
|
||||
|
||||
@@ -116,7 +117,7 @@ The only **required** field is `_class_name` — it must match a key in `_DIFFUS
|
||||
|
||||
## External Dependency Management
|
||||
|
||||
### Git clone + .pth injection
|
||||
### Git clone + .pth injection (DreamID-Omni pattern)
|
||||
|
||||
```python
|
||||
def download_dependency():
|
||||
@@ -162,7 +163,7 @@ The engine checks `isinstance(pipeline, SupportImageInput)` at startup to config
|
||||
|
||||
Diffusers models use `config.json` in each subfolder. Custom models often use:
|
||||
|
||||
**Module-level config dicts**:
|
||||
**Module-level config dicts** (DreamID-Omni):
|
||||
```python
|
||||
VIDEO_CONFIG = {
|
||||
"patch_size": [1, 2, 2], "model_type": "ti2v",
|
||||
@@ -180,7 +181,7 @@ vae_cfg = bagel_cfg.get("vae_config", {})
|
||||
|
||||
## Custom Architecture Patterns
|
||||
|
||||
### Split forward API
|
||||
### Split forward API (DreamID-Omni)
|
||||
|
||||
When a fusion model needs to interleave blocks from two backbones:
|
||||
|
||||
|
||||
@@ -1,151 +0,0 @@
|
||||
"""Benchmark CosyVoice3-style CUDA stream handoff strategies."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import statistics
|
||||
import time
|
||||
|
||||
import torch
|
||||
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
def run_case(
|
||||
*,
|
||||
device: torch.device,
|
||||
steps: int,
|
||||
warmup_steps: int,
|
||||
producer_cycles: int,
|
||||
estimator_cycles: int,
|
||||
use_stream_dependencies: bool,
|
||||
) -> dict[str, float | int | str]:
|
||||
caller_stream = torch.cuda.Stream(device=device)
|
||||
estimator_stream = torch.cuda.Stream(device=device)
|
||||
input_tensor = torch.zeros(1024, device=device)
|
||||
output_tensor = torch.empty_like(input_tensor)
|
||||
submission_ms = []
|
||||
|
||||
def step(measure: bool) -> None:
|
||||
with torch.cuda.stream(caller_stream):
|
||||
torch.cuda._sleep(producer_cycles)
|
||||
input_tensor.add_(1)
|
||||
|
||||
start = time.perf_counter_ns()
|
||||
if use_stream_dependencies:
|
||||
estimator_stream.wait_stream(caller_stream)
|
||||
else:
|
||||
caller_stream.synchronize()
|
||||
|
||||
with torch.cuda.stream(estimator_stream):
|
||||
torch.cuda._sleep(estimator_cycles)
|
||||
output_tensor.copy_(input_tensor)
|
||||
|
||||
if use_stream_dependencies:
|
||||
caller_stream.wait_stream(estimator_stream)
|
||||
else:
|
||||
estimator_stream.synchronize()
|
||||
|
||||
if measure:
|
||||
submission_ms.append((time.perf_counter_ns() - start) / 1e6)
|
||||
|
||||
for _ in range(warmup_steps):
|
||||
step(measure=False)
|
||||
caller_stream.synchronize()
|
||||
estimator_stream.synchronize()
|
||||
input_tensor.zero_()
|
||||
torch.accelerator.reset_peak_memory_stats(device)
|
||||
initial_allocated = torch.accelerator.memory_allocated(device)
|
||||
|
||||
start_event = torch.cuda.Event(enable_timing=True)
|
||||
end_event = torch.cuda.Event(enable_timing=True)
|
||||
wall_start = time.perf_counter()
|
||||
with torch.cuda.stream(caller_stream):
|
||||
start_event.record()
|
||||
for _ in range(steps):
|
||||
step(measure=True)
|
||||
with torch.cuda.stream(caller_stream):
|
||||
end_event.record()
|
||||
end_event.synchronize()
|
||||
wall_ms = (time.perf_counter() - wall_start) * 1e3
|
||||
peak_delta_bytes = torch.accelerator.max_memory_allocated(device) - initial_allocated
|
||||
|
||||
expected = torch.full_like(output_tensor, steps)
|
||||
torch.testing.assert_close(output_tensor, expected, rtol=0, atol=0)
|
||||
return {
|
||||
"mode": "wait_stream" if use_stream_dependencies else "synchronize",
|
||||
"steps": steps,
|
||||
"submission_median_ms": statistics.median(submission_ms),
|
||||
"submission_p95_ms": statistics.quantiles(submission_ms, n=20)[18],
|
||||
"submission_mean_ms": statistics.mean(submission_ms),
|
||||
"gpu_elapsed_ms": start_event.elapsed_time(end_event),
|
||||
"wall_elapsed_ms": wall_ms,
|
||||
"peak_allocation_delta_bytes": peak_delta_bytes,
|
||||
"final_value": int(output_tensor[0].item()),
|
||||
}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--device", default="cuda:0")
|
||||
parser.add_argument("--steps", type=int, default=100)
|
||||
parser.add_argument("--warmup-steps", type=int, default=10)
|
||||
parser.add_argument("--repeats", type=int, default=7)
|
||||
parser.add_argument("--producer-cycles", type=int, default=1_000_000)
|
||||
parser.add_argument("--estimator-cycles", type=int, default=2_000_000)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
device = torch.device(args.device)
|
||||
current_omni_platform.set_device(device)
|
||||
torch.manual_seed(0)
|
||||
|
||||
old_results = []
|
||||
new_results = []
|
||||
for repeat in range(args.repeats):
|
||||
common_args = {
|
||||
"device": device,
|
||||
"steps": args.steps,
|
||||
"warmup_steps": args.warmup_steps,
|
||||
"producer_cycles": args.producer_cycles,
|
||||
"estimator_cycles": args.estimator_cycles,
|
||||
}
|
||||
if repeat % 2 == 0:
|
||||
old_results.append(run_case(**common_args, use_stream_dependencies=False))
|
||||
new_results.append(run_case(**common_args, use_stream_dependencies=True))
|
||||
else:
|
||||
new_results.append(run_case(**common_args, use_stream_dependencies=True))
|
||||
old_results.append(run_case(**common_args, use_stream_dependencies=False))
|
||||
|
||||
old_submission_ms = statistics.median(result["submission_median_ms"] for result in old_results)
|
||||
new_submission_ms = statistics.median(result["submission_median_ms"] for result in new_results)
|
||||
old_gpu_ms = statistics.median(result["gpu_elapsed_ms"] for result in old_results)
|
||||
new_gpu_ms = statistics.median(result["gpu_elapsed_ms"] for result in new_results)
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"device": torch.cuda.get_device_name(device),
|
||||
"torch": torch.__version__,
|
||||
"cuda": torch.version.cuda,
|
||||
"repeats": args.repeats,
|
||||
"warmup_steps": args.warmup_steps,
|
||||
"producer_cycles": args.producer_cycles,
|
||||
"estimator_cycles": args.estimator_cycles,
|
||||
"old_submission_median_ms": old_submission_ms,
|
||||
"new_submission_median_ms": new_submission_ms,
|
||||
"submission_median_speedup": old_submission_ms / new_submission_ms,
|
||||
"old_gpu_elapsed_median_ms": old_gpu_ms,
|
||||
"new_gpu_elapsed_median_ms": new_gpu_ms,
|
||||
"gpu_elapsed_change_percent": (new_gpu_ms - old_gpu_ms) / old_gpu_ms * 100,
|
||||
"parity": all(result["final_value"] == args.steps for result in old_results + new_results),
|
||||
"old_runs": old_results,
|
||||
"new_runs": new_results,
|
||||
},
|
||||
indent=2,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,4 +1,4 @@
|
||||
# Target vllm version: v0.28.0 (tag v0.28.0 @ 2cf0a6915ce5).
|
||||
# Target vllm version: v0.27.0 (releases/v0.27.0 branch).
|
||||
# Keep VLLM_BASE_TAG aligned with the target vLLM release on every rebase.
|
||||
#
|
||||
# Release-boundary form: there is no precompiled-wheel pin here, because the
|
||||
@@ -14,14 +14,11 @@
|
||||
# core init), registering libnvJitLink.so.13 with ldconfig, and re-pinning numpy
|
||||
# under numba's ceiling. All of it was compensation for a base image and a vLLM
|
||||
# that disagreed. At a real release boundary they agree, so none of it is
|
||||
# needed. The 0.28 rebase repeated the cycle: the v0.28.0rc / final-tag wheels
|
||||
# rode on the v0.27.0 base image with the full reinstall block until the
|
||||
# official v0.28.0 image shipped, at which point this file dropped back to the
|
||||
# form you are reading. If you pin an untagged commit again, restore the
|
||||
# reinstall block — agent/tasks/20_update_wheel_url.sh fails the run rather
|
||||
# than letting the Dockerfile silently ship the wrong vLLM.
|
||||
# needed. If you pin an untagged commit again, restore the reinstall block —
|
||||
# agent/tasks/20_update_wheel_url.sh fails the run rather than letting the
|
||||
# Dockerfile silently ship the wrong vLLM.
|
||||
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
|
||||
ARG VLLM_BASE_TAG=v0.28.0
|
||||
ARG VLLM_BASE_TAG=v0.27.0
|
||||
|
||||
# We use registry-based layer caching for the dependencies,
|
||||
# where the tag is determined based on the hash of the contents
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai:v0.28.0
|
||||
ARG BASE_IMAGE=vllm/vllm-openai:v0.26.0
|
||||
FROM ${BASE_IMAGE}
|
||||
|
||||
ARG COMMON_WORKDIR=/app
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.28.0
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.27.0
|
||||
FROM ${BASE_IMAGE} AS base
|
||||
|
||||
# Declare a variable to know if we want to use the nightly build or the stable build.
|
||||
@@ -10,7 +10,7 @@ FROM ${BASE_IMAGE} AS base
|
||||
# we should swap over to use stable release ASAP.
|
||||
# We should avoid relying on custom commits.
|
||||
ARG USE_NIGHTLY_BUILD=0
|
||||
ARG VLLM_VERSION_OR_COMMIT_HASH=v0.28.0
|
||||
ARG VLLM_VERSION_OR_COMMIT_HASH=v0.27.0
|
||||
ARG ARG_PYTORCH_ROCM_ARCH
|
||||
ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
|
||||
|
||||
|
||||
+1
-1
@@ -24,6 +24,7 @@ nav:
|
||||
- Text-To-Speech: user_guide/examples/offline_inference/text_to_speech.md
|
||||
- Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
|
||||
- X-To-Text: user_guide/examples/offline_inference/x_to_text.md
|
||||
- X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
|
||||
- Online Serving:
|
||||
- Image-To-Image: user_guide/examples/online_serving/image_to_image.md
|
||||
- Image-To-Video: user_guide/examples/online_serving/image_to_video.md
|
||||
@@ -53,7 +54,6 @@ nav:
|
||||
- BitsAndBytes W4: user_guide/quantization/bitsandbytes.md
|
||||
- MXFP8 W8A8: user_guide/quantization/mxfp8.md
|
||||
- MXFP4 W4A4: user_guide/quantization/mxfp4.md
|
||||
- SVDQuant W4A4: user_guide/quantization/svdquant.md
|
||||
- ModelOpt: user_guide/quantization/modelopt.md
|
||||
- AutoRound: user_guide/quantization/autoround.md
|
||||
- msModelSlim: user_guide/quantization/msmodelslim.md
|
||||
|
||||
@@ -114,10 +114,6 @@ Configuration classes.
|
||||
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][]
|
||||
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][]
|
||||
- [vllm_omni.transformers_utils.configs.cosyvoice3.CosyVoice3Config][]
|
||||
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Config][]
|
||||
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLConfig][]
|
||||
- [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.voxcpm2.VoxCPM2Config][]
|
||||
|
||||
|
||||
@@ -82,8 +82,6 @@ depends on the installed kernels and model path.
|
||||
| --- | --- | --- | --- | --- |
|
||||
| `SPEAKER_SAMPLES_DIR` | Filesystem path; default `~/.cache/vllm-omni/speakers` | Speech server; read when speaker storage initializes | Environment-only setting. The directory is created; filesystem errors propagate. | Stable |
|
||||
| `SPEAKER_MAX_UPLOADED` | Integer; default `1000` | Speech server; read when speaker storage initializes | Environment-only setting. A non-integer logs a warning and uses `1000`; range is not otherwise validated. | Stable |
|
||||
| `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` | Float seconds; default `600` | Diffusion engine async-output wait in `step_streaming`; resolved per call on the request path, not at import | Environment-only setting. A non-float or `<=0` value warns once and uses the default. | Experimental |
|
||||
| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `1`, `true`, `yes` or `on` enables; default `0` (off) | Orchestration loop and the serving-side final-output drain; read once when the `Orchestrator` is constructed | Environment-only setting. Values are stripped and case-normalized; any unrecognized value leaves the legacy poll loop selected. | Experimental |
|
||||
| `VLLM_OMNI_INPUT_WAIT_TIMEOUT_S` | Float seconds; default `600`; `<=0` disables | Full-payload input coordinator, not async-chunk transfer; read when the scheduler module imports in each worker | Environment-only setting. A non-float logs a warning and uses `600`. | Stable operational control |
|
||||
| `VLLM_OMNI_ORCH_MONITOR_PATH` | Filesystem path; default `<current-working-directory>/vllm_omni_orch_monitor_<timestamp>.json` | Orchestrator monitor enabled by `--enable-orch-monitor`; read when the monitor is created | Environment-only path override. Parent directories are created; write errors are logged. | Diagnostic |
|
||||
| `VLLM_OMNI_VIDEO_SYNC_TIMEOUT` | Float seconds; default `600` | Synchronous Videos API; read when the API server module imports | Environment-only setting. A non-float raises `ValueError` during import. | Experimental |
|
||||
@@ -179,7 +177,7 @@ collection, examples, bug reports, or logs.
|
||||
|
||||
## Model-specific variables
|
||||
|
||||
The audit found 56 variables read by a single model or pipeline family. They are
|
||||
The audit found 57 variables read by a single model or pipeline family. They are
|
||||
not listed as public usage options here because doing so would turn implementation
|
||||
escape hatches into an accidental compatibility contract.
|
||||
|
||||
@@ -192,7 +190,7 @@ Every audited model-specific name has a migration disposition in the
|
||||
| Request scope | 6 | Move request-varying behavior into a declared request-option schema. |
|
||||
| External | 0 | Retain only when a supported third-party library owns the contract. |
|
||||
| Internalize | 11 | Keep a debug or diagnostic switch out of public documentation and configuration. |
|
||||
| Deprecate/remove | 6 | Remove a compatibility escape hatch that has no continuing contract. |
|
||||
| Deprecate/remove | 7 | Remove a compatibility escape hatch that has no continuing contract. |
|
||||
|
||||
The disposition is a migration target, not a statement that the existing
|
||||
environment switch is stable. Promote or request-scope work should land in
|
||||
|
||||
@@ -14,7 +14,7 @@ describe fixed model topology and are not accepted in deploy YAMLs.
|
||||
Common `PipelineConfig` fields include:
|
||||
|
||||
| Field | Description |
|
||||
| ------- | ------------- |
|
||||
|-------|-------------|
|
||||
| `model_type` | Pipeline identifier used during model and config resolution. |
|
||||
| `default_deploy_config_name` | Bundled deploy YAML loaded when the user does not pass `deploy_config`. |
|
||||
| `model_arch` | Default Hugging Face architecture for the pipeline. |
|
||||
@@ -26,7 +26,7 @@ Common `PipelineConfig` fields include:
|
||||
Common `StagePipelineConfig` fields include:
|
||||
|
||||
| Field | Description |
|
||||
| ------- | ------------- |
|
||||
|-------|-------------|
|
||||
| `stage_id` | Stable stage identifier. |
|
||||
| `model_stage` | Logical stage name used by runtime and strategy resolution. |
|
||||
| `execution_type` | `LLM_AR`, `LLM_GENERATION`, or `DIFFUSION`. |
|
||||
@@ -49,17 +49,17 @@ deployment overrides.
|
||||
The new deploy schema lives under `vllm_omni/deploy/` and is paired with a frozen `PipelineConfig` registered by the model's `pipeline.py`. Each deploy YAML has these top-level fields:
|
||||
|
||||
| Field | Type | Required | Default | Description |
|
||||
| ------- | ------ | ---------- | --------- | ------------- |
|
||||
|-------|------|----------|---------|-------------|
|
||||
| `base_config` | str (path) | optional | — | Overlay parent (relative or absolute). `stages:` / `platforms:` deep-merged by stage_id; other scalars overlay-wins. Intended for user-authored overlays; prod yamls stay flat. |
|
||||
| `async_chunk` | bool | optional | `true` | Enable chunked streaming between stages. Pin to `false` if the pipeline runs end-to-end. |
|
||||
| `session_mode` | str | optional | `"turn"` | Session behavior. MiniCPM-o 4.5 deploy YAMLs set `"duplex"` so `/v1/realtime` and chat share the same profile. |
|
||||
| `session_mode` | str | optional | `"turn"` | Session behavior. Use `"duplex"` only with a pipeline that enables duplex control. |
|
||||
| `active_stream_window` | int | optional | `0` | Number of active downstream stream slots; `0` preserves all-stream cycling. |
|
||||
| `duplex_session` | dict | optional | runtime defaults | Full-duplex session lifecycle, buffering, replay, and capacity limits. |
|
||||
| `connectors` | dict | optional | `null` | Named connector specs (`{name, extra}`). Referenced by each stage's `input_connectors` / `output_connectors`. See [Connector schema](#connector-schema). |
|
||||
| `edges` | list | optional | `null` | Explicit edge list for the KV transfer graph. Auto-derived from stage inputs if omitted. |
|
||||
| `stages` | list | optional | `[]` | Per-stage runtime overrides matched by `stage_id`. Pipeline stages are still created from `PipelineConfig` when this list is empty. |
|
||||
| `platforms` | dict | optional | `null` | Keyed by `npu` / `rocm` / `xpu`, each contains a `stages:` list with per-platform overrides applied on top of the CUDA defaults. |
|
||||
| `pipeline` | str | optional | `null` | Override the auto-detected pipeline registry key (used for structural variants like `qwen2_5_omni_thinker_only` / `qwen3_omni_moe_thinker_only`). |
|
||||
| `pipeline` | str | optional | `null` | Override the auto-detected pipeline registry key (used for structural variants like `qwen2_5_omni_thinker_only`). |
|
||||
| `trust_remote_code` | bool \| null | optional | `null` | **Pipeline-wide.** Trust HF remote code on model load; applies to every stage when specified. |
|
||||
| `distributed_executor_backend` | str \| null | optional | `null` | **Pipeline-wide.** Distributed executor backend forwarded to vLLM (`"mp"`, `"ray"`, `"external_launcher"`). If omitted, vLLM auto-selects backend from runtime topology. |
|
||||
| `dtype` | str \| null | optional | `null` | **Pipeline-wide.** Model dtype for every stage. |
|
||||
@@ -84,7 +84,7 @@ subprocess on one GPU. `ray` / `external_launcher` are not fully supported yet.
|
||||
Each entry under `stages:` accepts any `StageDeployConfig` field directly (no nested `engine_args:`). Only fields whose value legitimately varies across stages live here; pipeline-wide settings (trust_remote_code, distributed_executor_backend, dtype, quantization, prefix/chunked prefill, DP/PP sizes) are declared at the top level and applied to every stage. Unknown keys fall through to `engine_extras:` and are forwarded to the engine. Frequently used fields are listed below; the source-of-truth schema is `StageDeployConfig` in `vllm_omni/config/stage_config.py`.
|
||||
|
||||
| Field | Type | Required | Default | Description |
|
||||
| ------- | ------ | ---------- | --------- | ------------- |
|
||||
|-------|------|----------|---------|-------------|
|
||||
| `stage_id` | int | required | — | Stage identity; matched against `PipelineConfig.stages[*].stage_id`. |
|
||||
| `max_num_seqs` | int \| null | optional | `null` | Max concurrent sequences per stage. |
|
||||
| `gpu_memory_utilization` | float \| null | optional | `null` | Per-stage total memory target; used for automatic KV-cache sizing. |
|
||||
@@ -94,14 +94,12 @@ Each entry under `stages:` accepts any `StageDeployConfig` field directly (no ne
|
||||
| `max_num_batched_tokens` | int \| null | optional | `null` | Per-stage prefill/token budget; also contributes to the native maximum in-flight token limit. |
|
||||
| `max_model_len` | int \| null | optional | `null` | Per-sequence context or KV length; `-1` enables native cache-capacity auto-fitting, while values above the HF default auto-set `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1`. |
|
||||
| `async_scheduling` | bool \| null | optional | `null` | Per-stage async scheduling toggle. |
|
||||
| `devices` | str \| null | optional | `null` | Device list assigned to this stage. The number of device ids must equal this stage's local world size (`tensor_parallel_size` × local data-parallel size × `pipeline_parallel_size`, or `num_replicas` × that product for a replica pool); a mismatch fails early — see the note below. |
|
||||
| `devices` | str \| null | optional | `null` | Device list assigned to this stage. |
|
||||
| `output_connectors` | dict \| null | optional | `null` | Keyed by `to_stage_<n>`; values are names registered under top-level `connectors:`. |
|
||||
| `input_connectors` | dict \| null | optional | `null` | Keyed by `from_stage_<n>`; values are names registered under top-level `connectors:`. |
|
||||
| `default_sampling_params` | dict \| null | optional | `null` | Baseline sampling params. Deep-merged with pipeline `sampling_constraints` (pipeline wins). |
|
||||
| `engine_extras` | dict | optional | `{}` | Catch-all for engine fields not listed above; deep-merged across overlays and forwarded to the stage engine. |
|
||||
|
||||
**Note:** a stage's `devices` count must equal its local world size (`tensor_parallel_size` × `data_parallel_size_local` × `pipeline_parallel_size`, falling back to global `data_parallel_size` when the local size is unset), or `num_replicas` × that product for a replica pool. A mismatch fails early and names the offending stage. A top-level `--tensor-parallel-size` is broadcast to every stage, so it can make a single-GPU stage violate this contract ([issue #5003](gh-issue:5003)); fix that case with `--stage-overrides` (set `tensor_parallel_size` and `devices` together per stage) or set TP only on the multi-GPU stage.
|
||||
|
||||
### Connector schema
|
||||
|
||||
Each entry under top-level `connectors:` follows this shape:
|
||||
@@ -115,9 +113,9 @@ connectors:
|
||||
# Additional connector-specific options
|
||||
```
|
||||
|
||||
| Connector class | Use case | `extra` keys |
|
||||
|--------------------------|-----------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------|
|
||||
| `SharedMemoryConnector` | Same-host KV transfer between stages (default for bundled YAMLs). | None. All payloads use shared memory. |
|
||||
| Connector class | Use case | `extra` keys |
|
||||
|-----------------|----------|--------------|
|
||||
| `SharedMemoryConnector` | Same-host KV transfer between stages (default for bundled YAMLs). | None. All payloads use shared memory. |
|
||||
| `MooncakeStoreConnector` | Cross-host KV transfer over TCP. Required for multi-node deployments. | `host`, `metadata_server`, `master`, `segment` (int bytes), `localbuf` (int bytes), `proto` (`"tcp"` / `"rdma"`). |
|
||||
|
||||
A stage references a connector by name in its `input_connectors` / `output_connectors`:
|
||||
@@ -137,7 +135,7 @@ stages:
|
||||
### CLI flags
|
||||
|
||||
| Flag | Description |
|
||||
| ------ | ------------- |
|
||||
|------|-------------|
|
||||
| `--deploy-config PATH` | Load a deploy YAML. **Optional** — when omitted, the bundled `vllm_omni/deploy/<model_type>.yaml` is auto-loaded by the model registry. |
|
||||
| `--stage-overrides JSON` | Per-stage JSON overrides, e.g. `'{"0":{"gpu_memory_utilization":0.5}}'`. Per-stage values always win over global flags. |
|
||||
| `--async-chunk` / `--no-async-chunk` | Flip the deploy YAML's `async_chunk:` bool. Unset (default) leaves the YAML value in force. |
|
||||
@@ -251,7 +249,7 @@ CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
Effective config per stage after the merge:
|
||||
|
||||
| Stage | Field | Final value | Source |
|
||||
| ------- | ------- | ------------- | -------- |
|
||||
|-------|-------|-------------|--------|
|
||||
| 0 | `gpu_memory_utilization` | `0.9` | base YAML (overlay didn't touch stage 0) |
|
||||
| 0 | `max_num_seqs` | `8` | per-stage CLI (`--stage-overrides`) — wins over base `32` |
|
||||
| 0 | `max_model_len` | `16384` | global CLI |
|
||||
@@ -275,14 +273,12 @@ when loading a deploy configuration.
|
||||
Examples:
|
||||
|
||||
For offline inference (assuming the necessary dependencies have been imported):
|
||||
|
||||
```python
|
||||
model_name = "Qwen/Qwen2.5-Omni-7B"
|
||||
omni = Omni(model=model_name, deploy_config="/path/to/deploy_config.yaml")
|
||||
```
|
||||
|
||||
For online serving:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
|
||||
```
|
||||
|
||||
@@ -322,15 +322,6 @@ Each 1-second window records:
|
||||
|
||||
On shutdown the server also logs a short summary (`loop_active_pct`, per-replica queue averages/maxima).
|
||||
|
||||
`loop_idle` and `loop_active` count loop iterations, so their absolute scale is
|
||||
tied to which orchestration loop is running. Under the default poll loop an idle
|
||||
orchestrator records roughly one iteration per millisecond. Under the
|
||||
event-driven loop (`VLLM_OMNI_EVENT_DRIVEN_ORCH=1`, see
|
||||
[Speech API](../serving/speech_api.md#orchestration-loop-experimental)) an idle
|
||||
orchestrator wakes only once per 0.5 s reconcile timeout, while a busy one still
|
||||
records one iteration per routed output. Window counts and the `loop_active_pct`
|
||||
summary are therefore not comparable across the two modes.
|
||||
|
||||
### Relationship to other diagnostics
|
||||
|
||||
This monitor is intentionally separate from the existing profiling tools:
|
||||
|
||||
@@ -63,7 +63,7 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic.
|
||||
|
||||
### N-Branch CFG (3+ branches)
|
||||
|
||||
Some models require more than 2 CFG branches. For example, Bagel and OmniGen2 use 3 branches.
|
||||
Some models require more than 2 CFG branches. For example, Bagel and OmniGen2 use 3 branches, DreamID Omni uses 4 branches.
|
||||
|
||||
`predict_noise_with_multi_branch_cfg()` handles these by automatically dispatching N branches across M GPUs using round-robin (rule: branch `i` → rank `i % M`):
|
||||
|
||||
|
||||
@@ -14,12 +14,12 @@ specified in the [Host Weight Runtime module design](../module/host_weight_runti
|
||||
|
||||
## Status
|
||||
|
||||
The implementation provides contracts and a CPU local-filesystem store. The
|
||||
diffusion consumer additionally defines typed,
|
||||
The first implementation provides contracts and a CPU local-filesystem store.
|
||||
The initial diffusion consumer contract additionally defines typed,
|
||||
representation-independent final-layout identity/restoration mechanics plus a
|
||||
concrete BF16-with-preserved-FP32 policy for MiniMax H3 and
|
||||
`black-forest-labs/FLUX.2-klein-4B`. The opt-in no-AllGather DLO integration
|
||||
selects, publishes, restores, and transfers these artifacts.
|
||||
concrete BF16-with-preserved-FP32 policy for MiniMax H3. It remains
|
||||
library-only: no loader or DLO path selects, publishes, restores, or transports
|
||||
that artifact yet.
|
||||
|
||||
V1 includes:
|
||||
|
||||
@@ -33,12 +33,13 @@ V1 includes:
|
||||
|
||||
V1 does not include:
|
||||
|
||||
- default-on activation or consumers outside no-AllGather DLO;
|
||||
- online FP8, quantized, merged-adaptation, or additional model producers;
|
||||
- HWR interaction with DLO AllGather;
|
||||
- a public CLI, loader activation, or default-on model integration;
|
||||
- lease handoff to DLO or another transport consumer;
|
||||
- FP8, quantized, merged-adaptation, or additional model producers;
|
||||
- CUDA registration, pinned staging, H2D scheduling, or GPU kernels;
|
||||
- a remote artifact provider or cross-node coordination;
|
||||
- automatic eviction; or
|
||||
- a change to DLO collective or execution behavior.
|
||||
- a change to DLO AllGather or no-AllGather behavior.
|
||||
|
||||
## Motivation and use cases
|
||||
|
||||
@@ -254,8 +255,7 @@ The contract is intentionally separate from loader activation:
|
||||
- `FinalLayoutBF16Producer` accepts only the matching identity context and a
|
||||
finalized CPU model. It is `POST_LOAD_ONLY` and `SINGLE_PROCESS` per exact TP
|
||||
coordinate. Its BF16 policy preserves model-declared FP32 parameters and
|
||||
buffers, revalidates MiniMax H3 mixed-precision invariants, and revalidates
|
||||
FLUX.2-klein's two block stacks, packed QKV mapping, and BF16 base layout.
|
||||
buffers and revalidates MiniMax H3 mixed-precision invariants.
|
||||
|
||||
Other representations reuse source identity, typed parallel identity, tensor
|
||||
ownership, and exact restoration only when their policy proves those semantics.
|
||||
|
||||
@@ -29,6 +29,7 @@ Legend: ✅ supported, ⚠️ compatibility path or limited validation, ❌ unsu
|
||||
| **Other online quantization methods** | ❌ Rejected until runtime packing and scale layouts are validated. | ⚠️ Allowed through the ordinary loader; validation is method-specific. |
|
||||
| **Model-level or standard layerwise CPU offload** | ❌ Disabled because DLO takes priority. | ❌ Disabled because DLO takes priority. |
|
||||
| **Resident leading layers** | ❌ Rejected. | ✅ Requires eligible resident paths in the model's `OffloadPlan`. |
|
||||
| **Auxiliary component policy** | ✅ Encoders/VAEs may remain resident. | ✅ Encoders/VAEs may remain resident. |
|
||||
|
||||
See [Parallelism compatibility](#parallelism-compatibility) and
|
||||
[Request and loading constraints](#request-and-loading-constraints) for the
|
||||
@@ -220,15 +221,9 @@ model revision and parallel layout with `required`. TP coordinates have distinct
|
||||
identities, while equivalent DP replicas share them.
|
||||
|
||||
The current producer/consumer boundary is the model-declared final-layout BF16
|
||||
contract (MiniMax H3 and `black-forest-labs/FLUX.2-klein-4B`). The FLUX.2-klein
|
||||
contract covers both transformer block stacks, constructor-stable packed QKV
|
||||
mapping state, BF16 parameters, and any persistent loader-owned `beta`/`eps`
|
||||
buffers. The shared HWR machinery supports ordinary-loader final layouts for
|
||||
TP1 and TP2 rank identities plus SP layout identities. FLUX.2-klein evidence
|
||||
in this change is TP1-only; its TP2/SP layouts rely on the existing
|
||||
per-coordinate identity mechanics and remain unmeasured. Online quantization,
|
||||
HSDP, LoRA/adapted weights, and non-default load formats remain ineligible.
|
||||
HWR mode
|
||||
contract (MiniMax H3 today). It supports ordinary-loader final layouts for TP1
|
||||
and TP2 rank identities plus SP layout identities; online quantization, HSDP,
|
||||
LoRA/adapted weights, and non-default load formats remain ineligible. HWR mode
|
||||
`disabled`, DLO-disabled, and DLO AllGather configurations stop before source
|
||||
identity or store construction and retain the existing checkpoint-mmap or
|
||||
ordinary-loader path.
|
||||
@@ -425,6 +420,22 @@ failure.
|
||||
- **HSDP + DP or TP:** rejected independently by the diffusion parallel
|
||||
configuration.
|
||||
|
||||
### Auxiliary component policy
|
||||
|
||||
`dlo_offload_components` maps discovered encoder/VAE paths to booleans. A
|
||||
component set to `False` is materialized on the target device once and is not
|
||||
entered through the model's manual offload lifecycle. A `True` value permits
|
||||
the existing `OffloadPlan.on_demand_component_paths` and encoder block-hook
|
||||
behavior; it does not add support where the model declares none. The optional
|
||||
`default` key handles unmatched auxiliary components, and an empty map
|
||||
preserves the current policy.
|
||||
|
||||
The policy is resolved after component discovery and before DLO mutates DiT
|
||||
storage. Unknown names, including DiT paths, are rejected. DLO continues to
|
||||
own DiT streaming regardless of this map, which avoids an invalid state where
|
||||
the loader has skipped ordinary DiT materialization but the backend is asked
|
||||
to keep that DiT resident.
|
||||
|
||||
## Request and loading constraints
|
||||
|
||||
AllGather DP multi-concurrency requires:
|
||||
@@ -436,6 +447,15 @@ AllGather DP multi-concurrency requires:
|
||||
The no-AllGather path does not impose these DLO-specific synchronized-wave
|
||||
requirements.
|
||||
|
||||
MiniMax-H3 currently builds one global encoder TP group and broadcasts the
|
||||
resulting conditioning from global rank 0. Therefore an encoder TP group that
|
||||
spans multiple DP replicas is not valid for independent-prompt AllGather
|
||||
waves. Exact-recipe DP2/TP2 and DP2/SP2 experiments with encoder TP4 completed
|
||||
the collective schedule, but two distinct prompts with the same seed produced
|
||||
byte-identical video and audio in every measured wave. Enabling or disabling
|
||||
encoder offload produced paired-identical results, so this is a pre-existing
|
||||
model topology limitation rather than component-policy behavior.
|
||||
|
||||
Direct checkpoint mmap can back either transfer path. It is currently limited
|
||||
to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
|
||||
ordinary loader. Per-tensor online FP8 linears can use DLO AllGather after the
|
||||
@@ -500,6 +520,9 @@ AllGather, because DP replicas did not share checkpoint-backed runtime
|
||||
weights. This is a functional and memory smoke test, not a production-quality
|
||||
performance or output-quality benchmark.
|
||||
|
||||
These B300 DP waves intentionally reused the same prompt and seed, so they do
|
||||
not establish independent-prompt isolation across DP replicas.
|
||||
|
||||
### Host-memory measurement
|
||||
|
||||
A two-worker MiniMax-H3 FL2VA measurement on one L20X node compared the
|
||||
|
||||
+30
-76
@@ -6,7 +6,7 @@ This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage
|
||||
|
||||
- Expose pipeline-level request and latency metrics that span the full multi-stage execution (orchestrator scope).
|
||||
- Preserve all upstream vLLM per-engine metrics (`vllm:*`) for stages backed by an AR LLM engine, and reshape their `engine` label into `stage` + `replica` so multi-replica deployments gain per-replica visibility automatically.
|
||||
- Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss, plus image and diffusion workload, queueing, execution, and memory signals.
|
||||
- Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss.
|
||||
- Expose per-replica-edge cross-stage transfer metrics so the slack between E2E latency and the sum of per-stage `gen_time` (queueing, serialization, network) becomes attributable.
|
||||
- Keep the metrics collection overhead low enough that it does not regress TTFA or throughput.
|
||||
|
||||
@@ -20,7 +20,7 @@ vLLM's `unregister_vllm_metrics()` function strips every `prometheus_client` col
|
||||
|
||||
### The Problem
|
||||
|
||||
vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm_omni:` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm_omni:` / `vllm_omni`.
|
||||
vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm:omni_` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
|
||||
|
||||
Upstream per-engine metrics retain the `vllm:` prefix but are now registered by `OmniPrometheusStatLogger`, a thin subclass of upstream's `PrometheusStatLogger` that reshapes the single `engine` label into a `stage` + `replica` pair (see "OmniPrometheusStatLogger wrap" below).
|
||||
|
||||
@@ -38,7 +38,7 @@ Upstream per-engine metrics retain the `vllm:` prefix but are now registered by
|
||||
|
|
||||
+--------+--------+--------+--------+--------+
|
||||
| |
|
||||
vllm_omni:* vllm:*
|
||||
vllm:omni_* vllm:*
|
||||
collectors collectors
|
||||
| |
|
||||
+----+--------+ +-----------+ +----------+ +-----------+
|
||||
@@ -58,30 +58,17 @@ Upstream per-engine metrics retain the `vllm:` prefix but are now registered by
|
||||
|
||||
### Data Flow
|
||||
|
||||
There are five independent paths for metric collection.
|
||||
There are four independent paths for metric collection.
|
||||
|
||||
**Path 1: Pipeline-level metrics (`vllm_omni:*`)**
|
||||
**Path 1: Pipeline-level metrics (`vllm:omni_*`)**
|
||||
|
||||
`OmniPrometheusMetrics` registers the Gauge / Counter / Histogram collectors at import time. It is instantiated once per entrypoint, labeled with the model name. The entrypoint calls its methods as requests progress:
|
||||
|
||||
- `set_running(n)` / `set_waiting(n)` — updated after each request completes. The running count comes from `OmniRequestCounter`, a simple counter incremented/decremented by the Orchestrator as it tracks requests. Waiting is derived as `total - running`.
|
||||
- `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm_omni:requests_success_total{finished_reason}`.
|
||||
- `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` on the existing completion Counter.
|
||||
- `inc_requests_failed(reason)` — increments the dedicated failure-attribution Counter once per request. Reasons are normalized to the bounded set `client_abort`, `client_disconnect`, `stage_error`, and `unknown`.
|
||||
- `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm:omni_requests_success_total{finished_reason}`.
|
||||
- `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` so a single Counter family covers both natural and aborted completion.
|
||||
|
||||
**Path 2: Image and diffusion metrics**
|
||||
|
||||
Image and diffusion metrics use two cooperating collectors. `OmniPrometheusMetrics` owns pipeline/stage workload families, while `OmniModalityMetrics` owns per-replica diffusion breakdowns. Finished stage messages are consumed once in `OmniBase._process_single_result`, which prevents replayed terminal messages from incrementing Counters or Histograms twice.
|
||||
|
||||
- `stage_gen_time_s`, `stage_in_queue_s`, image workload, inference-step, and per-stage peak-memory values are emitted from `StageRequestStats` when a stage finishes. `image_pixels` and `image_count_total` are restricted to image outputs; `num_inference_steps` is restricted to diffusion stages.
|
||||
- Diffusion engine timings are accumulated per request by `OrchestratorAggregator.accumulate_diffusion_metrics`. Engine-side millisecond values are converted to seconds before `observe_modality_at_finalize(...)` dispatches the per-replica Histogram observations.
|
||||
- `request_queue_wait_s` is emitted once at final request completion from `pipeline_timings["queue_wait_ms"]`. A present value of `0` is a valid observation; a missing key produces no sample.
|
||||
- `stage_waiting_requests` is the sum of the latest scheduler snapshots for all live replicas in a stage. AR snapshots arrive through `SchedulerStats`; diffusion snapshots ride normal diffusion output metrics. A dead replica's snapshot is removed. Because this deliberately reuses existing message paths, the gauge can retain its last value while a replica produces no output.
|
||||
- `kv_wait_s` measures the AR scheduler interval from entering the KV-transfer-free wait state until extraction acknowledgement. The sample is carried to the Orchestrator through the existing engine output and labeled by connector type. Abort/error/removal paths discard incomplete timers rather than emitting partial waits.
|
||||
|
||||
Zero and missing values have different meanings. Valid zero-duration queue observations are recorded, while unavailable optional measurements are omitted instead of being exported as synthetic zeroes. `peak_memory_mb`, `diffusion_forward_s`, `diffusion_kv_load_s`, `vae_decode_s`, and `kv_wait_s` therefore appear only when their data source is active.
|
||||
|
||||
**Path 3: Audio modality metrics (`vllm_omni:audio_*`)**
|
||||
**Path 2: Audio modality metrics (`vllm:omni_audio_*`)**
|
||||
|
||||
`OmniModalityMetrics` registers seven audio families with `{model_name, stage, replica}` (plus an extra `threshold_ms` / `reason` label on the two extra-cardinality Counters). Three observation sites:
|
||||
|
||||
@@ -89,7 +76,7 @@ Zero and missing values have different meanings. Valid zero-duration queue obser
|
||||
- `observe_audio_first_packet(...)` — called from the OpenAI SSE audio branch in `serving_chat.py` on the first audio packet for a request. The once-per-request guard is held by `ClientRequestState.first_audio_ts`. The `request_arrival_ts` anchor is stored in `ClientRequestState` by `async_omni.generate()`, computed at request entry.
|
||||
- `observe_audio_streaming_finalize(...)` — called from `serving_chat.py` after the streaming chunk loop exhausts. It runs the per-chunk player simulation from `vllm_omni/benchmarks/audio_continuity.py` to compute the worst-case underrun and emits `audio_underrun_s` plus (when the request stayed below the threshold) `audio_continuity_ok_total{threshold_ms}`. Per-chunk PCM byte counts and arrival timestamps are recorded by the same audio branch that updates `first_audio_ts`.
|
||||
|
||||
**Path 4: Cross-stage transfer metrics (`vllm_omni:transfer_*`)**
|
||||
**Path 3: Cross-stage transfer metrics (`vllm:omni_transfer_*`)**
|
||||
|
||||
`OmniTransferMetrics` registers four Histogram families with `{model_name, from_stage, from_replica, to_stage, to_replica}` labels. Each observation corresponds to one physical transfer hop (one chunk between adjacent stages), not the per-request accumulated total — so the histograms track per-transfer distribution.
|
||||
|
||||
@@ -100,7 +87,7 @@ The hook lives in `OrchestratorAggregator.record_transfer_tx` and `record_transf
|
||||
|
||||
Defensive fail-safe: if `transfer_emitter` or `replica_resolver` is missing, or the resolver returns `None` for either side, the emit is skipped silently (the underlying `TransferEdgeStats` accumulation is unaffected).
|
||||
|
||||
**Path 5: Per-engine metrics (`vllm:*`, stage/replica wrap)**
|
||||
**Path 4: Per-engine metrics (`vllm:*`, stage/replica wrap)**
|
||||
|
||||
The Orchestrator instantiates `OmniPrometheusStatLogger` (a thin subclass of upstream `vllm.v1.metrics.loggers.PrometheusStatLogger`) and feeds it scheduler stats and iteration stats after processing each batch of engine outputs. This populates the standard ~37 vLLM metric families (TTFT, ITL, TPOT, KV cache usage, etc.) using the same upstream code path — but with the `engine` label reshaped into `stage` + `replica` so multi-replica deployments produce distinct series per replica. See the next section for the wrap mechanics.
|
||||
|
||||
@@ -110,9 +97,9 @@ The Orchestrator runs in a background thread. The API server (OmniBase) runs in
|
||||
|
||||
### Metric Registration and Lifecycle
|
||||
|
||||
All `vllm_omni:*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data.
|
||||
All `vllm:omni_*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data.
|
||||
|
||||
The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm_omni:*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime.
|
||||
The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm:omni_*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime.
|
||||
|
||||
## OmniPrometheusStatLogger Wrap
|
||||
|
||||
@@ -143,11 +130,11 @@ A reverse map `(stage_id, replica_id) -> flat_idx` is maintained on the Orchestr
|
||||
|
||||
## Gating: `--log-stats`
|
||||
|
||||
All metrics — both the omni-specific `vllm_omni:*` families and the upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`.
|
||||
All metrics — both the 15 `vllm:omni_*` families and the ~65 upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`.
|
||||
|
||||
Behavior with `--log-stats=off` (default):
|
||||
|
||||
- The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm_omni:*` stay materialized but never have data written to them.
|
||||
- The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm:omni_*` stay materialized but never have data written to them.
|
||||
- `OmniPrometheusStatLogger` is not constructed in `_init_metrics_state`, so the ~65 upstream `vllm:*` wrap families are not registered in the default registry at all.
|
||||
- The engine core's `Scheduler.make_stats()` also short-circuits inside upstream (`if not self.log_stats: return None`), so no `SchedulerStats` is produced per step — the per-iteration cost is bounded by the existing upstream gate.
|
||||
|
||||
@@ -175,43 +162,10 @@ second.
|
||||
|
||||
| Metric | Type | Labels | Description |
|
||||
|--------|------|--------|-------------|
|
||||
| `vllm_omni:num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages |
|
||||
| `vllm_omni:num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled |
|
||||
| `vllm_omni:requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
|
||||
| `vllm_omni:e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds |
|
||||
|
||||
### Image and diffusion service-level metrics
|
||||
|
||||
| Metric | Type | Labels | Description |
|
||||
|--------|------|--------|-------------|
|
||||
| `vllm_omni:stage_gen_time_s` | Histogram | `model_name`, `stage`, `stage_type` | Stage submit to finished output; includes in-stage queueing |
|
||||
| `vllm_omni:request_queue_wait_s` | Histogram | `model_name` | Orchestration-layer queue wait; a present zero is recorded |
|
||||
| `vllm_omni:stage_waiting_requests` | Gauge | `model_name`, `stage` | Sum of the latest waiting snapshots across live replicas in the stage |
|
||||
| `vllm_omni:stage_in_queue_s` | Histogram | `model_name`, `stage` | Diffusion scheduler wait from enqueue to initial execution admission |
|
||||
| `vllm_omni:num_inference_steps` | Histogram | `model_name` | Denoising-step distribution for diffusion stages |
|
||||
| `vllm_omni:image_count_total` | Counter | `model_name` | Cumulative number of generated image outputs |
|
||||
| `vllm_omni:image_pixels` | Histogram | `model_name` | Pixel-count distribution for image outputs |
|
||||
| `vllm_omni:peak_memory_mb` | Gauge | `model_name`, `stage` | Positive per-stage peak-memory samples reported by engine outputs |
|
||||
| `vllm_omni:requests_failed_total` | Counter | `model_name`, `reason` | Request failures with bounded reason values (`client_abort`, `client_disconnect`, `stage_error`, `unknown`) |
|
||||
| `vllm_omni:kv_wait_s` | Histogram | `model_name`, `connector_type` | Completed AR KV-transfer wait intervals, grouped by connector backend |
|
||||
|
||||
### Diffusion execution breakdown
|
||||
|
||||
Labels: `{model_name, stage, replica}`.
|
||||
|
||||
| Metric | Type | Description |
|
||||
|--------|------|-------------|
|
||||
| `vllm_omni:diffusion_exec_s` | Histogram | Core diffusion-step execution time per request |
|
||||
| `vllm_omni:diffusion_exec_per_step_s` | Histogram | Core diffusion execution divided by inference-step count |
|
||||
| `vllm_omni:diffusion_preprocess_s` | Histogram | Diffusion input preprocessing time |
|
||||
| `vllm_omni:diffusion_postprocess_s` | Histogram | Diffusion output postprocessing time |
|
||||
| `vllm_omni:vae_decode_s` | Histogram | VAE decode time when the engine reports it |
|
||||
| `vllm_omni:diffusion_forward_s` | Histogram | Forward-only denoise-loop time when pipeline profiling is enabled |
|
||||
| `vllm_omni:diffusion_kv_load_s` | Histogram | Diffusion-side KV receive/load time when upstream KV is used |
|
||||
| `vllm_omni:image_ttfp_s` | Histogram | Image-stage submit to first materialized image output |
|
||||
| `vllm_omni:denoise_step_latency_s` | Histogram | Mean diffusion forward time per denoise step for image output stages |
|
||||
|
||||
The optional breakdown and memory families are sparse by design: when an engine does not produce the source measurement, no sample is emitted. This differs from queue timings, where an explicitly measured zero is meaningful and is retained.
|
||||
| `vllm:omni_num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages |
|
||||
| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled |
|
||||
| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
|
||||
| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds |
|
||||
|
||||
### Audio (7)
|
||||
|
||||
@@ -219,13 +173,13 @@ Labels: `{model_name, stage, replica}` plus the listed extra label.
|
||||
|
||||
| Metric | Type | Extra label | Description |
|
||||
|--------|------|-------------|-------------|
|
||||
| `vllm_omni:audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
|
||||
| `vllm_omni:audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
|
||||
| `vllm_omni:audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` |
|
||||
| `vllm_omni:audio_frames_total` | Counter | — | Cumulative audio frames generated |
|
||||
| `vllm_omni:audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
|
||||
| `vllm_omni:audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
|
||||
| `vllm_omni:audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
|
||||
| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
|
||||
| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
|
||||
| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` |
|
||||
| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frames generated |
|
||||
| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
|
||||
| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
|
||||
| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
|
||||
|
||||
### Cross-stage transfer (4)
|
||||
|
||||
@@ -233,10 +187,10 @@ Labels: `{model_name, from_stage, from_replica, to_stage, to_replica}`.
|
||||
|
||||
| Metric | Type | Description |
|
||||
|--------|------|-------------|
|
||||
| `vllm_omni:transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
|
||||
| `vllm_omni:transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
|
||||
| `vllm_omni:transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
|
||||
| `vllm_omni:transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
|
||||
| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
|
||||
| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
|
||||
| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
|
||||
| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
|
||||
|
||||
### LLM stage-level (wrapped `vllm:*`)
|
||||
|
||||
@@ -249,7 +203,7 @@ After the wrap, every upstream `vllm:*` family — TTFT, ITL, TPOT, e2e latency,
|
||||
- `SECONDS_FAST_BUCKETS` (0.001 s – 60 s) for fine-grained cross-stage transfer and audio-underrun values that need millisecond-level resolution.
|
||||
- Counters use the `_total` suffix (auto-appended by `prometheus_client`).
|
||||
- Sizes use the `_bytes` suffix.
|
||||
- All omni-specific families are prefixed `vllm_omni:`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm_omni:` / `vllm_omni`.
|
||||
- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
|
||||
|
||||
## Logging vs. Prometheus
|
||||
|
||||
|
||||
@@ -66,12 +66,6 @@ the boundary affected by the in-flight stage client/process refactor in
|
||||
[#5441](https://github.com/vllm-project/vllm-omni/pull/5441). Names and
|
||||
responsibilities proposed only by that PR are not current contracts.
|
||||
|
||||
The orchestration loop also has an opt-in event-driven mode
|
||||
(`VLLM_OMNI_EVENT_DRIVEN_ORCH=1`, default off) proposed in
|
||||
[#5221](https://github.com/vllm-project/vllm-omni/pull/5221). It changes poll
|
||||
cadence only: the routing, ordering, and terminal-state contracts below hold
|
||||
identically on both loops.
|
||||
|
||||
## Ownership boundary
|
||||
|
||||
This document owns `AsyncOmniEngine`, `Orchestrator`, request-state creation,
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# Installation
|
||||
|
||||
!!! important
|
||||
vLLM-Omni is released against the matching upstream vLLM major/minor version. vLLM-Omni 0.28.x requires vLLM 0.28.x; the stable v0.28.0 instructions pin vLLM 0.28.0.
|
||||
vLLM-Omni is released against the matching upstream vLLM major/minor version. vLLM-Omni 0.26.x requires vLLM 0.26.x; the stable v0.26.0 instructions pin vLLM 0.26.0.
|
||||
|
||||
vLLM-Omni supports the following hardware platforms:
|
||||
|
||||
|
||||
@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
|
||||
|
||||
### Pre-built wheels
|
||||
|
||||
Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, and 0.26.0; the 0.28.0 wheels are published together with the v0.28.0 release this branch produces. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
|
||||
Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, and 0.26.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
|
||||
|
||||
=== "NVIDIA CUDA"
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
|
||||
|
||||
vLLM-Omni is built based on vLLM. Please install it with command below.
|
||||
```bash
|
||||
uv pip install vllm==0.28.0 --torch-backend=auto
|
||||
uv pip install vllm==0.26.0 --torch-backend=auto
|
||||
```
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
|
||||
# --8<-- [start:build-wheel-from-source]
|
||||
|
||||
#### Installation of vLLM
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.28.0 release version of the library
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.26.0 release version of the library
|
||||
|
||||
```bash
|
||||
uv pip install vllm==0.28.0 --torch-backend=auto
|
||||
uv pip install vllm==0.26.0 --torch-backend=auto
|
||||
```
|
||||
|
||||
The 0.28.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
The 0.26.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
Since vllm-omni is rapidly evolving, it's recommended to install it from source
|
||||
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm.git
|
||||
cd vllm
|
||||
git checkout v0.28.0
|
||||
git checkout v0.26.0
|
||||
```
|
||||
Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
|
||||
```bash
|
||||
#For CUDA 13.0 (the default for v0.28.0; the wheel filename has no `+cu130` suffix)
|
||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.28.0/vllm-0.28.0-cp38-abi3-manylinux_2_28_x86_64.whl
|
||||
#For CUDA 13.0 (the default for v0.26.0; the wheel filename has no `+cu130` suffix)
|
||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.26.0/vllm-0.26.0-cp38-abi3-manylinux_2_28_x86_64.whl
|
||||
```
|
||||
Install vllm with command below (If you have no existing PyTorch).
|
||||
```bash
|
||||
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
-p 8091:8091 \
|
||||
--ipc=host \
|
||||
vllm/vllm-omni:v0.28.0 \
|
||||
vllm/vllm-omni:v0.26.0 \
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
```
|
||||
|
||||
@@ -124,7 +124,7 @@ If you want to specify the base vLLM version:
|
||||
```bash
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
-f docker/Dockerfile.cuda \
|
||||
--build-arg BASE_IMAGE=vllm/vllm-openai:v0.28.0 \
|
||||
--build-arg BASE_IMAGE=vllm/vllm-openai:v0.26.0 \
|
||||
-t vllm-omni-cuda .
|
||||
```
|
||||
|
||||
|
||||
@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
|
||||
|
||||
vLLM-Omni is built based on vLLM. Please install it with command below.
|
||||
```bash
|
||||
uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
|
||||
uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
|
||||
```
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
|
||||
# --8<-- [start:build-wheel-from-source]
|
||||
|
||||
#### Installation of vLLM
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.28.0 release version of the library
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.26.0 release version of the library
|
||||
|
||||
```bash
|
||||
uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
|
||||
uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
|
||||
```
|
||||
|
||||
The pre-built 0.28.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
The pre-built 0.26.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
Since vllm-omni is rapidly evolving, it's recommended to install it from source
|
||||
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm.git
|
||||
cd vllm
|
||||
git checkout v0.28.0
|
||||
git checkout v0.26.0
|
||||
python3 -m pip install -r requirements/rocm.txt
|
||||
python3 setup.py develop
|
||||
```
|
||||
@@ -130,7 +130,7 @@ docker run --rm \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
-p 8091:8091 \
|
||||
vllm/vllm-omni-rocm:v0.28.0 \
|
||||
vllm/vllm-omni-rocm:v0.26.0 \
|
||||
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
```
|
||||
|
||||
@@ -149,7 +149,7 @@ docker run --rm -it \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
--entrypoint bash \
|
||||
vllm/vllm-omni-rocm:v0.28.0
|
||||
vllm/vllm-omni-rocm:v0.26.0
|
||||
```
|
||||
|
||||
# --8<-- [end:pre-built-images]
|
||||
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
|
||||
source .venv/bin/activate
|
||||
|
||||
# On CUDA
|
||||
uv pip install vllm==0.28.0 --torch-backend=auto
|
||||
uv pip install vllm==0.26.0 --torch-backend=auto
|
||||
|
||||
# On ROCm
|
||||
uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
|
||||
uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
|
||||
|
||||
git clone https://github.com/vllm-project/vllm-omni.git
|
||||
cd vllm-omni
|
||||
@@ -34,7 +34,7 @@ For additional installation methods — please see the [installation guide](inst
|
||||
!!! note
|
||||
It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
|
||||
|
||||
If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.28.0` and vLLM Omni `0.28.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
|
||||
If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.26.0` and vLLM Omni `0.26.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
|
||||
|
||||
## Offline Inference
|
||||
|
||||
|
||||
@@ -37,6 +37,7 @@ GENERAL_EXAMPLE_SLUGS = frozenset(
|
||||
"text_to_speech",
|
||||
"text_to_video",
|
||||
"x_to_text",
|
||||
"x_to_video_audio",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@@ -58,6 +58,7 @@ th {
|
||||
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Pi0Pipeline` | π0 (Pi-Zero) VLA | `lerobot/pi0_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi0.md) |
|
||||
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — |
|
||||
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | — |
|
||||
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
@@ -69,8 +70,6 @@ th {
|
||||
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — |
|
||||
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — |
|
||||
| `DotsTTSForConditionalGeneration` | dots.tts | `rednote-hilab/dots.tts-soar` | ✅︎ | | | | — |
|
||||
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — |
|
||||
| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — |
|
||||
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | | — |
|
||||
| `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
@@ -95,6 +94,7 @@ th {
|
||||
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — |
|
||||
| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — |
|
||||
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | | — |
|
||||
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | — |
|
||||
| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — |
|
||||
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — |
|
||||
|
||||
@@ -811,51 +811,6 @@ If you encounter OOM errors:
|
||||
|
||||
Use `/v1/audio/voices` to list available voices for the loaded model.
|
||||
|
||||
## Orchestration Loop (experimental)
|
||||
|
||||
Multi-stage omni deployments route stage outputs through a single orchestrator
|
||||
loop. By default that loop polls every stage replica on a 1 ms cadence. An
|
||||
opt-in event-driven mode replaces the poll with one reader task per live stage
|
||||
replica awaiting its client directly, and switches the serving-side
|
||||
final-output drain to a condition-variable wakeup at the same time.
|
||||
|
||||
**Configuration (environment variables):**
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `0` (off) | Switches the orchestration loop and the final-output drain from the legacy 1 ms poll to event-driven wakeups. Enabled by `1`, `true`, `yes`, or `on`, matched case-insensitively after surrounding whitespace is stripped; any other value leaves it off. |
|
||||
|
||||
Set it on the process that runs the orchestrator (stage 0 of an omni
|
||||
deployment) before starting the server:
|
||||
|
||||
```bash
|
||||
export VLLM_OMNI_EVENT_DRIVEN_ORCH=1
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
--omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
The server logs the selected loop mode and its reader/poller counts once at
|
||||
startup, so you can confirm which loop is live.
|
||||
|
||||
Routing, output ordering, and terminal-state behavior are identical on both
|
||||
loops; only the poll cadence changes. Leaving the variable unset keeps the
|
||||
legacy poll loop, which is the supported default.
|
||||
|
||||
**Known limitations:**
|
||||
|
||||
- The measured serving A/B (idle CPU 2.43% to 0.07%; TTFP p99 -32% at
|
||||
concurrency 8) predates the rebuild on the per-replica fault-isolation work
|
||||
in [#4285](https://github.com/vllm-project/vllm-omni/pull/4285). That work
|
||||
changed dead-replica handling and reader/poller lifecycle rather than the
|
||||
steady-state output path, and the parity suite covers it, but the serving
|
||||
A/B has not been re-run on the current head.
|
||||
- The diffusion-poller branch is covered by unit tests only. Deployments whose
|
||||
stages all run as standard engine cores never exercise it, including GLM-TTS,
|
||||
which deploys its DiT without `stage_type: diffusion`.
|
||||
- Concurrency 1 and 32 measured at parity with the legacy loop. At 32 the
|
||||
latency is admission-bound, which this mode does not address.
|
||||
|
||||
## Development
|
||||
|
||||
Enable debug logging:
|
||||
|
||||
@@ -36,7 +36,8 @@ vllm serve /path/to/model --omni \
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--data-parallel-size 4 \
|
||||
--dlo-no-use-allgather
|
||||
--dlo-no-use-allgather \
|
||||
--dlo-offload-components '{"text_encoder": false}'
|
||||
|
||||
# Sequence parallel deployment
|
||||
vllm serve /path/to/model --omni \
|
||||
@@ -51,6 +52,7 @@ omni = Omni(
|
||||
model="/path/to/model",
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=True,
|
||||
dlo_offload_components={"text_encoder": False},
|
||||
)
|
||||
```
|
||||
|
||||
@@ -63,6 +65,7 @@ omni = Omni(
|
||||
| `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
|
||||
| `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
|
||||
| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
|
||||
| `--dlo-offload-components JSON` | Permit or disable offload for discovered encoder/VAE components | `{}` |
|
||||
| `--host-weight-runtime-mode {disabled,preferred,required}` | HWR policy: no interaction, populate on a miss, or require an exact hit | `disabled` |
|
||||
| `--host-weight-runtime-root PATH` | Writable node-local HWR store shared by workers in one storage domain; required for `preferred` and `required` | unset |
|
||||
| `--dlo-host-registration-limit-gib N` | Optional per-worker ceiling for registering an HWR mmap; zero adds no ceiling | `0` |
|
||||
@@ -112,12 +115,6 @@ perform a collective and uses the same rank-local transfer behavior.
|
||||
HWR is an opt-in startup optimization for models that declare the final-layout
|
||||
BF16 restore contract. Use it only with no-AllGather DLO:
|
||||
|
||||
The validated BF16 model contracts currently cover MiniMax H3 and
|
||||
`black-forest-labs/FLUX.2-klein-4B`. FLUX.2-klein-9B shares the same model
|
||||
class but has not been validated against this contract. FLUX.2-dev, online
|
||||
FP8, HSDP, LoRA/adapted weights, and non-default load formats remain outside
|
||||
the validated scope.
|
||||
|
||||
```bash
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
@@ -250,6 +247,25 @@ class MyPipeline(nn.Module):
|
||||
When no plan exists, discovery falls back to
|
||||
`_layerwise_offload_blocks_attrs` and then heuristic attribute lookup.
|
||||
|
||||
## Auxiliary component policy
|
||||
|
||||
`dlo_offload_components` follows the same component-map style as diffusion
|
||||
quantization. Keys are discovered encoder or VAE paths, values are booleans,
|
||||
and `default` is the optional fallback:
|
||||
|
||||
```python
|
||||
dlo_offload_components = {
|
||||
"text_encoder": False,
|
||||
"default": True,
|
||||
}
|
||||
```
|
||||
|
||||
`True` permits the model's existing DLO behavior for that component. `False`
|
||||
keeps the whole component resident, including its child blocks and non-block
|
||||
modules. An empty map preserves current behavior. Unknown keys are rejected;
|
||||
the map does not control DiT blocks because enabling DLO gives the backend
|
||||
ownership of DiT streaming.
|
||||
|
||||
## Data-parallel concurrency
|
||||
|
||||
With `data_parallel_size > 1` and AllGather enabled, the scheduler can process
|
||||
@@ -274,6 +290,13 @@ must enter each collective.
|
||||
are validated.
|
||||
- Resident leading layers require `--dlo-no-use-allgather` and a model
|
||||
`OffloadPlan` that declares eligible `resident_dit_paths`.
|
||||
- Component policy applies only to discovered auxiliary encoders and VAEs;
|
||||
it cannot disable DLO for a DiT.
|
||||
- MiniMax-H3's encoder TP group is global rather than DP-replica-local. An
|
||||
encoder TP size that spans multiple DP replicas (for example TP4 with
|
||||
DP2/TP2 or DP2/SP2 on four ranks) is not supported for independent-prompt
|
||||
AllGather waves: the current pipeline broadcasts rank 0's conditioning to
|
||||
both replicas. This limitation is independent of the component policy.
|
||||
- DP concurrency requires an explicit, identical inference-step count.
|
||||
|
||||
Sharing quantized or otherwise unvalidated transformed layouts through a
|
||||
|
||||
@@ -129,7 +129,7 @@ The following tables show which models support each feature:
|
||||
| **Krea 2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **OmniGen2** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
@@ -147,7 +147,7 @@ The following tables show which models support each feature:
|
||||
|
||||
> Notes:
|
||||
> 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
|
||||
> 2. `Tongyi-MAI/Z-Image-Turbo` is a distilled model with minimal NFEs; CFG-Parallel is not necessary.
|
||||
> 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
|
||||
> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload swaps the nested UND reasoner and GEN generator pathways; layerwise offload remains available for blockwise GEN/UND offload.
|
||||
> 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`.
|
||||
> 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
|
||||
@@ -165,6 +165,7 @@ The following tables show which models support each feature:
|
||||
| **SANA-Video-2B T2V I2V** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ |
|
||||
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* |
|
||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||
| **LongCat-Video-Avatar-1.5** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **MiniMax-H3** | ✅ (FL2VA) | ✅ | ✅ | ❌ | ✅ (DiT/TE) | ❌ | ✅ | ✅ | ✅ (tile) | ✅ (DiT) | ❌ |
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
# X-To-Video-Audio
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/x_to_video_audio>.
|
||||
|
||||
|
||||
The `DreamID-Omni` pipeline generates short videos from text, image and video.
|
||||
|
||||
## Local CLI Usage
|
||||
### Download the Model locally
|
||||
Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally.
|
||||
|
||||
```bash
|
||||
python download_dreamid_omni.py --output-dir ./dreamid_omni
|
||||
```
|
||||
After download, the model directory will look like this:
|
||||
|
||||
```
|
||||
dreamid_omni/
|
||||
├── DreamID-Omni/
|
||||
│ ├── dreamid_omni.safetensors
|
||||
├── MMAudio/
|
||||
│ ├── ext_weights/
|
||||
│ │ ├── best_netG.pt
|
||||
│ │ ├── v1-16.pth
|
||||
├── Wan2.2-TI2V-5B/
|
||||
│ ├── google/*
|
||||
│ ├── models_t5_umt5-xxl-enc-bf16.pth
|
||||
│ ├── Wan2.2_VAE.pth
|
||||
│
|
||||
├── model_index.json # create by download_dreamid_omni.py
|
||||
```
|
||||
|
||||
### Run the Inference
|
||||
```python
|
||||
python x_to_video_audio.py \
|
||||
--model /path/to/dreamid_omni \
|
||||
--prompt "Two people walking together and singing happily" \
|
||||
--image-path ./example0.png ./example1.png \
|
||||
--audio-path ./example0.wav ./example1.wav \
|
||||
--video-negative-prompt "jitter, bad hands, blur, distortion" \
|
||||
--audio-negative-prompt "robotic, muffled, echo, distorted" \
|
||||
--cfg-parallel-size 2 \
|
||||
--num-inference-steps 45 \
|
||||
--height 704 \
|
||||
--width 1280 \
|
||||
--output out_dreamid_omni_twoip.mp4
|
||||
```
|
||||
In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
|
||||
The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
|
||||
For multi-GPU memory reduction on the fused DreamID-Omni transformer, you can also enable HSDP:
|
||||
|
||||
```python
|
||||
python x_to_video_audio.py \
|
||||
--model /path/to/dreamid_omni \
|
||||
--prompt "Two people walking together and singing happily" \
|
||||
--image-path ./example0.png ./example1.png \
|
||||
--audio-path ./example0.wav ./example1.wav \
|
||||
--video-negative-prompt "jitter, bad hands, blur, distortion" \
|
||||
--audio-negative-prompt "robotic, muffled, echo, distorted" \
|
||||
--cfg-parallel-size 2 \
|
||||
--num-inference-steps 45 \
|
||||
--height 704 \
|
||||
--width 1280 \
|
||||
--use-hsdp \
|
||||
--hsdp-shard-size 2 \
|
||||
--output out_dreamid_omni_twoip.mp4
|
||||
```
|
||||
|
||||
|
||||
You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni
|
||||
|
||||
For example, single IP ref resources can be found under https://github.com/Guoxu1233/DreamID-Omni/tree/main/test_case/oneip, you could download them correspondingly to your local and use them for testing.
|
||||
|
||||
```python
|
||||
# Example usage for oneip, ref media from the official repo DreamID-Omni
|
||||
python x_to_video_audio.py \
|
||||
--model /path/to/dreamid_omni \
|
||||
--prompt "<img1>: In the frame, a woman with black long hair is identified as <sub1>.\n**Overall Environment/Scene**: A lively open-kitchen café at night; stove flames flare, steam rises, and warm pendant lights swing slightly as staff move behind her. The shot is an upper-body close-up.\n**Main Characters/Subjects Appearance**: <sub1> is a young woman with thick dark wavy hair and a side part. She wears a fitted black top under a light apron, a thin gold chain necklace, and small stud earrings.\n**Main Characters/Subjects Actions**: <sub1> tastes the sauce with a spoon, then turns her face toward the camera while still holding the spoon, her expression shifting from focused to conflicted.\n<sub1> maintains eye contact, swallows as if choosing her words, and says, <S>I keep telling myself I’m fine,but some nights it feels like I’m just performing calm.<E>" \
|
||||
--image-path 9.png \
|
||||
--audio-path 9.wav \
|
||||
--video-negative-prompt "jitter, bad hands, blur, distortion" \
|
||||
--audio-negative-prompt "robotic, muffled, echo, distorted" \
|
||||
--cfg-parallel-size 2 \
|
||||
--num-inference-steps 45 \
|
||||
--height 704 \
|
||||
--width 1280 \
|
||||
--output out_dreamid_omni_oneip.mp4
|
||||
```
|
||||
|
||||
|
||||
Key arguments:
|
||||
- `--prompt`: text description (string).
|
||||
- `--model`: path to the model local directory.
|
||||
- `--height/--width`: output resolution (defaults 704 * 1024).
|
||||
- `--image-path`: path to the input image list.
|
||||
- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
|
||||
- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
|
||||
- `--use-hsdp`: enable HSDP weight sharding for DreamID-Omni fused blocks.
|
||||
- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
|
||||
- `--hsdp-replicate-size`: number of HSDP replica groups.
|
||||
- `--num-inference-steps`: number of denoising steps (defaults 45).
|
||||
- `--video-negative-prompt`: negative prompt for video generation.
|
||||
- `--audio-negative-prompt`: negative prompt for audio generation.
|
||||
- `--enable-cpu-offload`: enable CPU offload (defaults False).
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "download_dreamid_omni.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/x_to_video_audio/download_dreamid_omni.py"
|
||||
``````
|
||||
??? abstract "x_to_video_audio.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/x_to_video_audio/x_to_video_audio.py"
|
||||
``````
|
||||
@@ -28,18 +28,6 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--deploy-config /path/to/deploy_config_file
|
||||
```
|
||||
|
||||
To serve **thinker-only** (text output, no talker / code2wav loaded) with Instruct
|
||||
weights, pass the bundled thinker-only deploy YAML. The `pipeline:` key
|
||||
`qwen3_omni_moe_thinker_only` overrides the HF `enable_audio_output` resolver:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--deploy-config vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
|
||||
```
|
||||
|
||||
Captioner / Thinking checkpoints (`enable_audio_output=false`) still auto-select
|
||||
the same single-stage pipeline without `--deploy-config`.
|
||||
|
||||
### Launch individual stages (stage-based CLI)
|
||||
|
||||
Adopt the stage-based CLI architecture to independently instantiate execution processes per functional stage.
|
||||
|
||||
@@ -1,50 +0,0 @@
|
||||
# SVDQuant W4A4
|
||||
|
||||
## Overview
|
||||
|
||||
[SVDQuant](https://arxiv.org/abs/2411.05007) combines four-bit weights
|
||||
and activations with a small low-rank branch that corrects part of the
|
||||
quantization error. vLLM-Omni consumes an offline-quantized checkpoint;
|
||||
it does not calibrate the model while loading.
|
||||
|
||||
## Checkpoint contract
|
||||
|
||||
Place the following entry in the diffusion transformer's `config.json`:
|
||||
|
||||
```json
|
||||
{
|
||||
"quantization_config": {
|
||||
"quant_method": "svdquant",
|
||||
"rank": 32,
|
||||
"precision": "nvfp4",
|
||||
"act_unsigned": false,
|
||||
"modules_to_not_convert": []
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
For a quantized linear with input size `K`, output size `N`, and correction
|
||||
rank `R`, the checkpoint stores:
|
||||
|
||||
| Suffix | Shape | dtype |
|
||||
| --- | --- | --- |
|
||||
| `qweight` | `(N, K / 2)` | `int8` (two packed FP4 values per byte) |
|
||||
| `wscales` | `(K / 16, N)` | `float8_e4m3fn` |
|
||||
| `proj_down` | `(K, R)` | `bfloat16` |
|
||||
| `proj_up` | `(N, R)` | `bfloat16` |
|
||||
| `smooth_factor` | `(K,)` | `bfloat16` |
|
||||
| `wcscales` | `(N,)` | `bfloat16` |
|
||||
| `wtscale` | `(1,)` | `bfloat16` |
|
||||
|
||||
`K` must be divisible by 16 on every tensor-parallel rank. Set `wcscales` to
|
||||
ones when no per-output correction is needed. Modules listed in
|
||||
`modules_to_not_convert` keep their checkpoint precision.
|
||||
|
||||
## Runtime support
|
||||
|
||||
The compatibility path accepts BF16 inputs and executes an NVFP4 GEMM followed
|
||||
by the BF16 rank correction. It supports vLLM's FlashInfer, CUTLASS, and FBGEMM
|
||||
NVFP4 tensor layouts; incompatible forced backends fail during model loading.
|
||||
SM103 is the currently validated and enabled hardware target. Native fusion of
|
||||
the NVFP4 GEMM and rank correction is separate from this checkpoint-loading
|
||||
contract.
|
||||
@@ -1,9 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""vLLM-Omni example scripts, importable as a package.
|
||||
|
||||
Keep this ``__init__.py``: with the upstream vLLM installed editable (which also
|
||||
places an ``examples`` package on ``sys.path``), a regular package at this repo
|
||||
root wins over the upstream ``examples`` namespace during ``import examples``.
|
||||
"""
|
||||
@@ -1,117 +0,0 @@
|
||||
"""
|
||||
Offline inference example: MammothModa2 image summarization (single AR stage).
|
||||
|
||||
Example:
|
||||
uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
|
||||
--model path/to/MammothModa2-Preview \
|
||||
--deploy-config vllm_omni/deploy/mammoth_moda2_ar.yaml \
|
||||
--image /path/to/input.jpg \
|
||||
--question "Please summarize the content of this image."
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
|
||||
from PIL import Image
|
||||
from vllm import SamplingParams
|
||||
from vllm.multimodal.image import convert_image_mode
|
||||
|
||||
from vllm_omni import Omni
|
||||
|
||||
DEFAULT_SYSTEM = "You are a helpful assistant."
|
||||
DEFAULT_QUESTION = "Please summarize the content of this image."
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).")
|
||||
parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.")
|
||||
parser.add_argument("--deploy-config", type=str, required=True, help="Path to deploy config yaml.")
|
||||
parser.add_argument("--image", type=str, required=True, help="Path to input image.")
|
||||
parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.")
|
||||
parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.")
|
||||
parser.add_argument(
|
||||
"--max-tokens",
|
||||
type=int,
|
||||
default=512,
|
||||
help="Max new tokens to generate.",
|
||||
)
|
||||
parser.add_argument("--temperature", type=float, default=0.2)
|
||||
parser.add_argument("--top-p", type=float, default=0.9)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--trust-remote-code", action="store_true")
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def build_prompt(system: str, question: str) -> str:
|
||||
return (
|
||||
f"<|im_start|>system\n{system}<|im_end|>\n"
|
||||
"<|im_start|>user\n"
|
||||
"<|vision_start|><|image_pad|><|vision_end|>"
|
||||
f"{question}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
if not os.path.exists(args.image):
|
||||
raise FileNotFoundError(f"Image file not found: {args.image}")
|
||||
|
||||
pil_image = Image.open(args.image)
|
||||
image_data = convert_image_mode(pil_image, "RGB")
|
||||
prompt = build_prompt(args.system, args.question)
|
||||
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
deploy_config=args.deploy_config,
|
||||
trust_remote_code=args.trust_remote_code,
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
)
|
||||
try:
|
||||
sp = SamplingParams(
|
||||
temperature=float(args.temperature),
|
||||
top_p=float(args.top_p),
|
||||
top_k=-1,
|
||||
max_tokens=int(args.max_tokens),
|
||||
seed=int(args.seed),
|
||||
detokenize=True,
|
||||
)
|
||||
# NOTE: omni.generate() returns a Generator[OmniRequestOutput, None, None].
|
||||
# Consume it inside the try block so the worker isn't closed early.
|
||||
outputs = list(
|
||||
omni.generate(
|
||||
[
|
||||
{
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"image": image_data},
|
||||
"additional_information": {"omni_task": ["chat"]},
|
||||
}
|
||||
],
|
||||
[sp],
|
||||
)
|
||||
)
|
||||
finally:
|
||||
omni.close()
|
||||
|
||||
lines: list[str] = []
|
||||
for stage_outputs in outputs:
|
||||
ro = stage_outputs
|
||||
text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
|
||||
lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
|
||||
lines.append("answer:\n")
|
||||
lines.append(text.strip() + "\n")
|
||||
lines.append("\n")
|
||||
|
||||
print("\n".join(lines))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Validate MiniMax-H3 request-mode and DLO lifecycle behavior.
|
||||
|
||||
The DLO modes accept any positive data-parallel size. With AllGather enabled,
|
||||
@@ -85,12 +85,20 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument("--height", type=int, default=768)
|
||||
parser.add_argument("--batch-wait-ms", type=float, default=500.0)
|
||||
parser.add_argument("--init-timeout", type=float, default=1800.0)
|
||||
parser.add_argument(
|
||||
"--dlo-offload-components",
|
||||
type=json.loads,
|
||||
default=None,
|
||||
help="Component policy JSON, e.g. '{\"text_encoder\": false}'",
|
||||
)
|
||||
parser.add_argument("--output", type=Path)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
|
||||
is_dlo = args.mode in {"dlo", "dlo-no-allgather"}
|
||||
if args.dlo_offload_components and not is_dlo:
|
||||
raise ValueError("--dlo-offload-components requires a DLO mode")
|
||||
common: dict[str, Any] = {
|
||||
"model": args.model,
|
||||
"trust_remote_code": True,
|
||||
@@ -114,6 +122,7 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=args.mode == "dlo",
|
||||
dlo_resident_layers=0,
|
||||
dlo_offload_components=({} if args.dlo_offload_components is None else args.dlo_offload_components),
|
||||
)
|
||||
else:
|
||||
common.update(
|
||||
|
||||
@@ -599,17 +599,6 @@ def main():
|
||||
if args.seed is not None and hasattr(params, "seed"):
|
||||
params.seed = args.seed
|
||||
|
||||
# MammothModa2's AR stage emits one visual token per grid cell,
|
||||
# one EOL token per row, and one final look-ahead token whose hidden
|
||||
# state is unavailable. Size the first stage from the prompt metadata
|
||||
# instead of SamplingParams' default of 16 tokens.
|
||||
prompt_info = prompt_dict.get("additional_information", {})
|
||||
if idx == 0 and prompt_info.get("omni_task") == ["t2i"]:
|
||||
ar_width = int(prompt_info.get("ar_width", [0])[0])
|
||||
ar_height = int(prompt_info.get("ar_height", [0])[0])
|
||||
if ar_width > 0 and ar_height > 0:
|
||||
params.max_tokens = ar_height * (ar_width + 1) + 1
|
||||
|
||||
if not diffusion_replaced and len(sampling_params_list) == 1:
|
||||
sampling_params_list = [diffusion_params]
|
||||
|
||||
|
||||
@@ -10,12 +10,10 @@ Generate text from text or image inputs with vLLM-Omni's shared offline entrypoi
|
||||
| --- | --- | --- | --- |
|
||||
| `ByteDance-Seed/BAGEL-7B-MoT` | Yes | Yes | `vllm_omni/deploy/bagel.yaml` |
|
||||
| `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | `vllm_omni/deploy/hunyuan_image3_ar.yaml` |
|
||||
| `bytedance-research/MammothModa2-Preview` | Yes | Yes | `vllm_omni/deploy/mammoth_moda2_ar.yaml` |
|
||||
| `bytedance-research/MammothModa2-Dev` | Yes | Yes | `vllm_omni/deploy/mammoth_moda2_ar.yaml` |
|
||||
|
||||
The script recognizes these model families from `config.json`, applies the
|
||||
model-specific prompt format, and selects an AR-only deploy for HunyuanImage-3
|
||||
and MammothModa2. BAGEL uses its registered default deploy.
|
||||
model-specific prompt format, and selects an AR-only deploy for HunyuanImage-3.
|
||||
BAGEL uses its registered default deploy.
|
||||
|
||||
## Text-To-Text
|
||||
|
||||
@@ -35,22 +33,6 @@ python examples/offline_inference/x_to_text/x_to_text.py \
|
||||
--prompt "Explain multimodal inference in three concise sentences."
|
||||
```
|
||||
|
||||
### MammothModa2
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/x_to_text/x_to_text.py \
|
||||
--model bytedance-research/MammothModa2-Preview \
|
||||
--prompt "Explain multimodal inference in three concise sentences."
|
||||
```
|
||||
|
||||
The same command supports the Dev checkpoint:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/x_to_text/x_to_text.py \
|
||||
--model bytedance-research/MammothModa2-Dev \
|
||||
--prompt "Explain multimodal inference in three concise sentences."
|
||||
```
|
||||
|
||||
## Image-To-Text
|
||||
|
||||
Pass one image with `--image`. The default question can be replaced with any
|
||||
@@ -72,16 +54,6 @@ python examples/offline_inference/x_to_text/x_to_text.py \
|
||||
--prompt "Describe this image in detail."
|
||||
```
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/x_to_text/x_to_text.py \
|
||||
--model bytedance-research/MammothModa2-Preview \
|
||||
--image image.png \
|
||||
--prompt "Describe this image in detail."
|
||||
```
|
||||
|
||||
Use `bytedance-research/MammothModa2-Dev` in the same command to run I2T with
|
||||
the Dev checkpoint.
|
||||
|
||||
## Key Arguments
|
||||
|
||||
| Argument | Default | Description |
|
||||
@@ -102,4 +74,4 @@ the Dev checkpoint.
|
||||
|
||||
The committed HunyuanImage-3 AR-only default uses four GPUs with TP=4. To use a
|
||||
different topology, pass a custom AR-only YAML through `--deploy-config`.
|
||||
BAGEL and MammothModa2 use their repository deploy defaults unless overridden.
|
||||
BAGEL uses its repository deploy default unless overridden.
|
||||
@@ -15,7 +15,6 @@ from vllm_omni.model_extras import build_x_to_text_prompt, get_x_to_text_model_f
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
_HUNYUAN_AR_DEPLOY_CONFIG = _REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3_ar.yaml"
|
||||
_MAMMOTH_AR_DEPLOY_CONFIG = _REPO_ROOT / "vllm_omni" / "deploy" / "mammoth_moda2_ar.yaml"
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
@@ -62,7 +61,7 @@ def main() -> None:
|
||||
omni_kwargs: dict[str, Any] = {
|
||||
"model": args.model,
|
||||
"mode": "image-to-text" if image is not None else "text-to-text",
|
||||
"trust_remote_code": args.trust_remote_code or family in {"hunyuan_image3", "mammoth_moda2"},
|
||||
"trust_remote_code": args.trust_remote_code or family == "hunyuan_image3",
|
||||
"log_stats": args.log_stats,
|
||||
}
|
||||
if args.enforce_eager is not None:
|
||||
@@ -71,8 +70,6 @@ def main() -> None:
|
||||
omni_kwargs["deploy_config"] = args.deploy_config
|
||||
elif family == "hunyuan_image3":
|
||||
omni_kwargs["deploy_config"] = str(_HUNYUAN_AR_DEPLOY_CONFIG)
|
||||
elif family == "mammoth_moda2":
|
||||
omni_kwargs["deploy_config"] = str(_MAMMOTH_AR_DEPLOY_CONFIG)
|
||||
|
||||
omni = Omni(**omni_kwargs)
|
||||
try:
|
||||
|
||||
@@ -0,0 +1,108 @@
|
||||
import argparse
|
||||
import fcntl
|
||||
import os
|
||||
import site
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
DEPENDENCY_REPO = "https://github.com/bytedance/DreamID-V.git"
|
||||
DEPENDENCY_BRANCH = "omni"
|
||||
CACHE_DIR = Path(tempfile.gettempdir()) / "vllm-omni-dependency"
|
||||
LOCK_FILE = CACHE_DIR / ".install.lock"
|
||||
DEPENDENCY_DIR = CACHE_DIR / "DreamID-Omni"
|
||||
|
||||
|
||||
def download_dependency():
|
||||
CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with open(LOCK_FILE, "w") as f:
|
||||
fcntl.flock(f, fcntl.LOCK_EX)
|
||||
if not DEPENDENCY_DIR.exists():
|
||||
print(f"Downloading DreamID-Omni to {DEPENDENCY_DIR} ...")
|
||||
subprocess.run(
|
||||
["git", "clone", "--depth", "1", DEPENDENCY_REPO, "--branch", DEPENDENCY_BRANCH, str(DEPENDENCY_DIR)],
|
||||
check=True,
|
||||
)
|
||||
print("Download finished.")
|
||||
fcntl.flock(f, fcntl.LOCK_UN)
|
||||
|
||||
# write .pth to site-packages
|
||||
site_packages = Path(site.getsitepackages()[0])
|
||||
pth_file = site_packages / "vllm_omni_dependency.pth"
|
||||
pth_file.write_text(str(DEPENDENCY_DIR))
|
||||
print(f"Added {DEPENDENCY_DIR} to site-packages via {pth_file}")
|
||||
|
||||
|
||||
def timed_download(repo_id: str, local_dir: str, allow_patterns: list | None = None):
|
||||
"""Download files from HF repo and log time + destination."""
|
||||
if os.path.exists(local_dir):
|
||||
print(f"Directory {local_dir} already exists. Skipping download.")
|
||||
return
|
||||
print(f"Starting download from {repo_id} into {local_dir}")
|
||||
start_time = time.time()
|
||||
|
||||
snapshot_download(
|
||||
repo_id=repo_id,
|
||||
local_dir=local_dir,
|
||||
local_dir_use_symlinks=False,
|
||||
allow_patterns=allow_patterns,
|
||||
)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
print(f"✅ Finished downloading {repo_id} in {elapsed:.2f} seconds. Files saved at: {local_dir}")
|
||||
|
||||
|
||||
def main(output_dir: str):
|
||||
# Wan2.2
|
||||
wan_dir = os.path.join(output_dir, "Wan2.2-TI2V-5B")
|
||||
timed_download(
|
||||
repo_id="Wan-AI/Wan2.2-TI2V-5B",
|
||||
local_dir=wan_dir,
|
||||
allow_patterns=["google/*", "models_t5_umt5-xxl-enc-bf16.pth", "Wan2.2_VAE.pth"],
|
||||
)
|
||||
|
||||
# MMAudio
|
||||
mm_audio_dir = os.path.join(output_dir, "MMAudio")
|
||||
timed_download(
|
||||
repo_id="hkchengrex/MMAudio",
|
||||
local_dir=mm_audio_dir,
|
||||
allow_patterns=["ext_weights/best_netG.pt", "ext_weights/v1-16.pth"],
|
||||
)
|
||||
|
||||
dreamid_dir = os.path.join(output_dir, "DreamID-Omni")
|
||||
|
||||
timed_download(repo_id="XuGuo699/DreamID-Omni", local_dir=dreamid_dir)
|
||||
|
||||
# Now we construct the config file
|
||||
import json
|
||||
|
||||
data = {
|
||||
"_class_name": "DreamIDOmniPipeline",
|
||||
}
|
||||
|
||||
with open(os.path.join(output_dir, "model_index.json"), "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
|
||||
print(f"model_index.json created at {os.path.join(output_dir, 'model_index.json')}")
|
||||
|
||||
transformer_dir = os.path.join(output_dir, "transformer")
|
||||
os.makedirs(transformer_dir, exist_ok=True)
|
||||
with open(os.path.join(transformer_dir, "config.json"), "w", encoding="utf-8") as f:
|
||||
json.dump({"fusion": "DreamID-Omni/dreamid_omni.safetensors"}, f)
|
||||
print(f"transformer/config.json created at {os.path.join(transformer_dir, 'config.json')}")
|
||||
|
||||
# now we download the dependency code
|
||||
download_dependency()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Download models from Hugging Face")
|
||||
parser.add_argument(
|
||||
"--output-dir", type=str, default="./dreamid_omni", help="Base directory to save downloaded models"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
main(args.output_dir)
|
||||
@@ -0,0 +1,163 @@
|
||||
# X-To-Video-Audio
|
||||
|
||||
The `DreamID-Omni` pipeline generates short videos from text, image and video.
|
||||
|
||||
## Local CLI Usage
|
||||
### Download the Model locally
|
||||
Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally.
|
||||
|
||||
```bash
|
||||
python download_dreamid_omni.py --output-dir ./dreamid_omni
|
||||
```
|
||||
After download, the model directory will look like this:
|
||||
|
||||
```
|
||||
dreamid_omni/
|
||||
├── DreamID-Omni/
|
||||
│ ├── dreamid_omni.safetensors
|
||||
├── MMAudio/
|
||||
│ ├── ext_weights/
|
||||
│ │ ├── best_netG.pt
|
||||
│ │ ├── v1-16.pth
|
||||
├── Wan2.2-TI2V-5B/
|
||||
│ ├── google/*
|
||||
│ ├── models_t5_umt5-xxl-enc-bf16.pth
|
||||
│ ├── Wan2.2_VAE.pth
|
||||
│
|
||||
├── model_index.json
|
||||
└── transformer/
|
||||
└── config.json # create by download_dreamid_omni.py
|
||||
```
|
||||
|
||||
### Run the Inference
|
||||
```python
|
||||
python x_to_video_audio.py \
|
||||
--model /path/to/dreamid_omni \
|
||||
--prompt "Two people walking together and singing happily" \
|
||||
--image-path ./example0.png ./example1.png \
|
||||
--audio-path ./example0.wav ./example1.wav \
|
||||
--video-negative-prompt "jitter, bad hands, blur, distortion" \
|
||||
--audio-negative-prompt "robotic, muffled, echo, distorted" \
|
||||
--cfg-parallel-size 2 \
|
||||
--num-inference-steps 45 \
|
||||
--height 704 \
|
||||
--width 1280 \
|
||||
--output out_dreamid_omni_twoip.mp4
|
||||
```
|
||||
In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
|
||||
The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
|
||||
For multi-GPU memory reduction on the fused DreamID-Omni transformer, you can also enable HSDP:
|
||||
|
||||
```python
|
||||
python x_to_video_audio.py \
|
||||
--model /path/to/dreamid_omni \
|
||||
--prompt "Two people walking together and singing happily" \
|
||||
--image-path ./example0.png ./example1.png \
|
||||
--audio-path ./example0.wav ./example1.wav \
|
||||
--video-negative-prompt "jitter, bad hands, blur, distortion" \
|
||||
--audio-negative-prompt "robotic, muffled, echo, distorted" \
|
||||
--cfg-parallel-size 2 \
|
||||
--num-inference-steps 45 \
|
||||
--height 704 \
|
||||
--width 1280 \
|
||||
--use-hsdp \
|
||||
--hsdp-shard-size 2 \
|
||||
--output out_dreamid_omni_twoip.mp4
|
||||
```
|
||||
|
||||
|
||||
You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni
|
||||
|
||||
For example, single IP ref resources can be found under https://github.com/Guoxu1233/DreamID-Omni/tree/main/test_case/oneip, you could download them correspondingly to your local and use them for testing.
|
||||
|
||||
```python
|
||||
# Example usage for oneip, ref media from the official repo DreamID-Omni
|
||||
python x_to_video_audio.py \
|
||||
--model /path/to/dreamid_omni \
|
||||
--prompt "<img1>: In the frame, a woman with black long hair is identified as <sub1>.\n**Overall Environment/Scene**: A lively open-kitchen café at night; stove flames flare, steam rises, and warm pendant lights swing slightly as staff move behind her. The shot is an upper-body close-up.\n**Main Characters/Subjects Appearance**: <sub1> is a young woman with thick dark wavy hair and a side part. She wears a fitted black top under a light apron, a thin gold chain necklace, and small stud earrings.\n**Main Characters/Subjects Actions**: <sub1> tastes the sauce with a spoon, then turns her face toward the camera while still holding the spoon, her expression shifting from focused to conflicted.\n<sub1> maintains eye contact, swallows as if choosing her words, and says, <S>I keep telling myself I’m fine,but some nights it feels like I’m just performing calm.<E>" \
|
||||
--image-path 9.png \
|
||||
--audio-path 9.wav \
|
||||
--video-negative-prompt "jitter, bad hands, blur, distortion" \
|
||||
--audio-negative-prompt "robotic, muffled, echo, distorted" \
|
||||
--cfg-parallel-size 2 \
|
||||
--num-inference-steps 45 \
|
||||
--height 704 \
|
||||
--width 1280 \
|
||||
--output out_dreamid_omni_oneip.mp4
|
||||
```
|
||||
|
||||
|
||||
## MagiHuman (`--model-type magi-human`)
|
||||
|
||||
MagiHuman is a text → video+audio model with a DiT MoE backbone and a ~9B-param
|
||||
T5Gemma text encoder. A detailed text prompt is the only required input; an optional
|
||||
image and/or audio file may be supplied for conditioning. Natively supports Tensor
|
||||
Parallelism. For an 80GB node, `--tensor-parallel-size 4` is recommended to shard
|
||||
the MoE weights and the text encoder.
|
||||
|
||||
> Install [MagiCompiler](https://github.com/SandAI-org/MagiCompiler) for correct
|
||||
> attention-kernel behaviour (the pipeline otherwise falls back to stubs).
|
||||
|
||||
Text-only generation:
|
||||
|
||||
```bash
|
||||
python x_to_video_audio.py \
|
||||
--model-type magi-human \
|
||||
--model /path/to/daVinci-MagiHuman \
|
||||
--prompt "A young woman with long, wavy golden blonde hair... <dialogue and background sound>" \
|
||||
--tensor-parallel-size 4 \
|
||||
--height 256 --width 448 \
|
||||
--num-inference-steps 8 \
|
||||
--seed 52 \
|
||||
--extra-body '{"seconds": 5, "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \
|
||||
--output output_magihuman.mp4
|
||||
```
|
||||
|
||||
With optional image and audio conditioning:
|
||||
|
||||
```bash
|
||||
python x_to_video_audio.py \
|
||||
--model-type magi-human \
|
||||
--model /path/to/daVinci-MagiHuman \
|
||||
--prompt "A young woman..." \
|
||||
--tensor-parallel-size 4 \
|
||||
--height 256 --width 448 \
|
||||
--num-inference-steps 8 \
|
||||
--seed 52 \
|
||||
--extra-body '{"image_path": "/path/to/ref.jpg", "audio_path": "/path/to/ref.wav", "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \
|
||||
--output output_magihuman.mp4
|
||||
```
|
||||
|
||||
MagiHuman-specific arguments are passed as a JSON dict via `--extra-body` (declared
|
||||
in `vllm_omni/model_extras/magi_human.py`, routed via `extra_args`):
|
||||
|
||||
- `seconds`: output duration in seconds (default 10; ignored when `audio_path` is set,
|
||||
because the audio length then determines the number of frames).
|
||||
- `audio_path`: path to an audio file for audio-to-video conditioning. When provided,
|
||||
the audio drives the frame count and is encoded as a latent condition into the DiT
|
||||
(`is_a2v` mode). Omit for pure text-to-video+audio generation.
|
||||
- `image_path`: path to an image file for visual conditioning. Applied at both the
|
||||
base-resolution (BR) and super-resolution (SR) stages when SR is enabled.
|
||||
- `sr_height` / `sr_width`: super-resolution output resolution. SR stage is skipped
|
||||
when these are omitted.
|
||||
- `sr_num_inference_steps`: denoising steps for the SR stage.
|
||||
|
||||
## DreamID-Omni arguments
|
||||
|
||||
Key arguments:
|
||||
- `--prompt`: text description (string).
|
||||
- `--model`: path to the model local directory.
|
||||
- `--height/--width`: output resolution (defaults 704 * 1024).
|
||||
- `--image-path`: path to the input image list.
|
||||
- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
|
||||
- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
|
||||
- `--use-hsdp`: enable HSDP weight sharding for DreamID-Omni fused blocks.
|
||||
- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
|
||||
- `--hsdp-replicate-size`: number of HSDP replica groups.
|
||||
- `--num-inference-steps`: number of denoising steps (defaults 45).
|
||||
- `--video-negative-prompt`: negative prompt for video generation.
|
||||
- `--audio-negative-prompt`: negative prompt for audio generation.
|
||||
- `--enable-cpu-offload`: enable CPU offload (defaults False).
|
||||
- `--cache-backend`: enable `cache_dit` for acceleration.
|
||||
- `--quantization`: online (dynamic) quantization method — `fp8` or `int8`.
|
||||
(VAEs, the T5 text encoder, norms and modulation stay bf16.)
|
||||
@@ -0,0 +1,279 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
|
||||
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".gif"}
|
||||
AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg"}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Offline inference for X -> video+audio models.")
|
||||
parser.add_argument("--model", required=True, help="Model ckpt root directory.")
|
||||
parser.add_argument(
|
||||
"--model-type",
|
||||
default="dreamid-omni",
|
||||
choices=["dreamid-omni", "magi-human"],
|
||||
help="Model type.",
|
||||
)
|
||||
parser.add_argument("--prompt", default=None, help="Text prompt.")
|
||||
|
||||
parser.add_argument("--image-path", type=str, nargs="+", help="list of image-path")
|
||||
parser.add_argument("--audio-path", type=str, nargs="+", help="list of audio-path")
|
||||
parser.add_argument("--prompt-file", type=str, default=None, help="Text prompt in json format.")
|
||||
parser.add_argument(
|
||||
"--extra-body",
|
||||
type=str,
|
||||
default=None,
|
||||
help="[magi-human] JSON dict of model-specific extra params (declared in vllm_omni/model_extras/), "
|
||||
'merged into sampling extra_args. Example: \'{"image_path": "/path/to/img.jpg", "seconds": 5}\'.',
|
||||
)
|
||||
|
||||
parser.add_argument("--height", type=int, default=704, help="Video height.")
|
||||
parser.add_argument("--width", type=int, default=1280, help="Video width.")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=45, help="Sampling steps.")
|
||||
parser.add_argument("--solver-name", default="unipc", help="Solver name: unipc|dpm++|euler.")
|
||||
parser.add_argument("--shift", type=float, default=5.0, help="Scheduler shift.")
|
||||
parser.add_argument("--seed", type=int, default=103, help="Random seed for reproducible generation.")
|
||||
parser.add_argument(
|
||||
"--cfg-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
choices=[1, 2, 3, 4],
|
||||
help="Number of GPUs used for classifier free guidance parallel size (max 4 branches).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tensor-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help=("Number of GPUs used for tensor parallelism (TP) inside the DiT."),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video-negative-prompt",
|
||||
default="jitter, bad hands, blur, distortion",
|
||||
help="Negative prompt for video.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-negative-prompt",
|
||||
default="robotic, muffled, echo, distorted",
|
||||
help="Negative prompt for audio.",
|
||||
)
|
||||
parser.add_argument("--output", default="dreamid_output.mp4", help="Output video path.")
|
||||
parser.add_argument(
|
||||
"--quantization",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["fp8", "int8"],
|
||||
help=("Online (dynamic) quantization method for the DreamID-Omni transformer. "),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-cpu-offload",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable CPU offloading for diffusion models.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-layerwise-offload",
|
||||
action="store_true",
|
||||
help="Enable layerwise (blockwise) offloading on DiT modules.",
|
||||
)
|
||||
parser.add_argument("--cache-backend", type=str, default=None, choices=["cache_dit"], help="Cache backend.")
|
||||
parser.add_argument(
|
||||
"--use-hsdp",
|
||||
action="store_true",
|
||||
help="Enable HSDP for DreamID-Omni fused transformer blocks.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hsdp-shard-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of GPUs used for HSDP sharding when HSDP is enabled.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hsdp-replicate-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of HSDP replica groups. Default 1 means pure sharding.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def load_image_and_audio(image_paths, audio_paths):
|
||||
image = []
|
||||
audio = []
|
||||
|
||||
for path in image_paths:
|
||||
with Image.open(path) as img:
|
||||
img = img.convert("RGB")
|
||||
image.append(img)
|
||||
|
||||
for path in audio_paths:
|
||||
audio_array, sr = load_audio(path, sr=16000)
|
||||
audio_array = audio_array[int(sr * 1) : int(sr * 3)]
|
||||
audio.append(audio_array)
|
||||
return image, audio
|
||||
|
||||
|
||||
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 = result
|
||||
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() -> None:
|
||||
args = parse_args()
|
||||
if args.prompt is None and args.prompt_file is None:
|
||||
raise ValueError("Either --prompt or --prompt-file must be provided.")
|
||||
|
||||
text_prompt = args.prompt
|
||||
if args.prompt_file:
|
||||
with open(args.prompt_file) as f:
|
||||
text_prompt = json.load(f)
|
||||
text_prompt = re.sub(
|
||||
r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]", "", text_prompt, flags=re.DOTALL
|
||||
).strip()
|
||||
text_prompt = re.sub(
|
||||
r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]", "", text_prompt, flags=re.DOTALL
|
||||
).strip()
|
||||
text_prompt = re.sub(r"\[[A-Z_]+\]", "", text_prompt)
|
||||
text_prompt = re.sub(r"\n\s*\n", "\n", text_prompt).strip()
|
||||
|
||||
if args.model_type == "magi-human":
|
||||
prompt = text_prompt
|
||||
extra_args = json.loads(args.extra_body) if args.extra_body else {}
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
seed=args.seed,
|
||||
extra_args=extra_args,
|
||||
)
|
||||
else:
|
||||
image, audio = load_image_and_audio(args.image_path, args.audio_path)
|
||||
|
||||
prompt = {
|
||||
"prompt": text_prompt,
|
||||
"video_negative_prompt": args.video_negative_prompt,
|
||||
"audio_negative_prompt": args.audio_negative_prompt,
|
||||
"multi_modal_data": {"image": image, "audio": audio},
|
||||
}
|
||||
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
seed=args.seed,
|
||||
extra_args={
|
||||
"solver_name": args.solver_name,
|
||||
"shift": args.shift,
|
||||
},
|
||||
)
|
||||
|
||||
cache_config = None
|
||||
if args.cache_backend == "cache_dit":
|
||||
cache_config = {
|
||||
"Fn_compute_blocks": 1,
|
||||
"Bn_compute_blocks": 0,
|
||||
"max_warmup_steps": 4,
|
||||
"max_cached_steps": 20,
|
||||
"residual_diff_threshold": 0.24,
|
||||
"max_continuous_cached_steps": 3,
|
||||
"enable_taylorseer": False,
|
||||
"taylorseer_order": 1,
|
||||
"scm_steps_mask_policy": None,
|
||||
"scm_steps_policy": "dynamic",
|
||||
}
|
||||
|
||||
omni_kwargs = dict(
|
||||
model=args.model,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
use_hsdp=args.use_hsdp,
|
||||
hsdp_shard_size=args.hsdp_shard_size,
|
||||
hsdp_replicate_size=args.hsdp_replicate_size,
|
||||
model_type=args.model_type,
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
enable_layerwise_offload=args.enable_layerwise_offload,
|
||||
cache_backend=args.cache_backend,
|
||||
cache_config=cache_config,
|
||||
)
|
||||
if args.quantization is not None:
|
||||
omni_kwargs["quantization"] = args.quantization
|
||||
omni = Omni(**omni_kwargs)
|
||||
start = time.perf_counter()
|
||||
outputs = omni.generate(prompt, sampling_params)
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
if not outputs:
|
||||
raise RuntimeError("No output returned from the model.")
|
||||
result = outputs[0]
|
||||
if not result.images:
|
||||
raise RuntimeError("No video frames found in model output.")
|
||||
generated_video = result.images[0]
|
||||
mm = result.multimodal_output or {}
|
||||
generated_audio = mm.get("audio")
|
||||
|
||||
if args.model_type == "magi-human":
|
||||
# MagiHuman returns frames already as (F, H, W, C) uint8, ready to mux.
|
||||
fps = float(mm.get("fps", 25))
|
||||
sample_rate = int(mm.get("audio_sample_rate", 24000))
|
||||
frames = generated_video
|
||||
else:
|
||||
fps = float(mm.get("fps", 24))
|
||||
sample_rate = int(mm.get("audio_sample_rate", 16000))
|
||||
# DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1].
|
||||
# mux_video_audio_bytes expects (F, H, W, C) uint8.
|
||||
if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4:
|
||||
raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}")
|
||||
frames = generated_video.transpose(1, 2, 3, 0)
|
||||
frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8)
|
||||
|
||||
audio_np = None
|
||||
if generated_audio is not None:
|
||||
audio_np = np.squeeze(np.asarray(generated_audio)).astype(np.float32)
|
||||
|
||||
output_path = args.output
|
||||
video_bytes = mux_video_audio_bytes(
|
||||
frames,
|
||||
audio_np,
|
||||
fps=float(fps),
|
||||
audio_sample_rate=sample_rate,
|
||||
)
|
||||
with open(output_path, "wb") as f:
|
||||
f.write(video_bytes)
|
||||
print(f"Saved generated video to {output_path}")
|
||||
print(f"Total time: {elapsed:.2f}s")
|
||||
peak_mb = _extract_peak_memory_mb(outputs)
|
||||
if peak_mb:
|
||||
print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -32,14 +32,12 @@ GPU 0 (90%) and colocates the Talker (55%) and Code2Wav (35%) on GPU 1. That
|
||||
profile admits at most four concurrent sequences per stage.
|
||||
|
||||
| deploy config | GPUs | Notes |
|
||||
| --- | --- | --- |
|
||||
|---|---|---|
|
||||
| `minicpmo_4_5.yaml` (default) | 1 | Memory-constrained compatibility layout. |
|
||||
| `minicpmo_4_5_2gpu.yaml` | 2 | Recommended continuous-batching layout; Talker and Code2Wav share GPU 1. |
|
||||
| `minicpmo_4_5_3gpu.yaml` | 3 | One GPU per stage. |
|
||||
| `minicpmo_4_5_8x4090.yaml` | 8 | Full 8x4090 layout. |
|
||||
|
||||
Every profile sets `session_mode: duplex`, so native full-duplex is served
|
||||
from the same process as `/v1/chat/completions`.
|
||||
| `minicpmo_4_5_duplex.yaml` | 1 | Experimental native full-duplex overlay. |
|
||||
|
||||
The split pipeline preserves native-duplex epoch/turn identity, segment text,
|
||||
turn completion, reference voice, and terminal-audio metadata through
|
||||
@@ -57,8 +55,8 @@ vllm-omni serve openbmb/MiniCPM-o-4_5 \
|
||||
```
|
||||
|
||||
For local ModelScope checkpoints, replace `openbmb/MiniCPM-o-4_5` with the
|
||||
checkpoint path. For native full-duplex, connect `/v1/realtime?duplex=1` on
|
||||
this same server.
|
||||
checkpoint path. To start the experimental native duplex backend, use
|
||||
`vllm_omni/deploy/minicpmo_4_5_duplex.yaml`.
|
||||
|
||||
### Per-stage overrides
|
||||
|
||||
@@ -147,7 +145,7 @@ assistant template, so it is not an apples-to-apples accuracy run.
|
||||
|
||||
## Run the Realtime duplex CLI demo
|
||||
|
||||
After the server is running, stream one WAV through the Realtime
|
||||
After the duplex backend is running, stream one WAV through the Realtime
|
||||
WebSocket endpoint:
|
||||
|
||||
```bash
|
||||
@@ -192,64 +190,6 @@ stronger `response-required` mode is diagnostic: it requires a purpose-built
|
||||
two-response WAV, its `--input-sha256`, and an
|
||||
`--expect-second-response-substring` value.
|
||||
|
||||
## Run Omni-DuplexEval
|
||||
|
||||
Omni-DuplexEval generation uses the vLLM-Omni native MiniCPM-o duplex
|
||||
endpoint. Evaluation uses a separate OpenAI-compatible multimodal judge
|
||||
served from the same vLLM-Omni environment. This validates the vLLM-Omni
|
||||
duplex client and local judge integration; it does not reproduce the paper's
|
||||
original MiniCPM-o inference implementation.
|
||||
|
||||
Start the duplex generation server:
|
||||
|
||||
```bash
|
||||
vllm serve openbmb/MiniCPM-o-4_5 --omni --trust-remote-code \
|
||||
--deploy-config vllm_omni/deploy/minicpmo_4_5_duplex.yaml \
|
||||
--served-model-name openbmb/MiniCPM-o-4_5 \
|
||||
--host 0.0.0.0 --port 8099
|
||||
```
|
||||
|
||||
Start a separate multimodal judge. The allowed path must contain any local
|
||||
videos passed with `--judge-video-mode video_url`:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen2.5-VL-7B-Instruct \
|
||||
--served-model-name Qwen/Qwen2.5-VL-7B-Instruct \
|
||||
--allowed-local-media-path /data/omni-duplex-eval \
|
||||
--host 0.0.0.0 --port 8000
|
||||
```
|
||||
|
||||
Generate, evaluate, and summarize:
|
||||
|
||||
```bash
|
||||
vllm bench omni-duplex-eval --omni generate \
|
||||
--url ws://127.0.0.1:8099/v1/realtime?duplex=1 \
|
||||
--model openbmb/MiniCPM-o-4_5 \
|
||||
--ref-audio /data/ref.wav \
|
||||
--dataset Hothan/Omni-DuplexEval \
|
||||
--concurrency 2 \
|
||||
--response-root /data/omni-duplex-eval/responses
|
||||
|
||||
vllm bench omni-duplex-eval --omni evaluate \
|
||||
--dataset Hothan/Omni-DuplexEval \
|
||||
--response-root /data/omni-duplex-eval/responses \
|
||||
--score-root /data/omni-duplex-eval/scores \
|
||||
--judge-base-url http://127.0.0.1:8000 \
|
||||
--judge-model Qwen/Qwen2.5-VL-7B-Instruct \
|
||||
--judge-video-mode video_url \
|
||||
--eval-workers 4
|
||||
|
||||
vllm bench omni-duplex-eval --omni summarize \
|
||||
--score-root /data/omni-duplex-eval/scores
|
||||
```
|
||||
|
||||
The dataset names such as `RTD_OCR` and `PR_correction` are Hugging Face
|
||||
splits, not dataset configurations. Pass one with `--split`, or omit it to
|
||||
run all nine splits. `--limit 1` is useful for a smoke test. Realtime
|
||||
generation records `clock=media`; artifacts generated with
|
||||
`--pace as-fast-as-possible` record `clock=invalid` and evaluation rejects
|
||||
them unless `--allow-invalid-clock` is explicit.
|
||||
|
||||
## Related examples
|
||||
|
||||
- [Offline MiniCPM-o inference](../../offline_inference/minicpmo/)
|
||||
|
||||
@@ -24,18 +24,6 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--deploy-config /path/to/your_deploy_config.yaml
|
||||
```
|
||||
|
||||
To serve **thinker-only** (text output, no talker / code2wav loaded) with Instruct
|
||||
weights, pass the bundled thinker-only deploy YAML. The `pipeline:` key
|
||||
`qwen3_omni_moe_thinker_only` overrides the HF `enable_audio_output` resolver:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--deploy-config vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
|
||||
```
|
||||
|
||||
Captioner / Thinking checkpoints (`enable_audio_output=false`) still auto-select
|
||||
the same single-stage pipeline without `--deploy-config`.
|
||||
|
||||
For a 3x-GPU multi-replica layout (talker/code2wav scale-out on cuda:1,2),
|
||||
use `--stage-overrides` on top of the default config:
|
||||
|
||||
|
||||
@@ -1,212 +0,0 @@
|
||||
# MammothModa2
|
||||
|
||||
> MammothModa2-Preview and MammothModa2-Dev unified understanding and generation
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: ByteDance Research
|
||||
- Models: `bytedance-research/MammothModa2-Preview`, `bytedance-research/MammothModa2-Dev`
|
||||
- Tasks: Preview and Dev text-to-image (AR → DiT); Dev text/image understanding
|
||||
- Mode: Offline inference
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe to run MammothModa2-Preview through shared task-oriented
|
||||
offline examples. Text-to-image uses the shared image example (`text_to_image.py`) instead of a model-specific script.
|
||||
The generic example formats the AR prompt, drives the AR → DiT stage pipeline,
|
||||
and forwards MammothModa2-specific generation parameters through the
|
||||
pipeline-declared `extra_body` contract.
|
||||
|
||||
MammothModa2's DiT stage consumes its inputs through the multi-stage kwargs
|
||||
interface (not `OmniDiffusionRequest`), so its generation knobs
|
||||
(`text_guidance_scale`, `cfg_range`, `num_inference_steps`) are passed via
|
||||
`--extra-body` rather than the standard `--num-inference-steps` / `--cfg-scale`
|
||||
flags. Image size uses the standard `--height` / `--width` flags.
|
||||
|
||||
## References
|
||||
|
||||
- Upstream model:
|
||||
[`bytedance-research/MammothModa2-Preview`](https://huggingface.co/bytedance-research/MammothModa2-Preview)
|
||||
- Dev model:
|
||||
[`bytedance-research/MammothModa2-Dev`](https://huggingface.co/bytedance-research/MammothModa2-Dev)
|
||||
- Related offline example:
|
||||
[`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py)
|
||||
- Related T2T/I2T example:
|
||||
[`examples/offline_inference/x_to_text/x_to_text.py`](../../examples/offline_inference/x_to_text/x_to_text.py)
|
||||
- Declared parameters:
|
||||
[`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)
|
||||
- Deploy config:
|
||||
[`vllm_omni/deploy/mammoth_moda2.yaml`](../../vllm_omni/deploy/mammoth_moda2.yaml)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
The default deploy config runs both the AR and DiT stages on a single GPU
|
||||
(`devices: "0"`). The committed `gpu_memory_utilization` split (stage-0 AR `0.5`,
|
||||
stage-1 DiT `0.3`) is sized for an ~80 GB GPU. The model also fits on a 48 GB GPU
|
||||
after rebalancing the split so the AR weights (~23 GB) leave room for the KV
|
||||
cache — see the note under *1x L40S 48GB*.
|
||||
|
||||
## GPU
|
||||
|
||||
### 1x L40S 48GB
|
||||
|
||||
> **48 GB config adjustment:** the committed
|
||||
> `vllm_omni/deploy/mammoth_moda2.yaml` uses
|
||||
> `gpu_memory_utilization` 0.5 / 0.3 (sized for ~80 GB). To fit on a 48 GB L40S,
|
||||
> set the stage-0 (AR) value to `0.8` and the stage-1 (DiT) value to `0.16`
|
||||
> before running. (On an ~80 GB GPU, leave the defaults unchanged.)
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux
|
||||
- Python: Match the repository requirements for your checkout
|
||||
- Driver / runtime: NVIDIA CUDA environment with one L40S 48 GB (verified) or an
|
||||
~80 GB GPU for the default config
|
||||
- vLLM version: Match the repository requirements for your checkout
|
||||
- vLLM-Omni version or commit: Use the commit you are deploying from
|
||||
|
||||
#### Offline Commands
|
||||
|
||||
Download the model:
|
||||
|
||||
```bash
|
||||
hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview
|
||||
```
|
||||
|
||||
Run text-to-image with the shared offline example from the repository root. The
|
||||
deploy config sets `trust_remote_code`, so no extra flag is needed. Forward the
|
||||
MammothModa2 generation parameters as a JSON object through `--extra-body`:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model ./MammothModa2-Preview \
|
||||
--deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
|
||||
--prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
|
||||
--output mammoth_t2i.png
|
||||
```
|
||||
|
||||
The `--extra-body` JSON forwards MammothModa2-specific parameters into
|
||||
`OmniDiffusionSamplingParams.extra_args`. Keys are filtered against the model's
|
||||
declared `extra_body_params` (see
|
||||
[`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)),
|
||||
so unknown keys for MammothModa2 are silently dropped:
|
||||
|
||||
- `text_guidance_scale` — classifier-free guidance scale for the DiT stage
|
||||
(default `9.0`; CFG is active only when `> 1.0`).
|
||||
- `cfg_range` — relative step range `[start, end]` over which CFG is applied
|
||||
(default `[0.0, 1.0]`).
|
||||
- `num_inference_steps` — number of DiT denoising steps (default `50`).
|
||||
|
||||
`--height` and `--width` must be multiples of 16.
|
||||
|
||||
Run text-to-text through the shared understanding example. It recognizes the
|
||||
MammothModa2 checkpoint and automatically selects `mammoth_moda2_ar.yaml`:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/x_to_text/x_to_text.py \
|
||||
--model ./MammothModa2-Preview \
|
||||
--prompt "Explain multimodal generation in three sentences."
|
||||
```
|
||||
|
||||
Add an image for image-to-text or image summarization. The shared example
|
||||
uses MammothModa2's chat and vision-token template:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/x_to_text/x_to_text.py \
|
||||
--model ./MammothModa2-Preview \
|
||||
--image /path/to/input.jpg \
|
||||
--prompt "Please summarize the content of this image."
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
The example writes the generated image to the `--output` path. Confirm the file
|
||||
exists and is a valid image:
|
||||
|
||||
```bash
|
||||
ls -lh mammoth_t2i.png
|
||||
python -c "from PIL import Image; print(Image.open('mammoth_t2i.png').size)"
|
||||
```
|
||||
|
||||
### 1x AMD MI300X, MammothModa2 Preview
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux 6.8.0-134-generic, x86_64
|
||||
- Container: official ROCm image built from `docker/Dockerfile.rocm`
|
||||
- Python: 3.12.13
|
||||
- PyTorch: 2.11.0+gitd0c8b1f
|
||||
- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
|
||||
- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
|
||||
- vLLM version: 0.27.0+rocm723
|
||||
- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
|
||||
- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
|
||||
|
||||
#### Offline Commands
|
||||
|
||||
The checked run used the committed stage split, with `gpu_memory_utilization` set to 0.5 for AR and 0.3 for DiT:
|
||||
|
||||
```bash
|
||||
python3 examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model bytedance-research/MammothModa2-Preview \
|
||||
--deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
|
||||
--prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--seed 42 \
|
||||
--extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
|
||||
--enable-diffusion-pipeline-profiler \
|
||||
--log-stats \
|
||||
--output mammoth_t2i.png
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
The first request took 85.224 seconds. The AR stage generated 4,161 visual tokens in 72.996 seconds, and the DiT stage took 12.163 seconds. AR weight loading used 21.4 GiB and took 8.250 seconds. DiT weight loading used 5.49 GiB and took 1.824 seconds. The largest one second whole device memory sample was 106.57 GiB, including the AR KV cache reserved by the 0.5 memory setting.
|
||||
|
||||
The output was a valid 1024 by 1024 RGB PNG.
|
||||
|
||||
## MammothModa2-Dev unified inference
|
||||
|
||||
MammothModa2-Dev uses a Qwen3-VL AR backbone, while MammothModa2-Preview uses
|
||||
Qwen2.5-VL. vLLM-Omni selects the matching implementation from the nested
|
||||
`llm_config.model_type`; no checkpoint edits or `trust_remote_code` flag are
|
||||
required.
|
||||
|
||||
Text-to-text and image-to-text use the AR-only deploy. Text-to-image loads the
|
||||
Qwen3 generation experts (`gen_mlp`), extra visual vocabulary and image head,
|
||||
then sends the generated visual tokens and hidden states to the DiT stage.
|
||||
|
||||
Download the checkpoint:
|
||||
|
||||
```bash
|
||||
hf download bytedance-research/MammothModa2-Dev --local-dir ./MammothModa2-Dev
|
||||
```
|
||||
|
||||
Run text-to-text through the shared understanding example. It recognizes the
|
||||
Dev checkpoint as MammothModa2 and automatically selects
|
||||
`mammoth_moda2_ar.yaml`:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/x_to_text/x_to_text.py \
|
||||
--model ./MammothModa2-Dev \
|
||||
--prompt "Explain multimodal generation in three sentences."
|
||||
```
|
||||
|
||||
Add an image for image-to-text or image summarization:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/x_to_text/x_to_text.py \
|
||||
--model ./MammothModa2-Dev \
|
||||
--image ./image.png \
|
||||
--prompt "Please summarize the content of this image."
|
||||
```
|
||||
|
||||
The Dev checkpoint is approximately 47.55 GiB on disk. In the verified AR-only
|
||||
run, loaded model weights used approximately 16.97 GiB of GPU memory before KV
|
||||
and encoder caches. Allow additional GPU memory for those caches and the input
|
||||
image.
|
||||
@@ -839,7 +839,7 @@ hf download lightx2v/Minimax-h3-Turbo "${TURBO_FILE}" --local-dir "${TURBO_DIR}"
|
||||
export TURBO_LORA="${TURBO_DIR}/${TURBO_FILE}"
|
||||
```
|
||||
|
||||
Start from a non-offloaded or DLO FL2VA server command and add
|
||||
Start from a non-offloaded FL2VA server command and add
|
||||
`--task-type fl2va --lora-backend peft --lora-path "${TURBO_LORA}"`.
|
||||
`--lora-path` preloads the adapter; each request still activates it and uses
|
||||
the published sampling settings:
|
||||
@@ -853,10 +853,9 @@ the published sampling settings:
|
||||
|
||||
For FL2VA, change `task` and add `input_reference` as shown above. The 8-step,
|
||||
ComfyUI, Ref2VA, and v1.1 artifacts are not supported. This integration is
|
||||
dynamic-only and does not support prefusion or LoRA composition. DLO is
|
||||
supported by keeping the request-switchable LoRA A/B buffers resident on the
|
||||
accelerator while DLO streams only the base blocks; budget for this additional
|
||||
fixed HBM usage. Model-level and standard layerwise offload remain unsupported.
|
||||
dynamic-only and does not support prefusion, DLO, or LoRA composition.
|
||||
It also rejects model-level, layerwise, and distributed layerwise offload;
|
||||
the legacy dynamic LoRA tensors do not participate in those weight lifecycles.
|
||||
The five requested sigma points produce the four denoiser evaluations expected
|
||||
by the Turbo artifact.
|
||||
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# NemotronLabs VoiceChat 11B
|
||||
|
||||
> Speech-to-speech on a 3-stage vLLM-Omni pipeline: offline single-turn
|
||||
> inference and experimental model-native, frame-locked Realtime serving.
|
||||
> Offline speech-to-speech: a 16 kHz user utterance in, the agent's spoken reply
|
||||
> (text + 22.05 kHz WAV) out, on a 3-stage vLLM-Omni pipeline.
|
||||
|
||||
## Summary
|
||||
|
||||
@@ -11,8 +11,8 @@
|
||||
timeline: a Conformer + NemotronH hybrid-Mamba thinker emits one text token
|
||||
per acoustic frame, a Gemma3-1B EAR-TTS talker turns the text timeline into
|
||||
31-quantizer RVQ code stacks, and an RVQ-VAE codec decodes them to audio.
|
||||
- Mode: Offline single-turn inference (batch=1), plus experimental native
|
||||
duplex Realtime serving at one 80 ms microphone frame per scheduler wake.
|
||||
- Mode: Offline single-turn inference (batch=1). Online/duplex serving,
|
||||
function calling, and batch>1 are documented follow-ups.
|
||||
- Maintainer: [`@yuekaizhang`](https://github.com/yuekaizhang)
|
||||
|
||||
## When to use this recipe
|
||||
@@ -33,11 +33,6 @@ NeMo modules (`nemo_vendored/`), so no `nemo_toolkit` install is needed.
|
||||
[`vllm_omni/model_executor/models/nemotron_voicechat/`](../../vllm_omni/model_executor/models/nemotron_voicechat/)
|
||||
- Staged pipeline config:
|
||||
[`vllm_omni/deploy/nemotron_labs_voicechat.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat.yaml)
|
||||
- Native duplex config and E2E probe:
|
||||
[`vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml),
|
||||
[`tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py`](../../tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py)
|
||||
- Nightly model-level gate:
|
||||
[`tests/e2e/online_serving/test_nemotron_voicechat_duplex.py`](../../tests/e2e/online_serving/test_nemotron_voicechat_duplex.py)
|
||||
- Upstream model card:
|
||||
[`nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B)
|
||||
- Reference implementation: NVIDIA-NeMo/Speech, branch `nemotron-labs-voicechat`
|
||||
@@ -65,7 +60,7 @@ output stayed within one word of the reference in testing.
|
||||
|
||||
- OS: Linux
|
||||
- Python: 3.12
|
||||
- vLLM version: 0.27.0
|
||||
- vLLM version: 0.26.0
|
||||
- vLLM-Omni version or commit: this PR / current `main`
|
||||
|
||||
#### Command
|
||||
@@ -93,55 +88,6 @@ The reply text should read as a coherent spoken-style answer to the question in
|
||||
the input WAV, and the WAV should transcribe to (approximately) the same text
|
||||
with any ASR model.
|
||||
|
||||
#### Native duplex serving
|
||||
|
||||
The checkpoint does not contain the underlying Nemotron text tokenizer. Set
|
||||
`NEMOTRON_VOICECHAT_LLM_PATH` to a local snapshot of
|
||||
`nvidia/NVIDIA-Nemotron-Nano-9B-v2` before starting an offline deployment.
|
||||
|
||||
```bash
|
||||
export HF_HUB_OFFLINE=1
|
||||
export TRANSFORMERS_OFFLINE=1
|
||||
export NEMOTRON_VOICECHAT_LLM_PATH=/path/to/NVIDIA-Nemotron-Nano-9B-v2
|
||||
|
||||
vllm-omni serve /path/to/NVIDIA-NemotronLabs-VoiceChat-11B \
|
||||
--omni \
|
||||
--served-model-name nemotron-voicechat \
|
||||
--deploy-config vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml
|
||||
```
|
||||
|
||||
The single-GPU profile runs three eager fp32 engine processes on one device.
|
||||
|
||||
In another shell, stream only the user channel of one of NVIDIA's bundled
|
||||
stereo conversation fixtures:
|
||||
|
||||
```bash
|
||||
python tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py \
|
||||
--model nemotron-voicechat \
|
||||
--input-wav /path/to/NVIDIA-NemotronLabs-VoiceChat-11B/turn_taking.wav \
|
||||
--input-channel 0 \
|
||||
--output-dir results/nemotron_voicechat_duplex \
|
||||
--timeout-s 600
|
||||
```
|
||||
|
||||
The probe requires a completed response, non-silent 22.05 kHz output, and the
|
||||
advertised native 80 ms append capabilities. Add `--expect-function-call` and
|
||||
use `tool_call.wav` to validate the function-call channel. To validate the
|
||||
Realtime tool round trip, also provide the expected arguments and return a
|
||||
tool result:
|
||||
|
||||
```bash
|
||||
python tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py \
|
||||
--model nemotron-voicechat \
|
||||
--input-wav /path/to/NVIDIA-NemotronLabs-VoiceChat-11B/tool_call.wav \
|
||||
--output-dir results/nemotron_voicechat_tool_call \
|
||||
--expect-function-call \
|
||||
--expected-function-name generate_random_number \
|
||||
--expected-function-arguments '{"min":1,"max":50}' \
|
||||
--function-output 20 \
|
||||
--expected-post-tool-text "random number"
|
||||
```
|
||||
|
||||
#### Notes
|
||||
|
||||
- Memory usage: the shipped yaml runs all three stages on one GPU
|
||||
@@ -163,14 +109,5 @@ python tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py \
|
||||
instead of pausing.
|
||||
- The talker's `max_tokens` is 16383 (its stage prompt is one placeholder
|
||||
token, and the stage context is 16384).
|
||||
- The 80 ms frame period is a model/protocol contract, not a throughput
|
||||
guarantee. In an eager fp32 H200 profile, Stage 2 emitted 80 ms packets at
|
||||
roughly 0.87--0.88 s intervals (RTF about 10.8--10.9); placing the three
|
||||
stages on separate GPUs did not remove the serial 8-step EAR-TTS bottleneck.
|
||||
The current native duplex path is functionally streaming but not wall-clock
|
||||
realtime.
|
||||
- Known limitations: batch=1 only. The native duplex deployment allows one
|
||||
active session and does not support barge-in. Tool execution remains
|
||||
client-owned: the server emits function-call events, accepts a validated
|
||||
`function_call_output`, and resumes the live model with the returned result;
|
||||
it does not execute arbitrary tools itself.
|
||||
- Known limitations: batch=1 only; offline single turn (no online/duplex
|
||||
serving yet); the function-call channel is decoded but not yet acted on.
|
||||
+1
-2
@@ -85,9 +85,8 @@ recipes/
|
||||
| [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
|
||||
| [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
|
||||
| [`Wan-AI/Wan2.1-VACE.md`](./Wan-AI/Wan2.1-VACE.md) | Unified T2V, I2V, V2LF, FLF2V, inpaint, and R2V | 1x RTX 5090 (1.3B) / 1x L40S 48GB with layerwise offload (14B) |
|
||||
| [`XiaomiMiMo/MiMo-Audio.md`](./XiaomiMiMo/MiMo-Audio.md) | Offline + online omni audio (TTS, ASR, dialogue) | 1x RTX 5090 / 5090D 32GB |
|
||||
| [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB |
|
||||
| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB / 2x non-standard RTX 4090 48GB (49,140 MiB each) |
|
||||
| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
|
||||
| [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
|
||||
| [`JD/JoyAI-VL-Interaction.md`](./JD/JoyAI-VL-Interaction.md) | Real-time streaming video-language interaction (proactive speak/silence/delegate) | 1x GPU 24GB+ |
|
||||
|
||||
|
||||
@@ -1,161 +0,0 @@
|
||||
# MiMo-Audio for omni audio generation and understanding
|
||||
|
||||
> Offline + online TTS, ASR, audio understanding, and multi-turn dialogue
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Xiaomi MiMo
|
||||
- Model: `XiaomiMiMo/MiMo-Audio-7B-Instruct`
|
||||
- Task: Text-to-speech, speech-to-text, audio understanding, spoken dialogue
|
||||
- Mode: Offline end-to-end scripts and online OpenAI-compatible chat serving
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe when you want a known-good starting point for running
|
||||
MiMo-Audio on a single **32 GB Blackwell consumer GPU** (RTX 5090 / 5090D).
|
||||
The bundled deploy profile [`vllm_omni/deploy/mimo_audio_5090d.yaml`](../../vllm_omni/deploy/mimo_audio_5090d.yaml)
|
||||
splits VRAM between the AR thinker (stage 0) and code2wav (stage 1), avoids
|
||||
FlashAttention-2 on SM 12.x, and keeps the audio tokenizer on CPU to reduce
|
||||
Stage-1 GPU pressure.
|
||||
|
||||
## References
|
||||
|
||||
- Upstream model card: <https://huggingface.co/XiaomiMiMo/MiMo-Audio-7B-Instruct>
|
||||
- Audio tokenizer: <https://huggingface.co/XiaomiMiMo/MiMo-Audio-Tokenizer>
|
||||
- Offline example:
|
||||
[`examples/offline_inference/mimo_audio/end2end.py`](../../examples/offline_inference/mimo_audio/end2end.py)
|
||||
- Online example:
|
||||
[`examples/online_serving/mimo_audio/README.md`](../../examples/online_serving/mimo_audio/README.md)
|
||||
- Default deploy config (generic single GPU):
|
||||
[`vllm_omni/deploy/mimo_audio.yaml`](../../vllm_omni/deploy/mimo_audio.yaml)
|
||||
- 5090D-tuned deploy config:
|
||||
[`vllm_omni/deploy/mimo_audio_5090d.yaml`](../../vllm_omni/deploy/mimo_audio_5090d.yaml)
|
||||
- Pipeline / deploy docs:
|
||||
[Stage configs](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/)
|
||||
- Related issue or discussion:
|
||||
[RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645)
|
||||
- Tokenizer device alignment fix (required for CPU tokenizer on stage 1):
|
||||
[PR #6539](https://github.com/vllm-project/vllm-omni/pull/6539)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
## GPU
|
||||
|
||||
### 1x NVIDIA GeForce RTX 5090 / 5090D 32GB
|
||||
|
||||
Two-stage async-chunk pipeline (thinker + code2wav) on a single GPU. Tested
|
||||
offline with `tts_sft`.
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Ubuntu (Linux x86_64)
|
||||
- Python: 3.12
|
||||
- PyTorch: 2.11.0+cu129
|
||||
- Driver / runtime: NVIDIA driver 570.x, CUDA 12.9 runtime (Blackwell SM 12.x)
|
||||
- GPU: NVIDIA GeForce RTX 5090 D, 32 GB
|
||||
- vLLM: 0.26.0+cu129
|
||||
- vLLM-Omni: `main` (tested while [#6539](https://github.com/vllm-project/vllm-omni/pull/6539) was open for CPU-tokenizer support)
|
||||
|
||||
Install vLLM-Omni from source and align the vLLM wheel with your CUDA stack.
|
||||
On Blackwell, prefer the **cu129** vLLM build over cu13 defaults when using
|
||||
PyTorch 2.11+cu129.
|
||||
|
||||
Download or cache:
|
||||
|
||||
- `XiaomiMiMo/MiMo-Audio-7B-Instruct`
|
||||
- `XiaomiMiMo/MiMo-Audio-Tokenizer` (or a local checkout of the tokenizer repo)
|
||||
|
||||
Export the runtime knobs below before offline or online runs:
|
||||
|
||||
```bash
|
||||
export MIMO_AUDIO_TOKENIZER_PATH="${MIMO_AUDIO_TOKENIZER_PATH:-XiaomiMiMo/MiMo-Audio-Tokenizer}"
|
||||
export MIMO_AUDIO_TOKENIZER_DEVICE=cpu
|
||||
export MIMO_AUDIO_TOKENIZER_CUDA_GRAPH=0
|
||||
export VLLM_USE_FLASHINFER_SAMPLER=0
|
||||
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
|
||||
```
|
||||
|
||||
Optional (mirror / download stability):
|
||||
|
||||
```bash
|
||||
export HF_HUB_DISABLE_XET=1
|
||||
export HF_ENDPOINT=https://hf-mirror.com # if you use a Hugging Face mirror
|
||||
```
|
||||
|
||||
Do **not** set `VLLM_ATTENTION_BACKEND` on vLLM 0.26.x; configure
|
||||
`attention_backend: TRITON_ATTN` in the deploy YAML instead.
|
||||
|
||||
#### Command
|
||||
|
||||
Offline TTS validation from the repository root:
|
||||
|
||||
```bash
|
||||
cd examples/offline_inference/mimo_audio
|
||||
|
||||
python3 -u end2end.py \
|
||||
--deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft \
|
||||
--text "The weather is so nice today."
|
||||
```
|
||||
|
||||
Online serving (requires MiMo chat template):
|
||||
|
||||
```bash
|
||||
export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
|
||||
|
||||
vllm serve XiaomiMiMo/MiMo-Audio-7B-Instruct --omni \
|
||||
--deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml \
|
||||
--served-model-name "MiMo-Audio-7B-Instruct" \
|
||||
--port 18091 \
|
||||
--chat-template ./examples/online_serving/mimo_audio/chat_template.jinja
|
||||
```
|
||||
|
||||
Pass the same environment exports to the server process when using a CPU
|
||||
tokenizer or disabling FlashInfer sampling.
|
||||
|
||||
#### Verification
|
||||
|
||||
Offline:
|
||||
|
||||
```bash
|
||||
ls -lh output_audio/tts_sft/*.wav
|
||||
```
|
||||
|
||||
Expect a non-empty `.wav` under `output_audio/tts_sft/` and a companion `.txt`
|
||||
with the decoded text. A successful run completes both stages without a
|
||||
Stage-1 `decode_vq` device mismatch.
|
||||
|
||||
Online:
|
||||
|
||||
```bash
|
||||
curl -s -o /tmp/mimo_health.txt -w '%{http_code}' http://127.0.0.1:18091/health
|
||||
```
|
||||
|
||||
Then use the client under
|
||||
[`examples/online_serving/mimo_audio/`](../../examples/online_serving/mimo_audio/).
|
||||
|
||||
#### Notes
|
||||
|
||||
- Memory usage (observed offline): Stage-0 weights ~16.6 GiB; KV cache ~4.9 GiB
|
||||
with `gpu_memory_utilization: 0.78` on stage 0 and `0.12` on stage 1.
|
||||
- Key flags:
|
||||
- `--deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml` for 5090D memory
|
||||
and `TRITON_ATTN`.
|
||||
- `MIMO_AUDIO_TOKENIZER_PATH` is mandatory.
|
||||
- `MIMO_AUDIO_TOKENIZER_DEVICE=cpu` saves VRAM; requires
|
||||
[#6539](https://github.com/vllm-project/vllm-omni/pull/6539) (or equivalent
|
||||
fix) so code2wav decode uses the tokenizer device.
|
||||
- `VLLM_USE_FLASHINFER_SAMPLER=0` avoids FlashInfer top-p/top-k sampler
|
||||
issues on consumer GPUs.
|
||||
- Known limitations:
|
||||
- Blackwell may log `SM 12.x requires CUDA >= 12.9` from some optional
|
||||
tooling; the tested path uses TRITON attention and cu129 wheels.
|
||||
- Without a working `flash-attn` build, audio quality may be slightly
|
||||
metallic; install FlashAttention when a compatible wheel is available.
|
||||
- First request triggers Triton JIT for rotary / attention kernels (latency
|
||||
spike); subsequent requests are faster.
|
||||
- For longer outputs or batching, you may need to raise `max_model_len` in
|
||||
the deploy YAML and match `max_position_embeddings` in the model config
|
||||
(see the offline example README).
|
||||
+1
-135
@@ -1,4 +1,4 @@
|
||||
# GLM-TTS for Chinese/English TTS
|
||||
# GLM-TTS for Chinese/English TTS on 1x GPU
|
||||
|
||||
## Summary
|
||||
|
||||
@@ -103,137 +103,3 @@ python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
|
||||
- Key flags: `--omni` is required; `--trust-remote-code` is needed for the GLM-TTS phoneme tokenizer; the DiT stage enables bucketed model-internal CUDA graphs with eager fallback.
|
||||
- Voice cloning: requires `ref_audio` + `ref_text` together. Reference audio should be 3-10 seconds. Feature extraction (WhisperVQ tokenizer, CampPlus ONNX, mel) runs on the model side.
|
||||
- Known limitations: First request may be slow due to lazy model loading (WhisperVQ, CampPlus ONNX). Warm-cache RTF is approximately 0.6-0.7x on A40.
|
||||
|
||||
### 2x non-standard RTX 4090 48GB
|
||||
|
||||
This configuration was personally validated on two non-standard RTX 4090 cards
|
||||
reporting 49,140 MiB each in `nvidia-smi`. Stock 24GB RTX 4090 cards were not
|
||||
validated.
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Ubuntu 22.04.1 LTS
|
||||
- NVIDIA driver: 580.76.05
|
||||
- Python: 3.12.13
|
||||
- PyTorch: 2.11.0+cu130
|
||||
- CUDA runtime and JIT toolkit: 13.0
|
||||
- vLLM: 0.26.0
|
||||
- vLLM-Omni: commit `66a9b0c8`, version
|
||||
`0.26.0rc2.dev5+g66a9b0c84`
|
||||
- FlashInfer: 0.6.14
|
||||
|
||||
The host's default `nvcc` was 11.8. The working CUDA toolkit came from the
|
||||
Python environment, so `CUDA_HOME` was set to
|
||||
`$VIRTUAL_ENV/lib/python3.12/site-packages/nvidia/cu13`.
|
||||
|
||||
#### Command
|
||||
|
||||
From the repository root, start the async-chunk server with:
|
||||
|
||||
```bash
|
||||
export CUDA_HOME="$VIRTUAL_ENV/lib/python3.12/site-packages/nvidia/cu13"
|
||||
export PATH="$CUDA_HOME/bin:$PATH"
|
||||
export LD_LIBRARY_PATH="$CUDA_HOME/lib:${LD_LIBRARY_PATH:-}"
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0,1 vllm-omni serve zai-org/GLM-TTS \
|
||||
--omni \
|
||||
--trust-remote-code \
|
||||
--port 8091 \
|
||||
--deploy-config vllm_omni/deploy/glm_tts.yaml \
|
||||
--stage-overrides '{"1":{"devices":"1"}}' \
|
||||
--allowed-local-media-path "${PWD}/tests/assets"
|
||||
```
|
||||
|
||||
Stage 0 runs on GPU 0, while the override places stage 1 on GPU 1. For the sync
|
||||
variant, add `--no-async-chunk` to the same command.
|
||||
|
||||
#### Verification
|
||||
|
||||
Run the Chinese voice-cloning check:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
|
||||
--model zai-org/GLM-TTS \
|
||||
--text "今天天气很好,我们一起学习语音模型推理。" \
|
||||
--ref-audio "file://${PWD}/tests/assets/glm_tts/jiayan_zh.wav" \
|
||||
--ref-text "他当时还跟线下其他的站姐吵架,然后,打架进局子了。" \
|
||||
--response-format wav \
|
||||
--output glm_tts_zh.wav
|
||||
```
|
||||
|
||||
Run the English voice-cloning check with the same Chinese reference:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
|
||||
--model zai-org/GLM-TTS \
|
||||
--text "Hello, this is a voice cloning test with GLM-TTS." \
|
||||
--ref-audio "file://${PWD}/tests/assets/glm_tts/jiayan_zh.wav" \
|
||||
--ref-text "他当时还跟线下其他的站姐吵架,然后,打架进局子了。" \
|
||||
--response-format wav \
|
||||
--output glm_tts_en.wav
|
||||
```
|
||||
|
||||
Inspect either output:
|
||||
|
||||
```bash
|
||||
ffprobe -v error -select_streams a:0 \
|
||||
-show_entries stream=codec_name,sample_rate,channels \
|
||||
-of default=noprint_wrappers=1 glm_tts_zh.wav
|
||||
```
|
||||
|
||||
Expected output:
|
||||
|
||||
```text
|
||||
codec_name=pcm_s16le
|
||||
sample_rate=24000
|
||||
channels=1
|
||||
```
|
||||
|
||||
The English smoke test produced a valid WAV with an actual duration of 5.020
|
||||
seconds.
|
||||
|
||||
#### Observed warm-cache results
|
||||
|
||||
These observations are from one controlled run, not a hardware or performance
|
||||
guarantee. Cold-start and JIT costs were excluded. Values are mean +/- standard
|
||||
deviation.
|
||||
|
||||
##### Measurement procedure
|
||||
|
||||
- For the single-GPU server, use the documented command with
|
||||
`CUDA_VISIBLE_DEVICES=0` and omit `--stage-overrides`; both bundled stages
|
||||
remain on GPU 0. For the dual-GPU server, use the command as shown.
|
||||
- Async chunking is the default. Add `--no-async-chunk` for sync mode.
|
||||
- For non-streaming measurements, use the exact Chinese client command above.
|
||||
For streaming, add `--stream` to that command and change the output filename
|
||||
to `glm_tts_zh.pcm`; streaming output is raw PCM, not WAV.
|
||||
- The fixed Chinese prompt and reference were exactly those shown in
|
||||
Verification and yielded 3.900 seconds of audio. Each case used one warm-up
|
||||
followed by three measured requests at concurrency 1.
|
||||
- Wall time starts immediately before the HTTP POST and stops after the complete
|
||||
response body. TTFA runs from immediately before the POST to the first
|
||||
non-empty audio bytes. RTF is wall time divided by audio duration.
|
||||
- Peak total device memory used was sampled every 0.2 seconds and reported by
|
||||
NVML via `nvmlDeviceGetMemoryInfo.used`. The readings include background and
|
||||
baseline use and are not baseline-subtracted.
|
||||
|
||||
| GPUs | Mode | Response | Wall time (s) | TTFA (s) | RTF |
|
||||
|---:|---|---|---:|---:|---:|
|
||||
| 1 | Async | Non-streaming | 2.144 +/- 0.076 | -- | 0.550 +/- 0.019 |
|
||||
| 2 | Async | Non-streaming | 1.969 +/- 0.130 | -- | 0.505 +/- 0.033 |
|
||||
| 1 | Async | Streaming | 2.135 +/- 0.093 | 1.245 +/- 0.178 | 0.548 +/- 0.024 |
|
||||
| 2 | Async | Streaming | 1.891 +/- 0.006 | 1.170 +/- 0.042 | 0.485 +/- 0.002 |
|
||||
| 1 | Sync | Non-streaming | 1.822 +/- 0.016 | -- | 0.467 +/- 0.004 |
|
||||
| 2 | Sync | Non-streaming | 1.792 +/- 0.023 | -- | 0.459 +/- 0.006 |
|
||||
|
||||
Moving stage 1 to GPU 1 reduced async non-streaming wall time by 8.2%. It
|
||||
reduced async streaming wall time by 11.4% and TTFA by 6.0%, and reduced sync
|
||||
wall time by 1.6%. Single- and dual-GPU outputs were bitwise identical only
|
||||
within matching execution and response modes.
|
||||
|
||||
| Configuration | GPU 0 | GPU 1 |
|
||||
|---|---:|---:|
|
||||
| Single-GPU async | 31.52 GiB | -- |
|
||||
| Dual-GPU async | 29.23 GiB | 2.91 GiB |
|
||||
| Single-GPU sync | 31.47 GiB | -- |
|
||||
| Dual-GPU sync | 29.23 GiB | 2.91 GiB |
|
||||
@@ -1,143 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.benchmarks.duplex import omni_duplex_eval_eval as eval_module
|
||||
from vllm_omni.benchmarks.duplex.omni_duplex_eval_clock import (
|
||||
extract_timed_sentences,
|
||||
normalize_response_items,
|
||||
split_text,
|
||||
validate_clock,
|
||||
)
|
||||
from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import (
|
||||
DuplexSample,
|
||||
canonical_task_type,
|
||||
family_for_split,
|
||||
load_samples,
|
||||
task_type_for_split,
|
||||
)
|
||||
from vllm_omni.benchmarks.duplex.omni_duplex_eval_metrics import (
|
||||
PROTOCOL_PIN,
|
||||
build_content_prompt,
|
||||
build_reminder_prompt,
|
||||
build_temporal_prompt,
|
||||
parse_judge_json,
|
||||
reminder_window,
|
||||
summarize_pr_results,
|
||||
temporal_window,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_dataset_mapping_and_manifest(tmp_path):
|
||||
manifest = tmp_path / "samples.json"
|
||||
manifest.write_text(json.dumps([{"id": "a", "split": "RTD_OCR", "video": "clip.mp4"}]), encoding="utf-8")
|
||||
sample = load_samples(manifest, media_root=tmp_path)[0]
|
||||
assert sample.family == "rtd"
|
||||
assert sample.video == str(tmp_path / "clip.mp4")
|
||||
assert family_for_split("PR_event_reminder") == "pr"
|
||||
assert canonical_task_type("post-event-reminder") == "post_event_reminder"
|
||||
assert task_type_for_split("PR_correction") == "correction"
|
||||
assert load_samples([{"id": "pr", "split": "PR_correction"}])[0].task_type == "correction"
|
||||
|
||||
|
||||
def test_hugging_face_subset_names_are_splits(monkeypatch):
|
||||
calls = []
|
||||
|
||||
def fake_load_dataset(name, **kwargs):
|
||||
calls.append((name, kwargs))
|
||||
return [{"id": "pr", "question_text": "Correct this."}]
|
||||
|
||||
monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
|
||||
sample = load_samples("Hothan/Omni-DuplexEval", split="PR_correction")[0]
|
||||
assert calls == [("Hothan/Omni-DuplexEval", {"split": "PR_correction"})]
|
||||
assert sample.split == "PR_correction"
|
||||
assert sample.task_type == "correction"
|
||||
|
||||
|
||||
def test_response_aliases_and_clock_guard():
|
||||
assert split_text("One. Two!") == ["One.", "Two!"]
|
||||
assert normalize_response_items({"chunks": [{"text": "x", "current_time": 800}]}) == [
|
||||
{"sentence": "x", "start": 800.0, "end": 800.0}
|
||||
]
|
||||
with pytest.raises(ValueError, match="clock=invalid"):
|
||||
validate_clock({"clock": "invalid"})
|
||||
validate_clock({"clock": "invalid"}, allow_invalid=True)
|
||||
timed = extract_timed_sentences([{"type": "response.output_text.delta", "delta": "Done.", "_media_clock_ms": 800}])
|
||||
assert [item.as_dict() for item in timed] == [{"sentence": "Done.", "start": 0.8, "end": 0.8}]
|
||||
with pytest.raises(ValueError, match="only clock=media"):
|
||||
extract_timed_sentences([], clock="wall")
|
||||
|
||||
|
||||
def test_protocol_windows_and_parsing():
|
||||
assert PROTOCOL_PIN == "ca3c122b4d4bf67afd6b18ea5e724b4561bdde48"
|
||||
assert temporal_window(4, 5, 10) == (2.0, 3.0)
|
||||
assert temporal_window(1, 1.2, 10) == (0.0, 0.5)
|
||||
assert reminder_window(5) == (5.0, 15.0)
|
||||
assert parse_judge_json('noise {"content_score": 2.345, "is_relevant": 1}') == {
|
||||
"content_score": 2.35,
|
||||
"is_relevant": 1,
|
||||
}
|
||||
assert parse_judge_json('noise "success_score": 1') == {"success_score": 1}
|
||||
assert summarize_pr_results([{"task_type": "correction", "all_success": 1}])["mean_all_success"] == 1.0
|
||||
|
||||
|
||||
def test_protocol_prompt_wording_is_pinned():
|
||||
prompts = [
|
||||
build_temporal_prompt(1, 2, "response", "question"),
|
||||
build_content_prompt("response", "question", ["one", "two"]),
|
||||
build_reminder_prompt("instruction", "response", "proactive_reminder"),
|
||||
build_reminder_prompt("instruction", "response", "post_event_reminder"),
|
||||
build_reminder_prompt("instruction", "response", "correction", "answer"),
|
||||
]
|
||||
assert [hashlib.sha256(prompt.encode()).hexdigest() for prompt in prompts] == [
|
||||
"5f7f5fcb4a37e4d2eea8ba23beabf74630e6fc2fb5b2661ceb94403254bb04d5",
|
||||
"1fccbed6c1ed8ea1c76d5564d7289ed6adca3678f6a0da2cc735942123d64457",
|
||||
"d0a0e3a9d873d2f368c5c99b9fe46681f683cc36a9fe0e09b99d6cbca1d6dd0c",
|
||||
"08f9951243c859d271f82ee53da22cd0bdab1fe93c6ca04f4870d3b959c0803a",
|
||||
"3de131a7e36cb4b8ac369bb05e87ede518be20221e214df3644211c179131a73",
|
||||
]
|
||||
|
||||
|
||||
def test_frame_sample_content_passes_frames(tmp_path, monkeypatch):
|
||||
response = tmp_path / "response.json"
|
||||
response.write_text(json.dumps([{"sentence": "A person moves.", "start": 1, "end": 2}]), encoding="utf-8")
|
||||
response.with_name("response.meta.json").write_text(json.dumps({"clock": "media"}), encoding="utf-8")
|
||||
video = tmp_path / "video.mp4"
|
||||
video.write_bytes(b"video")
|
||||
monkeypatch.setattr(eval_module, "extract_jpeg", lambda *args, **kwargs: b"\xff\xd8jpeg")
|
||||
|
||||
class Judge:
|
||||
def temporal(self, *args, **kwargs):
|
||||
return '{"temporal_score": 3, "is_relevant": 1}'
|
||||
|
||||
def content(self, prompt, video, frames=None, *, mode="video_url"):
|
||||
assert mode == "frame-sample"
|
||||
assert frames == [b"\xff\xd8jpeg", b"\xff\xd8jpeg"]
|
||||
return '{"content_score": 3}'
|
||||
|
||||
sample = DuplexSample(
|
||||
id="rtd",
|
||||
split="RTD_OCR",
|
||||
family="rtd",
|
||||
task_type=None,
|
||||
video=video,
|
||||
video_duration=4,
|
||||
)
|
||||
score = eval_module.evaluate_sample(
|
||||
sample,
|
||||
response,
|
||||
tmp_path / "score.json",
|
||||
Judge(),
|
||||
judge_video_mode="frame-sample",
|
||||
)
|
||||
assert score["content"]["frame_count"] == 2
|
||||
@@ -534,7 +534,6 @@ class TestPipelineDiscovery:
|
||||
"""Check that specific models are in OMNI_PIPELINES."""
|
||||
assert "qwen2_5_omni" in OMNI_PIPELINES
|
||||
assert "qwen3_omni_moe" in OMNI_PIPELINES
|
||||
assert "qwen3_omni_moe_thinker_only" in OMNI_PIPELINES
|
||||
assert "qwen3_tts" in OMNI_PIPELINES
|
||||
|
||||
def test_registry_resolver_qwen3_omni_all_stages(self):
|
||||
@@ -559,18 +558,6 @@ class TestPipelineDiscovery:
|
||||
assert pipeline.default_deploy_config_name is None
|
||||
assert len(pipeline.stages) == 1 # thinker only
|
||||
|
||||
def test_registry_qwen3_omni_thinker_only_key_is_static(self):
|
||||
"""Explicit registry key must not re-enter the HF-config resolver."""
|
||||
pipeline = resolve_pipeline_config(
|
||||
"qwen3_omni_moe_thinker_only",
|
||||
Q3_OMNI_ALL_STAGES_HF_CONFIG,
|
||||
)
|
||||
assert pipeline is OMNI_PIPELINES["qwen3_omni_moe_thinker_only"]
|
||||
assert pipeline.model_type == "qwen3_omni_moe_thinker_only"
|
||||
assert len(pipeline.stages) == 1
|
||||
assert pipeline.stages[0].engine_output_type == "text"
|
||||
assert pipeline.stages[0].final_output_type == "text"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"pipeline",
|
||||
[registered for registered in OMNI_PIPELINES.values() if isinstance(registered, PipelineConfig)],
|
||||
@@ -1122,35 +1109,8 @@ class TestDeployConfigLoading:
|
||||
with pytest.raises(ValueError, match=r"stage_args.*PipelineConfig.*stages"):
|
||||
load_deploy_config(deploy_path)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("filename", "max_sessions"),
|
||||
[
|
||||
("minicpmo_4_5.yaml", 4),
|
||||
("minicpmo_4_5_2gpu.yaml", 4),
|
||||
("minicpmo_4_5_3gpu.yaml", 4),
|
||||
("minicpmo_4_5_8x4090.yaml", 1),
|
||||
("minicpmo_4_5_3gpu_stage1_replicas.yaml", 4),
|
||||
("minicpmo_4_5_4gpu_stage1_replicas.yaml", 4),
|
||||
("minicpmo_4_5_8x4090_stage1_replicas.yaml", 1),
|
||||
],
|
||||
)
|
||||
def test_minicpmo_deploy_configs_enable_duplex(self, filename: str, max_sessions: int):
|
||||
"""Every MiniCPM-o profile serves /v1/realtime, so the retired
|
||||
minicpmo_4_5_duplex.yaml overlay has no replacement to fall back on.
|
||||
"""
|
||||
deploy = load_deploy_config(Path(get_deploy_config_path(filename)))
|
||||
pipeline = resolve_pipeline_config("minicpmo_4_5")
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
assert deploy.session_mode == "duplex"
|
||||
assert deploy.active_stream_window == 1
|
||||
assert deploy.duplex_session.max_sessions == max_sessions
|
||||
assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
|
||||
assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
|
||||
|
||||
def test_load_minicpmo_default_deploy_config(self):
|
||||
deploy_path = Path(get_deploy_config_path("minicpmo_4_5.yaml"))
|
||||
def test_load_minicpmo_duplex_deploy_config(self):
|
||||
deploy_path = Path(get_deploy_config_path("minicpmo_4_5_duplex.yaml"))
|
||||
|
||||
deploy = load_deploy_config(deploy_path)
|
||||
pipeline = resolve_pipeline_config("minicpmo_4_5")
|
||||
@@ -1158,15 +1118,17 @@ class TestDeployConfigLoading:
|
||||
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
assert deploy.session_mode == "duplex"
|
||||
assert deploy.async_chunk is True
|
||||
assert stages[0].yaml_engine_args["async_scheduling"] is True
|
||||
assert stages[1].yaml_engine_args["async_scheduling"] is True
|
||||
assert "Async" in (stages[0].scheduler_cls or "")
|
||||
assert "Async" in (stages[1].scheduler_cls or "")
|
||||
assert deploy.active_stream_window == 1
|
||||
assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
|
||||
assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
|
||||
assert [stage.yaml_engine_args["async_scheduling"] for stage in stages] == [False, False, False]
|
||||
assert all("Async" not in (stage.scheduler_cls or "") for stage in stages)
|
||||
assert [stage.devices for stage in deploy.stages] == ["0", "0", "0"]
|
||||
assert deploy.stages[1].enforce_eager is False
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 4096
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["min_tokens"] == 50
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["min_tokens"] == 0
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["temperature"] == 0.8
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [6561]
|
||||
assert "codec_sampling_params" not in stages[1].yaml_engine_args
|
||||
@@ -1238,22 +1200,14 @@ stages:
|
||||
assert deploy.connectors is not None
|
||||
assert deploy.platforms is not None
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("hf_config", "model"),
|
||||
[
|
||||
(Q3_OMNI_THINKER_HF_CONFIG, "Qwen/Qwen3-Omni-30B-A3B-Thinking"),
|
||||
(Q3_OMNI_ALL_STAGES_HF_CONFIG, "Qwen/Qwen3-Omni-30B-A3B-Instruct"),
|
||||
],
|
||||
ids=["thinker_checkpoint", "instruct_forced_thinker_only"],
|
||||
)
|
||||
def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self, hf_config, model):
|
||||
def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self):
|
||||
deploy_path = Path(get_deploy_config_path("qwen3_omni_moe_thinking.yaml"))
|
||||
with patch(
|
||||
"vllm_omni.config.config_factory.get_config",
|
||||
return_value=hf_config,
|
||||
return_value=Q3_OMNI_THINKER_HF_CONFIG,
|
||||
):
|
||||
pipeline = StageConfigFactory.get_pipeline_config(
|
||||
model,
|
||||
"Qwen/Qwen3-Omni-30B-A3B-Thinking",
|
||||
trust_remote_code=False,
|
||||
deploy_config_path=str(deploy_path),
|
||||
)
|
||||
@@ -1262,7 +1216,6 @@ stages:
|
||||
assert pipeline.model_type == "qwen3_omni_moe_thinker_only"
|
||||
|
||||
deploy = load_deploy_config(deploy_path)
|
||||
assert deploy.pipeline == "qwen3_omni_moe_thinker_only"
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
assert len(stages) == 1
|
||||
@@ -1314,8 +1267,6 @@ stages:
|
||||
@pytest.mark.parametrize(
|
||||
("deploy_name", "pipeline_name", "stage_count", "final_output_type", "declared_pipeline"),
|
||||
[
|
||||
("mammoth_moda2.yaml", "mammoth_moda2", 2, "image", "mammoth_moda2"),
|
||||
("mammoth_moda2_ar.yaml", "mammoth_moda2_ar", 1, "text", "mammoth_moda2_ar"),
|
||||
("omnivoice.yaml", "omnivoice", 1, "audio", "omnivoice"),
|
||||
("mimo_audio.yaml", "mimo_audio", 2, "audio", None),
|
||||
("step_audio_2.yaml", "step_audio_2", 2, "audio", None),
|
||||
@@ -2165,8 +2116,7 @@ class TestBaseConfigInheritance:
|
||||
def test_minicpmo_overlays_inherit_talker_sampling_params(self):
|
||||
"""Overlays must keep the base Talker codec Sampler knobs."""
|
||||
for filename in (
|
||||
"minicpmo_4_5_2gpu.yaml",
|
||||
"minicpmo_4_5_3gpu.yaml",
|
||||
"minicpmo_4_5_duplex.yaml",
|
||||
"minicpmo_4_5_3gpu_stage1_replicas.yaml",
|
||||
"minicpmo_4_5_4gpu_stage1_replicas.yaml",
|
||||
"minicpmo_4_5_8x4090_stage1_replicas.yaml",
|
||||
@@ -2178,6 +2128,9 @@ class TestBaseConfigInheritance:
|
||||
sampling = deploy.stages[1].default_sampling_params
|
||||
assert sampling is not None, f"{filename} stage 1 lost default_sampling_params"
|
||||
assert sampling["temperature"] == 0.8, filename
|
||||
# Upstream utils.TTSSamplingParams. min_tokens is not checked here:
|
||||
# the duplex overlay drops it to 0 because the model budgets each
|
||||
# generate_chunk itself.
|
||||
assert sampling["top_k"] == 25, filename
|
||||
assert sampling["top_p"] == 0.85, filename
|
||||
assert sampling["repetition_penalty"] == 1.05, filename
|
||||
|
||||
@@ -66,8 +66,6 @@ def test_npu_platform_defers_minimax_h3_encoder_patch() -> None:
|
||||
runtime_src = src[runtime_start:]
|
||||
assert "apply_minimax_h3_qwen3vl_patch" not in init_src
|
||||
assert "apply_minimax_h3_qwen3vl_patch" in runtime_src
|
||||
assert "apply_minimax_h3_qwen3vl_sdpa_patch" not in init_src
|
||||
assert "apply_minimax_h3_qwen3vl_sdpa_patch" in runtime_src
|
||||
assert "apply_minimax_h3_qwen3vl_swiglu_patch" in runtime_src
|
||||
|
||||
|
||||
|
||||
@@ -164,10 +164,10 @@ def test_inventory_matches_reviewed_snapshot_counts():
|
||||
"""Make an inventory expansion an explicit review decision."""
|
||||
category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values())
|
||||
assert category_counts == {
|
||||
EnvironmentVariableCategory.PUBLIC_OMNI: 24,
|
||||
EnvironmentVariableCategory.PUBLIC_OMNI: 22,
|
||||
EnvironmentVariableCategory.INHERITED_VLLM: 20,
|
||||
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
|
||||
EnvironmentVariableCategory.MODEL_SPECIFIC: 56,
|
||||
EnvironmentVariableCategory.MODEL_SPECIFIC: 57,
|
||||
EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
|
||||
EnvironmentVariableCategory.INTERNAL: 2,
|
||||
}
|
||||
@@ -182,7 +182,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
|
||||
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
|
||||
ModelEnvironmentVariableDisposition.EXTERNAL: 0,
|
||||
ModelEnvironmentVariableDisposition.INTERNALIZE: 11,
|
||||
ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 6,
|
||||
ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 7,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -100,7 +100,6 @@ def _make_drain_sched(session: Request) -> MagicMock:
|
||||
sched.pending_stop_after_extraction = set()
|
||||
sched.connector = None
|
||||
sched.kv_cache_manager.take_events.return_value = None
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
sched.finished_req_ids_dict = {}
|
||||
sched.make_stats.return_value = None
|
||||
return sched
|
||||
|
||||
@@ -93,7 +93,6 @@ def _run_resumable_segment_stop(
|
||||
sched.pending_stop_after_extraction = set()
|
||||
sched.connector = None
|
||||
sched.kv_cache_manager.take_events.return_value = None
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
sched.finished_req_ids_dict = {}
|
||||
sched.make_stats.return_value = None
|
||||
|
||||
@@ -199,7 +198,6 @@ def test_resumable_segment_boundary_keeps_pre_transition_send_watermark() -> Non
|
||||
inter_stage_output,
|
||||
session,
|
||||
True,
|
||||
new_token_ids=[42],
|
||||
confirmed_num_computed_tokens=26,
|
||||
)
|
||||
|
||||
@@ -228,7 +226,6 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
|
||||
sched.pending_stop_after_extraction = set()
|
||||
sched.connector = None
|
||||
sched.kv_cache_manager.take_events.return_value = None
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
sched.finished_req_ids_dict = {}
|
||||
sched.make_stats.return_value = None
|
||||
|
||||
@@ -255,37 +252,6 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
|
||||
sched.chunk_transfer_adapter.save_async.assert_not_called()
|
||||
|
||||
|
||||
def test_queued_streaming_update_on_async_stop_fences_in_flight_once() -> None:
|
||||
"""Native duplex + async: the next unit is usually already queued when the
|
||||
current one stops, so _handle_stopped_request applies the update and the
|
||||
stop site below it fences the same in-flight decode a second time. An
|
||||
accumulated residue outlives the drain and swallows the next unit's
|
||||
listen/speak, which the client sees as silence until its timeout.
|
||||
"""
|
||||
session = _make_request()
|
||||
session.status = RequestStatus.RUNNING
|
||||
session.resumable = True
|
||||
session.append_output_token_ids([7])
|
||||
session.num_computed_tokens = 4
|
||||
session.num_output_placeholders = 1
|
||||
# This step's frame (settled to 0 by update_from_output) plus one extra
|
||||
# async decode still in flight.
|
||||
session.num_in_flight_tokens = 2
|
||||
|
||||
sched = _make_scheduler(stage_id=0)
|
||||
sched._enqueue_waiting_request = MagicMock()
|
||||
|
||||
def handle_stopped(request: Request) -> bool:
|
||||
sched._update_request_as_session(request, _make_update([10, 20]))
|
||||
return False
|
||||
|
||||
_run_resumable_segment_stop(session, handle_stopped=handle_stopped)
|
||||
|
||||
# Exactly the one unreported decode, so the drain reaches zero before the
|
||||
# new segment's first frame arrives.
|
||||
assert session.num_stale_output_tokens == session.num_in_flight_tokens == 1
|
||||
|
||||
|
||||
def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() -> None:
|
||||
sched = _make_scheduler(stage_id=0)
|
||||
session = _make_request()
|
||||
@@ -430,33 +396,3 @@ def test_ready_async_chunk_prompt_replacement_releases_stale_kv_once() -> None:
|
||||
assert sched.chunk_transfer_adapter.replaced_streaming_prompt_ids == set()
|
||||
assert sched.chunk_transfer_adapter.requests_with_ready_chunks == {session.request_id}
|
||||
assert sched.chunk_transfer_adapter.requests_num_chunks_sent == {}
|
||||
|
||||
|
||||
def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:
|
||||
"""A WAITING_FOR_CHUNK stop must not delete its newly parked session."""
|
||||
session = _make_request()
|
||||
session.resumable = True
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
|
||||
def queue(*requests):
|
||||
result = MagicMock()
|
||||
result.requests = set(requests)
|
||||
result.add_request.side_effect = result.requests.add
|
||||
result.remove_requests.side_effect = result.requests.difference_update
|
||||
return result
|
||||
|
||||
sched = OmniARScheduler.__new__(OmniARScheduler)
|
||||
sched.running = []
|
||||
sched.waiting = queue()
|
||||
sched.skipped_waiting = queue(session)
|
||||
sched.num_waiting_for_streaming_input = 1
|
||||
sched.chunk_transfer_adapter = SimpleNamespace(segment_finished_requests={session.request_id})
|
||||
|
||||
sched._resume_downstream_chunk_receiver(session)
|
||||
sched._remove_stopped_requests_from_queues(set(), {session})
|
||||
|
||||
assert session.status == RequestStatus.WAITING
|
||||
assert session in sched.waiting.requests
|
||||
assert session not in sched.skipped_waiting.requests
|
||||
assert sched.num_waiting_for_streaming_input == 0
|
||||
assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
|
||||
@@ -138,7 +138,6 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
|
||||
sched.recompute_kv_load_failures = False
|
||||
sched.connector = None
|
||||
sched.kv_cache_manager.take_events.return_value = None
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
sched.finished_req_ids_dict = {}
|
||||
sched.make_stats.return_value = None
|
||||
|
||||
|
||||
@@ -1,168 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""v0.28 deferred KV-block free fence in the Omni schedulers (#6606 review).
|
||||
|
||||
Upstream v0.28 (defer_block_free) advances ``sched_step_seq`` for every
|
||||
non-empty scheduled step and, at the top of ``update_from_output()``,
|
||||
advances ``processed_step_seq`` and drains the deferred frees. The Omni
|
||||
schedulers each carried only half of that contract:
|
||||
|
||||
- ``OmniARScheduler.schedule()`` delegates to ``super().schedule()`` (fence
|
||||
advances) but its reimplemented ``update_from_output()`` never advanced
|
||||
``processed_step_seq`` nor drained -> deferred blocks accumulate without
|
||||
bound under async/PP KV-consumer configurations.
|
||||
- ``OmniGenerationScheduler`` reimplements both: the fast-path
|
||||
``schedule()`` never advanced ``sched_step_seq`` (so finished requests
|
||||
freed blocks an in-flight GPU step could still be writing) and its
|
||||
``update_from_output()`` never drained.
|
||||
|
||||
These tests pin the restored contract at the call boundary; the drain and
|
||||
free logic itself is upstream-tested.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
|
||||
# before Request / RequestStatus are bound in this module.
|
||||
# isort: off
|
||||
import vllm_omni # noqa: F401 - import for side effects (patch vLLM)
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.core.sched.interface import PauseState
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.core.sched.request_queue import SchedulingPolicy
|
||||
from vllm.v1.outputs import ModelRunnerOutput
|
||||
from vllm.v1.request import Request, RequestStatus
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
|
||||
|
||||
# isort: on
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_update_sched(scheduler_cls) -> MagicMock:
|
||||
"""MagicMock scheduler driving the real ``update_from_output`` with an
|
||||
empty batch; only the fence state carries real values."""
|
||||
sched = MagicMock()
|
||||
sched.requests = {}
|
||||
sched.perf_metrics = None
|
||||
sched.chunk_transfer_adapter = None
|
||||
sched.defer_block_free = True
|
||||
sched.processed_step_seq = 0
|
||||
return sched
|
||||
|
||||
|
||||
def _empty_frame(*, total_num_scheduled_tokens: int):
|
||||
scheduler_output = MagicMock(spec=SchedulerOutput)
|
||||
scheduler_output.num_scheduled_tokens = {}
|
||||
scheduler_output.total_num_scheduled_tokens = total_num_scheduled_tokens
|
||||
model_runner_output = MagicMock(spec=ModelRunnerOutput)
|
||||
model_runner_output.sampled_token_ids = []
|
||||
model_runner_output.logprobs = None
|
||||
model_runner_output.prompt_logprobs_dict = {}
|
||||
model_runner_output.pooler_output = None
|
||||
model_runner_output.num_nans_in_logits = None
|
||||
model_runner_output.kv_connector_output = None
|
||||
model_runner_output.cudagraph_stats = None
|
||||
return scheduler_output, model_runner_output
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
|
||||
def test_update_from_output_advances_fence_and_drains(scheduler_cls) -> None:
|
||||
sched = _make_update_sched(scheduler_cls)
|
||||
scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=3)
|
||||
|
||||
scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
|
||||
|
||||
assert sched.processed_step_seq == 1
|
||||
sched._drain_deferred_frees.assert_called_once_with()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
|
||||
def test_update_from_output_zero_token_step_leaves_fence(scheduler_cls) -> None:
|
||||
"""0-token steps do not advance the fence upstream (their seq was never
|
||||
advanced in schedule()); draining on them would run ahead of the GPU."""
|
||||
sched = _make_update_sched(scheduler_cls)
|
||||
scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=0)
|
||||
|
||||
scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
|
||||
|
||||
assert sched.processed_step_seq == 0
|
||||
sched._drain_deferred_frees.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
|
||||
def test_update_from_output_defer_disabled_leaves_fence(scheduler_cls) -> None:
|
||||
sched = _make_update_sched(scheduler_cls)
|
||||
sched.defer_block_free = False
|
||||
scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=3)
|
||||
|
||||
scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
|
||||
|
||||
assert sched.processed_step_seq == 0
|
||||
sched._drain_deferred_frees.assert_not_called()
|
||||
|
||||
|
||||
def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request]:
|
||||
"""MagicMock scheduler that drives the real generation fast-path
|
||||
``schedule()`` over one running one-shot request."""
|
||||
request = Request(
|
||||
request_id="req-fence",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
sampling_params=SamplingParams(max_tokens=4),
|
||||
pooling_params=None,
|
||||
arrival_time=100.0,
|
||||
block_hasher=None,
|
||||
)
|
||||
request.status = RequestStatus.RUNNING
|
||||
|
||||
sched = MagicMock()
|
||||
sched.max_num_scheduled_tokens = 8
|
||||
sched._pause_state = PauseState.UNPAUSED
|
||||
sched.requests = {request.request_id: request}
|
||||
sched.running = [request]
|
||||
sched.waiting = [] # falsy: the waiting loop never runs
|
||||
sched.policy = SchedulingPolicy.FCFS
|
||||
sched.chunk_transfer_adapter = None
|
||||
sched.scheduler_config = MagicMock(enable_chunked_prefill=True)
|
||||
sched.num_lookahead_tokens = 0
|
||||
sched.log_stats = False
|
||||
sched.kv_cache_config = MagicMock(kv_cache_groups=[])
|
||||
sched.use_v2_model_runner = False
|
||||
sched.prev_step_scheduled_req_ids = set()
|
||||
sched.needs_kv_cache_zeroing = False
|
||||
sched.connector = None
|
||||
sched.ec_connector = None
|
||||
sched.defer_block_free = defer_block_free
|
||||
sched.sched_step_seq = 0
|
||||
# last_sched_seq stamping order: _update_after_schedule must observe the
|
||||
# already-advanced fence value, as upstream schedule() guarantees.
|
||||
sched.seq_at_update_after_schedule = None
|
||||
|
||||
def _capture(_scheduler_output):
|
||||
sched.seq_at_update_after_schedule = sched.sched_step_seq
|
||||
|
||||
sched._update_after_schedule.side_effect = _capture
|
||||
return sched, request
|
||||
|
||||
|
||||
def test_fast_path_schedule_advances_fence_before_update_after_schedule() -> None:
|
||||
sched, request = _make_fast_path_sched(defer_block_free=True)
|
||||
|
||||
OmniGenerationScheduler.schedule(sched)
|
||||
|
||||
assert sched.sched_step_seq == 1
|
||||
assert sched.seq_at_update_after_schedule == 1
|
||||
|
||||
|
||||
def test_fast_path_schedule_defer_disabled_leaves_fence() -> None:
|
||||
sched, request = _make_fast_path_sched(defer_block_free=False)
|
||||
|
||||
OmniGenerationScheduler.schedule(sched)
|
||||
|
||||
assert sched.sched_step_seq == 0
|
||||
@@ -1,180 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""v0.28 EC request-finish contract in the Omni schedulers (#6606 review).
|
||||
|
||||
Upstream v0.28 ``_free_request()`` fires ``ec_connector.request_finished()``
|
||||
BEFORE the encoder cache is freed (so the connector can inspect per-request
|
||||
state), honors its delayed-free result, and returns the generated
|
||||
``ec_transfer_params``; ``update_from_output()`` then carries them into the
|
||||
``EngineCoreOutput``. ``OmniARScheduler`` previously skipped the hook and
|
||||
always returned ``None``; both Omni schedulers discarded the second tuple
|
||||
item, so a remote encoder-cache handle could never reach the frontend.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
|
||||
# before Request / RequestStatus are bound in this module.
|
||||
# isort: off
|
||||
import vllm_omni # noqa: F401 - import for side effects (patch vLLM)
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.outputs import ModelRunnerOutput
|
||||
from vllm.v1.request import Request, RequestStatus
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
|
||||
|
||||
# isort: on
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
EC_PARAMS = {"remote_encoder_cache_handle": "ec-1"}
|
||||
|
||||
|
||||
def _make_free_request_sched(*, ec_connector) -> tuple[OmniARScheduler, MagicMock]:
|
||||
"""Minimal real-`_free_request` harness; `order` records the EC hook
|
||||
firing relative to the encoder-cache free."""
|
||||
order = MagicMock()
|
||||
sched = OmniARScheduler.__new__(OmniARScheduler)
|
||||
sched._omits_kv_transfer_cache = {}
|
||||
sched._connector_finished = lambda request: (False, None)
|
||||
sched.ec_connector = ec_connector
|
||||
if ec_connector is not None:
|
||||
order.attach_mock(ec_connector.request_finished, "request_finished")
|
||||
sched.encoder_cache_manager = MagicMock()
|
||||
order.attach_mock(sched.encoder_cache_manager.free, "encoder_cache_free")
|
||||
sched.finished_req_ids = set()
|
||||
sched._new_prompt_len_snapshot = {}
|
||||
sched.finished_req_ids_dict = None
|
||||
sched._should_transfer_kv_for_request = lambda req_id: False
|
||||
sched._free_blocks = MagicMock()
|
||||
sched._free_input_coordinator_request = MagicMock()
|
||||
sched.chunk_transfer_adapter = None
|
||||
return sched, order
|
||||
|
||||
|
||||
class _FakeFinishedRequest:
|
||||
def __init__(self, request_id: str) -> None:
|
||||
self.request_id = request_id
|
||||
|
||||
def is_finished(self) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
def test_free_request_fires_ec_hook_before_encoder_cache_free() -> None:
|
||||
ec_connector = MagicMock()
|
||||
ec_connector.request_finished.return_value = (False, EC_PARAMS)
|
||||
sched, order = _make_free_request_sched(ec_connector=ec_connector)
|
||||
request = _FakeFinishedRequest("req-ec")
|
||||
|
||||
kv_params, ec_params = sched._free_request(request)
|
||||
|
||||
assert ec_params == EC_PARAMS
|
||||
ec_connector.request_finished.assert_called_once_with(request)
|
||||
hook_calls = [name for name, *_ in order.mock_calls]
|
||||
assert hook_calls.index("request_finished") < hook_calls.index("encoder_cache_free")
|
||||
sched._free_blocks.assert_called_once_with(request)
|
||||
|
||||
|
||||
def test_free_request_honors_ec_delayed_free() -> None:
|
||||
ec_connector = MagicMock()
|
||||
ec_connector.request_finished.return_value = (True, EC_PARAMS)
|
||||
sched, _ = _make_free_request_sched(ec_connector=ec_connector)
|
||||
request = _FakeFinishedRequest("req-ec-delay")
|
||||
|
||||
kv_params, ec_params = sched._free_request(request)
|
||||
|
||||
assert ec_params == EC_PARAMS
|
||||
sched._free_blocks.assert_not_called()
|
||||
|
||||
|
||||
def test_free_request_without_ec_connector_returns_none() -> None:
|
||||
sched, _ = _make_free_request_sched(ec_connector=None)
|
||||
request = _FakeFinishedRequest("req-no-ec")
|
||||
|
||||
kv_params, ec_params = sched._free_request(request)
|
||||
|
||||
assert ec_params is None
|
||||
sched._free_blocks.assert_called_once_with(request)
|
||||
|
||||
|
||||
def _make_finishing_request() -> Request:
|
||||
request = Request(
|
||||
request_id="req-ec-finish",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
sampling_params=SamplingParams(max_tokens=4),
|
||||
pooling_params=None,
|
||||
arrival_time=100.0,
|
||||
block_hasher=None,
|
||||
)
|
||||
request.status = RequestStatus.RUNNING
|
||||
return request
|
||||
|
||||
|
||||
def _make_finish_sched(request: Request) -> MagicMock:
|
||||
"""MagicMock scheduler driving the real ``update_from_output`` through a
|
||||
finishing request; ``_free_request`` reports EC transfer params."""
|
||||
sched = MagicMock()
|
||||
sched.requests = {request.request_id: request}
|
||||
sched.perf_metrics = None
|
||||
sched.defer_block_free = False
|
||||
sched.structured_output_manager.should_advance.return_value = False
|
||||
sched._process_kv_transfer_trigger.return_value = False
|
||||
sched._maybe_decode_pooling_output.return_value = None
|
||||
sched._handle_stopped_request.return_value = True
|
||||
sched._free_request.return_value = (None, EC_PARAMS)
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
sched.chunk_transfer_adapter = None
|
||||
sched.finished_req_ids_dict = None
|
||||
sched._new_prompt_len_snapshot = {}
|
||||
return sched
|
||||
|
||||
|
||||
def test_ar_update_from_output_carries_ec_params_to_engine_output() -> None:
|
||||
request = _make_finishing_request()
|
||||
sched = _make_finish_sched(request)
|
||||
sched._update_request_with_output.return_value = ([42], True)
|
||||
|
||||
engine_core_outputs = OmniARScheduler.update_from_output(sched, *_finish_frame_parts(request))
|
||||
|
||||
(eco,) = engine_core_outputs[request.client_index].outputs
|
||||
assert eco.ec_transfer_params == EC_PARAMS
|
||||
sched._free_request.assert_called_once_with(request)
|
||||
|
||||
|
||||
def _finish_frame_parts(request: Request):
|
||||
scheduler_output = MagicMock(spec=SchedulerOutput)
|
||||
scheduler_output.num_scheduled_tokens = {request.request_id: 1}
|
||||
scheduler_output.total_num_scheduled_tokens = 1
|
||||
scheduler_output.scheduled_spec_decode_tokens = {}
|
||||
scheduler_output.num_invalid_spec_tokens = 0
|
||||
|
||||
model_runner_output = MagicMock(spec=ModelRunnerOutput)
|
||||
model_runner_output.sampled_token_ids = [[42]]
|
||||
model_runner_output.logprobs = None
|
||||
model_runner_output.prompt_logprobs_dict = {}
|
||||
model_runner_output.pooler_output = None
|
||||
model_runner_output.num_nans_in_logits = None
|
||||
model_runner_output.kv_connector_output = None
|
||||
model_runner_output.cudagraph_stats = None
|
||||
model_runner_output.req_id_to_index = {request.request_id: 0}
|
||||
model_runner_output.routed_experts = None
|
||||
return scheduler_output, model_runner_output
|
||||
|
||||
|
||||
def test_generation_update_from_output_carries_ec_params_to_engine_output() -> None:
|
||||
request = _make_finishing_request()
|
||||
# One-shot generation finish: the whole prompt is computed.
|
||||
request.num_computed_tokens = len(request.prompt_token_ids)
|
||||
sched = _make_finish_sched(request)
|
||||
|
||||
engine_core_outputs = OmniGenerationScheduler.update_from_output(sched, *_finish_frame_parts(request))
|
||||
|
||||
(eco,) = engine_core_outputs[request.client_index].outputs
|
||||
assert eco.ec_transfer_params == EC_PARAMS
|
||||
sched._free_request.assert_called_once_with(request)
|
||||
@@ -1,126 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""v0.28 prefill-stats finalization in the Omni schedulers (#6606 review).
|
||||
|
||||
Upstream v0.28 finalizes prefill statistics with
|
||||
``PrefillStats.finalize(kv_cache_manager.estimate_cached_tokens(request))``
|
||||
before emitting them, which computes ``num_cache_creation_tokens`` (tokens
|
||||
computed and written to the prefix cache). Both Omni schedulers previously
|
||||
emitted ``take_prefill_stats()`` raw, so creation usage was always 0.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
|
||||
# before Request / RequestStatus are bound in this module.
|
||||
# isort: off
|
||||
import vllm_omni # noqa: F401 - import for side effects (patch vLLM)
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.outputs import ModelRunnerOutput
|
||||
from vllm.v1.request import Request, RequestStatus
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
|
||||
|
||||
# isort: on
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_request() -> Request:
|
||||
request = Request(
|
||||
request_id="req-prefill-stats",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
sampling_params=SamplingParams(max_tokens=4),
|
||||
pooling_params=None,
|
||||
arrival_time=100.0,
|
||||
block_hasher=None,
|
||||
)
|
||||
request.status = RequestStatus.RUNNING
|
||||
# A scheduled prefill: 3 prompt tokens, 1 from the local prefix cache.
|
||||
request.prefill_stats.set(
|
||||
num_prompt_tokens=3,
|
||||
num_local_cached_tokens=1,
|
||||
num_external_cached_tokens=0,
|
||||
)
|
||||
return request
|
||||
|
||||
|
||||
def _make_sched(request: Request) -> MagicMock:
|
||||
sched = MagicMock()
|
||||
sched.requests = {request.request_id: request}
|
||||
sched.perf_metrics = None
|
||||
sched.defer_block_free = False
|
||||
sched.structured_output_manager.should_advance.return_value = False
|
||||
sched._process_kv_transfer_trigger.return_value = False
|
||||
sched._maybe_decode_pooling_output.return_value = None
|
||||
sched._handle_stopped_request.return_value = True
|
||||
# The whole 3-token prompt ends up cached: 3 - 1 pre-cached = 2 created.
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 3
|
||||
|
||||
def _free_and_invalidate(req):
|
||||
# Lifecycle guard: freeing releases the KV blocks, after which the
|
||||
# manager reports 0 cached tokens. If finalize ran after
|
||||
# _free_request (the reviewed bug), creation count would compute
|
||||
# against 0 and the assertions below would see 0, not 2.
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
return (None, None)
|
||||
|
||||
sched._free_request.side_effect = _free_and_invalidate
|
||||
sched.chunk_transfer_adapter = None
|
||||
sched.finished_req_ids_dict = None
|
||||
sched._new_prompt_len_snapshot = {}
|
||||
return sched
|
||||
|
||||
|
||||
def _frame(request: Request):
|
||||
scheduler_output = MagicMock(spec=SchedulerOutput)
|
||||
scheduler_output.num_scheduled_tokens = {request.request_id: 1}
|
||||
scheduler_output.total_num_scheduled_tokens = 1
|
||||
scheduler_output.scheduled_spec_decode_tokens = {}
|
||||
scheduler_output.num_invalid_spec_tokens = 0
|
||||
|
||||
model_runner_output = MagicMock(spec=ModelRunnerOutput)
|
||||
model_runner_output.sampled_token_ids = [[42]]
|
||||
model_runner_output.logprobs = None
|
||||
model_runner_output.prompt_logprobs_dict = {}
|
||||
model_runner_output.pooler_output = None
|
||||
model_runner_output.num_nans_in_logits = None
|
||||
model_runner_output.kv_connector_output = None
|
||||
model_runner_output.cudagraph_stats = None
|
||||
model_runner_output.req_id_to_index = {request.request_id: 0}
|
||||
model_runner_output.routed_experts = None
|
||||
return scheduler_output, model_runner_output
|
||||
|
||||
|
||||
def test_ar_scheduler_finalizes_prefill_stats_before_emitting() -> None:
|
||||
request = _make_request()
|
||||
sched = _make_sched(request)
|
||||
sched._update_request_with_output.return_value = ([42], True)
|
||||
|
||||
engine_core_outputs = OmniARScheduler.update_from_output(sched, *_frame(request))
|
||||
|
||||
(eco,) = engine_core_outputs[request.client_index].outputs
|
||||
assert eco.prefill_stats is not None
|
||||
assert eco.prefill_stats.num_cached_tokens == 1
|
||||
assert eco.prefill_stats.num_cache_creation_tokens == 2
|
||||
sched.kv_cache_manager.estimate_cached_tokens.assert_called_once_with(request)
|
||||
|
||||
|
||||
def test_generation_scheduler_finalizes_prefill_stats_before_emitting() -> None:
|
||||
request = _make_request()
|
||||
request.num_computed_tokens = len(request.prompt_token_ids)
|
||||
sched = _make_sched(request)
|
||||
|
||||
engine_core_outputs = OmniGenerationScheduler.update_from_output(sched, *_frame(request))
|
||||
|
||||
(eco,) = engine_core_outputs[request.client_index].outputs
|
||||
assert eco.prefill_stats is not None
|
||||
assert eco.prefill_stats.num_cached_tokens == 1
|
||||
assert eco.prefill_stats.num_cache_creation_tokens == 2
|
||||
sched.kv_cache_manager.estimate_cached_tokens.assert_called_once_with(request)
|
||||
@@ -18,7 +18,7 @@
|
||||
"model": "openbmb/MiniCPM-o-4_5",
|
||||
"extra_cli_args": [
|
||||
"--deploy-config",
|
||||
"vllm_omni/deploy/minicpmo_4_5.yaml",
|
||||
"vllm_omni/deploy/minicpmo_4_5_duplex.yaml",
|
||||
"--trust-remote-code"
|
||||
]
|
||||
},
|
||||
|
||||
@@ -308,91 +308,6 @@ def test_flash_attn_func_preferred_over_varlen():
|
||||
print("✓ flash_attn_func forward works correctly!")
|
||||
|
||||
|
||||
@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires CUDA or XPU")
|
||||
@pytest.mark.parametrize("k_len", [40, 256])
|
||||
def test_cross_attn_key_padding_vs_sdpa(k_len):
|
||||
"""
|
||||
Case 3: Cross-attention with a key-padding mask.
|
||||
|
||||
The mask covers only the key sequence; every query position must attend
|
||||
to the valid keys. FA and SDPA outputs should be very close. k_len=256
|
||||
covers equal Q/K lengths, which must not be mistaken for self-attention,
|
||||
and valid key lengths differ per batch element.
|
||||
"""
|
||||
device = torch.device(current_omni_platform.device_type)
|
||||
dtype = torch.bfloat16
|
||||
|
||||
batch_size = 2
|
||||
q_len = 256
|
||||
valid_lens = (17, 29)
|
||||
num_heads = 8
|
||||
head_dim = 64
|
||||
|
||||
fa_impl = FlashAttentionImpl(
|
||||
num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False, role="cross"
|
||||
)
|
||||
sdpa_impl = SDPAImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False)
|
||||
|
||||
torch.manual_seed(7)
|
||||
query = torch.randn(batch_size, q_len, num_heads, head_dim, device=device, dtype=dtype)
|
||||
key = torch.randn(batch_size, k_len, num_heads, head_dim, device=device, dtype=dtype)
|
||||
value = torch.randn(batch_size, k_len, num_heads, head_dim, device=device, dtype=dtype)
|
||||
|
||||
attn_mask = torch.zeros(batch_size, k_len, dtype=torch.bool, device=device)
|
||||
for i, valid_len in enumerate(valid_lens):
|
||||
attn_mask[i, :valid_len] = True
|
||||
attn_metadata = AttentionMetadata(attn_mask=attn_mask)
|
||||
|
||||
output_fa = fa_impl.forward(query=query, key=key, value=value, attn_metadata=attn_metadata)
|
||||
output_sdpa = sdpa_impl.forward(query=query, key=key, value=value, attn_metadata=attn_metadata)
|
||||
|
||||
max_diff = torch.max(torch.abs(output_fa - output_sdpa)).item()
|
||||
mean_diff = torch.mean(torch.abs(output_fa - output_sdpa)).item()
|
||||
|
||||
print("\n=== Case 3: Cross-Attention Key-Padding FA vs SDPA ===")
|
||||
print(f"Max absolute difference: {max_diff:.6f}")
|
||||
print(f"Mean absolute difference: {mean_diff:.6f}")
|
||||
|
||||
assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
|
||||
assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
|
||||
|
||||
print("✓ Case 3 PASSED: FA and SDPA cross-attention outputs are very close!")
|
||||
|
||||
|
||||
def test_varlen_masked_routing_by_role(monkeypatch):
|
||||
"""The unpad route is picked by role, not Q/K length equality; runs without a GPU."""
|
||||
calls = []
|
||||
|
||||
def fake_varlen_func(q, k, v, **kwargs):
|
||||
calls.append((q.shape[0], kwargs["cu_seqlens_q"], kwargs["cu_seqlens_k"]))
|
||||
return q
|
||||
|
||||
monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
|
||||
monkeypatch.setattr(fa, "flash_attn_varlen_func", fake_varlen_func)
|
||||
|
||||
batch_size, seq_len, num_heads, head_dim = 2, 4, 2, 8
|
||||
query = torch.randn(batch_size, seq_len, num_heads, head_dim)
|
||||
metadata = AttentionMetadata(attn_mask=torch.tensor([[True, True, False, False], [True, True, True, False]]))
|
||||
|
||||
cross_impl = FlashAttentionImpl(
|
||||
num_heads=num_heads, head_size=head_dim, softmax_scale=0.5, causal=False, role="cross"
|
||||
)
|
||||
output = cross_impl.forward_cuda(query, query, query, metadata)
|
||||
assert output.shape == query.shape
|
||||
num_q_rows, cu_seqlens_q, cu_seqlens_k = calls[-1]
|
||||
assert num_q_rows == batch_size * seq_len
|
||||
assert cu_seqlens_q.tolist() == [0, 4, 8]
|
||||
assert cu_seqlens_k.tolist() == [0, 2, 5]
|
||||
|
||||
self_impl = FlashAttentionImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=0.5, causal=False)
|
||||
output = self_impl.forward_cuda(query, query, query, metadata)
|
||||
assert output.shape == query.shape
|
||||
num_q_rows, cu_seqlens_q, cu_seqlens_k = calls[-1]
|
||||
assert num_q_rows == 5
|
||||
assert cu_seqlens_q.tolist() == [0, 2, 5]
|
||||
assert cu_seqlens_k.tolist() == [0, 2, 5]
|
||||
|
||||
|
||||
def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch):
|
||||
calls = []
|
||||
|
||||
|
||||
+6
-3
@@ -365,10 +365,13 @@ class TestCacheDiTBackend:
|
||||
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
|
||||
def test_enable_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter):
|
||||
"""Generic cache path should pick up ``transformer.fused_blocks``."""
|
||||
def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter):
|
||||
"""Test DreamID uses pipeline.transformer for cache enable/refresh.
|
||||
|
||||
NOTE: DreamID no longer has a custom enabler, so this tests against the generic path.
|
||||
"""
|
||||
mock_pipeline = Mock()
|
||||
mock_pipeline.__class__.__name__ = "FakeFusedBlocksPipeline"
|
||||
mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline"
|
||||
mock_pipeline.transformer = Mock()
|
||||
mock_pipeline.transformer.fused_blocks = Mock()
|
||||
mock_pipeline.transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
|
||||
|
||||
+10
@@ -6,6 +6,7 @@ Model specific tests for CacheDiT enablement.
|
||||
"""
|
||||
|
||||
import ast
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
@@ -24,9 +25,18 @@ from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import L
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
# NOTE: We patch DreamID Omni's modules here with mocks so that we can import and inspect
|
||||
# the class even though the dependency may not be set up correctly; this is ok for these
|
||||
# tests because we just inspect it and never initialize the model.
|
||||
for mod in ("dreamid_omni", "dreamid_omni.modules", "dreamid_omni.modules.model"):
|
||||
sys.modules.setdefault(mod, Mock())
|
||||
# isort: split
|
||||
from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel as DreamIdOmniModel # noqa: E402
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
SEPARATE_CFG_TRANSFORMERS = [
|
||||
DreamIdOmniModel,
|
||||
HeliosTransformer3DModel,
|
||||
LongCatImageTransformer2DModel,
|
||||
Cosmos3VFMTransformer,
|
||||
|
||||
@@ -1,6 +1,3 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
# ruff: noqa: N803, E741
|
||||
"""Layer-level correctness & performance test for MoT parallel linear layers.
|
||||
|
||||
@@ -240,53 +237,6 @@ def _reference_forward(x, text_indices, vae_indices, text_linear, vae_linear):
|
||||
return output
|
||||
|
||||
|
||||
def _assert_gen_expert_parallel_metadata(layer):
|
||||
"""The secondary expert must satisfy vLLM's parallel-linear contract."""
|
||||
for attr in (
|
||||
"tp_rank",
|
||||
"tp_size",
|
||||
"input_size",
|
||||
"input_size_per_partition",
|
||||
"output_size",
|
||||
"output_size_per_partition",
|
||||
):
|
||||
assert getattr(layer.gen_exp, attr) == getattr(layer, attr)
|
||||
|
||||
|
||||
def test_mot_gen_expert_parallel_metadata():
|
||||
"""Keep secondary experts compatible with vLLM online quantizers."""
|
||||
from vllm.model_executor.layers.quantization.online.fp8 import (
|
||||
_is_tp_sharded,
|
||||
)
|
||||
|
||||
with set_current_vllm_config(VllmConfig()):
|
||||
layers = (
|
||||
MoTQKVParallelLinear(
|
||||
hidden_size=16,
|
||||
head_size=4,
|
||||
total_num_heads=4,
|
||||
total_num_kv_heads=2,
|
||||
bias=False,
|
||||
vae_bias=False,
|
||||
params_dtype=torch.bfloat16,
|
||||
disable_tp=True,
|
||||
),
|
||||
MoTRowParallelLinear(
|
||||
16,
|
||||
8,
|
||||
bias=False,
|
||||
vae_bias=False,
|
||||
input_is_parallel=True,
|
||||
params_dtype=torch.bfloat16,
|
||||
disable_tp=True,
|
||||
),
|
||||
)
|
||||
|
||||
for layer in layers:
|
||||
_assert_gen_expert_parallel_metadata(layer)
|
||||
assert _is_tp_sharded(layer.gen_exp) is False
|
||||
|
||||
|
||||
def _check_and_report(ref: torch.Tensor, mot: torch.Tensor, tag: str):
|
||||
"""Compare outputs, print metrics, and assert correctness.
|
||||
|
||||
@@ -385,8 +335,6 @@ def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str, bias: bool
|
||||
disable_tp=True,
|
||||
).cuda()
|
||||
|
||||
_assert_gen_expert_parallel_metadata(mot_linear)
|
||||
|
||||
assert text_linear.output_size_per_partition == N, (
|
||||
f"Expected output_size_per_partition={N}, "
|
||||
f"got {text_linear.output_size_per_partition}. "
|
||||
@@ -475,8 +423,6 @@ def test_mot_o_proj(
|
||||
disable_tp=True,
|
||||
).cuda()
|
||||
|
||||
_assert_gen_expert_parallel_metadata(mot_linear)
|
||||
|
||||
_sync_weights(text_linear, vae_linear, mot_linear)
|
||||
|
||||
text_idx, vae_idx, M = _make_indices(image_num, _VAE_CHUNK_SIZE)
|
||||
@@ -555,13 +501,6 @@ def test_mot_qkv_und_mode(K: int, N: int, dtype: str, bias: bool):
|
||||
).cuda()
|
||||
|
||||
with torch.no_grad():
|
||||
# vLLM linear layers leave weights uninitialized (torch.empty), so
|
||||
# a freshly zeroed GPU page would make both outputs zero and the
|
||||
# cosine-similarity gate degenerate. Seed deterministic values so
|
||||
# und-mode is compared against a meaningful reference.
|
||||
ref_linear.weight.normal_(mean=0.0, std=0.02)
|
||||
if ref_linear.bias is not None:
|
||||
ref_linear.bias.normal_(mean=0.0, std=0.02)
|
||||
mot_linear.weight.copy_(ref_linear.weight)
|
||||
if bias and ref_linear.bias is not None:
|
||||
mot_linear.bias.copy_(ref_linear.bias)
|
||||
@@ -610,13 +549,6 @@ def test_mot_row_und_mode(K: int, N: int, dtype: str, bias: bool):
|
||||
).cuda()
|
||||
|
||||
with torch.no_grad():
|
||||
# vLLM linear layers leave weights uninitialized (torch.empty), so
|
||||
# a freshly zeroed GPU page would make both outputs zero and the
|
||||
# cosine-similarity gate degenerate. Seed deterministic values so
|
||||
# und-mode is compared against a meaningful reference.
|
||||
ref_linear.weight.normal_(mean=0.0, std=0.02)
|
||||
if ref_linear.bias is not None:
|
||||
ref_linear.bias.normal_(mean=0.0, std=0.02)
|
||||
mot_linear.weight.copy_(ref_linear.weight)
|
||||
if bias and ref_linear.bias is not None:
|
||||
mot_linear.bias.copy_(ref_linear.bias)
|
||||
|
||||
@@ -340,97 +340,6 @@ def test_rmsnorm_npu_receives_gamma_with_configured_dtype(monkeypatch: pytest.Mo
|
||||
assert captured["epsilon"].item() == pytest.approx(1e-5)
|
||||
|
||||
|
||||
# ── RMSNorm residual tests ──
|
||||
|
||||
|
||||
def test_rmsnorm_matches_native_residual_contract():
|
||||
"""The non-NPU path preserves separate residual-add and RMSNorm math."""
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
|
||||
eps = 1e-6
|
||||
norm = RMSNorm(4, eps=eps, dtype=torch.bfloat16)
|
||||
norm.weight.data.copy_(torch.tensor([1.0, 0.5, 1.5, 2.0], dtype=torch.bfloat16))
|
||||
x = torch.tensor([[1.0, -2.0, 3.0, -4.0]], dtype=torch.bfloat16)
|
||||
residual = torch.tensor([[0.00390625, 0.0078125, -0.0078125, 0.015625]], dtype=torch.bfloat16)
|
||||
|
||||
expected_residual = residual + x
|
||||
expected_fp32 = expected_residual.float()
|
||||
expected_output = expected_fp32 * torch.rsqrt(expected_fp32.square().mean(-1, keepdim=True) + eps)
|
||||
expected_output = (expected_output * norm.weight.float()).to(x.dtype)
|
||||
output, updated_residual = norm.forward_native(x, residual)
|
||||
|
||||
torch.testing.assert_close(updated_residual, expected_residual, atol=0, rtol=0)
|
||||
torch.testing.assert_close(output, expected_output, atol=0, rtol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"forward_name",
|
||||
["forward_cuda", "forward_hip", "forward_musa", "forward_xpu"],
|
||||
)
|
||||
def test_rmsnorm_non_npu_residual_delegates_to_native(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
forward_name: str,
|
||||
):
|
||||
"""Platform wrappers keep residual arithmetic in forward_native."""
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
|
||||
norm = RMSNorm(4, eps=1e-5, dtype=torch.bfloat16)
|
||||
x = torch.randn(2, 4, dtype=torch.bfloat16)
|
||||
residual = torch.randn_like(x)
|
||||
expected_output = torch.randn_like(x)
|
||||
expected_residual = torch.randn_like(x)
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def forward_native(native_x: torch.Tensor, native_residual: torch.Tensor | None = None):
|
||||
captured.update(x=native_x, residual=native_residual)
|
||||
return expected_output, expected_residual
|
||||
|
||||
monkeypatch.setattr(norm, "forward_native", forward_native)
|
||||
|
||||
output, updated_residual = getattr(norm, forward_name)(x, residual)
|
||||
|
||||
assert output is expected_output
|
||||
assert updated_residual is expected_residual
|
||||
assert captured["x"] is x
|
||||
assert captured["residual"] is residual
|
||||
|
||||
|
||||
def test_rmsnorm_npu_residual_uses_torch_npu_fused_op(monkeypatch: pytest.MonkeyPatch):
|
||||
"""NPU keeps the direct fused AddRMSNorm path instead of IR fallback."""
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
fused_output = torch.randn(2, 4, dtype=torch.bfloat16)
|
||||
fused_residual = torch.randn(2, 4, dtype=torch.bfloat16)
|
||||
|
||||
def npu_add_rms_norm(
|
||||
x: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
gamma: torch.Tensor,
|
||||
epsilon: float,
|
||||
):
|
||||
captured.update(x=x, residual=residual, gamma=gamma, epsilon=epsilon)
|
||||
return fused_output, torch.empty(0), fused_residual
|
||||
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"torch_npu",
|
||||
types.SimpleNamespace(npu_add_rms_norm=npu_add_rms_norm),
|
||||
)
|
||||
norm = RMSNorm(4, eps=1e-5, dtype=torch.bfloat16)
|
||||
x = torch.randn(2, 4, dtype=torch.bfloat16)
|
||||
residual = torch.randn(2, 4, dtype=torch.bfloat16)
|
||||
|
||||
output, updated_residual = norm.forward_npu(x, residual)
|
||||
|
||||
assert output is fused_output
|
||||
assert updated_residual is fused_residual
|
||||
assert captured["x"] is x
|
||||
assert captured["residual"] is residual
|
||||
assert captured["gamma"] is norm.weight
|
||||
assert captured["epsilon"] == 1e-5
|
||||
|
||||
|
||||
# ── RMSNorm compile-path regression tests ──
|
||||
|
||||
|
||||
|
||||
@@ -1,34 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.core_model,
|
||||
pytest.mark.diffusion,
|
||||
pytest.mark.cuda,
|
||||
pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="CUDA platform required"),
|
||||
]
|
||||
|
||||
|
||||
def test_rmsnorm_cuda_residual_matches_native_contract() -> None:
|
||||
eps = 1e-6
|
||||
device = torch.device(current_omni_platform.device_type)
|
||||
torch.manual_seed(0)
|
||||
norm = RMSNorm(64, eps=eps, dtype=torch.bfloat16).to(device)
|
||||
x = torch.randn(2, 4, 64, device=device, dtype=torch.bfloat16)
|
||||
residual = torch.randn_like(x)
|
||||
x_before = x.clone()
|
||||
residual_before = residual.clone()
|
||||
|
||||
expected_output, expected_residual = norm.forward_native(x_before, residual_before)
|
||||
output, updated_residual = norm.forward_cuda(x, residual)
|
||||
|
||||
torch.testing.assert_close(updated_residual, expected_residual, atol=0, rtol=0)
|
||||
torch.testing.assert_close(output, expected_output, atol=0, rtol=0)
|
||||
torch.testing.assert_close(x, x_before, atol=0, rtol=0)
|
||||
torch.testing.assert_close(residual, residual_before, atol=0, rtol=0)
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -89,29 +89,6 @@ def test_diffusion_base_linear_apply_multi_slice():
|
||||
assert torch.allclose(out, expected)
|
||||
|
||||
|
||||
def test_diffusion_lora_buffer_device_persists_across_reallocation(monkeypatch):
|
||||
from vllm.lora.layers.base_linear import BaseLinearLayerWithLoRA
|
||||
|
||||
layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
|
||||
torch.nn.Module.__init__(layer)
|
||||
layer.base_layer = torch.nn.Module()
|
||||
layer.n_slices = 1
|
||||
layer.lora_a_stacked = (torch.zeros(2),)
|
||||
layer.lora_b_stacked = (torch.zeros(3),)
|
||||
|
||||
layer._set_diffusion_lora_buffer_device(torch.device("meta"))
|
||||
assert all(tensor.is_meta for tensor in (*layer.lora_a_stacked, *layer.lora_b_stacked))
|
||||
|
||||
def recreate_on_cpu(self, max_loras, lora_config, model_config):
|
||||
self.lora_a_stacked = (torch.zeros(4),)
|
||||
self.lora_b_stacked = (torch.zeros(5),)
|
||||
|
||||
monkeypatch.setattr(BaseLinearLayerWithLoRA, "create_lora_weights", recreate_on_cpu)
|
||||
layer.create_lora_weights(max_loras=1, lora_config=object())
|
||||
|
||||
assert all(tensor.is_meta for tensor in (*layer.lora_a_stacked, *layer.lora_b_stacked))
|
||||
|
||||
|
||||
def test_diffusion_base_linear_reset_lora_disables_fast_path(monkeypatch):
|
||||
# Verify that after reset_lora(), apply() skips LoRA matmuls even if the
|
||||
# LoRA tensors are still allocated and non-empty.
|
||||
|
||||
@@ -3,8 +3,6 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.lora.lora_weights import LoRALayerWeights
|
||||
@@ -145,37 +143,6 @@ def test_lora_manager_replace_layers_does_not_rewrap_base_layer(monkeypatch):
|
||||
assert replace_calls == ["foo"]
|
||||
|
||||
|
||||
def test_lora_manager_keeps_dlo_lora_buffers_on_compute_device(monkeypatch):
|
||||
import vllm_omni.diffusion.lora.manager as manager_mod
|
||||
|
||||
class _DLOLoRALayer(DummyBaseLayerWithLoRA):
|
||||
lora_buffer_device = None
|
||||
|
||||
def _set_diffusion_lora_buffer_device(self, device):
|
||||
self.lora_buffer_device = device
|
||||
|
||||
monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DLOLoRALayer)
|
||||
monkeypatch.setattr(
|
||||
manager_mod,
|
||||
"from_layer_diffusion",
|
||||
lambda *, layer, **_kwargs: _DLOLoRALayer(layer),
|
||||
)
|
||||
|
||||
pipeline = torch.nn.Module()
|
||||
pipeline.od_config = SimpleNamespace(enable_distributed_layerwise_offload=True)
|
||||
pipeline.transformer = torch.nn.Module()
|
||||
pipeline.transformer.foo = _FakeLinearBase()
|
||||
manager = DiffusionLoRAManager(
|
||||
pipeline=pipeline,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
manager._replace_layers_with_lora(type("_PH", (), {"r": 1})())
|
||||
|
||||
assert pipeline.transformer.foo.lora_buffer_device == torch.device("cpu")
|
||||
|
||||
|
||||
def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch):
|
||||
import vllm_omni.diffusion.lora.manager as manager_mod
|
||||
|
||||
|
||||
@@ -249,30 +249,6 @@ def test_hwr_commit_failure_discards_model_and_reloads_without_hwr_or_mmap(tmp_p
|
||||
assert loader.take_host_weight_plan() is None
|
||||
|
||||
|
||||
def test_required_hwr_miss_fails_before_ordinary_loading_or_publication(
|
||||
tmp_path: Path,
|
||||
monkeypatch,
|
||||
):
|
||||
canonical_root = tmp_path / "canonical"
|
||||
canonical_root.mkdir()
|
||||
save_file(
|
||||
{"weight": torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)},
|
||||
str(canonical_root / "model.safetensors"),
|
||||
)
|
||||
loader = DiffusersPipelineLoader(
|
||||
LoadConfig(),
|
||||
_hwr_config(canonical_root, tmp_path / "empty-store", mode="required"),
|
||||
)
|
||||
pipeline = _HWRPipeline(canonical_root)
|
||||
monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: pipeline)
|
||||
|
||||
with pytest.raises(RuntimeError, match="Host Weight Runtime resolution failed"):
|
||||
loader.load_model(load_device="cpu", device=torch.device("cpu"))
|
||||
|
||||
assert pipeline.load_count == 0
|
||||
assert loader.take_host_weight_plan() is None
|
||||
|
||||
|
||||
def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
|
||||
return OmniDiffusionConfig(
|
||||
model="",
|
||||
@@ -285,11 +261,10 @@ def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("dist_offload", "use_allgather", "mode"),
|
||||
("dist_offload", "use_allgather"),
|
||||
[
|
||||
(False, False, "preferred"),
|
||||
(True, True, "preferred"),
|
||||
(True, False, "disabled"),
|
||||
(False, False),
|
||||
(True, True),
|
||||
],
|
||||
)
|
||||
def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction(
|
||||
@@ -297,13 +272,12 @@ def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction(
|
||||
tmp_path,
|
||||
dist_offload,
|
||||
use_allgather,
|
||||
mode,
|
||||
):
|
||||
"""Disabled and AllGather paths must never construct or probe HWR."""
|
||||
from vllm_omni.host_weight_runtime import HostWeightRuntime
|
||||
|
||||
root = tmp_path / "must-not-be-touched"
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config("dummy-model", root, mode=mode))
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config("dummy-model", root))
|
||||
model = _DummyPipelineModel(source_prefix="transformer.")
|
||||
modules = SimpleNamespace(dit_names=("transformer",), dits=(model.transformer,))
|
||||
|
||||
|
||||
@@ -303,83 +303,6 @@ class _FakeH3Attention(nn.Module):
|
||||
super().__init__()
|
||||
|
||||
|
||||
class _FakeFluxLinear(_FakeH3Linear):
|
||||
def __init__(
|
||||
self,
|
||||
*args: object,
|
||||
bias: bool = False,
|
||||
params_dtype: torch.dtype | None = None,
|
||||
total_num_heads: int | None = None,
|
||||
total_num_kv_heads: int | None = None,
|
||||
**kwargs: object,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
*args,
|
||||
bias=bias,
|
||||
params_dtype=params_dtype,
|
||||
total_num_heads=total_num_heads,
|
||||
total_num_kv_heads=total_num_kv_heads,
|
||||
**kwargs,
|
||||
)
|
||||
self.weight.weight_loader = self._load_weight
|
||||
|
||||
@staticmethod
|
||||
def _load_weight(
|
||||
parameter: nn.Parameter,
|
||||
loaded_weight: torch.Tensor,
|
||||
shard_id: str | None = None,
|
||||
) -> None:
|
||||
del shard_id
|
||||
parameter.data.copy_(loaded_weight.to(parameter).reshape_as(parameter))
|
||||
|
||||
|
||||
class _FakeFluxAttention(nn.Module):
|
||||
def __init__(self, **kwargs: object) -> None:
|
||||
del kwargs
|
||||
super().__init__()
|
||||
self.register_buffer("beta", torch.tensor([0.5]))
|
||||
self.register_buffer("eps", torch.tensor([-1e-6]))
|
||||
|
||||
|
||||
class _Flux2KleinPipeline(nn.Module):
|
||||
def __init__(self, transformer: nn.Module) -> None:
|
||||
super().__init__()
|
||||
self.transformer = transformer
|
||||
self.text_encoder = nn.Linear(2, 2, dtype=torch.bfloat16)
|
||||
self.vae = nn.Module()
|
||||
self.vae.register_buffer("gain", torch.tensor([9.0], dtype=torch.float32))
|
||||
|
||||
|
||||
def _reduced_flux2_klein(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
*,
|
||||
device: torch.device | str = "cpu",
|
||||
) -> nn.Module:
|
||||
from vllm_omni.diffusion.models.flux2_klein import flux2_klein_transformer as flux2
|
||||
|
||||
monkeypatch.setattr(flux2, "ColumnParallelLinear", _FakeFluxLinear)
|
||||
monkeypatch.setattr(flux2, "MergedColumnParallelLinear", _FakeFluxLinear)
|
||||
monkeypatch.setattr(flux2, "QKVParallelLinear", _FakeFluxLinear)
|
||||
monkeypatch.setattr(flux2, "RowParallelLinear", _FakeFluxLinear)
|
||||
monkeypatch.setattr(flux2, "Attention", _FakeFluxAttention)
|
||||
|
||||
transformer = flux2.Flux2Transformer2DModel(
|
||||
patch_size=1,
|
||||
in_channels=8,
|
||||
out_channels=8,
|
||||
num_layers=1,
|
||||
num_single_layers=1,
|
||||
attention_head_dim=8,
|
||||
num_attention_heads=1,
|
||||
joint_attention_dim=8,
|
||||
timestep_guidance_channels=8,
|
||||
mlp_ratio=1.0,
|
||||
axes_dims_rope=(2, 2, 2, 2),
|
||||
guidance_embeds=False,
|
||||
)
|
||||
return transformer.to(device=device, dtype=torch.bfloat16)
|
||||
|
||||
|
||||
class _H3Pipeline(nn.Module):
|
||||
def __init__(self, transformer: nn.Module) -> None:
|
||||
super().__init__()
|
||||
@@ -1277,152 +1200,3 @@ def test_minimax_h3_declares_the_final_layout_restore_contract() -> None:
|
||||
assert contract.implementation_id == "minimax-h3-dit"
|
||||
assert callable(MiniMaxH3DiTModel.validate_restored_host_weights)
|
||||
assert "data_parallel" not in FinalLayoutRequest.__dataclass_fields__
|
||||
|
||||
|
||||
def test_reduced_flux2_klein_declares_complete_final_layout_contract(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
|
||||
Flux2Transformer2DModel,
|
||||
)
|
||||
|
||||
transformer = _reduced_flux2_klein(monkeypatch)
|
||||
pipeline = _Flux2KleinPipeline(transformer)
|
||||
source = _prepared_source(tmp_path, directory="flux2-klein-source")
|
||||
context = _identity(pipeline, source)
|
||||
targets = collect_final_layout_targets(
|
||||
pipeline,
|
||||
(("transformer", transformer),),
|
||||
policy=FINAL_LAYOUT_BF16_POLICY,
|
||||
require_materialized=True,
|
||||
)
|
||||
target_names = {target.name for target in targets}
|
||||
|
||||
contract = Flux2Transformer2DModel.host_weight_restore_contract
|
||||
assert isinstance(contract, FinalLayoutModelContract)
|
||||
assert contract.schema == FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA
|
||||
assert contract.implementation_id == "flux2-klein-dit"
|
||||
assert contract.version == "1"
|
||||
assert context.model_contracts == (contract,)
|
||||
assert transformer.stacked_params_mapping == [
|
||||
(".to_qkv.", ".to_q.", "q"),
|
||||
(".to_qkv.", ".to_k.", "k"),
|
||||
(".to_qkv.", ".to_v.", "v"),
|
||||
(".add_kv_proj", ".add_q_proj", "q"),
|
||||
(".add_kv_proj", ".add_k_proj", "k"),
|
||||
(".add_kv_proj", ".add_v_proj", "v"),
|
||||
]
|
||||
loaded = transformer.load_weights(
|
||||
[
|
||||
("transformer_blocks.0.attn.to_q.weight", torch.tensor([1.0], dtype=torch.bfloat16)),
|
||||
("transformer_blocks.0.attn.add_q_proj.weight", torch.tensor([2.0], dtype=torch.bfloat16)),
|
||||
(
|
||||
"single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight",
|
||||
torch.tensor([3.0], dtype=torch.bfloat16),
|
||||
),
|
||||
]
|
||||
)
|
||||
assert loaded == {
|
||||
"transformer_blocks.0.attn.to_qkv.weight",
|
||||
"transformer_blocks.0.attn.add_kv_proj.weight",
|
||||
"single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight",
|
||||
}
|
||||
assert any(name.startswith("transformer.transformer_blocks.0.") for name in target_names)
|
||||
assert any(name.startswith("transformer.single_transformer_blocks.0.") for name in target_names)
|
||||
assert "transformer.transformer_blocks.0.attn.to_qkv.weight" in target_names
|
||||
assert "transformer.transformer_blocks.0.attn.add_kv_proj.weight" in target_names
|
||||
assert "transformer.single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight" in target_names
|
||||
assert "transformer.transformer_blocks.0.attn.attn.beta" in target_names
|
||||
assert "transformer.transformer_blocks.0.attn.attn.eps" in target_names
|
||||
assert all("text_encoder" not in name and "vae" not in name for name in target_names)
|
||||
transformer.validate_restored_host_weights()
|
||||
|
||||
|
||||
def test_flux2_klein_restore_validator_rejects_non_ready_state(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
transformer = _reduced_flux2_klein(monkeypatch)
|
||||
|
||||
mapping = transformer.stacked_params_mapping
|
||||
transformer.stacked_params_mapping = None
|
||||
with pytest.raises(ValueError, match="packed QKV mapping"):
|
||||
transformer.validate_restored_host_weights()
|
||||
transformer.stacked_params_mapping = mapping
|
||||
|
||||
weight = transformer.x_embedder.weight
|
||||
weight.data = weight.data.float()
|
||||
with pytest.raises(ValueError, match="must stay bf16"):
|
||||
transformer.validate_restored_host_weights()
|
||||
weight.data = weight.data.to(torch.bfloat16)
|
||||
|
||||
weight.data = weight.data.T
|
||||
with pytest.raises(ValueError, match="contiguous strided layout"):
|
||||
transformer.validate_restored_host_weights()
|
||||
weight.data = weight.data.contiguous()
|
||||
|
||||
attention = transformer.transformer_blocks[0].attn.attn
|
||||
attention._non_persistent_buffers_set.add("beta")
|
||||
with pytest.raises(ValueError, match="must be persistent"):
|
||||
transformer.validate_restored_host_weights()
|
||||
attention._non_persistent_buffers_set.remove("beta")
|
||||
|
||||
attention.beta.data = attention.beta.data.float()
|
||||
with pytest.raises(ValueError, match="scalar bf16"):
|
||||
transformer.validate_restored_host_weights()
|
||||
|
||||
|
||||
def test_flux2_klein_publication_restore_fallback_and_teardown(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
source = _prepared_source(tmp_path, directory="flux2-klein-source")
|
||||
cold_transformer = _reduced_flux2_klein(monkeypatch)
|
||||
cold_model = _Flux2KleinPipeline(cold_transformer)
|
||||
with torch.no_grad():
|
||||
for index, parameter in enumerate(cold_transformer.parameters(), start=1):
|
||||
parameter.fill_(index)
|
||||
for name, buffer in cold_transformer.named_buffers():
|
||||
if name.endswith((".beta", ".eps")):
|
||||
buffer.fill_(0.5 if name.endswith(".beta") else -1e-6)
|
||||
|
||||
context = _identity(cold_model, source)
|
||||
runtime = _runtime(tmp_path / "flux2-klein-store")
|
||||
miss = runtime.resolve(context.identity)
|
||||
assert miss.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK
|
||||
|
||||
producer = FinalLayoutBF16Producer(
|
||||
context,
|
||||
cold_model,
|
||||
(("transformer", cold_transformer),),
|
||||
max_shard_bytes=64,
|
||||
)
|
||||
publication = runtime.publish_after_load(context.identity, producer=producer)
|
||||
assert publication.outcome is PostLoadPublicationOutcome.PUBLISHED
|
||||
|
||||
warm_transformer = _reduced_flux2_klein(monkeypatch, device="meta")
|
||||
warm_model = _Flux2KleinPipeline(warm_transformer)
|
||||
warm_context = _identity(warm_model, source)
|
||||
assert warm_context.identity == context.identity
|
||||
hit = runtime.resolve(warm_context.identity)
|
||||
assert hit.report.outcome is ResolutionOutcome.LOCAL_HIT
|
||||
assert hit.lease is not None
|
||||
|
||||
plan = FinalLayoutTensorRestorer(warm_context).plan_restore(warm_model, hit.lease)
|
||||
assert all(tensor.is_meta for name, tensor in warm_model.named_parameters() if name.startswith("transformer."))
|
||||
plan.commit()
|
||||
warm_transformer.validate_restored_host_weights()
|
||||
|
||||
cold_tensors = dict(cold_model.named_parameters()) | dict(cold_model.named_buffers())
|
||||
warm_tensors = dict(warm_model.named_parameters()) | dict(warm_model.named_buffers())
|
||||
for name in warm_context.tensor_names:
|
||||
assert torch.equal(warm_tensors[name], cold_tensors[name])
|
||||
assert warm_tensors[name].untyped_storage().data_ptr() == hit.lease.tensors[name].untyped_storage().data_ptr()
|
||||
assert warm_transformer.stacked_params_mapping == cold_transformer.stacked_params_mapping
|
||||
|
||||
del warm_tensors, warm_model, warm_transformer, plan
|
||||
gc.collect()
|
||||
hit.lease.close()
|
||||
assert hit.lease.closed
|
||||
assert isinstance(runtime.store, FilesystemHostWeightStore)
|
||||
assert runtime.store.cleanup(context.identity) is None
|
||||
@@ -0,0 +1,27 @@
|
||||
import pytest
|
||||
import torch.nn as nn
|
||||
|
||||
pytest.importorskip("dreamid_omni.modules.model")
|
||||
|
||||
from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_dreamid_omni_exposes_hsdp_shard_conditions_for_fused_blocks():
|
||||
model = object.__new__(FusionModel)
|
||||
nn.Module.__init__(model)
|
||||
model.fused_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
|
||||
model.attn = nn.Identity()
|
||||
|
||||
conditions = getattr(model, "_hsdp_shard_conditions", None)
|
||||
|
||||
assert conditions is not None
|
||||
assert len(conditions) == 1
|
||||
|
||||
matched = []
|
||||
for name, module in model.named_modules():
|
||||
if any(cond(name, module) for cond in conditions):
|
||||
matched.append(name)
|
||||
|
||||
assert matched == ["fused_blocks.0", "fused_blocks.1"]
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -245,7 +245,7 @@ else:
|
||||
[sys.executable, "-c", script, str(_MODULE_PATH), str(tmp_path)],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
timeout=3,
|
||||
check=False,
|
||||
)
|
||||
|
||||
|
||||
@@ -11,31 +11,6 @@ import torch
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch):
|
||||
import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
|
||||
|
||||
class FakeBaseVocoder(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.input_dtype = None
|
||||
|
||||
def forward(self, value):
|
||||
self.input_dtype = value.dtype
|
||||
return value
|
||||
|
||||
monkeypatch.setattr(
|
||||
ltx_runtime.torch,
|
||||
"autocast",
|
||||
lambda **_kwargs: (_ for _ in ()).throw(AssertionError("base vocoder must not use autocast")),
|
||||
)
|
||||
vocoder = FakeBaseVocoder()
|
||||
|
||||
output = ltx_runtime._run_ltx_vocoder(vocoder, torch.ones(1, dtype=torch.bfloat16))
|
||||
|
||||
assert vocoder.input_dtype == torch.bfloat16
|
||||
assert output.dtype == torch.bfloat16
|
||||
|
||||
|
||||
class TestLTXOutputRank:
|
||||
@pytest.mark.parametrize(
|
||||
("distributed_vae_state", "expected_decode_calls"),
|
||||
|
||||
@@ -1,48 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""CUDA regression tests for LTX-2 BWE vocoder precision."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
||||
|
||||
|
||||
class _BWEConvVocoder(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.conv = torch.nn.Conv1d(
|
||||
1,
|
||||
1,
|
||||
kernel_size=1,
|
||||
bias=False,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
self.bwe_generator = torch.nn.Identity()
|
||||
self.input_dtype = None
|
||||
self.conv_output_dtype = None
|
||||
|
||||
def forward(self, value):
|
||||
self.input_dtype = value.dtype
|
||||
output = self.conv(value)
|
||||
self.conv_output_dtype = output.dtype
|
||||
return output
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_ltx_bwe_vocoder_runs_real_cuda_autocast_in_fp32():
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_runtime import _run_ltx_vocoder
|
||||
|
||||
vocoder = _BWEConvVocoder()
|
||||
generated_mel = torch.ones((1, 1, 4), device="cuda", dtype=torch.bfloat16)
|
||||
|
||||
output = _run_ltx_vocoder(vocoder, generated_mel)
|
||||
|
||||
assert next(vocoder.parameters()).dtype == torch.bfloat16
|
||||
assert vocoder.input_dtype == torch.float32
|
||||
assert vocoder.conv_output_dtype == torch.float32
|
||||
assert output.dtype == torch.bfloat16
|
||||
@@ -964,80 +964,6 @@ def test_minimax_h3_advertises_the_official_ref2va_image_limit():
|
||||
assert get_diffusion_model_metadata("MiniMaxH3Pipeline").max_multimodal_image_inputs == 9
|
||||
|
||||
|
||||
def test_text_attention_routes_local_gqa_heads_through_sdpa_helper(monkeypatch):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
|
||||
|
||||
class FakeEncoderGroup:
|
||||
rank_in_group = 0
|
||||
world_size = 2
|
||||
|
||||
def all_reduce(self, tensor):
|
||||
del tensor
|
||||
|
||||
config = SimpleNamespace(
|
||||
hidden_size=8,
|
||||
num_attention_heads=4,
|
||||
num_key_value_heads=2,
|
||||
head_dim=2,
|
||||
rms_norm_eps=1e-6,
|
||||
)
|
||||
with (
|
||||
patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0),
|
||||
patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1),
|
||||
):
|
||||
attention = encoder_module.MiniMaxH3Qwen3VLTextAttention(FakeEncoderGroup(), config, torch.float32)
|
||||
# This CPU contract test only needs projection shapes and head reshaping.
|
||||
# Avoid dispatching platform-specific norm and linear kernels (for example,
|
||||
# rocm_unquantized_gemm when the test suite is collected on ROCm).
|
||||
attention.q_norm = nn.Identity()
|
||||
attention.k_norm = nn.Identity()
|
||||
with torch.no_grad():
|
||||
attention.qkv_proj.weight.zero_()
|
||||
attention.o_proj = nn.Linear(
|
||||
attention.qkv_proj.local_num_heads * attention.head_dim,
|
||||
config.hidden_size,
|
||||
bias=False,
|
||||
)
|
||||
attn_call = {}
|
||||
|
||||
def fake_attention(query, key, value):
|
||||
attn_call.update(query_shape=query.shape, key_shape=key.shape, value_shape=value.shape)
|
||||
return query
|
||||
|
||||
monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", fake_attention)
|
||||
hidden_states = torch.randn(1, 3, config.hidden_size)
|
||||
cos = torch.ones(1, 3, config.head_dim)
|
||||
sin = torch.zeros_like(cos)
|
||||
|
||||
output = attention(hidden_states, (cos, sin))
|
||||
|
||||
assert attn_call["query_shape"] == (1, config.num_attention_heads // 2, 3, config.head_dim)
|
||||
assert attn_call["key_shape"] == (1, config.num_key_value_heads // 2, 3, config.head_dim)
|
||||
assert attn_call["value_shape"] == (1, config.num_key_value_heads // 2, 3, config.head_dim)
|
||||
assert output.shape == hidden_states.shape
|
||||
|
||||
|
||||
def test_text_attention_sdpa_helper_preserves_expanded_kv_fallback(monkeypatch):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
|
||||
|
||||
sdpa_call = {}
|
||||
|
||||
def fake_sdpa(query, key, value, **kwargs):
|
||||
sdpa_call.update(query_shape=query.shape, key_shape=key.shape, value_shape=value.shape, kwargs=kwargs)
|
||||
return query
|
||||
|
||||
monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa)
|
||||
output = encoder_module._scaled_dot_product_attention(
|
||||
torch.randn(1, 4, 3, 8),
|
||||
torch.randn(1, 2, 3, 8),
|
||||
torch.randn(1, 2, 3, 8),
|
||||
)
|
||||
|
||||
assert sdpa_call["query_shape"] == sdpa_call["key_shape"] == sdpa_call["value_shape"] == (1, 4, 3, 8)
|
||||
assert sdpa_call["kwargs"] == {"dropout_p": 0.0, "is_causal": True}
|
||||
assert output.shape == (1, 4, 3, 8)
|
||||
|
||||
|
||||
def test_encoder_forward_uses_hook_compatible_encode_entrypoint():
|
||||
from vllm_omni.diffusion.models.minimax_h3.encoder import (
|
||||
MiniMaxH3Qwen3VLEncoder,
|
||||
|
||||
@@ -99,3 +99,22 @@ def test_manual_component_offload_failure_forces_retained_cache_release(mocker):
|
||||
assert component.load_to_device.call_count == 2
|
||||
assert component.offload_to_cpu.call_count == 2
|
||||
pipeline._dlo_component_cache.release_if_needed.assert_called_once_with(force=True)
|
||||
|
||||
|
||||
def test_resident_component_skips_manual_dlo_lifecycle(mocker):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.od_config = mocker.Mock()
|
||||
pipeline.od_config.enable_layerwise_offload = False
|
||||
pipeline.od_config.enable_distributed_layerwise_offload = True
|
||||
pipeline._model_cpu_offload_modules = []
|
||||
component = mocker.Mock()
|
||||
component._omni_dlo_offload_enabled = False
|
||||
|
||||
with pipeline._component_on_device(component):
|
||||
pass
|
||||
|
||||
component.load_to_device.assert_not_called()
|
||||
component.offload_to_cpu.assert_not_called()
|
||||
@@ -197,6 +197,7 @@ def test_h3_turbo_rejects_wrong_alpha_and_ref2va(tmp_path):
|
||||
[
|
||||
"model-level CPU offload (--enable-cpu-offload)",
|
||||
"layerwise offload (--enable-layerwise-offload)",
|
||||
"distributed layerwise offload (--enable-distributed-layerwise-offload)",
|
||||
],
|
||||
)
|
||||
def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode):
|
||||
@@ -213,38 +214,6 @@ def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode):
|
||||
)
|
||||
|
||||
|
||||
def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
from vllm_omni.diffusion.models.minimax_h3 import (
|
||||
pipeline_minimax_h3 as pipeline_module,
|
||||
)
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.partition = "fl2va"
|
||||
pipeline._turbo_lora_adapter_ids = set()
|
||||
pipeline.od_config = SimpleNamespace(
|
||||
enable_cpu_offload=False,
|
||||
enable_layerwise_offload=False,
|
||||
enable_distributed_layerwise_offload=True,
|
||||
)
|
||||
captured = {}
|
||||
|
||||
def load_turbo(**kwargs):
|
||||
captured.update(kwargs)
|
||||
return object(), object()
|
||||
|
||||
monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", load_turbo)
|
||||
loaded = pipeline._load_diffusion_lora_adapter(
|
||||
lora_request=_request("turbo"),
|
||||
lora_path="turbo",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
assert loaded is not None
|
||||
assert captured["unsupported_offload_mode"] is None
|
||||
|
||||
|
||||
def test_h3_turbo_accepts_only_the_declared_v1_artifact(tmp_path):
|
||||
unsupported = tmp_path / "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors"
|
||||
_write_tiny_turbo(unsupported)
|
||||
|
||||
@@ -1,132 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
|
||||
def test_minimax_h3_npu_sdpa_patch_is_independent_and_idempotent(monkeypatch) -> None:
|
||||
from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
|
||||
from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
|
||||
|
||||
def sdpa_sentinel(query, key, value):
|
||||
del key, value
|
||||
return query
|
||||
|
||||
def rope_sentinel(q, k, cos, sin):
|
||||
del cos, sin
|
||||
return q, k
|
||||
|
||||
monkeypatch.setattr(npu_minimax_h3, "_ROPE_PATCHED", False)
|
||||
monkeypatch.setattr(npu_minimax_h3, "_SDPA_PATCHED", False)
|
||||
monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel)
|
||||
monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel)
|
||||
|
||||
npu_minimax_h3.apply_minimax_h3_qwen3vl_sdpa_patch()
|
||||
|
||||
assert encoder_module._scaled_dot_product_attention is npu_minimax_h3._scaled_dot_product_attention_npu
|
||||
assert encoder_module._apply_rotary_pos_emb is rope_sentinel
|
||||
assert not npu_minimax_h3._ROPE_PATCHED
|
||||
assert npu_minimax_h3._SDPA_PATCHED
|
||||
|
||||
monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel)
|
||||
npu_minimax_h3.apply_minimax_h3_qwen3vl_sdpa_patch()
|
||||
|
||||
assert encoder_module._scaled_dot_product_attention is sdpa_sentinel
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
|
||||
@pytest.mark.parametrize(
|
||||
("num_heads", "num_kv_heads", "expected_enable_gqa"),
|
||||
[(2, 2, False), (4, 2, True)],
|
||||
)
|
||||
def test_minimax_h3_npu_sdpa_preserves_compressed_heads_and_selects_gqa(
|
||||
monkeypatch,
|
||||
num_heads: int,
|
||||
num_kv_heads: int,
|
||||
expected_enable_gqa: bool,
|
||||
) -> None:
|
||||
from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_sdpa(query, key, value, **kwargs):
|
||||
captured.update(query=query, key=key, value=value, kwargs=kwargs)
|
||||
return query
|
||||
|
||||
monkeypatch.setattr(F, "scaled_dot_product_attention", fake_sdpa)
|
||||
query = torch.randn(1, num_heads, 3, 8)
|
||||
key = torch.randn(1, num_kv_heads, 3, 8)
|
||||
value = torch.randn_like(key)
|
||||
|
||||
output = npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value)
|
||||
|
||||
assert captured["query"] is query
|
||||
assert captured["key"] is key
|
||||
assert captured["value"] is value
|
||||
assert captured["kwargs"] == {
|
||||
"dropout_p": 0.0,
|
||||
"is_causal": True,
|
||||
"enable_gqa": expected_enable_gqa,
|
||||
}
|
||||
assert output is query
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
|
||||
@pytest.mark.parametrize(
|
||||
("num_heads", "num_key_heads", "num_value_heads", "message"),
|
||||
[
|
||||
(4, 2, 1, "key and value to have the same number of heads"),
|
||||
(4, 0, 0, "at least one KV head"),
|
||||
(3, 2, 2, "query heads to be a multiple of KV heads"),
|
||||
],
|
||||
)
|
||||
def test_minimax_h3_npu_sdpa_rejects_invalid_head_configurations(
|
||||
num_heads: int,
|
||||
num_key_heads: int,
|
||||
num_value_heads: int,
|
||||
message: str,
|
||||
) -> None:
|
||||
from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
|
||||
|
||||
query = torch.randn(1, num_heads, 3, 8)
|
||||
key = torch.randn(1, num_key_heads, 3, 8)
|
||||
value = torch.randn(1, num_value_heads, 3, 8)
|
||||
|
||||
with pytest.raises(ValueError, match=message):
|
||||
npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
|
||||
def test_minimax_h3_npu_causal_gqa_matches_expanded_kv() -> None:
|
||||
"""Exercise the real torch-npu causal GQA operator."""
|
||||
from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
|
||||
|
||||
torch.manual_seed(0)
|
||||
batch_size, seq_len = 1, 16
|
||||
num_heads, num_kv_heads, head_size = 8, 2, 128
|
||||
device = torch.device("npu")
|
||||
dtype = torch.bfloat16
|
||||
|
||||
query = torch.randn(batch_size, num_heads, seq_len, head_size, device=device, dtype=dtype)
|
||||
key = torch.randn(batch_size, num_kv_heads, seq_len, head_size, device=device, dtype=dtype)
|
||||
value = torch.randn_like(key)
|
||||
|
||||
actual = npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value)
|
||||
|
||||
repeat_num = num_heads // num_kv_heads
|
||||
expected = F.scaled_dot_product_attention(
|
||||
query,
|
||||
key.repeat_interleave(repeat_num, dim=1),
|
||||
value.repeat_interleave(repeat_num, dim=1),
|
||||
dropout_p=0.0,
|
||||
is_causal=True,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
|
||||
@@ -6,38 +6,7 @@ import torch
|
||||
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
|
||||
def test_qwen3_vl_npu_rope_patch_is_independent_and_idempotent(monkeypatch) -> None:
|
||||
from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
|
||||
from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
|
||||
|
||||
def rope_sentinel(q, k, cos, sin):
|
||||
del cos, sin
|
||||
return q, k
|
||||
|
||||
def sdpa_sentinel(query, key, value):
|
||||
del key, value
|
||||
return query
|
||||
|
||||
monkeypatch.setattr(npu_minimax_h3, "_ROPE_PATCHED", False)
|
||||
monkeypatch.setattr(npu_minimax_h3, "_SDPA_PATCHED", False)
|
||||
monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel)
|
||||
monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel)
|
||||
|
||||
npu_minimax_h3.apply_minimax_h3_qwen3vl_patch()
|
||||
|
||||
assert encoder_module._apply_rotary_pos_emb is npu_minimax_h3._apply_rotary_pos_emb_npu
|
||||
assert encoder_module._scaled_dot_product_attention is sdpa_sentinel
|
||||
assert npu_minimax_h3._ROPE_PATCHED
|
||||
assert not npu_minimax_h3._SDPA_PATCHED
|
||||
|
||||
monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel)
|
||||
npu_minimax_h3.apply_minimax_h3_qwen3vl_patch()
|
||||
|
||||
assert encoder_module._apply_rotary_pos_emb is rope_sentinel
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
|
||||
|
||||
@@ -1,20 +1,10 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
from vllm_omni.diffusion.models.minimax_h3.encoder import (
|
||||
MiniMaxH3Qwen3VLRMSNorm,
|
||||
MiniMaxH3Qwen3VLTextDecoderLayer,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
from vllm_omni.diffusion.models.minimax_h3.encoder import MiniMaxH3Qwen3VLRMSNorm
|
||||
|
||||
|
||||
def test_qwen3_vl_rmsnorm_uses_common_fused_rmsnorm_contract() -> None:
|
||||
@@ -34,43 +24,3 @@ def test_qwen3_vl_rmsnorm_uses_common_fused_rmsnorm_contract() -> None:
|
||||
assert norm.weight.dtype == torch.bfloat16
|
||||
assert set(norm.state_dict()) == {"weight"}
|
||||
torch.testing.assert_close(norm.forward_native(x), expected, atol=0, rtol=0)
|
||||
|
||||
|
||||
def test_qwen3_vl_decoder_preserves_post_attention_residual_contract() -> None:
|
||||
class ConstantAttention(nn.Module):
|
||||
def forward(self, hidden_states, position_embeddings):
|
||||
del position_embeddings
|
||||
return torch.full_like(hidden_states, 2.0)
|
||||
|
||||
config = SimpleNamespace(
|
||||
hidden_size=8,
|
||||
intermediate_size=16,
|
||||
num_attention_heads=4,
|
||||
num_key_value_heads=2,
|
||||
head_dim=2,
|
||||
rms_norm_eps=1e-6,
|
||||
)
|
||||
group = SimpleNamespace(rank_in_group=0, world_size=1)
|
||||
with (
|
||||
patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0),
|
||||
patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1),
|
||||
):
|
||||
layer = MiniMaxH3Qwen3VLTextDecoderLayer(group, config, torch.float32)
|
||||
layer.input_layernorm = nn.Identity()
|
||||
layer.self_attn = ConstantAttention()
|
||||
layer.mlp = nn.Identity()
|
||||
|
||||
hidden_states = torch.arange(24, dtype=torch.float32).reshape(1, 3, 8)
|
||||
attention_output = torch.full_like(hidden_states, 2.0)
|
||||
expected_residual = hidden_states + attention_output
|
||||
expected_normalized = RMSNorm.forward_native(layer.post_attention_layernorm, expected_residual)
|
||||
|
||||
with patch.object(
|
||||
layer.post_attention_layernorm,
|
||||
"forward",
|
||||
wraps=layer.post_attention_layernorm.forward_native,
|
||||
):
|
||||
output = layer(hidden_states, (torch.empty(0), torch.empty(0)))
|
||||
|
||||
assert isinstance(layer.post_attention_layernorm, MiniMaxH3Qwen3VLRMSNorm)
|
||||
torch.testing.assert_close(output, expected_residual + expected_normalized)
|
||||
@@ -0,0 +1,142 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
|
||||
from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
|
||||
DistributedLayerwiseOffloadBackend,
|
||||
)
|
||||
from vllm_omni.diffusion.offloader.module_collector import PipelineModules
|
||||
from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
|
||||
|
||||
|
||||
def _od_config(**overrides):
|
||||
values = {
|
||||
"enable_cpu_offload": False,
|
||||
"enable_layerwise_offload": False,
|
||||
"enable_distributed_layerwise_offload": True,
|
||||
"dlo_use_allgather": False,
|
||||
"dlo_resident_layers": 0,
|
||||
"dlo_offload_components": {},
|
||||
"pin_cpu_memory": False,
|
||||
"parallel_config": SimpleNamespace(
|
||||
use_hsdp=False,
|
||||
data_parallel_size=1,
|
||||
sequence_parallel_size=1,
|
||||
),
|
||||
}
|
||||
values.update(overrides)
|
||||
return SimpleNamespace(**values)
|
||||
|
||||
|
||||
def _backend(policy: dict[str, bool]) -> DistributedLayerwiseOffloadBackend:
|
||||
backend = object.__new__(DistributedLayerwiseOffloadBackend)
|
||||
backend.config = OffloadConfig(
|
||||
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
|
||||
pin_cpu_memory=False,
|
||||
dlo_use_allgather=False,
|
||||
dlo_offload_components=policy,
|
||||
)
|
||||
backend.device = torch.device("cpu")
|
||||
return backend
|
||||
|
||||
|
||||
def _modules() -> PipelineModules:
|
||||
return PipelineModules(
|
||||
dits=[],
|
||||
dit_names=[],
|
||||
encoders=[nn.Module()],
|
||||
encoder_names=["text_encoder"],
|
||||
vaes=[nn.Module(), nn.Module()],
|
||||
vae_names=["video_vae", "audio_vae"],
|
||||
)
|
||||
|
||||
|
||||
def test_component_policy_resolves_exact_then_default():
|
||||
backend = _backend({"text_encoder": False, "default": True})
|
||||
|
||||
assert not backend._component_offload_enabled("text_encoder")
|
||||
assert backend._component_offload_enabled("video_vae")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("value", "error", "match"),
|
||||
[
|
||||
([], TypeError, "must be a mapping"),
|
||||
(["text_encoder"], TypeError, "must be a mapping"),
|
||||
(False, TypeError, "must be a mapping"),
|
||||
({1: False}, TypeError, "keys must be non-empty strings"),
|
||||
({"text_encoder": 0}, TypeError, "must be a boolean"),
|
||||
],
|
||||
)
|
||||
def test_component_policy_rejects_invalid_shapes(value, error, match):
|
||||
with pytest.raises(error, match=match):
|
||||
OffloadConfig.from_od_config(_od_config(dlo_offload_components=value))
|
||||
|
||||
|
||||
def test_component_policy_accepts_none_as_default():
|
||||
config = OffloadConfig.from_od_config(_od_config(dlo_offload_components=None))
|
||||
|
||||
assert config.dlo_offload_components == {}
|
||||
|
||||
|
||||
def test_component_policy_requires_dlo():
|
||||
with pytest.raises(ValueError, match="requires distributed layerwise offload"):
|
||||
OffloadConfig.from_od_config(
|
||||
_od_config(
|
||||
enable_distributed_layerwise_offload=False,
|
||||
dlo_offload_components={"text_encoder": False},
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_component_policy_rejects_unknown_or_dit_names():
|
||||
backend = _backend({"transformer": False})
|
||||
|
||||
with pytest.raises(ValueError, match="unknown auxiliary components.*transformer"):
|
||||
backend._validate_component_offload_policy(_modules())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("policy", "encoder_stage", "video_stage", "audio_stage", "try_encoder"),
|
||||
[
|
||||
({}, True, True, True, True),
|
||||
({"text_encoder": False}, False, True, True, False),
|
||||
({"default": False, "video_vae": True}, False, True, False, False),
|
||||
],
|
||||
)
|
||||
def test_prepare_auxiliary_components_applies_policy(
|
||||
mocker,
|
||||
policy,
|
||||
encoder_stage,
|
||||
video_stage,
|
||||
audio_stage,
|
||||
try_encoder,
|
||||
):
|
||||
backend = _backend(policy)
|
||||
backend._try_layerwise_offload_encoder = mocker.Mock()
|
||||
backend._register_on_demand_hook = mocker.Mock()
|
||||
modules = _modules()
|
||||
plan = OffloadPlan(
|
||||
encoder_block_attrs={"text_encoder": ("layers",)},
|
||||
on_demand_component_paths=frozenset({"text_encoder", "video_vae", "audio_vae"}),
|
||||
)
|
||||
|
||||
backend._prepare_auxiliary_components(modules, plan)
|
||||
|
||||
assert backend._try_layerwise_offload_encoder.called is try_encoder
|
||||
assert modules.encoders[0]._omni_dlo_offload_enabled is encoder_stage
|
||||
assert modules.vaes[0]._omni_dlo_offload_enabled is video_stage
|
||||
assert modules.vaes[1]._omni_dlo_offload_enabled is audio_stage
|
||||
assert backend._register_on_demand_hook.call_args_list == [
|
||||
mocker.call(modules.encoders[0], "text_encoder", stage_on_demand=encoder_stage),
|
||||
mocker.call(modules.vaes[0], "video_vae", stage_on_demand=video_stage),
|
||||
mocker.call(modules.vaes[1], "audio_vae", stage_on_demand=audio_stage),
|
||||
]
|
||||
@@ -35,10 +35,7 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
|
||||
DistributedLayerwiseOffloadHook,
|
||||
PinnedResidentLayerGroup,
|
||||
)
|
||||
from vllm_omni.diffusion.offloader.host_registration import (
|
||||
HostRegistrationCleanupError,
|
||||
HostRegistrationError,
|
||||
)
|
||||
from vllm_omni.diffusion.offloader.host_registration import HostRegistrationCleanupError
|
||||
from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
|
||||
from vllm_omni.diffusion.offloader.offload_plan import (
|
||||
OffloadPlan,
|
||||
@@ -810,10 +807,7 @@ def test_registered_mmap_resident_group_bypasses_host_staging(patched_offload_ru
|
||||
assert torch.equal(block.weight, expected)
|
||||
|
||||
|
||||
def test_hwr_registration_uses_transport_budget(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
patched_offload_runtime,
|
||||
):
|
||||
def test_hwr_registration_uses_transport_budget(monkeypatch: pytest.MonkeyPatch):
|
||||
dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
|
||||
plan, _, lease = _fake_hwr_plan()
|
||||
backend = DistributedLayerwiseOffloadBackend(
|
||||
@@ -1131,51 +1125,6 @@ class TestMmapWeightLoading:
|
||||
backend.disable()
|
||||
assert lease.closed
|
||||
|
||||
def test_hwr_registration_failure_falls_back_to_bounded_staging(
|
||||
self,
|
||||
monkeypatch,
|
||||
patched_offload_runtime,
|
||||
):
|
||||
plan, carrier, lease = _fake_hwr_plan("hwr-registration-fallback")
|
||||
backend = DistributedLayerwiseOffloadBackend(
|
||||
OffloadConfig(
|
||||
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
|
||||
pin_cpu_memory=True,
|
||||
dp_size=1,
|
||||
dlo_use_allgather=False,
|
||||
),
|
||||
torch.device("cpu"),
|
||||
host_weight_plan=plan,
|
||||
)
|
||||
pipeline = _HWRPipeline()
|
||||
original_staging_allocator = backend._allocate_shared_cpu_staging_buffers
|
||||
|
||||
def fail_registration(*args, **kwargs):
|
||||
del args, kwargs
|
||||
raise HostRegistrationError("registration unavailable in CPU test")
|
||||
|
||||
def allocate_unpinned_staging(hooks, resident_group=None):
|
||||
# Registration selection still observes pin_cpu_memory=True. The
|
||||
# CPU-only test avoids asking the host for CUDA-pinned allocations.
|
||||
for hook in hooks:
|
||||
hook.pin_memory = False
|
||||
return original_staging_allocator(hooks, resident_group)
|
||||
|
||||
monkeypatch.setattr(dist_backend_module, "register_host_mappings", fail_registration)
|
||||
monkeypatch.setattr(backend, "_allocate_shared_cpu_staging_buffers", allocate_unpinned_staging)
|
||||
|
||||
backend.enable(pipeline)
|
||||
|
||||
assert carrier.taken
|
||||
assert backend._using_rank_local_mmap
|
||||
assert not backend._using_registered_mmap
|
||||
assert backend._host_registration is None
|
||||
assert all(len(hook.cpu_staging_buffers) == 2 for group in backend._all_hook_groups for hook in group)
|
||||
assert not lease.closed
|
||||
|
||||
backend.disable()
|
||||
assert lease.closed
|
||||
|
||||
def test_hwr_backend_failure_drains_partial_setup_before_lease_close(
|
||||
self,
|
||||
monkeypatch,
|
||||
@@ -2183,9 +2132,8 @@ class TestDynamicSlotTracking:
|
||||
prev_idx = (i - 1) % num_blocks
|
||||
|
||||
# Dynamic slot tracking: read from prev's prefetched slot
|
||||
prefetched_slot = prefetched_slots[prev_idx]
|
||||
if prefetched_slot is not None:
|
||||
current_slots[i] = prefetched_slot
|
||||
if prefetched_slots[prev_idx] is not None:
|
||||
current_slots[i] = prefetched_slots[prev_idx]
|
||||
|
||||
read_slot = current_slots[i]
|
||||
|
||||
|
||||
@@ -1,146 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.model_executor.layers.linear import (
|
||||
LinearBase,
|
||||
UnquantizedLinearMethod,
|
||||
)
|
||||
|
||||
from vllm_omni.quantization import svdquant_config
|
||||
from vllm_omni.quantization.factory import (
|
||||
SUPPORTED_QUANTIZATION_METHODS,
|
||||
build_quant_config,
|
||||
)
|
||||
from vllm_omni.quantization.svdquant_config import (
|
||||
DiffusionSVDQuantConfig,
|
||||
DiffusionSVDQuantLinearMethod,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def test_config_from_checkpoint_dict() -> None:
|
||||
config = DiffusionSVDQuantConfig.from_config(
|
||||
{
|
||||
"rank": 32,
|
||||
"precision": "nvfp4",
|
||||
"act_unsigned": False,
|
||||
"modules_to_not_convert": ["blocks.0.adaln_proj.linear"],
|
||||
}
|
||||
)
|
||||
|
||||
assert config.rank == 32
|
||||
assert config.precision == "nvfp4"
|
||||
assert config.modules_to_not_convert == ["blocks.0.adaln_proj.linear"]
|
||||
assert config.get_supported_act_dtypes() == [torch.bfloat16]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kwargs", "message"),
|
||||
[
|
||||
({"rank": 0}, "rank must be positive"),
|
||||
({"precision": "int4"}, "NVFP4 checkpoints only"),
|
||||
({"act_unsigned": True}, "unsigned activations"),
|
||||
],
|
||||
)
|
||||
def test_config_rejects_unsupported_phase1_options(
|
||||
kwargs: dict[str, object],
|
||||
message: str,
|
||||
) -> None:
|
||||
with pytest.raises(ValueError, match=message):
|
||||
DiffusionSVDQuantConfig(**kwargs)
|
||||
|
||||
|
||||
def test_factory_builds_svdquant_with_extra_checkpoint_fields() -> None:
|
||||
config = build_quant_config(
|
||||
{
|
||||
"quant_method": "svdquant",
|
||||
"rank": 16,
|
||||
"precision": "nvfp4",
|
||||
"w4a16_modules": ["adaln_proj.linear"],
|
||||
"w4a16_group_size": 64,
|
||||
}
|
||||
)
|
||||
|
||||
assert isinstance(config, DiffusionSVDQuantConfig)
|
||||
assert config.rank == 16
|
||||
assert "svdquant" in SUPPORTED_QUANTIZATION_METHODS
|
||||
|
||||
|
||||
def test_skipped_and_non_linear_modules_remain_unquantized() -> None:
|
||||
config = DiffusionSVDQuantConfig(modules_to_not_convert=["blocks.0.adaln_proj.linear"])
|
||||
linear = Mock(spec=LinearBase)
|
||||
|
||||
assert isinstance(
|
||||
config.get_quant_method(linear, "blocks.0.adaln_proj.linear"),
|
||||
UnquantizedLinearMethod,
|
||||
)
|
||||
assert config.get_quant_method(torch.nn.ReLU(), "blocks.0.act") is None
|
||||
|
||||
|
||||
def test_linear_method_creates_canonical_partitioned_weights(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(svdquant_config, "_assert_supported", lambda: None)
|
||||
method = DiffusionSVDQuantLinearMethod(DiffusionSVDQuantConfig(rank=32))
|
||||
layer = torch.nn.Module()
|
||||
|
||||
method.create_weights(
|
||||
layer,
|
||||
input_size_per_partition=128,
|
||||
output_partition_sizes=[64, 64],
|
||||
input_size=128,
|
||||
output_size=128,
|
||||
params_dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
assert layer.qweight.shape == (128, 64)
|
||||
assert layer.qweight.dtype == torch.int8
|
||||
assert layer.wscales.shape == (8, 128)
|
||||
assert layer.wscales.dtype == torch.float8_e4m3fn
|
||||
assert layer.proj_down.shape == (128, 32)
|
||||
assert layer.proj_up.shape == (128, 32)
|
||||
assert layer.smooth_factor.shape == (128,)
|
||||
assert layer.wcscales.shape == (128,)
|
||||
assert layer.wtscale.shape == (1,)
|
||||
|
||||
assert layer.qweight.input_dim == 1
|
||||
assert layer.qweight.output_dim == 0
|
||||
assert layer.wscales.input_dim == 0
|
||||
assert layer.wscales.output_dim == 1
|
||||
assert layer.proj_down.input_dim == 0
|
||||
assert layer.proj_up.output_dim == 0
|
||||
assert layer.wcscales.output_dim == 0
|
||||
|
||||
|
||||
def test_linear_method_rejects_misaligned_input_partition(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(svdquant_config, "_assert_supported", lambda: None)
|
||||
method = DiffusionSVDQuantLinearMethod(DiffusionSVDQuantConfig())
|
||||
|
||||
with pytest.raises(ValueError, match="divisible by the block size 16"):
|
||||
method.create_weights(
|
||||
torch.nn.Module(),
|
||||
input_size_per_partition=24,
|
||||
output_partition_sizes=[32],
|
||||
input_size=24,
|
||||
output_size=32,
|
||||
params_dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
|
||||
def test_active_linear_uses_svdquant_method(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(svdquant_config, "_assert_supported", lambda: None)
|
||||
config = DiffusionSVDQuantConfig()
|
||||
linear = Mock(spec=LinearBase)
|
||||
|
||||
method = config.get_quant_method(linear, "blocks.0.attn.out_proj")
|
||||
|
||||
assert isinstance(method, DiffusionSVDQuantLinearMethod)
|
||||
@@ -1,168 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.quantization import svdquant_config as svdquant
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
class _FakeNvfp4Kernel:
|
||||
def __init__(self, base_weight: torch.Tensor | None = None) -> None:
|
||||
self.base_weight = base_weight
|
||||
self.processed = False
|
||||
self.last_input: torch.Tensor | None = None
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
del layer
|
||||
self.processed = True
|
||||
|
||||
def apply_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
del layer
|
||||
assert bias is None
|
||||
assert self.base_weight is not None
|
||||
self.last_input = x.clone()
|
||||
return torch.mm(x, self.base_weight)
|
||||
|
||||
|
||||
def _register_parameter(
|
||||
layer: torch.nn.Module,
|
||||
name: str,
|
||||
value: torch.Tensor,
|
||||
) -> None:
|
||||
layer.register_parameter(
|
||||
name,
|
||||
torch.nn.Parameter(value, requires_grad=False),
|
||||
)
|
||||
|
||||
|
||||
def test_supports_only_validated_datacenter_blackwell() -> None:
|
||||
assert not svdquant._supports_capability(SimpleNamespace(major=10, minor=0))
|
||||
assert svdquant._supports_capability(SimpleNamespace(major=10, minor=3))
|
||||
assert not svdquant._supports_capability(SimpleNamespace(major=11, minor=0))
|
||||
assert not svdquant._supports_capability(SimpleNamespace(major=12, minor=0))
|
||||
assert not svdquant._supports_capability(None)
|
||||
|
||||
|
||||
def test_rejects_nvfp4_backend_with_incompatible_layer_contract(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
svdquant._nvfp4_kernel.cache_clear()
|
||||
monkeypatch.setattr(
|
||||
svdquant,
|
||||
"init_nvfp4_linear_kernel",
|
||||
lambda: _FakeNvfp4Kernel(),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="incompatible with the SVDQuant checkpoint"):
|
||||
svdquant._nvfp4_kernel()
|
||||
|
||||
svdquant._nvfp4_kernel.cache_clear()
|
||||
|
||||
|
||||
def test_prepare_weights_uses_existing_nvfp4_layout_contract(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
kernel = _FakeNvfp4Kernel()
|
||||
monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
|
||||
monkeypatch.setattr(svdquant, "_nvfp4_kernel", lambda: kernel)
|
||||
method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig())
|
||||
layer = torch.nn.Module()
|
||||
qweight = torch.arange(16, dtype=torch.int8).reshape(4, 4)
|
||||
wscales = torch.arange(8, dtype=torch.float32).to(torch.float8_e4m3fn).reshape(2, 4)
|
||||
_register_parameter(layer, "qweight", qweight)
|
||||
_register_parameter(layer, "wscales", wscales)
|
||||
_register_parameter(layer, "wtscale", torch.tensor([2.0], dtype=torch.bfloat16))
|
||||
_register_parameter(
|
||||
layer,
|
||||
"wcscales",
|
||||
torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.bfloat16),
|
||||
)
|
||||
|
||||
method.process_weights_after_loading(layer)
|
||||
|
||||
assert kernel.processed
|
||||
assert not hasattr(layer, "qweight")
|
||||
assert not hasattr(layer, "wscales")
|
||||
assert layer.weight.dtype == torch.uint8
|
||||
assert torch.equal(layer.weight.view(torch.int8), qweight)
|
||||
assert torch.equal(layer.weight_scale, wscales.transpose(0, 1))
|
||||
torch.testing.assert_close(layer.alpha, torch.tensor([2.0]))
|
||||
assert layer.input_global_scale_inv.dtype == torch.float32
|
||||
assert layer.output_channel_scale.shape == (4,)
|
||||
|
||||
|
||||
def test_prepare_weights_drops_identity_output_scale(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
kernel = _FakeNvfp4Kernel()
|
||||
monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
|
||||
monkeypatch.setattr(svdquant, "_nvfp4_kernel", lambda: kernel)
|
||||
method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig())
|
||||
layer = torch.nn.Module()
|
||||
_register_parameter(layer, "qweight", torch.zeros(4, 4, dtype=torch.int8))
|
||||
_register_parameter(
|
||||
layer,
|
||||
"wscales",
|
||||
torch.ones(2, 4, dtype=torch.float8_e4m3fn),
|
||||
)
|
||||
_register_parameter(layer, "wtscale", torch.ones(1, dtype=torch.bfloat16))
|
||||
_register_parameter(layer, "wcscales", torch.ones(4, dtype=torch.bfloat16))
|
||||
|
||||
method.process_weights_after_loading(layer)
|
||||
|
||||
assert layer.output_channel_scale is None
|
||||
|
||||
|
||||
def test_compatibility_path_applies_qkv_scale_and_rank_correction(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
base_weight = torch.tensor(
|
||||
[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
kernel = _FakeNvfp4Kernel(base_weight)
|
||||
monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
|
||||
monkeypatch.setattr(svdquant, "_nvfp4_kernel", lambda: kernel)
|
||||
method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig())
|
||||
|
||||
layer = torch.nn.Module()
|
||||
layer.smooth_factor = torch.tensor([2.0, 4.0], dtype=torch.bfloat16)
|
||||
layer.proj_down = torch.tensor([[1.0], [2.0]], dtype=torch.bfloat16)
|
||||
layer.proj_up = torch.tensor([[1.0], [2.0], [3.0]], dtype=torch.bfloat16)
|
||||
layer.output_channel_scale = torch.tensor([1.0, 2.0, 4.0], dtype=torch.bfloat16)
|
||||
layer.output_size_per_partition = 3
|
||||
x = torch.tensor([[2.0, 4.0], [4.0, 8.0]], dtype=torch.bfloat16)
|
||||
bias = torch.tensor([1.0, 2.0, 3.0], dtype=torch.bfloat16)
|
||||
|
||||
actual = method.apply(layer, x, bias)
|
||||
|
||||
smoothed = x / layer.smooth_factor
|
||||
expected = torch.mm(smoothed, base_weight)
|
||||
expected = expected * layer.output_channel_scale
|
||||
expected = expected + torch.mm(torch.mm(x, layer.proj_down), layer.proj_up.T)
|
||||
expected = expected + bias
|
||||
torch.testing.assert_close(actual, expected)
|
||||
torch.testing.assert_close(kernel.last_input, smoothed)
|
||||
|
||||
|
||||
def test_compatibility_path_rejects_fp16_activations(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
|
||||
method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig())
|
||||
layer = torch.nn.Module()
|
||||
with pytest.raises(ValueError, match="requires BF16 activations"):
|
||||
method.apply(
|
||||
layer,
|
||||
torch.zeros(1, 16, dtype=torch.float16),
|
||||
)
|
||||
@@ -1,130 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.model_executor.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
QKVParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
|
||||
from vllm_omni.quantization import svdquant_config as svdquant
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def _linear_without_init(linear_type: type[torch.nn.Module], tp_rank: int):
|
||||
layer = object.__new__(linear_type)
|
||||
torch.nn.Module.__init__(layer)
|
||||
layer.tp_rank = tp_rank
|
||||
return layer
|
||||
|
||||
|
||||
def _values(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor:
|
||||
values = torch.arange(torch.Size(shape).numel(), dtype=torch.int64).remainder(127)
|
||||
return values.reshape(shape).to(dtype)
|
||||
|
||||
|
||||
def _method(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
|
||||
return svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig(rank=8))
|
||||
|
||||
|
||||
def test_column_parallel_loader_shards_svdquant_outputs(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
layer = _linear_without_init(ColumnParallelLinear, tp_rank=1)
|
||||
method = _method(monkeypatch)
|
||||
method.create_weights(
|
||||
layer,
|
||||
input_size_per_partition=128,
|
||||
output_partition_sizes=[32],
|
||||
input_size=128,
|
||||
output_size=64,
|
||||
params_dtype=torch.bfloat16,
|
||||
weight_loader=layer.weight_loader,
|
||||
)
|
||||
|
||||
loaded_tensors = {
|
||||
"qweight": _values((64, 64), torch.int8),
|
||||
"wscales": _values((8, 64), torch.float8_e4m3fn),
|
||||
"proj_up": _values((64, 8), torch.bfloat16),
|
||||
"wcscales": _values((64,), torch.bfloat16),
|
||||
}
|
||||
for name, loaded in loaded_tensors.items():
|
||||
parameter = getattr(layer, name)
|
||||
parameter.weight_loader(parameter, loaded)
|
||||
expected = loaded.narrow(parameter.output_dim, 32, 32)
|
||||
torch.testing.assert_close(parameter, expected)
|
||||
|
||||
|
||||
def test_row_parallel_loader_shards_svdquant_inputs(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
layer = _linear_without_init(RowParallelLinear, tp_rank=1)
|
||||
method = _method(monkeypatch)
|
||||
method.create_weights(
|
||||
layer,
|
||||
input_size_per_partition=64,
|
||||
output_partition_sizes=[64],
|
||||
input_size=128,
|
||||
output_size=64,
|
||||
params_dtype=torch.bfloat16,
|
||||
weight_loader=layer.weight_loader,
|
||||
)
|
||||
|
||||
loaded_tensors = {
|
||||
"qweight": _values((64, 64), torch.int8),
|
||||
"wscales": _values((8, 64), torch.float8_e4m3fn),
|
||||
"proj_down": _values((128, 8), torch.bfloat16),
|
||||
"smooth_factor": _values((128,), torch.bfloat16),
|
||||
}
|
||||
for name, loaded in loaded_tensors.items():
|
||||
parameter = getattr(layer, name)
|
||||
parameter.weight_loader(parameter, loaded)
|
||||
shard_size = parameter.shape[parameter.input_dim]
|
||||
expected = loaded.narrow(parameter.input_dim, shard_size, shard_size)
|
||||
torch.testing.assert_close(parameter, expected)
|
||||
|
||||
|
||||
def test_qkv_parallel_loader_shards_fused_svdquant_outputs(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
layer = _linear_without_init(QKVParallelLinear, tp_rank=1)
|
||||
layer.total_num_heads = 4
|
||||
layer.total_num_kv_heads = 2
|
||||
layer.num_heads = 2
|
||||
layer.num_kv_heads = 1
|
||||
layer.num_kv_head_replicas = 1
|
||||
layer.head_size = 8
|
||||
layer.v_head_size = 8
|
||||
method = _method(monkeypatch)
|
||||
method.create_weights(
|
||||
layer,
|
||||
input_size_per_partition=128,
|
||||
output_partition_sizes=[16, 8, 8],
|
||||
input_size=128,
|
||||
output_size=64,
|
||||
params_dtype=torch.bfloat16,
|
||||
weight_loader=layer.weight_loader,
|
||||
)
|
||||
|
||||
loaded_tensors = {
|
||||
"qweight": _values((64, 64), torch.int8),
|
||||
"wscales": _values((8, 64), torch.float8_e4m3fn),
|
||||
"proj_up": _values((64, 8), torch.bfloat16),
|
||||
"wcscales": _values((64,), torch.bfloat16),
|
||||
}
|
||||
expected_global_outputs = [slice(16, 32), slice(40, 48), slice(56, 64)]
|
||||
for name, loaded in loaded_tensors.items():
|
||||
parameter = getattr(layer, name)
|
||||
parameter.weight_loader(parameter, loaded)
|
||||
expected = torch.cat(
|
||||
[
|
||||
loaded.narrow(parameter.output_dim, part.start, part.stop - part.start)
|
||||
for part in expected_global_outputs
|
||||
],
|
||||
dim=parameter.output_dim,
|
||||
)
|
||||
torch.testing.assert_close(parameter, expected)
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Source-level regression tests for diffusion output/engine helpers.
|
||||
|
||||
These tests verify naming conventions and patterns by inspecting source code
|
||||
@@ -37,50 +37,10 @@ _FORMATTER_PATH = os.path.normpath(
|
||||
"output_formatter.py",
|
||||
)
|
||||
)
|
||||
_METRICS_PATH = os.path.normpath(
|
||||
os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
os.pardir,
|
||||
os.pardir,
|
||||
"vllm_omni",
|
||||
"metrics",
|
||||
"utils.py",
|
||||
)
|
||||
)
|
||||
_INLINE_CLIENT_PATH = os.path.normpath(
|
||||
os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
os.pardir,
|
||||
os.pardir,
|
||||
"vllm_omni",
|
||||
"diffusion",
|
||||
"inline_stage_diffusion_client.py",
|
||||
)
|
||||
)
|
||||
_STAGE_PROC_PATH = os.path.normpath(
|
||||
os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
os.pardir,
|
||||
os.pardir,
|
||||
"vllm_omni",
|
||||
"diffusion",
|
||||
"stage_diffusion_proc.py",
|
||||
)
|
||||
)
|
||||
_ORCHESTRATOR_PATH = os.path.normpath(
|
||||
os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
os.pardir,
|
||||
os.pardir,
|
||||
"vllm_omni",
|
||||
"engine",
|
||||
"orchestrator.py",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _read_source(path: str) -> str:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
with open(path) as f:
|
||||
return f.read()
|
||||
|
||||
|
||||
@@ -140,116 +100,9 @@ class TestMetricKeys:
|
||||
raise AssertionError("diffusion_engine_total_time_ms should be attached in step_streaming()")
|
||||
|
||||
assert '"diffusion_engine_exec_time_ms": exec_total_time * 1000' in step_streaming_source
|
||||
# step_total_ms is no longer emitted as a metric (no family consumes
|
||||
# it); it lives only in the debug log breakdown below.
|
||||
assert '"diffusion_engine_total_time_ms"' not in step_streaming_source
|
||||
assert '"diffusion_engine_total_time_ms": step_total_ms' in step_streaming_source
|
||||
assert '"postprocess_time_ms": postprocess_time * 1000' in step_streaming_source
|
||||
|
||||
def test_step_streaming_piggybacks_scheduler_waiting_snapshot(self) -> None:
|
||||
source = _read_source(_ENGINE_PATH)
|
||||
step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming")
|
||||
assert "diffusion_scheduler_waiting_metrics(" in step_streaming_source
|
||||
assert "_scheduler_num_waiting_reqs" in step_streaming_source
|
||||
|
||||
def test_step_streaming_preserves_snapshot_on_postprocess_error(self) -> None:
|
||||
source = _read_source(_ENGINE_PATH)
|
||||
step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming")
|
||||
assert "scheduler_metrics = diffusion_scheduler_waiting_metrics(" in step_streaming_source
|
||||
assert 'setattr(exc, "diffusion_metrics", scheduler_metrics)' in step_streaming_source
|
||||
assert step_streaming_source.index("scheduler_metrics =") < step_streaming_source.index(
|
||||
"self.postprocess_output("
|
||||
)
|
||||
|
||||
def test_abort_keeps_output_consumers_alive_for_terminal_snapshot(self) -> None:
|
||||
inline_source = _read_source(_INLINE_CLIENT_PATH)
|
||||
abort_source = _get_function_source(inline_source, "InlineStageDiffusionClient", "abort_requests_async")
|
||||
assert ".cancel(" not in abort_source
|
||||
|
||||
proc_source = _read_source(_STAGE_PROC_PATH)
|
||||
run_loop_source = _get_function_source(proc_source, "StageDiffusionProc", "run_loop")
|
||||
abort_branch = run_loop_source[run_loop_source.index('elif msg_type == "abort":') :]
|
||||
abort_branch = abort_branch[: abort_branch.index('elif msg_type == "collective_rpc":')]
|
||||
assert ".cancel(" not in abort_branch
|
||||
|
||||
def test_orchestrator_consumes_metrics_only_output_without_routing(self) -> None:
|
||||
"""A metrics-only sentinel contributes its queue depth and is not routed.
|
||||
|
||||
Absorption lives in ``_absorb_diffusion_metrics`` rather than inline in a
|
||||
loop, so both orchestration loops share one implementation. The ordering
|
||||
is asserted where it now lives: the snapshot inside the helper, and the
|
||||
absorb-before-route guard inside every loop that polls diffusion output.
|
||||
"""
|
||||
source = _read_source(_ORCHESTRATOR_PATH)
|
||||
|
||||
# Snapshot before the verdict, so a sentinel still reports its waiting
|
||||
# depth on the way out instead of being dropped unaccounted.
|
||||
absorb_source = _get_function_source(source, "Orchestrator", "_absorb_diffusion_metrics")
|
||||
snapshot_pos = absorb_source.index("_update_stage_replica_waiting(")
|
||||
sentinel_pos = absorb_source.index("diffusion_output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID")
|
||||
assert snapshot_pos < sentinel_pos
|
||||
|
||||
# Absorb before routing, or a sentinel reaches the downstream consumer
|
||||
# as if it were a real request output.
|
||||
for loop_name in ("_orchestration_loop", "_orchestration_loop_event_driven"):
|
||||
loop_source = _get_function_source(source, "Orchestrator", loop_name)
|
||||
absorb_pos = loop_source.index("self._absorb_diffusion_metrics(")
|
||||
route_pos = loop_source.index("record_output_timestamps(")
|
||||
assert absorb_pos < route_pos, loop_name
|
||||
|
||||
|
||||
class TestVaeDecodeEmit:
|
||||
"""VAE decode timing is sourced from ``DiffusionOutput.stage_durations``
|
||||
(populated by the diffusion pipeline profiler) and emitted as
|
||||
``vae_decode_time_ms`` so the ``_MS_TO_S`` accumulator picks it up.
|
||||
"""
|
||||
|
||||
def test_extract_vae_decode_ms_helper_exists(self) -> None:
|
||||
source = _read_source(_METRICS_PATH)
|
||||
helper_src = _get_function_source(source, None, "extract_diffusion_vae_decode_ms")
|
||||
assert ".vae.decode" in helper_src, "helper must key on the '.vae.decode' suffix"
|
||||
|
||||
def test_step_streaming_emits_vae_decode_time_ms(self) -> None:
|
||||
source = _read_source(_ENGINE_PATH)
|
||||
step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming")
|
||||
assert "extract_diffusion_vae_decode_ms(" in step_streaming_source, (
|
||||
"step_streaming must call extract_diffusion_vae_decode_ms to source VAE decode timing"
|
||||
)
|
||||
assert '"vae_decode_time_ms"' in step_streaming_source, (
|
||||
"step_streaming must emit vae_decode_time_ms key for the _MS_TO_S accumulator"
|
||||
)
|
||||
|
||||
|
||||
class TestDiffuseForwardEmit:
|
||||
"""Forward-only timing → forward_time_ms, mirroring TestVaeDecodeEmit."""
|
||||
|
||||
def test_extract_diffuse_ms_helper_exists(self) -> None:
|
||||
source = _read_source(_METRICS_PATH)
|
||||
helper_src = _get_function_source(source, None, "extract_diffusion_denoise_ms")
|
||||
assert ".diffuse" in helper_src, "helper must key on the '.diffuse' suffix"
|
||||
|
||||
def test_step_streaming_emits_forward_time_ms(self) -> None:
|
||||
source = _read_source(_ENGINE_PATH)
|
||||
step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming")
|
||||
assert "extract_diffusion_denoise_ms(" in step_streaming_source, (
|
||||
"step_streaming must call extract_diffusion_denoise_ms to source forward-only timing"
|
||||
)
|
||||
assert '"forward_time_ms"' in step_streaming_source, (
|
||||
"step_streaming must emit forward_time_ms key for the _MS_TO_S accumulator"
|
||||
)
|
||||
|
||||
|
||||
class TestKvRecvEmit:
|
||||
"""KV-recv timing → kv_recv_time_ms, sourced from DiffusionOutput.kv_recv_ms
|
||||
(set by the runner's _prepare_request_for_forward, not the profiler)."""
|
||||
|
||||
def test_step_streaming_emits_kv_recv_time_ms(self) -> None:
|
||||
source = _read_source(_ENGINE_PATH)
|
||||
step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming")
|
||||
assert "kv_recv_ms" in step_streaming_source, "step_streaming must read output.kv_recv_ms"
|
||||
assert '"kv_recv_time_ms"' in step_streaming_source, (
|
||||
"step_streaming must emit kv_recv_time_ms key for the _MS_TO_S accumulator"
|
||||
)
|
||||
|
||||
|
||||
class TestDummyRunAllocation:
|
||||
"""Verify _dummy_run generates exact-sized audio arrays."""
|
||||
|
||||
@@ -181,7 +181,7 @@ def test_diffusion_output_dict_tensors_round_trip_through_shm() -> None:
|
||||
|
||||
|
||||
def test_diffusion_output_tuple_tensors_round_trip_through_shm() -> None:
|
||||
# LTX2 returns (video, audio) tuples as DiffusionOutput.output.
|
||||
# LTX2 / DreamID return (video, audio) tuples as DiffusionOutput.output.
|
||||
video = torch.arange(300_000, dtype=torch.float32)
|
||||
audio = torch.arange(300_000, dtype=torch.float32) * 3
|
||||
output = DiffusionOutput(output=(video, audio))
|
||||
|
||||
@@ -773,19 +773,6 @@ class TestRequestScheduler:
|
||||
assert third.num_running_reqs == 1
|
||||
assert third.num_waiting_reqs == 0
|
||||
|
||||
def test_records_initial_scheduler_queue_wait(self, mocker: MockerFixture) -> None:
|
||||
perf_counter = mocker.patch(
|
||||
"vllm_omni.diffusion.sched.base_scheduler.time.perf_counter",
|
||||
side_effect=[10.0, 10.125],
|
||||
)
|
||||
request = _make_step_request("queue-wait", num_inference_steps=2)
|
||||
|
||||
self.scheduler.add_request(request)
|
||||
self.scheduler.schedule()
|
||||
|
||||
assert request.scheduler_queue_wait_ms == pytest.approx(125.0)
|
||||
assert perf_counter.call_count == 2
|
||||
|
||||
def test_batches_compatible_requests_up_to_max_num_seqs(self) -> None:
|
||||
scheduler = RequestScheduler()
|
||||
scheduler.initialize(SimpleNamespace(max_num_seqs=2))
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import asyncio
|
||||
import multiprocessing as mp
|
||||
@@ -1079,42 +1079,6 @@ class TestStageDiffusionClientErrorPropagation:
|
||||
|
||||
assert client.get_diffusion_output_nowait() is None
|
||||
|
||||
def test_error_response_preserves_scheduler_metrics(self):
|
||||
from vllm_omni.metrics import definitions as metric_defs
|
||||
|
||||
client = self._make_client()
|
||||
client._response_socket.recv.side_effect = [b"message", zmq.Again()]
|
||||
client._decoder.decode.return_value = {
|
||||
"type": "error",
|
||||
"request_id": "req-error",
|
||||
"error": "gpu fault",
|
||||
"metrics": {metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY: 0},
|
||||
}
|
||||
|
||||
output = client.get_diffusion_output_nowait()
|
||||
|
||||
assert output is not None
|
||||
assert output.error == "gpu fault"
|
||||
assert output.metrics[metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY] == 0
|
||||
|
||||
def test_metrics_response_creates_metrics_only_output(self):
|
||||
from vllm_omni.metrics import definitions as metric_defs
|
||||
from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID
|
||||
|
||||
client = self._make_client()
|
||||
client._response_socket.recv.side_effect = [b"message", zmq.Again()]
|
||||
client._decoder.decode.return_value = {
|
||||
"type": "metrics",
|
||||
"metrics": {metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY: 0},
|
||||
}
|
||||
|
||||
output = client.get_diffusion_output_nowait()
|
||||
|
||||
assert output is not None
|
||||
assert output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID
|
||||
assert output.error is None
|
||||
assert output.metrics[metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY] == 0
|
||||
|
||||
def test_check_health_raises_when_proc_dead(self):
|
||||
"""``check_health`` detects a dead subprocess via the manager's proc
|
||||
and raises ``EngineDeadError``, setting ``_engine_dead`` as a
|
||||
|
||||
@@ -1209,67 +1209,6 @@ def test_cleanup_preserves_pending_save(build_adapter):
|
||||
assert len(adapter._pending_save_reqs) == 1
|
||||
|
||||
|
||||
def test_abort_invalidates_queued_sender_task_before_external_id_reuse(build_adapter):
|
||||
adapter, connector = build_adapter(stage_id=1)
|
||||
request = _req("req-old", RequestStatus.WAITING, external_req_id="ext-reused")
|
||||
adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
|
||||
codes=CodesStruct(audio=torch.tensor([1], dtype=torch.long))
|
||||
)
|
||||
|
||||
adapter.save_async(multimodal_output=None, request=request)
|
||||
stale_task = adapter._pending_save_reqs.popleft()
|
||||
adapter.cleanup(request.request_id, request.external_req_id)
|
||||
adapter._send_single_request(stale_task)
|
||||
|
||||
connector.put.assert_not_called()
|
||||
assert request.external_req_id not in adapter.put_req_chunk
|
||||
|
||||
replacement = _req("req-new", RequestStatus.WAITING, external_req_id=request.external_req_id)
|
||||
adapter.save_async(multimodal_output=None, request=replacement)
|
||||
current_task = adapter._pending_save_reqs.popleft()
|
||||
adapter._send_single_request(current_task)
|
||||
|
||||
connector.put.assert_called_once()
|
||||
assert connector.put.call_args.kwargs["put_key"] == "ext-reused_1_0"
|
||||
|
||||
|
||||
def test_finish_requests_does_not_wait_for_inflight_send(build_adapter):
|
||||
adapter, connector = build_adapter(stage_id=1)
|
||||
first = _req("req-first", RequestStatus.WAITING, external_req_id="ext-first")
|
||||
adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct()
|
||||
put_started = threading.Event()
|
||||
release_put = threading.Event()
|
||||
|
||||
def blocking_put(**kwargs):
|
||||
put_started.set()
|
||||
release_put.wait(timeout=2)
|
||||
return True, 1, {}
|
||||
|
||||
connector.put.side_effect = blocking_put
|
||||
adapter.save_async(multimodal_output=None, request=first)
|
||||
task = adapter._pending_save_reqs.popleft()
|
||||
sender = threading.Thread(target=adapter._send_single_request, args=(task,))
|
||||
sender.start()
|
||||
assert put_started.wait(timeout=1)
|
||||
cleaner = threading.Thread(
|
||||
target=adapter.finish_requests,
|
||||
args=([first.request_id], RequestStatus.FINISHED_ABORTED, {first.request_id: first}),
|
||||
)
|
||||
cleaner.start()
|
||||
try:
|
||||
cleaner.join(timeout=0.5)
|
||||
assert not cleaner.is_alive()
|
||||
assert task["sender_token"].cancelled
|
||||
assert first.external_req_id in adapter._sender_tokens
|
||||
finally:
|
||||
release_put.set()
|
||||
sender.join(timeout=1)
|
||||
cleaner.join(timeout=1)
|
||||
assert not sender.is_alive()
|
||||
assert first.external_req_id not in adapter._sender_tokens
|
||||
assert first.external_req_id not in adapter.put_req_chunk
|
||||
|
||||
|
||||
def test_cleanup_only_affects_target_request(build_adapter):
|
||||
"""Cleanup for one request must not affect another request's state."""
|
||||
adapter, _ = build_adapter(stage_id=1)
|
||||
@@ -1568,7 +1507,6 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
|
||||
scheduler.structured_output_manager.should_advance.return_value = False
|
||||
scheduler.finished_req_ids_dict = {}
|
||||
scheduler.kv_cache_manager.take_events.return_value = None
|
||||
scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
scheduler.kv_event_publisher = mocker.MagicMock()
|
||||
|
||||
request = _HashableRequest(
|
||||
@@ -1649,7 +1587,6 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
|
||||
scheduler.finished_req_ids_dict = {}
|
||||
scheduler.kv_cache_manager = mocker.MagicMock()
|
||||
scheduler.kv_cache_manager.take_events.return_value = None
|
||||
scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
scheduler.kv_event_publisher = mocker.MagicMock()
|
||||
scheduler.waiting_for_transfer_free = set()
|
||||
scheduler.transfer_triggered_requests = set()
|
||||
@@ -1805,7 +1742,6 @@ def _build_deferred_finish_scheduler(mocker, *, running, pending_finish_reqs):
|
||||
scheduler.structured_output_manager.should_advance.return_value = False
|
||||
scheduler.finished_req_ids_dict = {}
|
||||
scheduler.kv_cache_manager.take_events.return_value = None
|
||||
scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
scheduler.kv_event_publisher = mocker.MagicMock()
|
||||
scheduler._pending_finish_reqs = list(pending_finish_reqs)
|
||||
|
||||
@@ -2280,43 +2216,3 @@ def test_expiry_is_per_request(build_adapter):
|
||||
|
||||
assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == {"stalled"}
|
||||
assert "healthy" in adapter._waiting_since
|
||||
|
||||
|
||||
def test_abort_clears_native_codec_state_before_external_id_reuse(build_adapter):
|
||||
from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
|
||||
talker2code2wav_async_chunk,
|
||||
)
|
||||
|
||||
adapter, _ = build_adapter(stage_id=1, connector_extra={"codec_chunk_frames": 1})
|
||||
external_req_id = "ext-native-abort"
|
||||
first = _req("req-native-abort", RequestStatus.WAITING, external_req_id=external_req_id)
|
||||
first.resumable = True
|
||||
first.model_intermediate_buffer = None
|
||||
first.additional_information = {"meta": {"codec_streaming": True}, "nvc_logical_prompt_len": 1}
|
||||
first_frame = torch.full((1, 31), 101, dtype=torch.long)
|
||||
|
||||
talker2code2wav_async_chunk(
|
||||
adapter,
|
||||
{"codes": {"audio": first_frame}, "meta": {"codec_streaming": True}},
|
||||
first,
|
||||
is_finished=False,
|
||||
)
|
||||
|
||||
adapter.requests_num_chunks_sent[external_req_id] = 1
|
||||
adapter.finish_requests([first.request_id], RequestStatus.FINISHED_ABORTED, {first.request_id: first})
|
||||
|
||||
assert external_req_id not in adapter.request_payload
|
||||
assert external_req_id not in adapter.requests_num_chunks_sent
|
||||
|
||||
replacement = _req("req-native-abort", RequestStatus.WAITING, external_req_id=external_req_id)
|
||||
replacement.resumable = True
|
||||
replacement.model_intermediate_buffer = None
|
||||
replacement.additional_information = first.additional_information
|
||||
replacement_payload = talker2code2wav_async_chunk(
|
||||
adapter,
|
||||
{"codes": {"audio": torch.full((1, 31), 202, dtype=torch.long)}, "meta": {"codec_streaming": True}},
|
||||
replacement,
|
||||
is_finished=False,
|
||||
)
|
||||
|
||||
assert torch.equal(replacement_payload.codes.audio, torch.full((1, 31), 202, dtype=torch.long))
|
||||
@@ -22,7 +22,6 @@ import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import asdict, dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
@@ -84,27 +83,35 @@ NEGATIVE_PROMPT = (
|
||||
# Release parity exercises the accelerated BF16 cuDNN path in both runtimes.
|
||||
# Degenerate one-token attention falls back to PyTorch SDPA MATH.
|
||||
ATTENTION_BACKEND = "torch_sdpa_cudnn_bf16"
|
||||
VIDEO_SSIM_MEAN_THRESHOLD = 0.95
|
||||
VIDEO_SSIM_MIN_THRESHOLD = 0.90
|
||||
VIDEO_PSNR_MEAN_THRESHOLD = 30.0
|
||||
AUDIO_RELATIVE_L2_THRESHOLD = 0.2
|
||||
AUDIO_COSINE_THRESHOLD = 0.95
|
||||
|
||||
# Decoded-output gates for the pinned official/Omni BF16 cuDNN comparison.
|
||||
# The metrics artifact preserves the exact observed values.
|
||||
LTX25_VIDEO_SSIM_MEAN_THRESHOLD = 0.99
|
||||
LTX25_I2V_VIDEO_SSIM_MEAN_THRESHOLD = 0.99
|
||||
LTX25_VIDEO_SSIM_MIN_THRESHOLD = 0.99
|
||||
LTX25_VIDEO_PSNR_MEAN_THRESHOLD = 40.0
|
||||
LTX25_VIDEO_LPIPS_MEAN_THRESHOLD = 0.01
|
||||
LTX25_AUDIO_RELATIVE_L2_THRESHOLD = 0.32
|
||||
LTX25_AUDIO_COSINE_THRESHOLD = 0.95
|
||||
|
||||
# Full/SFT uses independently converted dev weights, but the release contract
|
||||
# still requires decoded-video parity at the same 0.99 SSIM floor as the
|
||||
# distilled paths when both runtimes use BF16 cuDNN attention.
|
||||
LTX25_FULL_VIDEO_SSIM_MEAN_THRESHOLD = 0.99
|
||||
LTX25_FULL_VIDEO_SSIM_MIN_THRESHOLD = 0.99
|
||||
LTX25_FULL_VIDEO_PSNR_MEAN_THRESHOLD = 40.0
|
||||
LTX25_FULL_AUDIO_RELATIVE_L2_THRESHOLD = 0.3
|
||||
LTX25_FULL_VIDEO_LPIPS_MEAN_THRESHOLD = 0.01
|
||||
LTX25_FULL_I2V_VIDEO_LPIPS_MEAN_THRESHOLD = 0.02
|
||||
LTX25_FULL_AUDIO_COSINE_THRESHOLD = AUDIO_COSINE_THRESHOLD
|
||||
LTX25_STAGE_2_SIGMAS = [0.909375, 0.725, 0.421875, 0.0]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LTXAccuracyThresholds:
|
||||
video_ssim_mean: float
|
||||
video_ssim_min: float
|
||||
video_psnr_mean_db: float
|
||||
audio_relative_l2: float
|
||||
audio_cosine_similarity: float
|
||||
|
||||
|
||||
STRICT_THRESHOLDS = LTXAccuracyThresholds(
|
||||
video_ssim_mean=0.99,
|
||||
video_ssim_min=0.99,
|
||||
video_psnr_mean_db=40.0,
|
||||
audio_relative_l2=0.10,
|
||||
audio_cosine_similarity=0.99,
|
||||
)
|
||||
|
||||
|
||||
def test_ltx_reference_runner_unwraps_flattened_pipeline_output() -> None:
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
@@ -412,6 +419,8 @@ def test_ltx25_i2v_requests_pin_official_crf() -> None:
|
||||
|
||||
|
||||
def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float]:
|
||||
from benchmarks.diffusion.quantization_quality import compute_lpips_video
|
||||
|
||||
assert reference.shape == prediction.shape
|
||||
assert reference.ndim == 4 and reference.shape[-1] == 3
|
||||
ssim_scores: list[float] = []
|
||||
@@ -426,6 +435,7 @@ def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
|
||||
"ssim_mean": float(np.mean(ssim_scores)),
|
||||
"ssim_min": float(np.min(ssim_scores)),
|
||||
"psnr_mean_db": float(np.mean(psnr_scores)),
|
||||
"lpips_alex_mean": compute_lpips_video(reference, prediction),
|
||||
"max_abs": float(difference.max()),
|
||||
"mean_abs": float(difference.mean()),
|
||||
}
|
||||
@@ -453,17 +463,6 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
|
||||
}
|
||||
|
||||
|
||||
def _assert_strict_similarity(
|
||||
video_metrics: dict[str, float],
|
||||
audio_metrics: dict[str, float | bool],
|
||||
) -> None:
|
||||
assert video_metrics["ssim_mean"] >= STRICT_THRESHOLDS.video_ssim_mean
|
||||
assert video_metrics["ssim_min"] >= STRICT_THRESHOLDS.video_ssim_min
|
||||
assert video_metrics["psnr_mean_db"] >= STRICT_THRESHOLDS.video_psnr_mean_db
|
||||
assert audio_metrics["relative_l2"] <= STRICT_THRESHOLDS.audio_relative_l2
|
||||
assert audio_metrics["cosine_similarity"] >= STRICT_THRESHOLDS.audio_cosine_similarity
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.benchmark
|
||||
@pytest.mark.diffusion
|
||||
@@ -570,7 +569,6 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path
|
||||
"omni_model": omni_model_source,
|
||||
"omni_model_revision": omni_model_revision,
|
||||
"resolved_omni_model_path": str(omni_model),
|
||||
"thresholds": asdict(STRICT_THRESHOLDS),
|
||||
"peak_memory_mb": {
|
||||
"official_allocated": official_metadata["peak_memory_allocated_mb"],
|
||||
"official_reserved": official_metadata["peak_memory_reserved_mb"],
|
||||
@@ -584,7 +582,13 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path
|
||||
(output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
|
||||
print(json.dumps(result, indent=2))
|
||||
|
||||
_assert_strict_similarity(video_metrics, audio_metrics)
|
||||
ssim_mean_threshold = LTX25_VIDEO_SSIM_MEAN_THRESHOLD if task == "t2v" else LTX25_I2V_VIDEO_SSIM_MEAN_THRESHOLD
|
||||
assert video_metrics["ssim_mean"] >= ssim_mean_threshold
|
||||
assert video_metrics["ssim_min"] >= LTX25_VIDEO_SSIM_MIN_THRESHOLD
|
||||
assert video_metrics["psnr_mean_db"] >= LTX25_VIDEO_PSNR_MEAN_THRESHOLD
|
||||
assert video_metrics["lpips_alex_mean"] <= LTX25_VIDEO_LPIPS_MEAN_THRESHOLD
|
||||
assert audio_metrics["cosine_similarity"] >= LTX25_AUDIO_COSINE_THRESHOLD
|
||||
assert audio_metrics["relative_l2"] <= LTX25_AUDIO_RELATIVE_L2_THRESHOLD
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@@ -725,7 +729,6 @@ def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pi
|
||||
"omni_model": omni_model_source,
|
||||
"omni_model_revision": omni_model_revision,
|
||||
"resolved_omni_model_path": str(omni_model),
|
||||
"thresholds": asdict(STRICT_THRESHOLDS),
|
||||
"peak_memory_mb": {
|
||||
"official_allocated": official_metadata["peak_memory_allocated_mb"],
|
||||
"official_reserved": official_metadata["peak_memory_reserved_mb"],
|
||||
@@ -739,4 +742,12 @@ def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pi
|
||||
(output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
|
||||
print(json.dumps(result, indent=2))
|
||||
|
||||
_assert_strict_similarity(video_metrics, audio_metrics)
|
||||
assert video_metrics["ssim_mean"] >= LTX25_FULL_VIDEO_SSIM_MEAN_THRESHOLD
|
||||
assert video_metrics["ssim_min"] >= LTX25_FULL_VIDEO_SSIM_MIN_THRESHOLD
|
||||
assert video_metrics["psnr_mean_db"] >= LTX25_FULL_VIDEO_PSNR_MEAN_THRESHOLD
|
||||
lpips_threshold = (
|
||||
LTX25_FULL_VIDEO_LPIPS_MEAN_THRESHOLD if task == "t2v" else LTX25_FULL_I2V_VIDEO_LPIPS_MEAN_THRESHOLD
|
||||
)
|
||||
assert audio_metrics["relative_l2"] <= LTX25_FULL_AUDIO_RELATIVE_L2_THRESHOLD
|
||||
assert video_metrics["lpips_alex_mean"] <= lpips_threshold
|
||||
assert audio_metrics["cosine_similarity"] >= LTX25_FULL_AUDIO_COSINE_THRESHOLD
|
||||
@@ -3,8 +3,9 @@
|
||||
|
||||
"""E2E accuracy guard against a pinned Lightricks LTX pipeline revision.
|
||||
|
||||
Four complementary default-shape cases cover both model versions, both
|
||||
two-stage weight families, and T2V/I2V without duplicating one-stage coverage.
|
||||
The original reduced one-stage guards remain unchanged. Four complementary
|
||||
default-shape cases cover both model versions, both two-stage weight families,
|
||||
and T2V/I2V without expanding every combination into a duplicate golden.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -68,23 +69,25 @@ class LTXAccuracyThresholds:
|
||||
video_ssim_mean: float
|
||||
video_ssim_min: float
|
||||
video_psnr_mean_db: float
|
||||
# Zero disables gating for waveform-sensitive audio that is not expected
|
||||
# to align with the official output.
|
||||
audio_relative_l2: float
|
||||
audio_cosine_similarity: float
|
||||
|
||||
|
||||
STRICT_THRESHOLDS = LTXAccuracyThresholds(
|
||||
video_ssim_mean=0.99,
|
||||
video_ssim_min=0.99,
|
||||
video_psnr_mean_db=40.0,
|
||||
audio_relative_l2=0.10,
|
||||
audio_cosine_similarity=0.99,
|
||||
video_ssim_mean=0.95,
|
||||
video_ssim_min=0.90,
|
||||
video_psnr_mean_db=30.0,
|
||||
audio_relative_l2=0.20,
|
||||
audio_cosine_similarity=0.95,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LTXAccuracyCase:
|
||||
name: str
|
||||
pipeline_kind: Literal["distilled", "two_stage"]
|
||||
pipeline_kind: Literal["one_stage", "distilled", "two_stage"]
|
||||
model_id: str
|
||||
model_revision: str
|
||||
model_env: str
|
||||
@@ -96,6 +99,8 @@ class LTXAccuracyCase:
|
||||
num_inference_steps: int
|
||||
seed: int
|
||||
stg_block: int | None
|
||||
thresholds: LTXAccuracyThresholds
|
||||
prompt: str = PROMPT
|
||||
gemma_model_id: str | None = None
|
||||
gemma_model_revision: str | None = None
|
||||
gemma_model_env: str | None = None
|
||||
@@ -163,6 +168,59 @@ I2V_IMAGE = LTXArtifact(
|
||||
)
|
||||
|
||||
|
||||
LEGACY_CASES = (
|
||||
LTXAccuracyCase(
|
||||
name="ltx2",
|
||||
pipeline_kind="one_stage",
|
||||
model_id="Lightricks/LTX-2",
|
||||
model_revision=LTX2_REVISION,
|
||||
model_env="VLLM_TEST_LTX2_MODEL",
|
||||
model_class_name="LTX2Pipeline",
|
||||
checkpoint=LTX2_CHECKPOINT,
|
||||
width=512,
|
||||
height=384,
|
||||
num_frames=25,
|
||||
num_inference_steps=20,
|
||||
seed=42,
|
||||
stg_block=29,
|
||||
thresholds=STRICT_THRESHOLDS,
|
||||
),
|
||||
LTXAccuracyCase(
|
||||
name="ltx2_3",
|
||||
pipeline_kind="one_stage",
|
||||
model_id="diffusers/LTX-2.3-Diffusers",
|
||||
model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
|
||||
model_env="VLLM_TEST_LTX23_MODEL",
|
||||
model_class_name="LTX2Pipeline",
|
||||
checkpoint=LTX23_CHECKPOINT,
|
||||
width=512,
|
||||
height=384,
|
||||
num_frames=25,
|
||||
num_inference_steps=20,
|
||||
seed=42,
|
||||
stg_block=28,
|
||||
thresholds=STRICT_THRESHOLDS,
|
||||
),
|
||||
LTXAccuracyCase(
|
||||
name="ltx2_3_i2v",
|
||||
pipeline_kind="one_stage",
|
||||
model_id="diffusers/LTX-2.3-Diffusers",
|
||||
model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
|
||||
model_env="VLLM_TEST_LTX23_MODEL",
|
||||
model_class_name="LTX2Pipeline",
|
||||
checkpoint=LTX23_CHECKPOINT,
|
||||
width=512,
|
||||
height=384,
|
||||
num_frames=25,
|
||||
num_inference_steps=20,
|
||||
seed=42,
|
||||
stg_block=28,
|
||||
thresholds=STRICT_THRESHOLDS,
|
||||
image=I2V_IMAGE,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
TWO_STAGE_CASES = (
|
||||
LTXAccuracyCase(
|
||||
name="ltx2_distilled_t2v",
|
||||
@@ -182,6 +240,7 @@ TWO_STAGE_CASES = (
|
||||
gemma_model_id="Lightricks/LTX-2",
|
||||
gemma_model_revision=LTX2_REVISION,
|
||||
gemma_model_env="VLLM_TEST_LTX2_MODEL",
|
||||
thresholds=STRICT_THRESHOLDS,
|
||||
),
|
||||
LTXAccuracyCase(
|
||||
name="ltx23_distilled_i2v",
|
||||
@@ -201,6 +260,7 @@ TWO_STAGE_CASES = (
|
||||
gemma_model_id="diffusers/LTX-2.3-Diffusers",
|
||||
gemma_model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
|
||||
gemma_model_env="VLLM_TEST_LTX23_MODEL",
|
||||
thresholds=STRICT_THRESHOLDS,
|
||||
image=I2V_IMAGE,
|
||||
),
|
||||
LTXAccuracyCase(
|
||||
@@ -220,6 +280,7 @@ TWO_STAGE_CASES = (
|
||||
seed=10,
|
||||
stg_block=29,
|
||||
enable_layerwise_offload=True,
|
||||
thresholds=STRICT_THRESHOLDS,
|
||||
),
|
||||
LTXAccuracyCase(
|
||||
name="ltx23_two_stage_layer_fused_i2v",
|
||||
@@ -238,11 +299,12 @@ TWO_STAGE_CASES = (
|
||||
seed=10,
|
||||
stg_block=28,
|
||||
enable_layerwise_offload=True,
|
||||
thresholds=STRICT_THRESHOLDS,
|
||||
image=I2V_IMAGE,
|
||||
),
|
||||
)
|
||||
|
||||
WEEKLY_CASES = TWO_STAGE_CASES
|
||||
CASES = (*LEGACY_CASES, *TWO_STAGE_CASES)
|
||||
|
||||
|
||||
def _run(command: list[str], *, env: dict[str, str], timeout: int = 1800) -> None:
|
||||
@@ -398,49 +460,10 @@ def _resolve_image(case: LTXAccuracyCase) -> Path | None:
|
||||
return _resolve_artifact(case.image)
|
||||
|
||||
|
||||
def _omni_model_with_pinned_artifacts(
|
||||
model: Path,
|
||||
output_root: Path,
|
||||
artifacts: tuple[Path | None, ...],
|
||||
) -> Path:
|
||||
"""Expose the already-resolved sidecars to Omni without copying the model."""
|
||||
resolved_artifacts = tuple(artifact for artifact in artifacts if artifact is not None)
|
||||
if not resolved_artifacts:
|
||||
return model
|
||||
|
||||
overlay = output_root / "omni-model"
|
||||
overlay.mkdir()
|
||||
for source in model.iterdir():
|
||||
(overlay / source.name).symlink_to(source.resolve(), target_is_directory=source.is_dir())
|
||||
for artifact in resolved_artifacts:
|
||||
target = overlay / artifact.name
|
||||
if target.exists():
|
||||
assert target.samefile(artifact), f"Model sidecar does not match pinned artifact: {target} != {artifact}"
|
||||
else:
|
||||
target.symlink_to(artifact.resolve())
|
||||
return overlay
|
||||
|
||||
|
||||
def test_omni_model_overlay_pins_artifacts(tmp_path: Path) -> None:
|
||||
model = tmp_path / "model"
|
||||
model.mkdir()
|
||||
component = model / "model_index.json"
|
||||
component.write_text("{}\n")
|
||||
artifact = tmp_path / "sidecar.safetensors"
|
||||
artifact.write_bytes(b"pinned")
|
||||
output_root = tmp_path / "output"
|
||||
output_root.mkdir()
|
||||
|
||||
overlay = _omni_model_with_pinned_artifacts(model, output_root, (artifact,))
|
||||
|
||||
assert (overlay / component.name).samefile(component)
|
||||
assert (overlay / artifact.name).samefile(artifact)
|
||||
|
||||
|
||||
def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]:
|
||||
request: dict[str, object] = {
|
||||
"pipeline_kind": case.pipeline_kind,
|
||||
"prompt": PROMPT,
|
||||
"prompt": case.prompt,
|
||||
"negative_prompt": "" if case.pipeline_kind == "distilled" else NEGATIVE_PROMPT,
|
||||
"width": case.width,
|
||||
"height": case.height,
|
||||
@@ -450,17 +473,14 @@ def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]:
|
||||
"seed": case.seed,
|
||||
}
|
||||
if case.pipeline_kind != "distilled":
|
||||
assert case.stg_block is not None
|
||||
request.update(
|
||||
video_cfg_scale=3.0,
|
||||
audio_cfg_scale=7.0,
|
||||
video_modality_scale=3.0,
|
||||
audio_modality_scale=3.0,
|
||||
)
|
||||
if case.pipeline_kind == "two_stage":
|
||||
assert case.stg_block is not None
|
||||
request.update(
|
||||
video_stg_scale=1.0,
|
||||
audio_stg_scale=1.0,
|
||||
video_modality_scale=3.0,
|
||||
audio_modality_scale=3.0,
|
||||
video_rescale_scale=0.7,
|
||||
audio_rescale_scale=0.7,
|
||||
video_stg_blocks=[case.stg_block],
|
||||
@@ -523,7 +543,7 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", WEEKLY_CASES, ids=lambda case: case.name)
|
||||
@pytest.mark.parametrize("case", CASES, ids=lambda case: case.name)
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.benchmark
|
||||
@pytest.mark.diffusion
|
||||
@@ -541,11 +561,9 @@ def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Pat
|
||||
checkpoint = _resolve_artifact(case.checkpoint, model)
|
||||
spatial_upsampler = _resolve_artifact(case.spatial_upsampler, model) if case.spatial_upsampler is not None else None
|
||||
distilled_lora = _resolve_artifact(case.distilled_lora, model) if case.distilled_lora is not None else None
|
||||
omni_model = _omni_model_with_pinned_artifacts(model, output_root, (spatial_upsampler, distilled_lora))
|
||||
image = _resolve_image(case)
|
||||
request = _request(case, image)
|
||||
request_path = output_root / "request.json"
|
||||
request_path.write_text(json.dumps(request, indent=2) + "\n")
|
||||
request_path.write_text(json.dumps(_request(case, image), indent=2) + "\n")
|
||||
|
||||
runner = Path(__file__).with_name("run_ltx_reference.py")
|
||||
runner_args = [
|
||||
@@ -601,7 +619,7 @@ def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Pat
|
||||
"--output-dir",
|
||||
str(omni_output),
|
||||
"--model",
|
||||
str(omni_model),
|
||||
str(model),
|
||||
"--model-class-name",
|
||||
case.model_class_name,
|
||||
]
|
||||
@@ -633,16 +651,18 @@ def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Pat
|
||||
"spatial_upsampler_revision": (case.spatial_upsampler.revision if case.spatial_upsampler is not None else None),
|
||||
"distilled_lora_revision": case.distilled_lora.revision if case.distilled_lora is not None else None,
|
||||
"enable_layerwise_offload": enable_layerwise_offload,
|
||||
"thresholds": asdict(STRICT_THRESHOLDS),
|
||||
"request": request,
|
||||
"thresholds": asdict(case.thresholds),
|
||||
"request": _request(case, image),
|
||||
"video": video_metrics,
|
||||
"audio": audio_metrics,
|
||||
}
|
||||
(output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
|
||||
print(json.dumps(result, indent=2))
|
||||
|
||||
assert video_metrics["ssim_mean"] >= STRICT_THRESHOLDS.video_ssim_mean
|
||||
assert video_metrics["ssim_min"] >= STRICT_THRESHOLDS.video_ssim_min
|
||||
assert video_metrics["psnr_mean_db"] >= STRICT_THRESHOLDS.video_psnr_mean_db
|
||||
assert audio_metrics["relative_l2"] <= STRICT_THRESHOLDS.audio_relative_l2
|
||||
assert audio_metrics["cosine_similarity"] >= STRICT_THRESHOLDS.audio_cosine_similarity
|
||||
assert video_metrics["ssim_mean"] >= case.thresholds.video_ssim_mean
|
||||
assert video_metrics["ssim_min"] >= case.thresholds.video_ssim_min
|
||||
assert video_metrics["psnr_mean_db"] >= case.thresholds.video_psnr_mean_db
|
||||
if case.thresholds.audio_relative_l2 > 0:
|
||||
assert audio_metrics["relative_l2"] <= case.thresholds.audio_relative_l2
|
||||
if case.thresholds.audio_cosine_similarity > 0:
|
||||
assert audio_metrics["cosine_similarity"] >= case.thresholds.audio_cosine_similarity
|
||||
@@ -1,94 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.fullduplex.engine.duplex_runtime import duplex_resource_request_id
|
||||
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
|
||||
from vllm_omni.experimental.fullduplex.nemotron_voicechat.data_plane import (
|
||||
NemotronVoiceChatDataPlaneContext,
|
||||
NemotronVoiceChatDataPlaneSession,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
_RUNTIME = {
|
||||
"nvc_text_bos_id": 0,
|
||||
"nvc_text_eos_id": 1,
|
||||
"nvc_text_pad_id": 12,
|
||||
"nvc_function_sotc_id": 20,
|
||||
"nvc_function_eotc_id": 21,
|
||||
"nvc_tokenizer_ref": "test-tokenizer",
|
||||
}
|
||||
|
||||
|
||||
def _projector(encode_audio=lambda *_: None) -> NemotronVoiceChatDataPlaneSession:
|
||||
projector = NemotronVoiceChatDataPlaneSession(encode_audio)
|
||||
projector.configure_runtime(_RUNTIME, tokenizer=SimpleNamespace(decode=lambda *_args, **_kwargs: "text"))
|
||||
return projector
|
||||
|
||||
|
||||
def _stage0_output(request_id: str = "req-0", *, text_token: int = 12, function_token: int | None = None) -> object:
|
||||
metadata: dict[str, object] = {"nvc_text_token_ids": [text_token]}
|
||||
if function_token is not None:
|
||||
metadata["nvc_function_token"] = [function_token]
|
||||
completion = SimpleNamespace(multimodal_output=metadata)
|
||||
return SimpleNamespace(
|
||||
stage_id=0,
|
||||
request_output=SimpleNamespace(request_id=request_id, outputs=[completion]),
|
||||
)
|
||||
|
||||
|
||||
def test_function_channel_projects_completed_call_without_ending_speech() -> None:
|
||||
projector = _projector()
|
||||
projector._decode = lambda token_ids: (
|
||||
'[{"name":"weather","arguments":{"city":"Shanghai"}}]' if token_ids == [99] else ""
|
||||
)
|
||||
context = NemotronVoiceChatDataPlaneContext(epoch=0)
|
||||
|
||||
events = []
|
||||
for function_token in (20, 99, 21):
|
||||
events.extend(projector.project_output(_stage0_output(function_token=function_token), context=context))
|
||||
|
||||
listen = [event for event in events if event.get("is_listen") is True]
|
||||
function = [event for event in events if event.get("function_call") is True]
|
||||
assert len(listen) == 3
|
||||
assert len(function) == 1
|
||||
assert function[0]["name"] == "weather"
|
||||
assert function[0]["arguments"] == '{"city":"Shanghai"}'
|
||||
|
||||
projector._decode = lambda _token_ids: "not-json"
|
||||
bad_events = []
|
||||
for function_token in (20, 98, 21):
|
||||
bad_events.extend(projector.project_output(_stage0_output(function_token=function_token)))
|
||||
assert any(event.get("error_code") == "nemotron_function_call_parse_error" for event in bad_events)
|
||||
|
||||
|
||||
def test_new_epoch_request_does_not_inherit_partial_function_or_frame_state() -> None:
|
||||
projector = _projector(lambda *_: "audio")
|
||||
old_request = duplex_resource_request_id(DuplexFence("sid", epoch=1), "stage0")
|
||||
new_request = duplex_resource_request_id(DuplexFence("sid", epoch=2), "stage0")
|
||||
|
||||
projector.begin_request(old_request)
|
||||
list(projector.project_output(_stage0_output(old_request, text_token=42, function_token=20)))
|
||||
projector.begin_request(new_request)
|
||||
|
||||
# EOTC in the new epoch must not close the old epoch's partial function.
|
||||
new_events = list(projector.project_output(_stage0_output(new_request, text_token=1, function_token=21)))
|
||||
assert not [event for event in new_events if event.get("function_call") is True]
|
||||
assert not [event for event in new_events if event.get("end_of_turn") is True]
|
||||
|
||||
audio = SimpleNamespace(
|
||||
stage_id=2,
|
||||
request_output=SimpleNamespace(
|
||||
request_id=new_request,
|
||||
outputs=[
|
||||
SimpleNamespace(multimodal_output={"model_outputs": [np.ones(1764, dtype=np.float32)], "sr": [22050]})
|
||||
],
|
||||
),
|
||||
)
|
||||
audio_events = list(projector.project_output(audio))
|
||||
assert sum(event.get("end_of_turn") is True for event in audio_events) == 1
|
||||
@@ -1,73 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.fullduplex.nemotron_voicechat.input import (
|
||||
NemotronVoiceChatPcmAppendBuffer,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _payload(samples: np.ndarray, *, sample_rate_hz: int = 16000, fmt: str = "pcm_f32le") -> dict[str, object]:
|
||||
return {
|
||||
"type": "audio",
|
||||
"audio": base64.b64encode(np.asarray(samples, dtype=np.float32).tobytes()).decode("ascii"),
|
||||
"format": fmt,
|
||||
"sample_rate_hz": sample_rate_hz,
|
||||
}
|
||||
|
||||
|
||||
def _append(buffer, samples, operation_id):
|
||||
return buffer.prepare_append(
|
||||
_payload(np.asarray(samples, dtype=np.float32)),
|
||||
operation_id=operation_id,
|
||||
chunk_period_ms=80,
|
||||
allow_emit=True,
|
||||
)
|
||||
|
||||
|
||||
def test_irregular_browser_packets_emit_exact_ordered_80ms_frame() -> None:
|
||||
buffer = NemotronVoiceChatPcmAppendBuffer()
|
||||
assert _append(buffer, np.arange(1000), "packet-1") is None
|
||||
|
||||
reservation = _append(buffer, np.arange(1000, 1280), "packet-2")
|
||||
|
||||
assert reservation is not None
|
||||
assert reservation.operation_id == "packet-2"
|
||||
decoded = np.frombuffer(base64.b64decode(reservation.payload["audio"]), dtype=np.float32)
|
||||
np.testing.assert_array_equal(decoded, np.arange(1280, dtype=np.float32))
|
||||
assert buffer.pending_byte_count == 0
|
||||
|
||||
|
||||
def test_append_rejects_invalid_frame_contract() -> None:
|
||||
cases = (
|
||||
(1281, 16000, "pcm_f32le", 0.0, "at most 1280"),
|
||||
(1, 8000, "pcm_f32le", 0.0, "sample_rate_hz"),
|
||||
(1, 16000, "audio/pcm", 0.0, "format"),
|
||||
(1, 16000, "pcm_f32le", np.nan, "finite"),
|
||||
)
|
||||
for samples, sample_rate_hz, fmt, value, message in cases:
|
||||
with pytest.raises(ValueError, match=message):
|
||||
NemotronVoiceChatPcmAppendBuffer().prepare_append(
|
||||
_payload(np.full(samples, value, dtype=np.float32), sample_rate_hz=sample_rate_hz, fmt=fmt),
|
||||
operation_id="invalid",
|
||||
chunk_period_ms=80,
|
||||
allow_emit=True,
|
||||
)
|
||||
|
||||
|
||||
def test_terminal_commit_rejects_more_than_one_frame_tail() -> None:
|
||||
buffer = NemotronVoiceChatPcmAppendBuffer()
|
||||
for index, samples in enumerate((1000, 281)):
|
||||
buffer.prepare_append(
|
||||
_payload(np.zeros(samples, dtype=np.float32)),
|
||||
operation_id=f"tail-{index}",
|
||||
chunk_period_ms=80,
|
||||
allow_emit=False,
|
||||
)
|
||||
with pytest.raises(ValueError, match="more than one 80 ms frame"):
|
||||
buffer.prepare_commit(operation_id="commit", chunk_period_ms=80)
|
||||
@@ -1,115 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
from vllm_omni.experimental.fullduplex.engine.contracts import DuplexInputMode
|
||||
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
|
||||
from vllm_omni.experimental.fullduplex.nemotron_voicechat.runtime import (
|
||||
NemotronVoiceChatDuplexRuntimeExtension,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.nemotron_voicechat.serving_adapter import (
|
||||
NemotronVoiceChatServingRuntimeAdapter,
|
||||
_render_tool_response,
|
||||
)
|
||||
from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
|
||||
NemotronVoiceChatThinkerForConditionalGeneration,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _frame() -> dict[str, object]:
|
||||
raw = np.zeros(1280, dtype=np.float32).tobytes()
|
||||
return {
|
||||
"type": "audio",
|
||||
"audio": base64.b64encode(raw).decode("ascii"),
|
||||
"format": "pcm_f32le",
|
||||
"sample_rate_hz": 16000,
|
||||
}
|
||||
|
||||
|
||||
def _plan(extension, *, input_seq: int):
|
||||
return extension.plan_append(
|
||||
request_id="req",
|
||||
fence=DuplexFence("sid", incarnation=2, epoch=3),
|
||||
session_config={},
|
||||
runtime_config={
|
||||
"nvc_prompt_token_ids": [0, 42, 1],
|
||||
"nvc_text_pad_id": 12,
|
||||
"nvc_max_model_len": 8192,
|
||||
},
|
||||
seq=input_seq,
|
||||
turn_seq=input_seq,
|
||||
mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
|
||||
payload=_frame(),
|
||||
final=False,
|
||||
sampling_params=SamplingParams(),
|
||||
)
|
||||
|
||||
|
||||
def test_first_append_prefills_prompt_then_each_append_consumes_one_frame() -> None:
|
||||
extension = NemotronVoiceChatDuplexRuntimeExtension()
|
||||
|
||||
params = extension.configure_sampling_params(
|
||||
runtime_config={}, defaults=(SamplingParams(temperature=0.7, max_tokens=4),)
|
||||
)[0]
|
||||
assert params.temperature == 0.0 and params.max_tokens == 1 and params.top_p == 1.0 and params.top_k == 0
|
||||
|
||||
first = _plan(extension, input_seq=1)
|
||||
later = _plan(extension, input_seq=2)
|
||||
|
||||
assert first.prompt["prompt_token_ids"] == [0, 42, 1, 12]
|
||||
assert later.prompt["prompt_token_ids"] == [12]
|
||||
assert first.prompt["model_intermediate_buffer"]["duplex"]["source_input_seq"] == 1
|
||||
assert later.prompt["model_intermediate_buffer"]["duplex"]["source_input_seq"] == 2
|
||||
|
||||
|
||||
def test_append_rejects_stage0_context_overflow() -> None:
|
||||
with pytest.raises(ValueError, match="max_model_len"):
|
||||
_plan(NemotronVoiceChatDuplexRuntimeExtension(), input_seq=8190)
|
||||
|
||||
|
||||
def test_function_output_becomes_versioned_nvidia_channel_tokens() -> None:
|
||||
encoded: list[str] = []
|
||||
adapter = NemotronVoiceChatServingRuntimeAdapter(lambda *_: None)
|
||||
adapter._tokenizer = SimpleNamespace(
|
||||
encode=lambda text, **_kwargs: encoded.append(text) or [31, 32, 33],
|
||||
)
|
||||
|
||||
first = adapter.runtime_config_for_function_output(
|
||||
{"type": "function_call_output", "call_id": "call-1", "output": '{"result":20}'},
|
||||
{},
|
||||
)
|
||||
second = adapter.runtime_config_for_function_output(
|
||||
{"type": "function_call_output", "call_id": "call-2", "output": "plain text"},
|
||||
first,
|
||||
)
|
||||
|
||||
assert _render_tool_response('{"result":20}') == '<TOOL_RESPONSE>[{"result":20}]</TOOL_RESPONSE>'
|
||||
assert encoded == [
|
||||
'<TOOL_RESPONSE>[{"result":20}]</TOOL_RESPONSE>',
|
||||
'<TOOL_RESPONSE>["plain text"]</TOOL_RESPONSE>',
|
||||
]
|
||||
assert first["nvc_function_response_generation"] == 1
|
||||
assert second["nvc_function_response_generation"] == 2
|
||||
assert second["nvc_function_response_token_ids"] == [31, 32, 33]
|
||||
assert second["nvc_function_response_call_id"] == "call-2"
|
||||
assert second["nvc_function_response_batches"] == [
|
||||
{"generation": 1, "call_id": "call-1", "token_ids": [31, 32, 33]},
|
||||
{"generation": 2, "call_id": "call-2", "token_ids": [31, 32, 33]},
|
||||
]
|
||||
|
||||
session = {"function_response_generation": 0, "forced_function_tokens": []}
|
||||
NemotronVoiceChatThinkerForConditionalGeneration._sync_forced_function_response(session, second)
|
||||
|
||||
assert session["function_response_generation"] == 2
|
||||
assert session["forced_function_token"] == 31
|
||||
assert session["forced_function_tokens"] == [31, 32, 33, 31, 32, 33]
|
||||
@@ -18,10 +18,6 @@ from vllm_omni.experimental.fullduplex.engine.duplex_runtime import (
|
||||
DuplexInputMode,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
|
||||
from vllm_omni.experimental.fullduplex.openai.runtime_adapter import (
|
||||
ServingRuntimeConfigError,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.personaplex.config import DEFAULT_PERSONA
|
||||
from vllm_omni.experimental.fullduplex.personaplex.data_plane import (
|
||||
PersonaPlexDataPlaneContext,
|
||||
PersonaPlexDataPlaneSession,
|
||||
@@ -246,53 +242,6 @@ def test_personaplex_capabilities_are_honest() -> None:
|
||||
assert multi.supports_barge_in is False
|
||||
|
||||
|
||||
def test_personaplex_runtime_config_update_rejects_changed_persona() -> None:
|
||||
adapter = PersonaPlexServingRuntimeAdapter(lambda *_: None)
|
||||
short_persona = "You are a helpful assistant."
|
||||
longer_persona = (
|
||||
"You are now a swashbuckling pirate captain who speaks entirely in "
|
||||
"nautical slang, references buried treasure constantly, and never "
|
||||
"breaks character no matter what the user asks."
|
||||
)
|
||||
current = {"personaplex_persona": short_persona, "personaplex_voice_prompt": "NATF2.pt"}
|
||||
|
||||
with pytest.raises(ServingRuntimeConfigError, match="persona"):
|
||||
adapter.runtime_config_for_update(
|
||||
SimpleNamespace(instructions=longer_persona, extra_body={}),
|
||||
current,
|
||||
)
|
||||
|
||||
unchanged = adapter.runtime_config_for_update(
|
||||
SimpleNamespace(instructions=short_persona, extra_body={}),
|
||||
current,
|
||||
)
|
||||
assert unchanged["personaplex_persona"] == short_persona
|
||||
|
||||
|
||||
def test_personaplex_runtime_config_update_rejects_changed_voice() -> None:
|
||||
adapter = PersonaPlexServingRuntimeAdapter(lambda *_: None)
|
||||
current = {"personaplex_persona": DEFAULT_PERSONA, "personaplex_voice_prompt": "NATF2.pt"}
|
||||
|
||||
with pytest.raises(ServingRuntimeConfigError, match="voice"):
|
||||
adapter.runtime_config_for_update(
|
||||
SimpleNamespace(instructions=None, voice="OtherVoice.pt", extra_body={}),
|
||||
current,
|
||||
)
|
||||
|
||||
unchanged = adapter.runtime_config_for_update(
|
||||
SimpleNamespace(instructions=None, voice="NATF2.pt", extra_body={}),
|
||||
current,
|
||||
)
|
||||
assert unchanged["personaplex_voice_prompt"] == "NATF2.pt"
|
||||
|
||||
|
||||
def test_personaplex_rejects_minicpmo_only_opt_out_flag() -> None:
|
||||
adapter = PersonaPlexServingRuntimeAdapter(lambda *_: None)
|
||||
|
||||
with pytest.raises(ServingRuntimeConfigError, match="minicpmo45_native_duplex"):
|
||||
adapter.validate_client_extra_body({"minicpmo45_native_duplex": False})
|
||||
|
||||
|
||||
def test_pcm_buffer_emits_one_80ms_frame_transactionally() -> None:
|
||||
buffer = PersonaPlexPcmAppendBuffer()
|
||||
reservation = buffer.prepare_append(
|
||||
|
||||
@@ -1,194 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import threading
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.benchmarks.duplex import omni_duplex_eval_runner as runner
|
||||
from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import DuplexSample
|
||||
from vllm_omni.benchmarks.duplex.omni_duplex_eval_judge import DuplexJudge
|
||||
from vllm_omni.entrypoints.cli.benchmark import omni_duplex_eval as cli
|
||||
from vllm_omni.entrypoints.cli.benchmark.omni_duplex_eval import OmniDuplexEvalSubcommand
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
|
||||
|
||||
|
||||
def test_cli_generate_evaluate_summarize_flow(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys):
|
||||
manifest = tmp_path / "manifest.json"
|
||||
manifest.write_text(
|
||||
json.dumps(
|
||||
[
|
||||
{
|
||||
"id": "sample-1",
|
||||
"split": "PR_correction",
|
||||
"question_text": "What changed?",
|
||||
"answer1": "The object moved.",
|
||||
}
|
||||
]
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
response_root = tmp_path / "responses"
|
||||
score_root = tmp_path / "scores"
|
||||
|
||||
async def fake_generate(sample, *, output_root, **kwargs):
|
||||
output = Path(output_root) / sample.split / f"{sample.id}.json"
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
output.write_text(json.dumps([{"sentence": "It moved.", "start": 0, "end": 1}]), encoding="utf-8")
|
||||
output.with_name(output.stem + ".meta.json").write_text(json.dumps({"clock": "media"}), encoding="utf-8")
|
||||
return output
|
||||
|
||||
class FakeJudge:
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def chat(self, *args, **kwargs):
|
||||
return '{"success_score": 1, "is_relevant": 1}'
|
||||
|
||||
monkeypatch.setattr(cli, "generate_sample", fake_generate)
|
||||
monkeypatch.setattr(cli, "DuplexJudge", FakeJudge)
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
OmniDuplexEvalSubcommand.add_cli_args(parser)
|
||||
|
||||
common = ["--dataset", str(manifest), "--family", "pr"]
|
||||
generate = parser.parse_args(
|
||||
[
|
||||
"generate",
|
||||
*common,
|
||||
"--model",
|
||||
"mock",
|
||||
"--ref-audio",
|
||||
str(manifest),
|
||||
"--response-root",
|
||||
str(response_root),
|
||||
"--concurrency",
|
||||
"2",
|
||||
]
|
||||
)
|
||||
OmniDuplexEvalSubcommand.cmd(generate)
|
||||
evaluate = parser.parse_args(
|
||||
[
|
||||
"evaluate",
|
||||
*common,
|
||||
"--response-root",
|
||||
str(response_root),
|
||||
"--score-root",
|
||||
str(score_root),
|
||||
"--judge-model",
|
||||
"mock-judge",
|
||||
"--eval-workers",
|
||||
"2",
|
||||
]
|
||||
)
|
||||
OmniDuplexEvalSubcommand.cmd(evaluate)
|
||||
OmniDuplexEvalSubcommand.cmd(parser.parse_args(["summarize", "--score-root", str(score_root)]))
|
||||
|
||||
summary = json.loads(capsys.readouterr().out)
|
||||
assert summary["samples"] == 1
|
||||
assert summary["pr"]["mean_all_success"] == 1.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_exercises_realtime_socket_and_media_clock(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
import websockets
|
||||
|
||||
received = []
|
||||
|
||||
async def handler(websocket):
|
||||
response_started = False
|
||||
async for raw in websocket:
|
||||
event = json.loads(raw)
|
||||
received.append(event)
|
||||
if event["type"] == "session.update":
|
||||
await websocket.send(json.dumps({"type": "session.created"}))
|
||||
elif event["type"] == "input_audio_buffer.append" and not response_started:
|
||||
response_started = True
|
||||
await websocket.send(json.dumps({"type": "response.created", "response": {"id": "r1"}}))
|
||||
await websocket.send(
|
||||
json.dumps({"type": "response.output_text.delta", "response_id": "r1", "delta": "Done."})
|
||||
)
|
||||
await websocket.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "response.audio.delta",
|
||||
"response_id": "r1",
|
||||
"delta": base64.b64encode(b"\0\0" * 240).decode(),
|
||||
"sample_rate_hz": 24_000,
|
||||
}
|
||||
)
|
||||
)
|
||||
elif event["type"] == "input_audio_buffer.commit":
|
||||
await websocket.send(json.dumps({"type": "response.done", "response": {"id": "r1"}}))
|
||||
elif event["type"] == "session.close":
|
||||
await websocket.send(json.dumps({"type": "session.closed"}))
|
||||
return
|
||||
|
||||
monkeypatch.setattr(runner, "read_audio_pcm16", lambda path: b"\0\0" * (16_000 * 8 // 10))
|
||||
monkeypatch.setattr(runner, "video_duration", lambda path: 0.8)
|
||||
monkeypatch.setattr(runner, "iter_jpegs", lambda *args, **kwargs: iter([(0.0, b"jpeg")]))
|
||||
audio = tmp_path / "question.wav"
|
||||
video = tmp_path / "video.mp4"
|
||||
ref = tmp_path / "ref.wav"
|
||||
for path in (audio, video, ref):
|
||||
path.write_bytes(b"data")
|
||||
sample = DuplexSample("sample", "PR_correction", "pr", "correction", video, audio)
|
||||
|
||||
async with websockets.serve(handler, "127.0.0.1", 0) as server:
|
||||
port = server.sockets[0].getsockname()[1]
|
||||
output = await runner.generate_sample(
|
||||
sample,
|
||||
url=f"ws://127.0.0.1:{port}/v1/realtime?duplex=1",
|
||||
model="mock",
|
||||
ref_audio=ref,
|
||||
output_root=tmp_path / "responses",
|
||||
)
|
||||
|
||||
assert json.loads(output.read_text(encoding="utf-8")) == [{"sentence": "Done.", "start": 0.8, "end": 0.8}]
|
||||
meta = json.loads(output.with_name("sample.meta.json").read_text(encoding="utf-8"))
|
||||
assert meta["response_done"] is True
|
||||
assert meta["drain_timeout"] is None
|
||||
assert any(event["type"] == "playback.ack" for event in received)
|
||||
|
||||
|
||||
def test_judge_exercises_openai_http_schema():
|
||||
requests = []
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def do_POST(self):
|
||||
length = int(self.headers["Content-Length"])
|
||||
requests.append((self.path, self.headers["Authorization"], json.loads(self.rfile.read(length))))
|
||||
payload = json.dumps({"choices": [{"message": {"content": '{"success_score": 1}'}}]}).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(payload)))
|
||||
self.end_headers()
|
||||
self.wfile.write(payload)
|
||||
|
||||
def log_message(self, *args):
|
||||
pass
|
||||
|
||||
server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
|
||||
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
judge = DuplexJudge(f"http://127.0.0.1:{server.server_port}", "judge", api_key="token")
|
||||
assert judge.chat("prompt") == '{"success_score": 1}'
|
||||
finally:
|
||||
server.shutdown()
|
||||
thread.join()
|
||||
server.server_close()
|
||||
|
||||
path, authorization, payload = requests[0]
|
||||
assert path == "/v1/chat/completions"
|
||||
assert authorization == "Bearer token"
|
||||
assert payload["model"] == "judge"
|
||||
assert payload["messages"] == [{"role": "user", "content": "prompt"}]
|
||||
@@ -0,0 +1,290 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""E2E test for DreamID-Omni offline inference with Cache-Dit.
|
||||
|
||||
Uses the official DreamID-Omni ``test_case/twoip`` example assets instead of
|
||||
synthetic inputs so the coverage matches the upstream demo structure.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
from typing import NamedTuple
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import soundfile as sf
|
||||
from PIL import Image
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniRunner
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.cache, pytest.mark.slow]
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
DEFAULT_MODEL_PATH = REPO_ROOT / "dreamid_omni"
|
||||
MODEL_PATH_ENV = "DREAMID_OMNI_MODEL_PATH"
|
||||
EXAMPLE_ROOT_ENV = "DREAMID_OMNI_EXAMPLE_ROOT"
|
||||
NUM_INFERENCE_STEPS = 45
|
||||
|
||||
CACHE_CONFIG = {
|
||||
"Fn_compute_blocks": 1,
|
||||
"Bn_compute_blocks": 0,
|
||||
"max_warmup_steps": 4,
|
||||
"max_cached_steps": 20,
|
||||
"residual_diff_threshold": 0.24,
|
||||
"max_continuous_cached_steps": 3,
|
||||
"enable_taylorseer": False,
|
||||
"taylorseer_order": 1,
|
||||
"scm_steps_mask_policy": None,
|
||||
"scm_steps_policy": "dynamic",
|
||||
}
|
||||
|
||||
|
||||
class DreamIDGenerationResult(NamedTuple):
|
||||
video: np.ndarray
|
||||
audio: np.ndarray
|
||||
fps: int
|
||||
duration_s: float
|
||||
elapsed_s: float
|
||||
|
||||
|
||||
class DreamIDExampleCase(NamedTuple):
|
||||
prompt: str
|
||||
image_paths: tuple[Path, ...]
|
||||
audio_paths: tuple[Path, ...]
|
||||
|
||||
|
||||
def _require_model_path() -> str:
|
||||
model_path = os.environ.get(MODEL_PATH_ENV, str(DEFAULT_MODEL_PATH))
|
||||
if not Path(model_path).exists():
|
||||
pytest.skip(f"DreamID-Omni model bundle not found at {model_path}. Set {MODEL_PATH_ENV}.")
|
||||
|
||||
if importlib.util.find_spec("dreamid_omni") is None:
|
||||
pytest.skip(
|
||||
"dreamid_omni dependency package is not importable. "
|
||||
"Run examples/offline_inference/x_to_video_audio/download_dreamid_omni.py first."
|
||||
)
|
||||
|
||||
return model_path
|
||||
|
||||
|
||||
def _dreamid_repo_root_candidates() -> list[Path]:
|
||||
candidates: list[Path] = []
|
||||
example_root = os.environ.get(EXAMPLE_ROOT_ENV)
|
||||
if example_root:
|
||||
candidates.append(Path(example_root).expanduser().resolve())
|
||||
|
||||
spec = importlib.util.find_spec("dreamid_omni")
|
||||
if spec is not None and spec.origin:
|
||||
candidates.append(Path(spec.origin).resolve().parents[1] / "test_case" / "twoip")
|
||||
|
||||
return candidates
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _load_example_case() -> DreamIDExampleCase:
|
||||
for candidate_root in _dreamid_repo_root_candidates():
|
||||
caption_path = candidate_root / "captions" / "20.json"
|
||||
image_paths = (
|
||||
candidate_root / "imgs" / "20" / "0.png",
|
||||
candidate_root / "imgs" / "20" / "1.png",
|
||||
)
|
||||
audio_paths = (
|
||||
candidate_root / "audios" / "20" / "0.wav",
|
||||
candidate_root / "audios" / "20" / "1.wav",
|
||||
)
|
||||
if caption_path.exists() and all(path.exists() for path in image_paths + audio_paths):
|
||||
prompt = _load_prompt_from_json(caption_path)
|
||||
return DreamIDExampleCase(
|
||||
prompt=prompt,
|
||||
image_paths=image_paths,
|
||||
audio_paths=audio_paths,
|
||||
)
|
||||
|
||||
searched = ", ".join(str(path) for path in _dreamid_repo_root_candidates()) or "<none>"
|
||||
pytest.skip(
|
||||
"DreamID-Omni official example assets not found. "
|
||||
f"Set {EXAMPLE_ROOT_ENV} to the local 'test_case/twoip' directory. "
|
||||
f"Searched: {searched}"
|
||||
)
|
||||
|
||||
|
||||
def _load_prompt_from_json(path: Path) -> str:
|
||||
raw_text = path.read_text(encoding="utf-8").strip()
|
||||
parsed = json.loads(raw_text)
|
||||
if isinstance(parsed, str):
|
||||
prompt = parsed
|
||||
elif isinstance(parsed, dict) and isinstance(parsed.get("content"), str):
|
||||
prompt = parsed["content"]
|
||||
else:
|
||||
raise AssertionError(f"Unexpected DreamID caption format in {path}")
|
||||
|
||||
prompt = re.sub(
|
||||
r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]",
|
||||
"",
|
||||
prompt,
|
||||
flags=re.DOTALL,
|
||||
).strip()
|
||||
prompt = re.sub(
|
||||
r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]",
|
||||
"",
|
||||
prompt,
|
||||
flags=re.DOTALL,
|
||||
).strip()
|
||||
prompt = re.sub(r"\[[A-Z_]+\]", "", prompt)
|
||||
prompt = re.sub(r"\n\s*\n", "\n", prompt).strip()
|
||||
return prompt
|
||||
|
||||
|
||||
def _load_example_images(example_case: DreamIDExampleCase) -> list[Image.Image]:
|
||||
images: list[Image.Image] = []
|
||||
for path in example_case.image_paths:
|
||||
with Image.open(path) as image:
|
||||
images.append(image.convert("RGB"))
|
||||
return images
|
||||
|
||||
|
||||
def _load_example_audios(example_case: DreamIDExampleCase) -> list[np.ndarray]:
|
||||
audios: list[np.ndarray] = []
|
||||
for path in example_case.audio_paths:
|
||||
audio_array, sample_rate = sf.read(str(path), dtype="float32")
|
||||
audio_np = np.asarray(audio_array, dtype=np.float32)
|
||||
if audio_np.ndim == 2:
|
||||
audio_np = np.mean(audio_np, axis=1)
|
||||
audio_np = audio_np[int(sample_rate * 1) : int(sample_rate * 3)]
|
||||
audios.append(audio_np.reshape(-1))
|
||||
return audios
|
||||
|
||||
|
||||
def _sampling_params(num_inference_steps: int) -> OmniDiffusionSamplingParams:
|
||||
return OmniDiffusionSamplingParams(
|
||||
height=704,
|
||||
width=1280,
|
||||
num_inference_steps=num_inference_steps,
|
||||
seed=42,
|
||||
extra_args={
|
||||
"solver_name": "unipc",
|
||||
"shift": 5.0,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _request_payload(example_case: DreamIDExampleCase) -> dict[str, object]:
|
||||
return {
|
||||
"prompt": example_case.prompt,
|
||||
"video_negative_prompt": "jitter, bad hands, blur, distortion",
|
||||
"audio_negative_prompt": "robotic, muffled, echo, distorted",
|
||||
"multi_modal_data": {
|
||||
"image": _load_example_images(example_case),
|
||||
"audio": _load_example_audios(example_case),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _extract_generation(outputs: list[object], elapsed_s: float) -> DreamIDGenerationResult:
|
||||
"""Extract video and audio from raw Omni outputs"""
|
||||
assert outputs, "DreamID-Omni returned no outputs"
|
||||
|
||||
result = outputs[0]
|
||||
assert hasattr(result, "images") and result.images, "not return video frames"
|
||||
|
||||
video = result.images[0]
|
||||
assert isinstance(video, np.ndarray), f"Expected video ndarray, got {type(video)}"
|
||||
assert video.ndim == 4, f"Expected 4D DreamID video tensor, got {video.shape}"
|
||||
|
||||
multimodal_output = getattr(result, "multimodal_output", None) or {}
|
||||
assert "audio" in multimodal_output, "not return audio"
|
||||
audio = np.asarray(multimodal_output["audio"], dtype=np.float32).reshape(-1)
|
||||
assert audio.size > 0, "returned empty audio"
|
||||
|
||||
audio_sample_rate = int(multimodal_output.get("audio_sample_rate", 16000))
|
||||
fps = multimodal_output.get("fps", 24)
|
||||
duration_s = float(audio.size) / float(audio_sample_rate)
|
||||
|
||||
generation = DreamIDGenerationResult(
|
||||
video=video.astype(np.float32, copy=False),
|
||||
audio=audio,
|
||||
fps=fps,
|
||||
duration_s=duration_s,
|
||||
elapsed_s=elapsed_s,
|
||||
)
|
||||
return generation
|
||||
|
||||
|
||||
def _run_generation(
|
||||
*,
|
||||
model_path: str,
|
||||
example_case: DreamIDExampleCase,
|
||||
use_cache_dit: bool,
|
||||
) -> DreamIDGenerationResult:
|
||||
runner_kwargs = {
|
||||
"model_type": "dreamid-omni",
|
||||
"parallel_config": DiffusionParallelConfig(cfg_parallel_size=1),
|
||||
}
|
||||
if use_cache_dit:
|
||||
runner_kwargs["cache_backend"] = "cache_dit"
|
||||
runner_kwargs["cache_config"] = CACHE_CONFIG
|
||||
|
||||
with OmniRunner(model_path, **runner_kwargs) as runner:
|
||||
warmup_payload = _request_payload(example_case)
|
||||
_ = runner.omni.generate(warmup_payload, _sampling_params(2))
|
||||
|
||||
elapsed_samples: list[float] = []
|
||||
extracted_result: DreamIDGenerationResult | None = None
|
||||
|
||||
request_payload = _request_payload(example_case)
|
||||
start = time.perf_counter()
|
||||
outputs = runner.omni.generate(request_payload, _sampling_params(NUM_INFERENCE_STEPS))
|
||||
elapsed_s = time.perf_counter() - start
|
||||
elapsed_samples.append(elapsed_s)
|
||||
if extracted_result is None:
|
||||
extracted_result = _extract_generation(outputs, elapsed_s)
|
||||
|
||||
assert extracted_result is not None
|
||||
median_elapsed_s = float(np.median(np.asarray(elapsed_samples, dtype=np.float64)))
|
||||
return extracted_result._replace(elapsed_s=median_elapsed_s)
|
||||
|
||||
|
||||
def _video_mae(a: np.ndarray, b: np.ndarray) -> float:
|
||||
assert a.shape == b.shape, f"Video shapes differ: {a.shape} vs {b.shape}"
|
||||
return float(np.mean(np.abs(a - b)))
|
||||
|
||||
|
||||
def _audio_rms(audio: np.ndarray) -> float:
|
||||
audio = np.asarray(audio, dtype=np.float32).reshape(-1)
|
||||
return float(np.sqrt(np.mean(audio**2)))
|
||||
|
||||
|
||||
def _audio_rms_ratio(a: DreamIDGenerationResult, b: DreamIDGenerationResult) -> float:
|
||||
rms_a = _audio_rms(a.audio)
|
||||
rms_b = _audio_rms(b.audio)
|
||||
return rms_b / max(rms_a, 1e-6)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100"})
|
||||
def test_dreamid_omni_cache_dit_e2e() -> None:
|
||||
model_path = _require_model_path()
|
||||
example_case = _load_example_case()
|
||||
|
||||
baseline = _run_generation(model_path=model_path, example_case=example_case, use_cache_dit=False)
|
||||
cached = _run_generation(model_path=model_path, example_case=example_case, use_cache_dit=True)
|
||||
|
||||
video_mae = _video_mae(baseline.video, cached.video)
|
||||
audio_rms_ratio = _audio_rms_ratio(baseline, cached)
|
||||
speedup = baseline.elapsed_s / cached.elapsed_s if cached.elapsed_s > 0 else float("inf")
|
||||
|
||||
assert video_mae <= 0.15, f"video drift too much: {video_mae:.6f}"
|
||||
assert 0.4 <= audio_rms_ratio <= 2.5, f"audio quality regressed too much: rms_ratio={audio_rms_ratio:.6f}"
|
||||
assert speedup >= 1.05, f"cache-dit improve not enough: {speedup:.3f}x"
|
||||
@@ -76,7 +76,6 @@ def _send_inpaint_with_generator(
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.advanced_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_flux2_klein_inpaint_deterministic(offline_client: OfflineOmniClient):
|
||||
input_image, mask_image = _create_test_inputs()
|
||||
@@ -98,7 +97,6 @@ def test_flux2_klein_inpaint_deterministic(offline_client: OfflineOmniClient):
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.advanced_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_flux2_klein_inpaint_different_seeds_different_output(offline_client: OfflineOmniClient):
|
||||
input_image, mask_image = _create_test_inputs()
|
||||
|
||||
@@ -0,0 +1,162 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""End-to-end tests for MagiHuman pipeline via vLLM-Omni."""
|
||||
|
||||
import io
|
||||
|
||||
import av
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniRunner
|
||||
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
MODEL = "SII-GAIR/daVinci-MagiHuman-Base-1080p"
|
||||
|
||||
# (model, stage_config_path, extra_omni_kwargs) — see ``omni_runner`` in tests.helpers.fixtures.runtime
|
||||
_OMNI_RUNNER_PARAM = (
|
||||
MODEL,
|
||||
None,
|
||||
{
|
||||
"init_timeout": 1200,
|
||||
"tensor_parallel_size": 2,
|
||||
"vae_use_tiling": True,
|
||||
},
|
||||
)
|
||||
|
||||
_OMNI_RUNNER_PARAM_OFFLOAD = (
|
||||
MODEL,
|
||||
None,
|
||||
{
|
||||
"init_timeout": 1200,
|
||||
"tensor_parallel_size": 2,
|
||||
"enable_cpu_offload": True,
|
||||
"vae_use_tiling": True,
|
||||
},
|
||||
)
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.slow,
|
||||
pytest.mark.diffusion,
|
||||
pytest.mark.parametrize(
|
||||
"omni_runner",
|
||||
[_OMNI_RUNNER_PARAM, _OMNI_RUNNER_PARAM_OFFLOAD],
|
||||
indirect=True,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
|
||||
"""Validate that the MP4 contains meaningful video and audio tracks."""
|
||||
container = av.open(io.BytesIO(video_bytes))
|
||||
|
||||
v_streams = [s for s in container.streams if s.type == "video"]
|
||||
assert len(v_streams) >= 1, "No video stream found in MP4"
|
||||
|
||||
a_streams = [s for s in container.streams if s.type == "audio"]
|
||||
assert len(a_streams) >= 1, "No audio stream found in MP4"
|
||||
|
||||
v_stream = v_streams[0]
|
||||
assert v_stream.width >= 1080, f"Unexpected video width: {v_stream.width}"
|
||||
assert v_stream.height >= 1056, f"Unexpected video height: {v_stream.height}"
|
||||
|
||||
frame_count = 0
|
||||
for frame in container.decode(video=0):
|
||||
frame_count += 1
|
||||
if frame_count >= min_frames:
|
||||
break
|
||||
assert frame_count >= min_frames, f"Video has only {frame_count} frames (expected >= {min_frames})"
|
||||
|
||||
container.close()
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
||||
def test_magi_human_e2e(omni_runner: OmniRunner):
|
||||
"""End-to-end test for MagiHuman generating video and audio."""
|
||||
prompt = (
|
||||
"A young woman with long, wavy golden blonde hair and bright blue eyes, "
|
||||
"wearing a fitted ivory silk blouse with a delicate lace collar, sits "
|
||||
"stationary in front of a softly lit, blurred warm-toned interior. Her "
|
||||
"overall disposition is warm, composed, and gently confident. The camera "
|
||||
"holds a static medium close-up, framing her from the shoulders up, "
|
||||
"with shallow depth of field keeping her face in sharp focus. Soft "
|
||||
"directional key light falls from the upper left, casting a gentle "
|
||||
"highlight along her cheekbone and nose bridge. She draws a quiet breath, "
|
||||
"the levator labii superiors relaxing as her lips part. She speaks in "
|
||||
"clear, warm, unhurried American English: "
|
||||
"\"The most beautiful things in life aren't things at all — "
|
||||
"they're moments, feelings, and the people who make you feel truly alive.\" "
|
||||
"Her jaw descends smoothly on each stressed syllable; the orbicularis oris "
|
||||
"shapes each vowel with precision. A faint, genuine smile engages the "
|
||||
"zygomaticus major, lifting her lip corners fractionally. Her brows rest "
|
||||
"in a soft, neutral arch throughout. She maintains steady, forward-facing "
|
||||
"eye contact. Head position remains level; no torso displacement occurs.\n\n"
|
||||
"Dialogue:\n"
|
||||
"<Young blonde woman, American English>: "
|
||||
"\"The most beautiful things in life aren't things at all — "
|
||||
"they're moments, feelings, and the people who make you feel truly alive.\"\n\n"
|
||||
"Background Sound:\n"
|
||||
"<Soft, warm indoor ambience with a faint distant piano melody>"
|
||||
)
|
||||
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=256,
|
||||
width=448,
|
||||
num_inference_steps=8,
|
||||
seed=52,
|
||||
extra_args={
|
||||
"seconds": 5,
|
||||
"sr_height": 1080,
|
||||
"sr_width": 1920,
|
||||
"sr_num_inference_steps": 5,
|
||||
},
|
||||
)
|
||||
|
||||
omni = omni_runner.omni
|
||||
outputs = list(
|
||||
omni.generate(
|
||||
prompts=[prompt],
|
||||
sampling_params_list=[sampling_params],
|
||||
)
|
||||
)
|
||||
|
||||
assert len(outputs) > 0, "No outputs returned"
|
||||
first = outputs[0]
|
||||
|
||||
assert hasattr(first, "images") and first.images, "No video frames in output"
|
||||
video_frames = first.images[0]
|
||||
assert isinstance(video_frames, np.ndarray), f"Expected numpy array, got {type(video_frames)}"
|
||||
assert video_frames.ndim == 4, f"Expected 4D array (T,H,W,3), got shape {video_frames.shape}"
|
||||
|
||||
mm = first.multimodal_output
|
||||
assert mm, "multimodal_output is empty or missing"
|
||||
|
||||
audio_waveform = mm.get("audio")
|
||||
assert audio_waveform is not None, "No audio waveform in multimodal_output"
|
||||
|
||||
audio_sample_rate = mm.get("audio_sample_rate")
|
||||
assert audio_sample_rate is not None, (
|
||||
"audio_sample_rate not found in multimodal_output; model post-process must propagate it"
|
||||
)
|
||||
assert isinstance(audio_sample_rate, (int, float)), (
|
||||
f"audio_sample_rate should be numeric, got {type(audio_sample_rate)}"
|
||||
)
|
||||
assert int(audio_sample_rate) > 0, f"audio_sample_rate must be positive, got {audio_sample_rate}"
|
||||
|
||||
fps = mm.get("fps")
|
||||
assert fps is not None, "fps not found in multimodal_output; model post-process must propagate it"
|
||||
assert isinstance(fps, (int, float)), f"fps should be numeric, got {type(fps)}"
|
||||
assert int(fps) > 0, f"fps must be positive, got {fps}"
|
||||
|
||||
video_bytes = mux_video_audio_bytes(
|
||||
video_frames,
|
||||
audio_waveform,
|
||||
fps=float(fps),
|
||||
audio_sample_rate=int(audio_sample_rate),
|
||||
)
|
||||
assert isinstance(video_bytes, bytes), f"Expected MP4 bytes, got {type(video_bytes)}"
|
||||
assert len(video_bytes) > 1000, f"MP4 too small ({len(video_bytes)} bytes)"
|
||||
|
||||
_validate_mp4(video_bytes)
|
||||
Loaded 100 of 271 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user