Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a5ee63fb1f | ||
|
|
3335e76c97 | ||
|
|
9167599b3c | ||
|
|
99f946210b | ||
|
|
ec4e12f379 | ||
|
|
d272d57bef | ||
|
|
bc95e7415f | ||
|
|
4d617b6ed9 | ||
|
|
d4e0a5a5fd | ||
|
|
f879ad4983 | ||
|
|
f2e3d2e25b | ||
|
|
240e6655be | ||
|
|
9235b0ae65 | ||
|
|
1c2a81f6d8 | ||
|
|
b4581b293b | ||
|
|
9f978923f2 | ||
|
|
589afcef6c | ||
|
|
556c5f88a0 | ||
|
|
2eada79b99 | ||
|
|
7fb62db5c2 | ||
|
|
6510de344e | ||
|
|
740cb35a59 | ||
|
|
ae7bee0c43 | ||
|
|
6a604be610 | ||
|
|
f201b7171f | ||
|
|
0a7ac4c67e | ||
|
|
e788ef6e76 | ||
|
|
f24a6165ff | ||
|
|
bfaaa3c2c3 | ||
|
|
c9f2e5ad94 | ||
|
|
fbea788668 | ||
|
|
78c144f3a8 | ||
|
|
1c4aca7f6e |
No files matched your search
@@ -309,7 +309,7 @@ For Omni or custom models, create:
|
|||||||
|
|
||||||
Required updates:
|
Required updates:
|
||||||
1. `docs/user_guide/diffusion/parallelism/overview.md` — parallelism support overview/table
|
1. `docs/user_guide/diffusion/parallelism/overview.md` — parallelism support overview/table
|
||||||
2. `docs/user_guide/diffusion/cpu_offload_diffusion.md` — if CPU offload supported (add to supported models table)
|
2. `docs/user_guide/diffusion/cpu_offload.md` — if CPU offload supported (add to supported models table)
|
||||||
3. `docs/user_guide/diffusion/cache_acceleration/teacache.md` — if TeaCache supported
|
3. `docs/user_guide/diffusion/cache_acceleration/teacache.md` — if TeaCache supported
|
||||||
4. `docs/user_guide/diffusion/cache_acceleration/cache_dit.md` — if Cache-DiT supported
|
4. `docs/user_guide/diffusion/cache_acceleration/cache_dit.md` — if Cache-DiT supported
|
||||||
5. Offline example docs under `examples/offline_inference/<name>/` (`README.md` or category-specific `.md`)
|
5. Offline example docs under `examples/offline_inference/<name>/` (`README.md` or category-specific `.md`)
|
||||||
|
|||||||
+4
-4
@@ -16,15 +16,15 @@
|
|||||||
/.claude/ @hsliuustc0106
|
/.claude/ @hsliuustc0106
|
||||||
|
|
||||||
# Tests
|
# Tests
|
||||||
/tests/ @yenuo26
|
/tests/ @yenuo26 @NickCao
|
||||||
|
|
||||||
# CI
|
# CI
|
||||||
/.buildkite/ @yenuo26 @congw729
|
/.buildkite/ @yenuo26 @congw729 @NickCao
|
||||||
|
|
||||||
# Engine and model_executor
|
# Engine and model_executor
|
||||||
/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307
|
/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307
|
||||||
/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan @Sy0307
|
/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan @Sy0307
|
||||||
/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks
|
/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks @NickCao
|
||||||
/vllm_omni/model_executor/ @tzhouam @gcanlin
|
/vllm_omni/model_executor/ @tzhouam @gcanlin
|
||||||
/vllm_omni/worker/ @tzhouam @gcanlin
|
/vllm_omni/worker/ @tzhouam @gcanlin
|
||||||
|
|
||||||
@@ -37,7 +37,7 @@
|
|||||||
|
|
||||||
# Omni and TTS support
|
# Omni and TTS support
|
||||||
/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
|
/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
|
||||||
/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin
|
/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin @NickCao
|
||||||
|
|
||||||
# Benchmarks
|
# Benchmarks
|
||||||
/benchmarks/accuracy/ @david6666666
|
/benchmarks/accuracy/ @david6666666
|
||||||
|
|||||||
+44
-28
@@ -68,13 +68,12 @@ nav:
|
|||||||
- Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md
|
- Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md
|
||||||
- Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md
|
- Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md
|
||||||
- VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md
|
- VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md
|
||||||
- CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
|
- CPU Offloading: user_guide/diffusion/cpu_offload.md
|
||||||
- LoRA: user_guide/diffusion/lora.md
|
- LoRA: user_guide/diffusion/lora.md
|
||||||
- Custom Pipeline: features/custom_pipeline.md
|
- Custom Pipeline: features/custom_pipeline.md
|
||||||
- Request-Level Batching: user_guide/diffusion/request_batching.md
|
- Execution Modes: user_guide/diffusion/execution_modes.md
|
||||||
- Step Execution: user_guide/diffusion/step_execution.md
|
|
||||||
- Quantization:
|
- Quantization:
|
||||||
- Overview: user_guide/quantization/overview.md
|
- Quantization: user_guide/quantization/overview.md
|
||||||
- Online Quantization: user_guide/quantization/online.md
|
- Online Quantization: user_guide/quantization/online.md
|
||||||
- Quantized KV Cache: user_guide/quantization/quantized_kvcache.md
|
- Quantized KV Cache: user_guide/quantization/quantized_kvcache.md
|
||||||
- FP8 W8A8: user_guide/quantization/fp8.md
|
- FP8 W8A8: user_guide/quantization/fp8.md
|
||||||
@@ -102,32 +101,49 @@ nav:
|
|||||||
- CI: contributing/ci
|
- CI: contributing/ci
|
||||||
- Design Documents:
|
- Design Documents:
|
||||||
- design/index.md
|
- design/index.md
|
||||||
- design/architecture_overview.md
|
- Architecture Documents:
|
||||||
|
- Architecture Overview: design/architecture_overview.md
|
||||||
- Feature Design:
|
- Feature Design:
|
||||||
- design/feature/disaggregated_inference.md
|
- Runtime and Stage Execution:
|
||||||
- design/feature/ray_based_execution.md
|
- Disaggregated Inference: design/feature/disaggregated_inference.md
|
||||||
- design/feature/omni_connectors/
|
- Async Chunk: design/feature/async_chunk.md
|
||||||
- design/feature/prefix_caching.md
|
- Async Diffusion Output: design/feature/async_diffusion_output.md
|
||||||
- design/feature/cfg_parallel.md
|
- Async Omni Output Materialization: design/feature/omni_async_output_materialization.md
|
||||||
- design/feature/expert_parallel.md
|
- Automatic Prefix Caching: design/feature/prefix_caching.md
|
||||||
- design/feature/sequence_parallel.md
|
- Communication:
|
||||||
- design/feature/tensor_parallel.md
|
- OmniConnector Implementations:
|
||||||
- design/feature/vae_parallel.md
|
- Mooncake Store Connector: design/feature/omni_connectors/mooncake_store_connector.md
|
||||||
- design/feature/hsdp.md
|
- Mooncake Transfer Engine Connector: design/feature/omni_connectors/mooncake_transfer_engine_connector.md
|
||||||
- design/feature/cache_dit.md
|
- Mori Transfer Engine Connector: design/feature/omni_connectors/mori_transfer_engine_connector.md
|
||||||
- design/feature/teacache.md
|
- Shared Memory Connector: design/feature/omni_connectors/shared_memory_connector.md
|
||||||
- design/feature/skip_softmax.md
|
- Yuanrong Store Connector: design/feature/omni_connectors/yuanrong_connector.md
|
||||||
- design/feature/async_chunk.md
|
- Yuanrong Transfer Engine Connector: design/feature/omni_connectors/yuanrong_transfer_engine_connector.md
|
||||||
- design/feature/omni_async_output_materialization.md
|
- Diffusion Acceleration:
|
||||||
- design/feature/vae_parallel.md
|
- Parallelism:
|
||||||
- design/feature/diffusion_step_execution.md
|
- CFG-Parallel: design/feature/cfg_parallel.md
|
||||||
- design/feature/diffusion_request_level_batching.md
|
- Expert Parallel: design/feature/expert_parallel.md
|
||||||
- design/feature/diffusion_continuous_batching.md
|
- Hybrid Sharded Data Parallel (HSDP): design/feature/hsdp.md
|
||||||
|
- Pipeline Parallel: design/feature/pipeline_parallel.md
|
||||||
|
- Sequence Parallel: design/feature/sequence_parallel.md
|
||||||
|
- Tensor Parallel: design/feature/tensor_parallel.md
|
||||||
|
- VAE Patch Parallelism: design/feature/vae_parallel.md
|
||||||
|
- Attention Backends:
|
||||||
|
- Skip-Softmax: design/feature/skip_softmax.md
|
||||||
|
- Quantization: design/feature/quantization.md
|
||||||
|
- Cache-DiT: design/feature/cache_dit.md
|
||||||
|
- TeaCache: design/feature/teacache.md
|
||||||
|
- Diffusion Continuous Batching: design/feature/diffusion_continuous_batching.md
|
||||||
|
- Distributed Layerwise Offload: design/feature/distributed_layerwise_offload.md
|
||||||
|
- Infrastructure and Performance:
|
||||||
|
- Prometheus Metrics: design/metrics.md
|
||||||
|
- Speech Generation Performance Optimizations: design/qwen3_omni_tts_performance_optimization.md
|
||||||
- Module Design:
|
- Module Design:
|
||||||
- design/module/ar_module.md
|
- Runtime Modules:
|
||||||
- design/module/dit_module.md
|
- AR Module: design/module/ar_module.md
|
||||||
- design/module/entrypoint_module.md
|
- DIT Module: design/module/dit_module.md
|
||||||
- design/module/async_omni_architecture.md
|
- Orchestration:
|
||||||
|
- Entrypoint Module: design/module/entrypoint_module.md
|
||||||
|
- AsyncOmni Architecture: design/module/async_omni_architecture.md
|
||||||
- Docs Guide: contributing/DOCS_GUIDE.md
|
- Docs Guide: contributing/DOCS_GUIDE.md
|
||||||
- API Reference:
|
- API Reference:
|
||||||
- api/README.md
|
- api/README.md
|
||||||
|
|||||||
@@ -40,6 +40,7 @@ Committers have write access and merge rights. They typically have deep expertis
|
|||||||
- [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
|
- [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
|
||||||
- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
|
- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
|
||||||
- [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration
|
- [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration
|
||||||
|
- [@NickCao](https://github.com/NickCao): CI, Tests, Entrypoint, and TTS Model Support
|
||||||
- [@princepride](https://github.com/princepride): Diffusion and Omni Support
|
- [@princepride](https://github.com/princepride): Diffusion and Omni Support
|
||||||
- [@RuixiangMa](https://github.com/RuixiangMa): Diffusion models, parallel, cache, and docs
|
- [@RuixiangMa](https://github.com/RuixiangMa): Diffusion models, parallel, cache, and docs
|
||||||
- [@SamitHuang](https://github.com/SamitHuang): RL, Diffusion, and cache
|
- [@SamitHuang](https://github.com/SamitHuang): RL, Diffusion, and cache
|
||||||
|
|||||||
@@ -387,8 +387,8 @@ that support request-level batching with `step_execution` disabled.
|
|||||||
|
|
||||||
Use this together with `max_num_seqs > 1` for bursty serving traffic. `0`
|
Use this together with `max_num_seqs > 1` for bursty serving traffic. `0`
|
||||||
disables admission waiting and preserves the lowest first-request latency.
|
disables admission waiting and preserves the lowest first-request latency.
|
||||||
For diffusion request-level batching tuning, see
|
For diffusion execution and batching tuning, see
|
||||||
[Request-Level Batching](../user_guide/diffusion/request_batching.md).
|
[Diffusion Execution Modes](../user_guide/diffusion/execution_modes.md).
|
||||||
|
|
||||||
Default: `0.0`
|
Default: `0.0`
|
||||||
|
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
**Scope**
|
**Scope**
|
||||||
|
|
||||||
For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/quantization/overview/).
|
For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/quantization/overview/).
|
||||||
|
|
||||||
**Test Case Design**
|
**Test Case Design**
|
||||||
|
|
||||||
|
|||||||
@@ -44,9 +44,11 @@ Pass **`--test-config-file`** to run one JSON file as-is, or omit it for the bul
|
|||||||
"ignore_eos": true,
|
"ignore_eos": true,
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [500, 800],
|
"H100": {
|
||||||
"mean_audio_ttfp_ms": [2000, 3500],
|
"mean_ttft_ms": [500, 800],
|
||||||
"mean_audio_rtf": [0.25, 0.35]
|
"mean_audio_ttfp_ms": [2000, 3500],
|
||||||
|
"mean_audio_rtf": [0.25, 0.35]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
@@ -142,10 +144,6 @@ pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py \
|
|||||||
--test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
|
--test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
|
||||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \
|
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \
|
||||||
--test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
|
--test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
|
||||||
|
|
||||||
# Optional baseline assertion (default off)
|
|
||||||
pytest -sv tests/dfx/perf/scripts/run_diffusion_benchmark.py --assert-baseline \
|
|
||||||
--test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
|
|
||||||
```
|
```
|
||||||
|
|
||||||
See also [Markers for Tests](#markers-for-tests) for registered hardware markers (`H100`, `L4`, `cuda`, `distributed_cuda`, …).
|
See also [Markers for Tests](#markers-for-tests) for registered hardware markers (`H100`, `L4`, `cuda`, `distributed_cuda`, …).
|
||||||
@@ -202,4 +200,4 @@ You can add any benchmark running parameters you need here. For all optional par
|
|||||||
| num_prompts | int / array | Yes | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
|
| num_prompts | int / array | Yes | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
|
||||||
| request_rate | float / array | No | 0.5, [0.5, 1, inf] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts. |
|
| request_rate | float / array | No | 0.5, [0.5, 1, inf] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts. |
|
||||||
| max_concurrency | int / array | No | 1, [1, 2, 3] | Maximum concurrent in-flight requests. Same array / expansion rules as `request_rate` (mutually exclusive with QPS mode). |
|
| max_concurrency | int / array | No | 1, [1, 2, 3] | Maximum concurrent in-flight requests. Same array / expansion rules as `request_rate` (mutually exclusive with QPS mode). |
|
||||||
| baseline | object | No | see above | Optional per-metric thresholds; keys must match benchmark output fields. Scalar, list (per sweep step), or object (keyed by concurrency or QPS string).
|
| baseline | object | No | see below | Optional thresholds. Prefer hardware-nested form: `{"H100": {"throughput_qps": …}, "A3": {…}}` (selected via runtime hardware label). Legacy flat `{metric: threshold}` still works. Per metric: scalar, list (sweep step), or object keyed by concurrency / QPS.
|
||||||
@@ -755,7 +755,8 @@ omni = Omni(model="your-model", ulysses_degree=2, ring_degree=2)
|
|||||||
|
|
||||||
### Step Execution
|
### Step Execution
|
||||||
|
|
||||||
See detailed design guide: [How to add step execution support](../../design/feature/diffusion_step_execution.md)
|
See the detailed
|
||||||
|
[Diffusion Continuous Batching design guide](../../design/feature/diffusion_continuous_batching.md).
|
||||||
|
|
||||||
Use this only when your pipeline can be split into stable request-scoped and
|
Use this only when your pipeline can be split into stable request-scoped and
|
||||||
step-scoped phases. The reference implementation is
|
step-scoped phases. The reference implementation is
|
||||||
@@ -769,9 +770,8 @@ step-scoped phases. The reference implementation is
|
|||||||
Do not enable `step_execution=True` until those four methods are implemented
|
Do not enable `step_execution=True` until those four methods are implemented
|
||||||
and validated against the request-level path.
|
and validated against the request-level path.
|
||||||
|
|
||||||
If you want the pipeline to work with the experimental batched step-wise path
|
The same design guide covers the experimental batched step-wise path used when
|
||||||
(`max_num_seqs > 1`), also see:
|
`max_num_seqs > 1`.
|
||||||
[Continuous Batching for Step-Wise Diffusion](../../design/feature/diffusion_continuous_batching.md).
|
|
||||||
|
|
||||||
If you expose this in example scripts or recipes, keep it opt-in. Surface
|
If you expose this in example scripts or recipes, keep it opt-in. Surface
|
||||||
runtime features like `step_execution` as optional flags instead of silently
|
runtime features like `step_execution` as optional flags instead of silently
|
||||||
@@ -824,7 +824,7 @@ omni = Omni(model="your-model",
|
|||||||
|
|
||||||
### CPU Offload
|
### CPU Offload
|
||||||
|
|
||||||
See detailed guide: [CPU Offloading for Diffusion Models](../../user_guide/diffusion/cpu_offload_diffusion.md)
|
See detailed guide: [CPU Offloading for Diffusion Models](../../user_guide/diffusion/cpu_offload.md)
|
||||||
|
|
||||||
vLLM-Omni provides two offloading strategies to reduce GPU memory usage:
|
vLLM-Omni provides two offloading strategies to reduce GPU memory usage:
|
||||||
|
|
||||||
|
|||||||
@@ -1,165 +1,232 @@
|
|||||||
# Continuous Batching for Step-Wise Diffusion
|
# Diffusion Continuous Batching
|
||||||
|
|
||||||
!!! warning "Experimental Feature"
|
This document describes the unified diffusion execution architecture, including
|
||||||
This feature is experimental. It currently applies only to native
|
request execution, request-level batching, step execution, continuous batching,
|
||||||
diffusion pipelines running with `step_execution=True`.
|
and output delivery. For user-facing configuration and CLI examples, see
|
||||||
|
[Diffusion Execution Modes](../../user_guide/diffusion/execution_modes.md).
|
||||||
This document describes the batching extension built on top of
|
|
||||||
[Diffusion Step Execution](diffusion_step_execution.md). The base
|
|
||||||
step-execution contract is unchanged. The batching work is mainly in the
|
|
||||||
scheduler and runner layers.
|
|
||||||
|
|
||||||
## Why It Helps
|
|
||||||
|
|
||||||
Step-wise execution breaks a long denoise loop into scheduler-visible units.
|
|
||||||
That gives the runtime a place to admit other compatible requests between
|
|
||||||
steps instead of waiting for an entire request to finish.
|
|
||||||
|
|
||||||
This matters most in low-MFU or bursty serving scenarios:
|
|
||||||
|
|
||||||
- one request's denoise step may not fully saturate the GPU
|
|
||||||
- several compatible requests can share the same denoise forward pass
|
|
||||||
- throughput and device utilization can improve without changing request-local
|
|
||||||
scheduler state
|
|
||||||
|
|
||||||
This is **not** a guaranteed single-request latency win. The main benefit is
|
|
||||||
usually higher utilization and better throughput when the workload contains
|
|
||||||
multiple in-flight compatible requests.
|
|
||||||
|
|
||||||
## Overview
|
## Overview
|
||||||
|
|
||||||
With continuous batching enabled:
|
`DiffusionEngine` selects an execution mode from `step_execution` and uses
|
||||||
|
`max_num_seqs` as that mode's scheduler capacity:
|
||||||
|
|
||||||
- the scheduler may keep multiple compatible requests active at the same time
|
| Configuration | Engine mode | Scheduler | Execution |
|
||||||
- the runner packs request-local step state into one `InputBatch`
|
|---|---|---|---|
|
||||||
- `denoise_step()` runs on that batch
|
| `step_execution=False`, `max_num_seqs=1` | `REQUEST_BATCH` | `RequestScheduler` | One complete request-level `forward()` |
|
||||||
- `step_scheduler()` and `post_decode()` still run per request
|
| `step_execution=False`, `max_num_seqs>1` | `REQUEST_BATCH` | `RequestScheduler` | One fused `forward()` over compatible requests |
|
||||||
|
| `step_execution=True`, `max_num_seqs=1` | `STEP_BATCH` | `StepScheduler` | One request advanced one denoise step per scheduler tick |
|
||||||
|
| `step_execution=True`, `max_num_seqs>1` | `STEP_BATCH` | `StepScheduler` | Compatible requests advanced together in step waves |
|
||||||
|
|
||||||
The current implementation is conservative:
|
Serial request execution is handled inside `REQUEST_BATCH`; it is not a
|
||||||
|
separate engine mode. Similarly, continuous batching is the multi-request
|
||||||
|
configuration of `STEP_BATCH`, not an independent execution mode.
|
||||||
|
|
||||||
- only compatible requests are batched together
|
The engine performs the following initialization:
|
||||||
- per-request progress and completion remain independent
|
|
||||||
|
|
||||||
Here, "continuous batching" means the step-wise path enabled by
|
1. Resolve model-specific pre- and post-processing hooks.
|
||||||
`step_execution=True`. Request-mode `DiffusionRequestBatch` is static
|
2. Select `REQUEST_BATCH` or `STEP_BATCH`.
|
||||||
request-level batching for one full pipeline `forward()` call; it does not
|
3. Construct the configured `DiffusionExecutor`.
|
||||||
admit or remove requests between denoise steps.
|
4. Construct `RequestScheduler` or `StepScheduler`, unless a scheduler was
|
||||||
|
explicitly injected.
|
||||||
|
5. Bind `execute_batch` or `execute_step`.
|
||||||
|
|
||||||
## Enablement
|
`DiffusionEngine.make_engine()` owns startup warmup after engine construction.
|
||||||
|
This keeps construction, backend selection, and warmup failure cleanup in one
|
||||||
|
lifecycle.
|
||||||
|
|
||||||
Use `--step-execution` as the feature gate, then increase `--max-num-seqs`
|
## Unified Output Stream
|
||||||
above `1` if you want batching:
|
|
||||||
|
|
||||||
```bash
|
All asynchronous diffusion requests use the same per-request output stream,
|
||||||
vllm serve Qwen/Qwen-Image --omni \
|
regardless of whether the caller exposes streaming results to the user.
|
||||||
--port 8091 \
|
|
||||||
--step-execution \
|
```python
|
||||||
--max-num-seqs 8
|
async for output in engine.step_streaming(request):
|
||||||
|
...
|
||||||
```
|
```
|
||||||
|
|
||||||
`--max-num-seqs 1` keeps the step-wise path without enabling batching.
|
When a request is admitted, the engine creates an
|
||||||
|
`asyncio.Queue[DiffusionOutput]` keyed by its request ID. The busy loop sends
|
||||||
|
both intermediate chunks and final results through that queue:
|
||||||
|
|
||||||
For a reproducible replay flow using the bundled serving benchmark, see the
|
```text
|
||||||
Qwen-Image replay commands in
|
request
|
||||||
[`benchmarks/diffusion/README.md`](gh-file:benchmarks/diffusion/README.md)
|
-> scheduler.add_request()
|
||||||
and
|
-> scheduler.schedule()
|
||||||
[`benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md`](gh-file:benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md).
|
-> execute_batch() or execute_step()
|
||||||
|
-> scheduler.update_from_output()
|
||||||
|
-> per-request output queue
|
||||||
|
-> step_streaming()
|
||||||
|
```
|
||||||
|
|
||||||
## Scheduler
|
Streaming callers forward every yielded output. Non-streaming callers consume
|
||||||
|
the same generator to completion and return only its final output. This avoids
|
||||||
|
maintaining separate future-based and queue-based result paths.
|
||||||
|
|
||||||
The scheduler derives its batch capacity from `max_num_seqs` through
|
`DiffusionEngine.step()` and
|
||||||
`max_num_running_reqs`.
|
`async_add_req_and_wait_for_response()` remain deprecated compatibility
|
||||||
|
wrappers. New integrations should consume `step_streaming()` or
|
||||||
|
`async_add_req_and_stream_response()`.
|
||||||
|
|
||||||
Batch admission is gated by
|
### Completion and Cancellation
|
||||||
[`StepBatchSamplingParamsKey`](gh-file:vllm_omni/diffusion/sched/interface.py),
|
|
||||||
which is built from shape-sensitive and CFG-sensitive sampling fields. This is
|
|
||||||
the core correctness rule for batching: requests are only co-batched when they
|
|
||||||
share the same denoise tensor contract.
|
|
||||||
|
|
||||||
There are three important details:
|
Output delivery and scheduler lifecycle are related but separate:
|
||||||
|
|
||||||
- `num_inference_steps` is not part of the key, so requests with different
|
- A terminal scheduler result always finalizes request state.
|
||||||
total step counts can still share a batch
|
- A disconnected or cancelled consumer removes only its delivery queue.
|
||||||
- requests also do not need to be at the same current denoise progress; active
|
- The scheduler may still finish and clean up the request after its queue has
|
||||||
requests can continue batching even when their current step indices diverge
|
been removed.
|
||||||
- admission is still FIFO, so an incompatible request at the head of the
|
- Engine shutdown sends an error output to all remaining streams.
|
||||||
waiting queue blocks later compatible requests
|
|
||||||
|
|
||||||
Today that compatibility rule is still shape-sensitive. `height`, `width`,
|
This separation prevents cancelled consumers from leaving active request IDs in
|
||||||
`num_frames`, and CFG-related fields remain part of the key, so different
|
the scheduler.
|
||||||
resolutions or incompatible guidance settings do **not** co-batch yet. The
|
|
||||||
key also covers LoRA identity (`lora_int_id`, `lora_scale`), so requests
|
|
||||||
targeting different adapters or scales run in separate batches and the
|
|
||||||
worker can activate exactly one adapter per step.
|
|
||||||
|
|
||||||
The scheduler batching unit is one logical `OmniDiffusionRequest`. In the
|
## Request-Batch Execution
|
||||||
step-wise path, runtime tensor batching is represented as `StepInputBatch`. For
|
|
||||||
request-mode prompt semantics, see
|
|
||||||
[Request-Level Batching](../../user_guide/diffusion/request_batching.md).
|
|
||||||
|
|
||||||
## Runner
|
Request-batch execution runs the complete diffusion pipeline for each scheduler
|
||||||
|
wave. Each `OmniDiffusionRequest` remains an independent logical request with
|
||||||
|
its own prompt, sampling parameters, seed, request ID, output, error, and abort
|
||||||
|
state.
|
||||||
|
|
||||||
The runner keeps persistent per-request execution state in
|
With `max_num_seqs=1`, `execute_batch()` handles the serial request path. With a
|
||||||
[`StepRequestState`](gh-file:vllm_omni/diffusion/worker/utils.py),
|
larger value, a pipeline can opt into fused execution by declaring:
|
||||||
while the scheduler owns a separate lightweight
|
|
||||||
[`SchedulerRequestState`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
|
||||||
for queueing and lifecycle tracking.
|
|
||||||
|
|
||||||
For each step, the runner builds an
|
```python
|
||||||
[`InputBatch`](gh-file:vllm_omni/diffusion/worker/input_batch.py) from the
|
supports_request_batch = True
|
||||||
active request states:
|
```
|
||||||
|
|
||||||
- prompt embeddings and masks are normalized and padded
|
Its `forward()` method must accept `DiffusionRequestBatch` and return one
|
||||||
- dynamic tensors such as `latents` and `timesteps` are gathered each step
|
`DiffusionOutput` per request. The runner validates the result count and maps
|
||||||
- buffers are reused when batch composition stays the same
|
outputs back to their original request IDs with `BatchRunnerOutput`.
|
||||||
|
|
||||||
The step-wise batched path is:
|
Current in-tree request-batch implementations include the Flux, LTX-2, SD3, and
|
||||||
|
Qwen-Image pipelines. Treat the pipeline capability flag in the source as
|
||||||
|
authoritative because support continues to expand.
|
||||||
|
|
||||||
1. Run `prepare_encode()` for newly admitted requests.
|
### Request-Batch Data Flow
|
||||||
|
|
||||||
|
- `DiffusionSchedulerOutput` carries scheduled request IDs and payloads.
|
||||||
|
- `DiffusionRequestBatch` provides the pipeline-facing static batch.
|
||||||
|
- `BatchRunnerOutput` carries per-request results back to the engine.
|
||||||
|
|
||||||
|
Request-local setup, such as seed handling and output/error mapping, remains
|
||||||
|
per request. Cache refresh, homogeneous LoRA activation, and
|
||||||
|
`pipeline.forward(req_batch)` can be shared by the batch.
|
||||||
|
|
||||||
|
### Compatibility and Admission
|
||||||
|
|
||||||
|
`RequestBatchSamplingParamsKey` contains fields that affect the tensor contract,
|
||||||
|
including shape, guidance, output count, and LoRA identity. Only compatible
|
||||||
|
requests can share a fused forward call.
|
||||||
|
|
||||||
|
Admission is FIFO and conservative. An incompatible request at the front of the
|
||||||
|
waiting queue can prevent later compatible requests from joining the current
|
||||||
|
wave.
|
||||||
|
|
||||||
|
`request_batch_max_wait_ms` optionally delays the first schedule operation of a
|
||||||
|
new wave so bursty arrivals can accumulate. It applies only when:
|
||||||
|
|
||||||
|
- the pipeline supports request batching;
|
||||||
|
- step execution is disabled;
|
||||||
|
- no requests are currently running; and
|
||||||
|
- the configured wait is greater than zero.
|
||||||
|
|
||||||
|
The wait ends when the queue reaches capacity, stabilizes briefly, reaches its
|
||||||
|
deadline, or the engine stops.
|
||||||
|
|
||||||
|
## Step-Batch Execution
|
||||||
|
|
||||||
|
Step execution exposes denoising progress to the scheduler. A supporting
|
||||||
|
pipeline implements four stateful operations:
|
||||||
|
|
||||||
|
| Operation | Responsibility |
|
||||||
|
|---|---|
|
||||||
|
| `prepare_encode(state)` | Validate input, encode prompts, initialize latents and timesteps, and create request-local scheduler state |
|
||||||
|
| `denoise_step(input_batch, *, states=...)` | Run one denoise forward for the scheduler-provided request states |
|
||||||
|
| `step_scheduler(state, noise_pred)` | Update latents and advance request progress |
|
||||||
|
| `post_decode(state)` | Decode and postprocess a completed request or output boundary |
|
||||||
|
|
||||||
|
Persistent request state lives in `StepRequestState`. Pipeline-specific fields
|
||||||
|
that do not belong in the shared contract should be stored in `state.extra`.
|
||||||
|
Queueing and lifecycle metadata remain in the scheduler's request state.
|
||||||
|
|
||||||
|
Current native pipelines that explicitly enable step execution include
|
||||||
|
Qwen-Image, HunyuanImage3, and Helios. Step execution alone does not imply
|
||||||
|
continuous-batching support: Qwen-Image accepts batched step states.
|
||||||
|
HunyuanImage3 accepts batched step states only when its resolved self-attention
|
||||||
|
backend is `TORCH_SDPA`; otherwise it rejects groups larger than one request.
|
||||||
|
Configure `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` or
|
||||||
|
`diffusion_attention_config.default.backend=TORCH_SDPA` when
|
||||||
|
`max_num_seqs>1`. See the
|
||||||
|
[HunyuanImage-3.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md)
|
||||||
|
for its validated configuration. Helios supports only a single active step
|
||||||
|
request and must use `max_num_seqs=1`.
|
||||||
|
|
||||||
|
### Continuous Batching
|
||||||
|
|
||||||
|
When `max_num_seqs>1`, `StepScheduler` can keep multiple compatible requests
|
||||||
|
active. The runner gathers their state into `InputBatch`, performs one batched
|
||||||
|
denoise forward, then applies scheduler updates per request:
|
||||||
|
|
||||||
|
1. Receive transferred KV payloads, then run `prepare_encode()` for newly
|
||||||
|
admitted requests.
|
||||||
2. Build or refresh `InputBatch`.
|
2. Build or refresh `InputBatch`.
|
||||||
3. Run one batched `denoise_step(input_batch)`.
|
3. Run one batched `denoise_step(input_batch, states=states)`.
|
||||||
4. Slice the batched `noise_pred` back per request.
|
4. Slice noise predictions back to each request.
|
||||||
5. Run per-request `step_scheduler()`.
|
5. Run each request's `step_scheduler()`.
|
||||||
6. Run `post_decode()` only for requests that finished denoising.
|
6. Run `post_decode()` at a chunk boundary or request completion.
|
||||||
7. Scatter updated latents back into persistent request state with
|
7. Scatter updated latents back into persistent request state.
|
||||||
[`scatter_latents()`](gh-file:vllm_omni/diffusion/worker/input_batch.py).
|
|
||||||
|
|
||||||
This keeps the shared work limited to the denoise forward pass while preserving
|
`StepBatchSamplingParamsKey` protects the batched tensor contract. Requests can
|
||||||
request-local scheduler state and outputs.
|
have different total step counts and current step indices, but shape-sensitive,
|
||||||
|
CFG-sensitive, and LoRA fields must be compatible. FIFO head-of-line blocking
|
||||||
|
also applies here.
|
||||||
|
|
||||||
## Engine
|
Chunk-capable pipelines may emit intermediate results. Final-only step
|
||||||
|
pipelines emit when the request finishes even if `streaming_output=True`.
|
||||||
|
|
||||||
[`DiffusionEngine`](gh-file:vllm_omni/diffusion/diffusion_engine.py) provides
|
## Model Author Guidelines
|
||||||
the background loop and async add-request path needed for multiple requests to
|
|
||||||
accumulate in the scheduler.
|
|
||||||
|
|
||||||
When `step_execution=True`, the engine routes work through the step-wise
|
When converting a request-level pipeline:
|
||||||
executor path. The continuous batching behavior is defined by scheduler-side
|
|
||||||
compatibility gating and runner-side `StepInputBatch` packing.
|
|
||||||
|
|
||||||
## Current Limitations
|
- Reuse the same helpers as `forward()` to avoid behavior drift.
|
||||||
|
- Copy request-scoped scheduler state into `state.scheduler`.
|
||||||
|
- Advance `state.step_index` only in `step_scheduler()`.
|
||||||
|
- Keep model forward work in `denoise_step()`.
|
||||||
|
- Keep latent mutation in `step_scheduler()`.
|
||||||
|
- Keep final decoding equivalent to the tail of `forward()`.
|
||||||
|
- Store masks, condition latents, and other request-local tensors in the
|
||||||
|
request state rather than on the shared pipeline.
|
||||||
|
|
||||||
- Experimental feature; use `max_num_seqs=1` for the conservative single-request
|
Before setting `supports_step_execution = True`, validate output parity for the
|
||||||
step path.
|
same seed and sampling parameters, scheduler isolation across concurrent
|
||||||
- Only native pipelines that already support `step_execution=True`.
|
requests, abort cleanup, progress reporting, and all supported CFG paths.
|
||||||
- Only homogeneous batches keyed by `StepBatchSamplingParamsKey` are supported.
|
|
||||||
- `cache_backend`, KV transfer, and other request-mode extras are not wired
|
|
||||||
into the batched step-wise path yet.
|
|
||||||
- Future work can relax the current same-shape restriction with richer
|
|
||||||
heterogeneous batching policies such as bucketing or padded execution for
|
|
||||||
different resolutions.
|
|
||||||
|
|
||||||
## Related Files
|
Before setting `supports_request_batch = True`, validate the batch input
|
||||||
|
contract, one-output-per-request result shape, request identity and error
|
||||||
|
mapping, seeded concurrency, LoRA compatibility, and tensor IPC.
|
||||||
|
|
||||||
- Scheduler base:
|
## Limitations
|
||||||
[`vllm_omni/diffusion/sched/base_scheduler.py`](gh-file:vllm_omni/diffusion/sched/base_scheduler.py)
|
|
||||||
- Scheduler interface:
|
- Request batching requires an explicit pipeline capability when
|
||||||
[`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
`max_num_seqs>1`.
|
||||||
- Step scheduler:
|
- Step execution requires the four-method stateful pipeline contract.
|
||||||
[`vllm_omni/diffusion/sched/step_scheduler.py`](gh-file:vllm_omni/diffusion/sched/step_scheduler.py)
|
- Both batching paths currently require homogeneous compatibility keys.
|
||||||
- Runner:
|
- FIFO scheduling can reduce batching opportunities.
|
||||||
[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
|
- `request_batch_max_wait_ms` trades first-request latency for burst
|
||||||
- Input batch:
|
coalescing.
|
||||||
[`vllm_omni/diffusion/worker/input_batch.py`](gh-file:vllm_omni/diffusion/worker/input_batch.py)
|
- All diffusion cache backends are currently unsupported in step mode.
|
||||||
- Tests:
|
- KV transfer is supported for newly admitted step requests; some other
|
||||||
[`tests/diffusion/test_diffusion_scheduler.py`](gh-file:tests/diffusion/test_diffusion_scheduler.py)
|
request-mode extras remain unsupported.
|
||||||
|
- Step continuous batching remains experimental; use `max_num_seqs=1` for the
|
||||||
|
conservative step path.
|
||||||
|
|
||||||
|
## Related Implementation
|
||||||
|
|
||||||
|
- Engine: [`vllm_omni/diffusion/diffusion_engine.py`](gh-file:vllm_omni/diffusion/diffusion_engine.py)
|
||||||
|
- Model contract: [`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py)
|
||||||
|
- Scheduler interface: [`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
||||||
|
- Request scheduler: [`vllm_omni/diffusion/sched/request_scheduler.py`](gh-file:vllm_omni/diffusion/sched/request_scheduler.py)
|
||||||
|
- Step scheduler: [`vllm_omni/diffusion/sched/step_scheduler.py`](gh-file:vllm_omni/diffusion/sched/step_scheduler.py)
|
||||||
|
- Runner: [`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
|
||||||
|
- Request batch: [`vllm_omni/diffusion/worker/request_batch.py`](gh-file:vllm_omni/diffusion/worker/request_batch.py)
|
||||||
|
- Step input batch: [`vllm_omni/diffusion/worker/input_batch.py`](gh-file:vllm_omni/diffusion/worker/input_batch.py)
|
||||||
@@ -1,175 +0,0 @@
|
|||||||
# Request-Level Batching for Diffusion
|
|
||||||
|
|
||||||
This document describes the request-mode batching path for diffusion pipelines.
|
|
||||||
For end-user enablement and tuning, see
|
|
||||||
[Request-Level Batching](../../user_guide/diffusion/request_batching.md).
|
|
||||||
|
|
||||||
This is separate from
|
|
||||||
[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md).
|
|
||||||
Request-level batching runs one full pipeline `forward()` over a static batch of
|
|
||||||
compatible requests. Step-wise continuous batching admits work between denoise
|
|
||||||
steps when `step_execution=True`.
|
|
||||||
|
|
||||||
## Why It Helps
|
|
||||||
|
|
||||||
The request-level design avoids coupling several logical requests to one request
|
|
||||||
object. This keeps request identity, abort/error handling, and per-request
|
|
||||||
metadata unambiguous while still allowing one fused pipeline forward pass for
|
|
||||||
bursty or concurrent traffic.
|
|
||||||
|
|
||||||
## Overview
|
|
||||||
|
|
||||||
With request-level batching enabled:
|
|
||||||
|
|
||||||
- each `OmniDiffusionRequest` contains one `prompt` and one `request_id`
|
|
||||||
- the scheduler groups compatible waiting requests into one scheduler wave
|
|
||||||
- `DiffusionRequestBatch` wraps the scheduled requests for pipeline `forward()`
|
|
||||||
- batch-capable pipelines return `list[DiffusionOutput]`, one output per
|
|
||||||
request
|
|
||||||
- `BatchRunnerOutput` maps each result back to its original `request_id`
|
|
||||||
|
|
||||||
Pipelines opt in with `supports_request_batch = True` and a `forward()` method
|
|
||||||
that accepts `DiffusionRequestBatch` and returns `list[DiffusionOutput]`.
|
|
||||||
Pipelines that do not opt in keep the existing per-request execution path.
|
|
||||||
|
|
||||||
## Enablement
|
|
||||||
|
|
||||||
Request-level batching is the request-mode path, so `step_execution` must remain
|
|
||||||
disabled. Increase `max_num_seqs` above `1` to let the scheduler keep multiple
|
|
||||||
compatible requests active:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve Qwen/Qwen-Image --omni \
|
|
||||||
--port 8091 \
|
|
||||||
--max-num-seqs 4
|
|
||||||
```
|
|
||||||
|
|
||||||
For bursty online ingress, `request_batch_max_wait_ms` can add a bounded
|
|
||||||
admission wait before the first `schedule()` of a scheduler wave:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve Qwen/Qwen-Image --omni \
|
|
||||||
--port 8091 \
|
|
||||||
--max-num-seqs 4 \
|
|
||||||
--request-batch-max-wait-ms 20
|
|
||||||
```
|
|
||||||
|
|
||||||
`request_batch_max_wait_ms=0` disables this wait and is the default.
|
|
||||||
|
|
||||||
## Request Contract
|
|
||||||
|
|
||||||
`OmniDiffusionRequest` represents one logical request. It owns one prompt,
|
|
||||||
sampling parameters, request id, and request-local metadata. Runtime batches are
|
|
||||||
formed by the scheduler and represented separately from the request payload.
|
|
||||||
|
|
||||||
Runtime batching is represented by:
|
|
||||||
|
|
||||||
- [`DiffusionSchedulerOutput`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
|
||||||
for scheduled request ids and request payloads
|
|
||||||
- [`DiffusionRequestBatch`](gh-file:vllm_omni/diffusion/worker/request_batch.py)
|
|
||||||
for the pipeline-facing request batch
|
|
||||||
- [`BatchRunnerOutput`](gh-file:vllm_omni/diffusion/worker/utils.py) for
|
|
||||||
per-request results
|
|
||||||
|
|
||||||
`DiffusionRequestBatch` intentionally exposes compatibility properties such as
|
|
||||||
`prompts`, `sampling_params`, `request_id`, and `kv_sender_info` so migrated
|
|
||||||
pipelines can stay close to upstream code while using a batch-aware contract.
|
|
||||||
|
|
||||||
## Scheduler
|
|
||||||
|
|
||||||
The scheduler derives its capacity from `max_num_seqs` through
|
|
||||||
`max_num_running_reqs`. It exposes waiting/running queue counters so the engine
|
|
||||||
can decide whether admission wait is useful before scheduling a new wave.
|
|
||||||
|
|
||||||
Batch compatibility is controlled by
|
|
||||||
[`RequestBatchSamplingParamsKey`](gh-file:vllm_omni/diffusion/sched/interface.py). The key
|
|
||||||
contains shape-sensitive and guidance-sensitive fields, including output count
|
|
||||||
and LoRA identity. Requests with incompatible shapes, CFG settings, output
|
|
||||||
counts, LoRA adapters, or LoRA scales are kept in separate batches.
|
|
||||||
|
|
||||||
Admission is conservative:
|
|
||||||
|
|
||||||
- the scheduler only batches compatible requests
|
|
||||||
- FIFO ordering is preserved
|
|
||||||
- an incompatible request at the head of the waiting queue blocks later
|
|
||||||
compatible requests
|
|
||||||
|
|
||||||
## Engine
|
|
||||||
|
|
||||||
[`DiffusionEngine`](gh-file:vllm_omni/diffusion/diffusion_engine.py) resolves
|
|
||||||
request-batch capability during initialization from the configured pipeline
|
|
||||||
class, including custom pipeline classes.
|
|
||||||
|
|
||||||
The capability check uses the pipeline class attribute
|
|
||||||
`supports_request_batch = True`. Pipelines that set this attribute must implement
|
|
||||||
a request-batch-compatible `forward()` contract and return one
|
|
||||||
`DiffusionOutput` per request; the runner validates that return shape at runtime.
|
|
||||||
|
|
||||||
When the selected pipeline is batch-capable and `step_execution=False`, request
|
|
||||||
mode routes scheduler waves through the batch executor path. Otherwise it keeps
|
|
||||||
the per-request executor path.
|
|
||||||
|
|
||||||
The optional admission wait runs only when:
|
|
||||||
|
|
||||||
- request batching is supported
|
|
||||||
- `step_execution=False`
|
|
||||||
- `request_batch_max_wait_ms > 0`
|
|
||||||
- no requests are currently running
|
|
||||||
|
|
||||||
The wait exits early when the waiting queue reaches capacity, when the queue is
|
|
||||||
stable for a short window, when the deadline expires, or when the engine stops.
|
|
||||||
|
|
||||||
## Executor And Runner
|
|
||||||
|
|
||||||
The executor exposes two request-mode entries:
|
|
||||||
|
|
||||||
- `execute_request`: one worker call per scheduled request
|
|
||||||
- `execute_batch`: one worker call for the whole `DiffusionSchedulerOutput`
|
|
||||||
|
|
||||||
On the batch path, the worker builds a `DiffusionRequestBatch` and runs the
|
|
||||||
pipeline once. Request-local setup remains per request:
|
|
||||||
|
|
||||||
- KV transfer metadata
|
|
||||||
- random generator and seed handling
|
|
||||||
- request output/error/abort mapping
|
|
||||||
|
|
||||||
Shared batch setup happens once per batch when possible:
|
|
||||||
|
|
||||||
- cache refresh
|
|
||||||
- LoRA activation for the homogeneous adapter key
|
|
||||||
- pipeline `forward(req_batch)`
|
|
||||||
|
|
||||||
Large tensor IPC still uses the shared-memory packing path. The packer traverses
|
|
||||||
both normal `RunnerOutput.result` wrappers and nested batch results so batched
|
|
||||||
outputs do not fall back to pickle IPC for tensor payloads.
|
|
||||||
|
|
||||||
## Current Limitations
|
|
||||||
|
|
||||||
- Only pipelines that declare the request-batch contract use fused batch
|
|
||||||
execution.
|
|
||||||
- Batches are homogeneous under `RequestBatchSamplingParamsKey`; heterogeneous resolution or
|
|
||||||
incompatible guidance settings do not co-batch yet.
|
|
||||||
- FIFO scheduling can reduce batching opportunities when an incompatible
|
|
||||||
request is at the front of the queue.
|
|
||||||
- `request_batch_max_wait_ms` improves burst coalescing but can add latency to
|
|
||||||
the first request in a scheduler wave. Keep it small for latency-sensitive
|
|
||||||
serving.
|
|
||||||
- Step-wise continuous batching is documented separately and only applies when
|
|
||||||
`step_execution=True`.
|
|
||||||
|
|
||||||
## Related Files
|
|
||||||
|
|
||||||
- Request object and request batch:
|
|
||||||
[`vllm_omni/diffusion/request.py`](gh-file:vllm_omni/diffusion/request.py)
|
|
||||||
- Scheduler interface:
|
|
||||||
[`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
|
||||||
- Scheduler base:
|
|
||||||
[`vllm_omni/diffusion/sched/base_scheduler.py`](gh-file:vllm_omni/diffusion/sched/base_scheduler.py)
|
|
||||||
- Engine:
|
|
||||||
[`vllm_omni/diffusion/diffusion_engine.py`](gh-file:vllm_omni/diffusion/diffusion_engine.py)
|
|
||||||
- Worker runner:
|
|
||||||
[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
|
|
||||||
- Executor interface:
|
|
||||||
[`vllm_omni/diffusion/executor/abstract.py`](gh-file:vllm_omni/diffusion/executor/abstract.py)
|
|
||||||
- Tests:
|
|
||||||
[`tests/diffusion/test_diffusion_engine.py`](gh-file:tests/diffusion/test_diffusion_engine.py)
|
|
||||||
@@ -1,131 +0,0 @@
|
|||||||
# Adding Step Execution Support for Diffusion Pipelines
|
|
||||||
|
|
||||||
This guide documents vLLM-Omni's stepwise diffusion contract for model authors
|
|
||||||
and contributors implementing `step_execution=True` support for a diffusion
|
|
||||||
pipeline.
|
|
||||||
|
|
||||||
For end-user enablement, supported models, and current limitations, see
|
|
||||||
[Step Execution](../../user_guide/diffusion/step_execution.md).
|
|
||||||
|
|
||||||
This document describes the base step-execution contract only. For the
|
|
||||||
experimental batching policy layered on top of the step-wise path, see
|
|
||||||
[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md).
|
|
||||||
|
|
||||||
## Current Support Scope
|
|
||||||
|
|
||||||
`step_execution` is **not** a generic diffusion toggle. It only works for
|
|
||||||
pipelines that implement the segmented stateful contract in
|
|
||||||
[`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py).
|
|
||||||
|
|
||||||
This page is intentionally author-facing. Treat runtime enablement
|
|
||||||
(`step_execution=True` in Python or `--step-execution` in serving) as an
|
|
||||||
opt-in user knob layered on top of the implementation contract below.
|
|
||||||
|
|
||||||
Current in-tree support:
|
|
||||||
|
|
||||||
| Pipeline | Example models | Step execution |
|
|
||||||
|----------|----------------|----------------|
|
|
||||||
| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes |
|
|
||||||
| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No |
|
|
||||||
|
|
||||||
Current engine/runtime limitations:
|
|
||||||
|
|
||||||
- Continuous batching with `max_num_seqs > 1` is experimental and documented in
|
|
||||||
[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md).
|
|
||||||
Keep `max_num_seqs=1` for the conservative single-request step path.
|
|
||||||
- `cache_backend` is not supported in step mode.
|
|
||||||
- Request-mode extras such as KV transfer are not wired into step mode yet.
|
|
||||||
- Unsupported pipelines now fail early during model loading instead of failing on the first request.
|
|
||||||
|
|
||||||
## Execution Contract
|
|
||||||
|
|
||||||
Step mode is driven by four pipeline methods plus the shared mutable request
|
|
||||||
state object:
|
|
||||||
|
|
||||||
- `prepare_encode(state)`: one-time request preparation.
|
|
||||||
- `denoise_step(state)`: compute the noise prediction for the current step.
|
|
||||||
- `step_scheduler(state, noise_pred)`: mutate latents and advance step state.
|
|
||||||
- `post_decode(state)`: decode the final output after denoising is complete.
|
|
||||||
|
|
||||||
The state lives in
|
|
||||||
[`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py)
|
|
||||||
as `StepRequestState`. Store request-scoped tensors there, or use
|
|
||||||
`state.extra` for model-specific fields that do not justify extending the core
|
|
||||||
dataclass.
|
|
||||||
|
|
||||||
The worker-side step loop lives in
|
|
||||||
[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py):
|
|
||||||
|
|
||||||
1. `prepare_encode()` runs once for a new request.
|
|
||||||
2. `denoise_step()` runs every scheduler tick.
|
|
||||||
3. `step_scheduler()` mutates `state.latents` and advances `state.step_index`.
|
|
||||||
4. `post_decode()` runs exactly once after `state.denoise_completed` becomes true.
|
|
||||||
|
|
||||||
## Recommended Split
|
|
||||||
|
|
||||||
When converting an existing request-level `forward()` pipeline, keep the split
|
|
||||||
strict and mechanical:
|
|
||||||
|
|
||||||
| Request-level phase | Stepwise method | What belongs there |
|
|
||||||
|---------------------|-----------------|--------------------|
|
|
||||||
| Input validation, prompt encoding, latent init, timestep prep, per-request scheduler creation | `prepare_encode()` | Anything that should happen once per request |
|
|
||||||
| Transformer forward / noise prediction | `denoise_step()` | Pure denoise computation for the current timestep |
|
|
||||||
| `scheduler.step(...)` and `step_index += 1` | `step_scheduler()` | Only latent/state mutation for one step |
|
|
||||||
| VAE decode / postprocess | `post_decode()` | Final decode only |
|
|
||||||
|
|
||||||
Keep the stepwise path reusing the same helpers as the request-level path
|
|
||||||
whenever possible. Reimplementing the denoise loop from scratch is the easiest
|
|
||||||
way to introduce behavioral drift.
|
|
||||||
|
|
||||||
## Qwen-Image Reference
|
|
||||||
|
|
||||||
[`pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py)
|
|
||||||
is the reference implementation and is split correctly for the current
|
|
||||||
contract:
|
|
||||||
|
|
||||||
- `prepare_encode()` reuses `_prepare_generation_context()` so prompt encoding,
|
|
||||||
latent init, timestep creation, CFG setup, and shape bookkeeping stay aligned
|
|
||||||
with `forward()`.
|
|
||||||
- `prepare_encode()` deep-copies `self.scheduler` **after**
|
|
||||||
`prepare_timesteps()` so request-specific scheduler state is isolated.
|
|
||||||
- `denoise_step()` reuses `_build_denoise_kwargs()` plus
|
|
||||||
`predict_noise_maybe_with_cfg()`, so sequential CFG, CFG-parallel, and
|
|
||||||
non-CFG behavior stay identical to the request-level path.
|
|
||||||
- `step_scheduler()` only calls
|
|
||||||
`scheduler_step_maybe_with_cfg(..., per_request_scheduler=state.scheduler)`
|
|
||||||
and increments `state.step_index`.
|
|
||||||
- `post_decode()` reuses `_decode_latents()`, so the final image decode matches
|
|
||||||
the normal `forward()` path.
|
|
||||||
|
|
||||||
That decomposition is the target pattern for future models.
|
|
||||||
|
|
||||||
## Rules For New Pipelines
|
|
||||||
|
|
||||||
- Do not keep request-scoped scheduler state on `self.scheduler`. Copy it into
|
|
||||||
`state.scheduler` during `prepare_encode()`.
|
|
||||||
- Do not mutate `state.step_index` inside `denoise_step()`. Only
|
|
||||||
`step_scheduler()` should advance the step.
|
|
||||||
- Do not decode partial outputs in `denoise_step()` or `step_scheduler()`.
|
|
||||||
- If the request-level pipeline has condition latents, masks, or edit-specific
|
|
||||||
tensors, store them in `state` or `state.extra`, not in global pipeline
|
|
||||||
attributes.
|
|
||||||
- Preserve CFG behavior by sharing the same helper path used by `forward()`.
|
|
||||||
- Keep `post_decode()` equivalent to the tail of `forward()`.
|
|
||||||
|
|
||||||
## Validation Checklist
|
|
||||||
|
|
||||||
Before marking a pipeline as `supports_step_execution = True`, verify:
|
|
||||||
|
|
||||||
- Stepwise output matches request-level output for the same seed and sampling params.
|
|
||||||
- Per-request scheduler state is isolated across concurrent requests.
|
|
||||||
- Abort during denoise does not leak cached state.
|
|
||||||
- `step_index` reported by `RunnerOutput` matches the scheduler progress.
|
|
||||||
- CFG-parallel and non-CFG paths both work if the request-level pipeline supports them.
|
|
||||||
|
|
||||||
## Related Files
|
|
||||||
|
|
||||||
- Contract: [`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py)
|
|
||||||
- State: [`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py)
|
|
||||||
- Runner loop: [`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
|
|
||||||
- Scheduler transport: [`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
|
||||||
- Reference pipeline: [`vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py)
|
|
||||||
@@ -0,0 +1,156 @@
|
|||||||
|
# Distributed Layerwise Offload
|
||||||
|
|
||||||
|
This document describes distributed layerwise offload (DLO) for diffusion
|
||||||
|
models. DLO keeps only a small number of DiT blocks on the accelerator and
|
||||||
|
streams the remaining blocks from host memory. The distributed backend can
|
||||||
|
either shard those host-side weights across an existing parallel group or keep
|
||||||
|
the standard loader's rank-local weights and avoid an additional collective.
|
||||||
|
|
||||||
|
For user-facing commands, see the
|
||||||
|
[distributed layerwise offloading guide](../../user_guide/diffusion/cpu_offload.md)
|
||||||
|
and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md).
|
||||||
|
|
||||||
|
## Status
|
||||||
|
|
||||||
|
DLO is implemented for multi-device diffusion execution. The default
|
||||||
|
AllGather path is the primary path for DP and SP deployments. The
|
||||||
|
`--dlo-no-use-allgather` path is a rank-local compatibility mode: it is useful
|
||||||
|
for standard-loader sharding, workstation bring-up, and systems where an
|
||||||
|
additional DLO collective is undesirable, but it does not reduce host weight
|
||||||
|
storage across ranks.
|
||||||
|
|
||||||
|
The compatibility matrix below describes the current implementation. The
|
||||||
|
unit-level guards are covered, but not every parallelism combination has a
|
||||||
|
full model-and-hardware end-to-end test.
|
||||||
|
|
||||||
|
## Design
|
||||||
|
|
||||||
|
### DLO consumes the existing parallel topology
|
||||||
|
|
||||||
|
DLO does not create a new DP, TP, or SP topology. It reads the configured
|
||||||
|
`DiffusionParallelConfig` and attaches offload hooks to the DiT blocks after the
|
||||||
|
standard distributed groups have been initialized.
|
||||||
|
|
||||||
|
The DLO weight-sharding group is selected as follows:
|
||||||
|
|
||||||
|
1. Use the existing DP group when `data_parallel_size > 1`.
|
||||||
|
2. When DP is one and SP is greater than one, use the SP group.
|
||||||
|
3. Otherwise, run rank-locally without a DLO process group.
|
||||||
|
|
||||||
|
TP is deliberately not used as DLO's AllGather group. HSDP has its own
|
||||||
|
parameter-sharding lifecycle and is not allowed to be sharded a second time by
|
||||||
|
DLO's AllGather path.
|
||||||
|
|
||||||
|
### AllGather path
|
||||||
|
|
||||||
|
With the default `dlo_use_allgather=True`, each rank stores approximately
|
||||||
|
`1 / group_size` of each streamable block in pinned host memory. The next
|
||||||
|
block's shard is copied to a device buffer and reconstructed with
|
||||||
|
`all_gather_into_tensor` on a communication stream while the current block is
|
||||||
|
executing.
|
||||||
|
|
||||||
|
```text
|
||||||
|
Compute: [Block N] [Block N+1] [Block N+2]
|
||||||
|
H2D: [shard N+1] [shard N+2]
|
||||||
|
AllGather: [full N+1] [full N+2]
|
||||||
|
Buffers: [current slot] [prefetch slot] [current slot]
|
||||||
|
```
|
||||||
|
|
||||||
|

|
||||||
|
|
||||||
|
The backend uses two shared device buffers, so accelerator weight residency is
|
||||||
|
bounded by the largest streamed blocks rather than the complete model.
|
||||||
|
|
||||||
|
When DP is greater than one, the engine can process one request per DP rank in
|
||||||
|
the same denoising wave. Because AllGather is a collective, all participating
|
||||||
|
requests must take the same execution path at every denoising step.
|
||||||
|
|
||||||
|
### Rank-local path without DLO AllGather
|
||||||
|
|
||||||
|
With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
|
||||||
|
one. The regular model loader remains responsible for preparing each rank's
|
||||||
|
weights, including TP-local tensors or HSDP-managed parameters. DLO then
|
||||||
|
streams those rank-local tensors block by block using H2D copies only.
|
||||||
|
|
||||||
|
This mode means:
|
||||||
|
|
||||||
|
- DP still provides independent replicas, but DLO does not shard weights
|
||||||
|
across DP ranks.
|
||||||
|
- SP still performs its normal activation/attention collectives, but DLO does
|
||||||
|
not shard weights across SP ranks.
|
||||||
|
- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
|
||||||
|
DLO's additional weight AllGather is disabled.
|
||||||
|
- Pure DP deployments keep a full host-side model copy per rank, subject to
|
||||||
|
shared OS page-cache behavior.
|
||||||
|
- The scheduler does not require a synchronized DP request wave for DLO.
|
||||||
|
|
||||||
|
This path is intentionally not implemented by reusing the AllGather mmap
|
||||||
|
loader. It relies on the standard loader so model-specific TP/HSDP transforms
|
||||||
|
remain intact.
|
||||||
|
|
||||||
|
## Parallelism compatibility
|
||||||
|
|
||||||
|
| Parallelism | DLO + AllGather | DLO without AllGather |
|
||||||
|
|---|---|---|
|
||||||
|
| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight sharding or cross-DP DLO collective. |
|
||||||
|
| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
|
||||||
|
| **TP > 1** | Unsupported in the DLO mmap path. TP-aware loader callbacks are bypassed, so the backend rejects this configuration when it enters that path. | Standard loading is retained and TP-local tensors can be streamed. This is the intended compatibility path, but it still needs broader model and hardware validation. |
|
||||||
|
| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
|
||||||
|
|
||||||
|
### Combined dimensions
|
||||||
|
|
||||||
|
- **DP + SP:** DLO uses the DP group for weight sharding when DP is greater
|
||||||
|
than one; SP continues to use its own sequence-parallel group. If DP is one,
|
||||||
|
the SP group becomes DLO's sharding group in AllGather mode.
|
||||||
|
- **DP + TP/SP without AllGather:** standard model loading defines the
|
||||||
|
rank-local tensor layout. DLO adds no cross-DP, cross-TP, or cross-SP weight
|
||||||
|
collective.
|
||||||
|
- **HSDP + SP:** the general parallel configuration permits HSDP over SP, but
|
||||||
|
DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight
|
||||||
|
materialization and synchronization.
|
||||||
|
- **HSDP + DP or TP:** rejected independently by the diffusion parallel
|
||||||
|
configuration.
|
||||||
|
|
||||||
|
## Request and loading constraints
|
||||||
|
|
||||||
|
AllGather DP multi-concurrency requires:
|
||||||
|
|
||||||
|
- explicit `num_inference_steps`;
|
||||||
|
- the same `num_inference_steps` for all requests in a wave; and
|
||||||
|
- identical request arguments that affect the collective execution path.
|
||||||
|
|
||||||
|
The no-AllGather path does not impose these DLO-specific synchronized-wave
|
||||||
|
requirements.
|
||||||
|
|
||||||
|
The mmap loader is used only by the supported DLO+AllGather path when the
|
||||||
|
model has a compatible checkpoint layout. Online quantization is incompatible
|
||||||
|
with that sharded mmap path; use `--dlo-no-use-allgather` or disable online
|
||||||
|
quantization. Models that do not meet the mmap requirements use the regular
|
||||||
|
loader path.
|
||||||
|
|
||||||
|
## Validation coverage
|
||||||
|
|
||||||
|
Current source-level validation includes:
|
||||||
|
|
||||||
|
- HSDP + DLO + AllGather rejection;
|
||||||
|
- HSDP + DLO without AllGather acceptance at configuration level;
|
||||||
|
- TP rejection in the DLO+AllGather mmap path;
|
||||||
|
- resident-layer requests requiring no-AllGather;
|
||||||
|
- DP request-wave validation for denoising-step compatibility;
|
||||||
|
- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
|
||||||
|
tests.
|
||||||
|
|
||||||
|
The highest-value missing coverage is end-to-end numerical comparison against
|
||||||
|
ordinary layerwise offload for DP+SP, TP+no-AllGather, and HSDP+SP+no-AllGather
|
||||||
|
on the target CUDA/NCCL or CANN/HCCL hardware.
|
||||||
|
|
||||||
|
## Recommendations
|
||||||
|
|
||||||
|
- Use **DP + DLO AllGather** for the supported throughput and host-memory
|
||||||
|
scaling path.
|
||||||
|
- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
|
||||||
|
not the goal.
|
||||||
|
- Use **no-AllGather** to bring up TP or workstation PCIe configurations, with
|
||||||
|
the expectation of higher host-memory use and lower validation confidence.
|
||||||
|
- Prefer **HSDP alone** for production HSDP deployments until the combined
|
||||||
|
HSDP + DLO no-AllGather path has broader end-to-end coverage.
|
||||||
@@ -0,0 +1,249 @@
|
|||||||
|
# Quantization
|
||||||
|
|
||||||
|
This document describes the quantization architecture in vLLM-Omni and the
|
||||||
|
extension points for contributors. For user-facing configuration, supported
|
||||||
|
hardware, and method-specific instructions, see the
|
||||||
|
[quantization user guide](../../user_guide/quantization/overview.md).
|
||||||
|
|
||||||
|
## Goals and scope
|
||||||
|
|
||||||
|
The design has four goals:
|
||||||
|
|
||||||
|
1. Provide one `quantization_config` entry point for diffusion-only,
|
||||||
|
multi-stage omni/TTS, and multi-stage diffusion models.
|
||||||
|
2. Reuse vLLM's quantization registry and `QuantizationConfig` contract where
|
||||||
|
a vLLM backend already provides the required behavior.
|
||||||
|
3. Allow different components of one pipeline to use different quantization
|
||||||
|
configurations.
|
||||||
|
4. Preserve the distinction between online quantization of a BF16/FP16
|
||||||
|
checkpoint and loading a pre-quantized checkpoint with serialized scales.
|
||||||
|
|
||||||
|
This document covers configuration resolution, component routing, model
|
||||||
|
integration, backend extension, and validation. Kernel implementations and
|
||||||
|
user-facing support matrices are outside its scope.
|
||||||
|
|
||||||
|
## Architecture
|
||||||
|
|
||||||
|
The quantization path is split across the following layers:
|
||||||
|
|
||||||
|
| Layer | Responsibility | Source |
|
||||||
|
| --- | --- | --- |
|
||||||
|
| Public API | Exposes `build_quant_config`, component routing, and backend registration | `vllm_omni/quantization/__init__.py` |
|
||||||
|
| Factory | Normalizes method names, selects Omni overrides or vLLM registry backends, detects ModelOpt configs, and reconciles checkpoint metadata | `vllm_omni/quantization/factory.py` |
|
||||||
|
| Component router | Resolves a quantization config from a layer prefix using longest-prefix matching | `vllm_omni/quantization/component_config.py` |
|
||||||
|
| Method configs | Implement `QuantizationConfig` and select the platform-specific linear method | `vllm_omni/quantization/*_config.py` |
|
||||||
|
| Diffusion configuration | Stores the canonical runtime config and propagates checkpoint configuration | `vllm_omni/diffusion/data.py` |
|
||||||
|
| Worker and model integration | Passes the resolved config into vLLM's model configuration and model layers | `vllm_omni/diffusion/worker/diffusion_worker.py` and `vllm_omni/diffusion/models/` |
|
||||||
|
|
||||||
|
The runtime contract remains vLLM's `QuantizationConfig`:
|
||||||
|
|
||||||
|
```text
|
||||||
|
user/CLI or model config
|
||||||
|
|
|
||||||
|
v
|
||||||
|
OmniDiffusionConfig.quantization_config
|
||||||
|
|
|
||||||
|
+--> build_quant_config()
|
||||||
|
| |
|
||||||
|
| +--> Omni override
|
||||||
|
| +--> vLLM registry backend
|
||||||
|
| +--> ComponentQuantizationConfig
|
||||||
|
|
|
||||||
|
+--> checkpoint metadata reconciliation
|
||||||
|
|
|
||||||
|
+--> diffusion worker / vLLM config
|
||||||
|
|
|
||||||
|
+--> model layers call get_quant_method(layer, prefix)
|
||||||
|
```
|
||||||
|
|
||||||
|
The model layer receives either a method implementation or `None`. Returning
|
||||||
|
`None` leaves that layer unquantized.
|
||||||
|
|
||||||
|
## Configuration resolution
|
||||||
|
|
||||||
|
`build_quant_config()` accepts the following forms:
|
||||||
|
|
||||||
|
| Input | Resolution |
|
||||||
|
| --- | --- |
|
||||||
|
| `None` or `"none"` | Disable quantization. |
|
||||||
|
| Method string, such as `"fp8"` | Resolve an Omni override first, then the vLLM quantization registry. |
|
||||||
|
| Flat dictionary with `method` or `quant_method` | Normalize the method and construct its `QuantizationConfig`. |
|
||||||
|
| Per-component dictionary | Construct a `ComponentQuantizationConfig`. |
|
||||||
|
| Existing `QuantizationConfig` | Pass through without rebuilding it. |
|
||||||
|
|
||||||
|
Method aliases are normalized case-insensitively, with `-` and `_` treated as
|
||||||
|
equivalent. `auto-round` and `auto_round` use the Omni INC/AutoRound adapter.
|
||||||
|
ModelOpt checkpoint metadata is detected from its method, producer metadata, or
|
||||||
|
`quant_algo`, then mapped to the corresponding vLLM ModelOpt configuration.
|
||||||
|
|
||||||
|
The factory uses Omni-specific overrides for methods whose diffusion behavior
|
||||||
|
or platform dispatch is not provided by the vLLM registry. The current override
|
||||||
|
set includes INT8, BitsAndBytes, MXFP8, MXFP4, MXFP4 dual-scale, and INC/
|
||||||
|
AutoRound. Other methods, such as FP8, GGUF, and ModelOpt, are resolved through
|
||||||
|
vLLM's registry when their configuration is compatible.
|
||||||
|
|
||||||
|
## Per-component routing
|
||||||
|
|
||||||
|
Multi-stage models do not necessarily want one quantization policy for every
|
||||||
|
module. A component map expresses that policy without requiring each pipeline
|
||||||
|
to construct separate vLLM configurations:
|
||||||
|
|
||||||
|
```python
|
||||||
|
quantization_config = {
|
||||||
|
"transformer": {"method": "fp8"},
|
||||||
|
"vae": None,
|
||||||
|
"default": None,
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
`ComponentQuantizationConfig` applies the following rules:
|
||||||
|
|
||||||
|
1. Component keys are interpreted as runtime layer-prefixes.
|
||||||
|
2. The longest matching prefix wins.
|
||||||
|
3. `None` disables quantization for the matching component.
|
||||||
|
4. `default`, when present, handles prefixes with no explicit match.
|
||||||
|
5. A matching child config delegates to its normal
|
||||||
|
`get_quant_method(layer, prefix)` implementation.
|
||||||
|
|
||||||
|
Prefix matching happens after the model's weight-name mapping. New model
|
||||||
|
integrations must verify the actual runtime prefixes, especially when a
|
||||||
|
`WeightsMapper` changes checkpoint names. A prefix mismatch can silently route
|
||||||
|
a layer to the default configuration, so component-routing tests should cover
|
||||||
|
both positive matches and fall-through behavior.
|
||||||
|
|
||||||
|
## Online and pre-quantized checkpoints
|
||||||
|
|
||||||
|
The same method name can represent two different loading paths:
|
||||||
|
|
||||||
|
- **Online**: the source checkpoint contains BF16/FP16 weights; the method
|
||||||
|
creates quantized weights or scales while the model is loaded.
|
||||||
|
- **Pre-quantized**: the checkpoint contains quantized weights and the metadata
|
||||||
|
required by the method; the loader consumes those serialized tensors.
|
||||||
|
|
||||||
|
Method configuration owns the mode flag, for example
|
||||||
|
`is_checkpoint_int8_serialized` or `is_checkpoint_mxfp8_serialized`. The
|
||||||
|
configuration's `get_quant_method()` selects the corresponding linear method.
|
||||||
|
|
||||||
|
For pipelines with transformer-specific `config.json` files, such as cascade
|
||||||
|
models, `resolve_quant_config_from_disk()` reconciles the active configuration
|
||||||
|
with each transformer's checkpoint metadata. Its invariants are:
|
||||||
|
|
||||||
|
- If no active configuration exists, valid checkpoint metadata can be used for
|
||||||
|
auto-detection.
|
||||||
|
- A method mismatch is an error rather than an implicit conversion.
|
||||||
|
- Serialized-checkpoint flags rebuild an online configuration into the
|
||||||
|
matching offline configuration.
|
||||||
|
- Checkpoint-specific `ignored_layers` are preserved by rebuilding the active
|
||||||
|
configuration when they differ.
|
||||||
|
- AutoRound MXFP8 metadata (`data_type="mx_fp"`) selects its offline path.
|
||||||
|
|
||||||
|
This prevents a model from interpreting serialized weights with an online
|
||||||
|
loader or silently applying the wrong method to one transformer in a
|
||||||
|
multi-transformer pipeline.
|
||||||
|
|
||||||
|
## Model and layer integration
|
||||||
|
|
||||||
|
Diffusion model constructors receive the resolved configuration and pass it to
|
||||||
|
quantizable vLLM layers. Each method configuration is responsible for:
|
||||||
|
|
||||||
|
- identifying supported layers in `get_quant_method()`;
|
||||||
|
- creating the correct parameter shapes and loaders;
|
||||||
|
- processing weights after loading when conversion is required; and
|
||||||
|
- applying the platform-specific quantized operation at runtime.
|
||||||
|
|
||||||
|
Model code must keep precision-sensitive or unsupported modules unquantized.
|
||||||
|
The shared `safe_quant_config()` helper is used by model integrations for
|
||||||
|
normalization and modulation layers where a generic FP8 configuration would be
|
||||||
|
unsafe. Text encoders, audio/vision encoders, VAEs, schedulers, and tokenizers
|
||||||
|
are not implicitly quantized just because the main transformer is quantized;
|
||||||
|
they require explicit component routing and method support.
|
||||||
|
|
||||||
|
Weight quantization and quantized KV cache are separate paths. The former is
|
||||||
|
resolved through `quantization_config`; runtime attention/KV-cache quantization
|
||||||
|
is configured by the diffusion attention configuration and must be validated
|
||||||
|
independently.
|
||||||
|
|
||||||
|
## Platform and parallelism boundaries
|
||||||
|
|
||||||
|
Platform checks belong in the method implementation rather than in the common
|
||||||
|
factory. This keeps configuration parsing portable while allowing CUDA, NPU,
|
||||||
|
and XPU methods to select different kernels or reject unsupported execution at
|
||||||
|
model construction time. The user guide is the source of truth for the current
|
||||||
|
hardware support matrix.
|
||||||
|
|
||||||
|
Quantization methods must also respect the weight layout produced by the
|
||||||
|
configured parallelism strategy:
|
||||||
|
|
||||||
|
- Tensor-parallel layers may receive partitioned weight shapes and loaders.
|
||||||
|
- Sequence-parallel execution changes activation movement but does not by
|
||||||
|
itself change the quantization configuration.
|
||||||
|
- Data-parallel replicas build their own method state; quantization metadata
|
||||||
|
must be available consistently on every rank.
|
||||||
|
- `WeightsMapper` transformations must be applied before matching ignored-layer
|
||||||
|
lists or component prefixes.
|
||||||
|
|
||||||
|
Distributed layerwise offload has an additional boundary: its sharded mmap
|
||||||
|
AllGather path is not compatible with online quantization because the path
|
||||||
|
expects already materialized checkpoint weights. Use the rank-local
|
||||||
|
`--dlo-no-use-allgather` path or a compatible pre-quantized checkpoint when
|
||||||
|
combining these features. See the [DLO design document](distributed_layerwise_offload.md)
|
||||||
|
for the full compatibility matrix.
|
||||||
|
|
||||||
|
## Adding a quantization backend
|
||||||
|
|
||||||
|
Contributors should follow this sequence:
|
||||||
|
|
||||||
|
1. Check whether an existing vLLM backend can provide the required weight
|
||||||
|
format and kernel behavior. Prefer reuse over a duplicate Omni backend.
|
||||||
|
2. If Omni-specific behavior is required, implement a `QuantizationConfig` and
|
||||||
|
its `QuantizeMethodBase`/linear method in `vllm_omni/quantization/`.
|
||||||
|
3. Implement `get_name()`, supported activation dtypes, minimum capability,
|
||||||
|
checkpoint parsing, ignored-layer handling, and `get_quant_method()`.
|
||||||
|
4. Add a lazy factory override in `factory.py`. Keep optional platform
|
||||||
|
dependencies out of package import paths and normalize aliases to one
|
||||||
|
canonical method name.
|
||||||
|
5. If the method is pre-quantized, define the checkpoint metadata and the
|
||||||
|
serialized loading path before adding an online path.
|
||||||
|
6. Add model-specific component routing only when the model contains stages
|
||||||
|
that must use different policies.
|
||||||
|
7. Document user-visible configuration in the quantization user guide and keep
|
||||||
|
implementation rationale here.
|
||||||
|
|
||||||
|
The public factory registration hook,
|
||||||
|
`register_quantization_override(method, builder)`, is intended for integrations
|
||||||
|
that need to register an Omni-compatible builder without changing the vLLM
|
||||||
|
registry.
|
||||||
|
|
||||||
|
## Validation
|
||||||
|
|
||||||
|
Validation should cover the complete path, not only whether a config object can
|
||||||
|
be instantiated:
|
||||||
|
|
||||||
|
| Area | Required checks |
|
||||||
|
| --- | --- |
|
||||||
|
| Factory | Strings, aliases, flat dictionaries, invalid inputs, passthrough, and component maps |
|
||||||
|
| Method config | Online/offline selection, ignored layers, platform guards, and weight loading |
|
||||||
|
| Checkpoint metadata | Auto-detection, method mismatch errors, serialized flags, and per-transformer metadata |
|
||||||
|
| Model integration | Correct quantization scope, unquantized components, and parallel weight layouts |
|
||||||
|
| Numerical quality | Same-seed BF16 versus quantized output comparison at representative resolutions and step counts |
|
||||||
|
| Performance | Peak device memory, generation latency, and any backend-specific throughput metric |
|
||||||
|
|
||||||
|
The unit tests live under `tests/diffusion/quantization/`. The trajectory
|
||||||
|
similarity tool in
|
||||||
|
`vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py`
|
||||||
|
provides a repeatable quality and memory comparison for diffusion models.
|
||||||
|
Method support and feature combinations remain model- and hardware-dependent;
|
||||||
|
passing factory tests alone does not establish production support.
|
||||||
|
|
||||||
|
## Design invariants
|
||||||
|
|
||||||
|
- `quantization_config` is the single runtime source of truth after config
|
||||||
|
normalization.
|
||||||
|
- A checkpoint method mismatch fails loudly instead of risking weight
|
||||||
|
corruption.
|
||||||
|
- Components without an explicit supported quantization path remain in their
|
||||||
|
original precision.
|
||||||
|
- Optional backend dependencies are imported lazily.
|
||||||
|
- Method implementations own platform-specific capability checks and kernels.
|
||||||
|
- User-facing support claims require model, hardware, numerical, and memory
|
||||||
|
validation in addition to configuration-level tests.
|
||||||
@@ -1,59 +0,0 @@
|
|||||||
# Distributed utils
|
|
||||||
|
|
||||||
This directory (vllm_omni/distributed/ray_utils) contains utilities for distributed execution in vllm-omni, supporting both **Ray** and **Multiprocessing** backends.
|
|
||||||
## 1. Installation
|
|
||||||
```bash
|
|
||||||
pip install "ray[default]"
|
|
||||||
```
|
|
||||||
## 2. Ray Utils
|
|
||||||
|
|
||||||
The `ray_utils` module provides helper functions for managing Ray clusters and actors, which is essential for:
|
|
||||||
* **Multi-node deployment**: Running pipeline stages across different physical machines.
|
|
||||||
* **Resource management**: Efficient GPU/CPU allocation.
|
|
||||||
|
|
||||||
### 2.1 Basic Usage
|
|
||||||
|
|
||||||
To use the Ray backend, specify `worker_backend="ray"` when initializing the engine.
|
|
||||||
|
|
||||||
**Command Line Example:**
|
|
||||||
```bash
|
|
||||||
vllm serve Qwen/Qwen2.5-Omni-7B \
|
|
||||||
--omni \
|
|
||||||
--port 8091 \
|
|
||||||
--worker-backend ray \
|
|
||||||
--ray-address auto
|
|
||||||
```
|
|
||||||
|
|
||||||
### 2.2 Cluster Setup
|
|
||||||
|
|
||||||
**Step 1: Start Head Node**
|
|
||||||
Run this on your primary machine:
|
|
||||||
```bash
|
|
||||||
ray start --head --port=6399
|
|
||||||
```
|
|
||||||
|
|
||||||
**Step 2: Connect Worker Nodes**
|
|
||||||
Run this on each worker machine:
|
|
||||||
```bash
|
|
||||||
ray start --address=<HEAD_NODE_IP>:6399
|
|
||||||
```
|
|
||||||
|
|
||||||
> **Tip**: For a complete cluster setup script, refer to the vLLM example:
|
|
||||||
> [run_cluster.sh](https://github.com/vllm-project/vllm/blob/main/examples/online_serving/run_cluster.sh)
|
|
||||||
|
|
||||||
### 2.3 Distributed Connector Support
|
|
||||||
|
|
||||||
When running on Ray, the system automatically adapts its communication strategy:
|
|
||||||
|
|
||||||
* **Cross-Node**: Recommended to use `MooncakeTransferEngineConnector` (RDMA, fastest) or `MooncakeStoreConnector` (TCP fallback).
|
|
||||||
* **Same-Node**: Can still use `SharedMemoryConnector` for efficiency, or Ray's native object store (plasma).
|
|
||||||
* **SHM payload routing**: `SharedMemoryConnector` always stores payloads in `/dev/shm`.
|
|
||||||
|
|
||||||
### 2.4 Internal Helpers
|
|
||||||
|
|
||||||
* **`initialize_ray_cluster`**: Connects to an existing Ray cluster or starts a local one.
|
|
||||||
|
|
||||||
## 3. Troubleshooting
|
|
||||||
|
|
||||||
* **Connection Issues**: Ensure the Ray head node is accessible and ports (default 6399 in this example) are open.
|
|
||||||
* **Version Mismatch**: Ensure all nodes run the same version of Ray and Python.
|
|
||||||
Binary file not shown.
|
After Width: | Height: | Size: 195 KiB |
+54
-8
@@ -1,6 +1,8 @@
|
|||||||
# Design Documents
|
# Design Documents
|
||||||
|
|
||||||
This section contains design documents and architecture specifications for vLLM-Omni.
|
This section contains design documents and architecture specifications for
|
||||||
|
vLLM-Omni. The sidebar groups documents by the system concern they describe,
|
||||||
|
while the source paths remain stable so existing links continue to work.
|
||||||
|
|
||||||
## Architecture Documents
|
## Architecture Documents
|
||||||
|
|
||||||
@@ -8,19 +10,63 @@ This section contains design documents and architecture specifications for vLLM-
|
|||||||
|
|
||||||
## Feature Design Documents
|
## Feature Design Documents
|
||||||
|
|
||||||
- [Disaggregated Inference](feature/disaggregated_inference.md)
|
### Runtime and stage execution
|
||||||
- [Ray-based Execution](feature/ray_based_execution.md)
|
|
||||||
- [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
|
|
||||||
- [Adding Step Execution Support for Diffusion Pipelines](feature/diffusion_step_execution.md)
|
|
||||||
- [Request-Level Batching for Diffusion](feature/diffusion_request_level_batching.md)
|
|
||||||
- [Continuous Batching for Step-Wise Diffusion](feature/diffusion_continuous_batching.md)
|
|
||||||
|
|
||||||
## Infrastructure Design Documents
|
- [Disaggregated Inference](feature/disaggregated_inference.md)
|
||||||
|
- [Async Chunk](feature/async_chunk.md)
|
||||||
|
- [Async Diffusion Output](feature/async_diffusion_output.md)
|
||||||
|
- [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
|
||||||
|
- [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
|
||||||
|
|
||||||
|
### Communication
|
||||||
|
|
||||||
|
#### OmniConnector implementations
|
||||||
|
|
||||||
|
- [Mooncake Store Connector](feature/omni_connectors/mooncake_store_connector.md)
|
||||||
|
- [Mooncake Transfer Engine Connector](feature/omni_connectors/mooncake_transfer_engine_connector.md)
|
||||||
|
- [Mori Transfer Engine Connector](feature/omni_connectors/mori_transfer_engine_connector.md)
|
||||||
|
- [Shared Memory Connector](feature/omni_connectors/shared_memory_connector.md)
|
||||||
|
- [Yuanrong Store Connector](feature/omni_connectors/yuanrong_connector.md)
|
||||||
|
- [Yuanrong Transfer Engine Connector](feature/omni_connectors/yuanrong_transfer_engine_connector.md)
|
||||||
|
|
||||||
|
### Diffusion acceleration
|
||||||
|
|
||||||
|
#### Parallelism
|
||||||
|
|
||||||
|
- [CFG-Parallel](feature/cfg_parallel.md)
|
||||||
|
- [Expert Parallel](feature/expert_parallel.md)
|
||||||
|
- [Hybrid Sharded Data Parallel (HSDP)](feature/hsdp.md)
|
||||||
|
- [Pipeline Parallel](feature/pipeline_parallel.md)
|
||||||
|
- [Sequence Parallel](feature/sequence_parallel.md)
|
||||||
|
- [Tensor Parallel](feature/tensor_parallel.md)
|
||||||
|
- [VAE Patch Parallelism](feature/vae_parallel.md)
|
||||||
|
|
||||||
|
#### Attention Backends
|
||||||
|
|
||||||
|
- [Skip-Softmax](feature/skip_softmax.md)
|
||||||
|
|
||||||
|
#### Quantization
|
||||||
|
|
||||||
|
- [Quantization](feature/quantization.md)
|
||||||
|
|
||||||
|
- [Cache-DiT](feature/cache_dit.md)
|
||||||
|
- [TeaCache](feature/teacache.md)
|
||||||
|
- [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md)
|
||||||
|
- [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md)
|
||||||
|
|
||||||
|
## Infrastructure and Performance
|
||||||
|
|
||||||
- [Prometheus Metrics](metrics.md)
|
- [Prometheus Metrics](metrics.md)
|
||||||
|
- [Speech Generation Performance Optimizations](qwen3_omni_tts_performance_optimization.md)
|
||||||
|
|
||||||
## Module Design Documents
|
## Module Design Documents
|
||||||
|
|
||||||
|
### Runtime Modules
|
||||||
|
|
||||||
- [AR Module](module/ar_module.md)
|
- [AR Module](module/ar_module.md)
|
||||||
- [DIT Module](module/dit_module.md)
|
- [DIT Module](module/dit_module.md)
|
||||||
|
|
||||||
|
### Orchestration
|
||||||
|
|
||||||
- [Entrypoint Module](module/entrypoint_module.md)
|
- [Entrypoint Module](module/entrypoint_module.md)
|
||||||
|
- [AsyncOmni Architecture (Qwen3-Omni Example)](module/async_omni_architecture.md)
|
||||||
@@ -252,8 +252,7 @@ while True:
|
|||||||
batch of compatible independent requests for one full pipeline forward when
|
batch of compatible independent requests for one full pipeline forward when
|
||||||
the pipeline supports request-level batching; the step-wise path can also
|
the pipeline supports request-level batching; the step-wise path can also
|
||||||
admit or remove compatible requests between denoise steps. See
|
admit or remove compatible requests between denoise steps. See
|
||||||
[Request-Level Batching](../feature/diffusion_request_level_batching.md) and
|
[Diffusion Continuous Batching](../feature/diffusion_continuous_batching.md).
|
||||||
[Continuous Batching for Step-Wise Diffusion](../feature/diffusion_continuous_batching.md).
|
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
|||||||
@@ -87,9 +87,8 @@ if __name__ == "__main__":
|
|||||||
|
|
||||||
!!! info
|
!!! info
|
||||||
|
|
||||||
For diffusion request-level batching controls such as `max_num_seqs` and
|
For diffusion request batching, step execution, and streaming controls, see
|
||||||
`request_batch_max_wait_ms`, see
|
[Diffusion Execution Modes](../user_guide/diffusion/execution_modes.md).
|
||||||
[Request-Level Batching](../user_guide/diffusion/request_batching.md).
|
|
||||||
|
|
||||||
For more usages, please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md)
|
For more usages, please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md)
|
||||||
|
|
||||||
|
|||||||
@@ -44,7 +44,7 @@ th {
|
|||||||
| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
||||||
| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
|
| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
|
||||||
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | |
|
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | |
|
||||||
| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, FL2VA, and Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | | | |
|
| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | | | |
|
||||||
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||||
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
|
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
|
||||||
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
|
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
|
||||||
|
|||||||
+29
-12
@@ -184,10 +184,15 @@ class Flux2Transformer2DModel(nn.Module):
|
|||||||
### How It Works
|
### How It Works
|
||||||
|
|
||||||
Distributed layerwise offloading extends single-GPU layerwise offloading to
|
Distributed layerwise offloading extends single-GPU layerwise offloading to
|
||||||
multi-device deployments. Each DP rank stores only **1/dp_size** of the model
|
multi-device deployments. In the default mode, each rank stores only
|
||||||
weights in host memory; full layer weights are reconstructed at runtime via
|
**1/dp_size** of the model weights in host memory; full layer weights are
|
||||||
**AllGather** on a dedicated communication stream, overlapped with computation
|
reconstructed at runtime via **AllGather** on a dedicated communication stream,
|
||||||
via a fixed double-buffer scheme.
|
overlapped with computation via a fixed double-buffer scheme. The
|
||||||
|
`--dlo-no-use-allgather` mode instead streams standard-loader rank-local
|
||||||
|
weights independently and does not shard weights across ranks.
|
||||||
|
|
||||||
|
For the implementation design and compatibility matrix, see
|
||||||
|
[Distributed Layerwise Offload](../../design/feature/distributed_layerwise_offload.md).
|
||||||
|
|
||||||
**Key features:**
|
**Key features:**
|
||||||
- **Weight sharding + AllGather**: each rank stores 1/dp_size of weights,
|
- **Weight sharding + AllGather**: each rank stores 1/dp_size of weights,
|
||||||
@@ -219,7 +224,7 @@ vllm serve /path/to/model --omni \
|
|||||||
--enable-distributed-layerwise-offload \
|
--enable-distributed-layerwise-offload \
|
||||||
--data-parallel-size 4
|
--data-parallel-size 4
|
||||||
|
|
||||||
# Without AllGather (each rank loads full weights, no sharding)
|
# Without DLO AllGather (each rank streams standard-loader rank-local weights)
|
||||||
vllm serve /path/to/model --omni \
|
vllm serve /path/to/model --omni \
|
||||||
--enable-distributed-layerwise-offload \
|
--enable-distributed-layerwise-offload \
|
||||||
--data-parallel-size 4 \
|
--data-parallel-size 4 \
|
||||||
@@ -247,9 +252,9 @@ m = Omni(
|
|||||||
| Flag | Description | Default |
|
| Flag | Description | Default |
|
||||||
|------|-------------|---------|
|
|------|-------------|---------|
|
||||||
| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
|
| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
|
||||||
| `--data-parallel-size N` | Number of DP ranks for weight sharding | `1` |
|
| `--data-parallel-size N` | Number of DP ranks; DLO weight-sharding group when AllGather is enabled | `1` |
|
||||||
| `--dlo-use-allgather` | Shard + AllGather (saves CPU, requires concurrent requests) | `true` |
|
| `--dlo-use-allgather` | Shard + AllGather (saves CPU, requires concurrent requests) | `true` |
|
||||||
| `--dlo-no-use-allgather` | Full weights per rank (no sharding, no AllGather) | `false` |
|
| `--dlo-no-use-allgather` | Use standard-loader rank-local weights; no DLO sharding or AllGather | `false` |
|
||||||
|
|
||||||
### How model weights are loaded (mmap path)
|
### How model weights are loaded (mmap path)
|
||||||
|
|
||||||
@@ -303,9 +308,15 @@ request while AllGather synchronizes only weight shards (request-independent).
|
|||||||
|
|
||||||
### Limitations
|
### Limitations
|
||||||
- Online quantization (FP8) is incompatible with mmap loading — falls back to
|
- Online quantization (FP8) is incompatible with mmap loading — falls back to
|
||||||
regular `load_weights()` automatically
|
regular loading only when the model is outside the DLO+AllGather mmap path;
|
||||||
- Tensor Parallel is not supported (DLO uses DP-based sharding)
|
online quantization with DLO+AllGather is rejected. Use
|
||||||
- HSDP + AllGather is rejected (would double-shard weights)
|
`--dlo-no-use-allgather` or disable online quantization.
|
||||||
|
- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path because TP-aware
|
||||||
|
weight-loader callbacks are bypassed. The no-AllGather path retains the
|
||||||
|
standard TP loader and is experimental.
|
||||||
|
- HSDP + AllGather is rejected because it would double-shard parameters. HSDP
|
||||||
|
with no-AllGather is accepted at configuration level, but has limited
|
||||||
|
end-to-end validation.
|
||||||
- `num_inference_steps=None` is not allowed in DP multi-concurrency mode
|
- `num_inference_steps=None` is not allowed in DP multi-concurrency mode
|
||||||
|
|
||||||
**Module Discovery**
|
**Module Discovery**
|
||||||
@@ -338,7 +349,7 @@ OffloadBackend (base class)
|
|||||||
├── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
|
├── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
|
||||||
│ (single-GPU, full weights on host)
|
│ (single-GPU, full weights on host)
|
||||||
└── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook
|
└── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook
|
||||||
(multi-GPU, 1/dp_size sharded weights + AllGather)
|
(multi-GPU, sharded AllGather or rank-local mode)
|
||||||
```
|
```
|
||||||
|
|
||||||
Factory function `get_offload_backend()` selects the appropriate backend based on
|
Factory function `get_offload_backend()` selects the appropriate backend based on
|
||||||
@@ -369,4 +380,10 @@ reasoner/generator components inside the model forward pass.
|
|||||||
**Notes:**
|
**Notes:**
|
||||||
- Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
|
- Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
|
||||||
- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks
|
- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks
|
||||||
- Distributed Layerwise Offloading works with any model that supports Layerwise Offloading — no additional model changes required. See [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) for usage examples.
|
- Distributed Layerwise Offloading reuses the layerwise block topology, but the
|
||||||
|
model and parallelism combination still needs validation. AllGather mode
|
||||||
|
additionally requires a compatible checkpoint/mmap layout; no-AllGather mode
|
||||||
|
retains the standard loader's rank-local tensors. See the
|
||||||
|
[DLO design](../../design/feature/distributed_layerwise_offload.md) and
|
||||||
|
[Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
|
||||||
|
for usage examples.
|
||||||
@@ -0,0 +1,252 @@
|
|||||||
|
# Diffusion Execution Modes
|
||||||
|
|
||||||
|
vLLM-Omni supports complete-request execution and step-wise diffusion
|
||||||
|
execution. Both modes use the same asynchronous engine output stream; the
|
||||||
|
configuration controls how work is scheduled and whether intermediate outputs
|
||||||
|
are exposed.
|
||||||
|
|
||||||
|
## Choose a Mode
|
||||||
|
|
||||||
|
| Goal | CLI configuration |
|
||||||
|
|---|---|
|
||||||
|
| Serial request execution | `--max-num-seqs 1` |
|
||||||
|
| Fused request-level batching | `--max-num-seqs N` |
|
||||||
|
| Single-request step execution | `--step-execution --max-num-seqs 1` |
|
||||||
|
| Step-wise continuous batching | `--step-execution --max-num-seqs N` |
|
||||||
|
| Chunked diffusion output | `--diffusion-streaming-output` |
|
||||||
|
|
||||||
|
`N` must be greater than `1` to allow batching. The selected pipeline must
|
||||||
|
support request-level batching or batched-step execution as appropriate;
|
||||||
|
single-request step support alone is not sufficient for step-wise continuous
|
||||||
|
batching.
|
||||||
|
|
||||||
|
## Request Execution
|
||||||
|
|
||||||
|
Request execution is the default when `--step-execution` is omitted. The
|
||||||
|
pipeline performs a complete `forward()` for each scheduler wave.
|
||||||
|
|
||||||
|
### Serial Requests
|
||||||
|
|
||||||
|
Use the conservative serial path with any request-mode pipeline:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
vllm serve MODEL --omni \
|
||||||
|
--port 8091 \
|
||||||
|
--max-num-seqs 1
|
||||||
|
```
|
||||||
|
|
||||||
|
### Request-Level Batching
|
||||||
|
|
||||||
|
Set `max_num_seqs` above one to combine compatible independent requests into a
|
||||||
|
single pipeline forward:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
vllm serve Qwen/Qwen-Image --omni \
|
||||||
|
--port 8091 \
|
||||||
|
--max-num-seqs 4
|
||||||
|
```
|
||||||
|
|
||||||
|
For bursty traffic, add a small admission window:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
vllm serve Qwen/Qwen-Image --omni \
|
||||||
|
--port 8091 \
|
||||||
|
--max-num-seqs 4 \
|
||||||
|
--request-batch-max-wait-ms 20
|
||||||
|
```
|
||||||
|
|
||||||
|
`--request-batch-max-wait-ms 0` is the default. A nonzero value can improve
|
||||||
|
batch formation but adds up to that much latency before a new scheduler wave.
|
||||||
|
It has no effect in step mode.
|
||||||
|
|
||||||
|
Each prompt remains a separate logical request. Do not submit a top-level list
|
||||||
|
as one packed prompt. Concurrent serving requests are batched internally when
|
||||||
|
their shapes, guidance settings, output counts, and LoRA settings are
|
||||||
|
compatible.
|
||||||
|
|
||||||
|
Only pipelines that declare request-batch support accept
|
||||||
|
`max_num_seqs>1` in request mode. Unsupported pipelines fail during engine
|
||||||
|
initialization; use `max_num_seqs=1` instead.
|
||||||
|
|
||||||
|
## Step Execution
|
||||||
|
|
||||||
|
Step execution lets the scheduler advance and abort requests between denoise
|
||||||
|
steps:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
vllm serve Qwen/Qwen-Image --omni \
|
||||||
|
--port 8091 \
|
||||||
|
--step-execution \
|
||||||
|
--max-num-seqs 1
|
||||||
|
```
|
||||||
|
|
||||||
|
Set a larger capacity to allow compatible requests to share denoise waves:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
vllm serve Qwen/Qwen-Image --omni \
|
||||||
|
--port 8091 \
|
||||||
|
--step-execution \
|
||||||
|
--max-num-seqs 8
|
||||||
|
```
|
||||||
|
|
||||||
|
Step continuous batching is experimental. Start with `max_num_seqs=1` when
|
||||||
|
validating a model or debugging correctness, then increase it for
|
||||||
|
multi-request throughput.
|
||||||
|
|
||||||
|
Step execution is capability-based, not a generic switch for every diffusion
|
||||||
|
model. Qwen-Image supports step-wise continuous batching. HunyuanImage3 also
|
||||||
|
supports it, but only when its resolved self-attention backend is `TORCH_SDPA`;
|
||||||
|
set `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` or configure
|
||||||
|
`diffusion_attention_config.default.backend=TORCH_SDPA` before using
|
||||||
|
`--max-num-seqs >1`. See the
|
||||||
|
[HunyuanImage-3.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md)
|
||||||
|
for its validated configuration. Helios supports single-request step execution only: use
|
||||||
|
`--step-execution --max-num-seqs 1` for Helios. Consult the selected pipeline's
|
||||||
|
documentation and source for the latest support status.
|
||||||
|
|
||||||
|
## Streaming Output
|
||||||
|
|
||||||
|
Use `--diffusion-streaming-output` for a pipeline that can produce intermediate
|
||||||
|
diffusion outputs:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
vllm serve BestWishYsh/Helios-Distilled --omni \
|
||||||
|
--port 8000 \
|
||||||
|
--diffusion-streaming-output
|
||||||
|
```
|
||||||
|
|
||||||
|
Streaming output requires step execution. If
|
||||||
|
`--diffusion-streaming-output` is set without `--step-execution`, the engine
|
||||||
|
enables step execution automatically. Model initialization fails if the
|
||||||
|
pipeline does not implement step execution.
|
||||||
|
|
||||||
|
Chunk-capable pipelines emit intermediate and final outputs through the same
|
||||||
|
request stream. Final-only step pipelines emit only the final result even when
|
||||||
|
streaming output is enabled.
|
||||||
|
|
||||||
|
The normal non-streaming serving path uses the same internal stream but drains
|
||||||
|
it and returns only the final output. Users do not need to enable
|
||||||
|
`--diffusion-streaming-output` to benefit from the unified engine lifecycle.
|
||||||
|
|
||||||
|
## Send Requests
|
||||||
|
|
||||||
|
Execution-mode flags configure the server. They do not change the
|
||||||
|
OpenAI-compatible client request format. For example, after starting any of
|
||||||
|
the non-streaming configurations above:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
curl -X POST http://localhost:8091/v1/images/generations \
|
||||||
|
-H "Content-Type: application/json" \
|
||||||
|
-d '{
|
||||||
|
"prompt": "a cat sitting on a windowsill",
|
||||||
|
"size": "1024x1024",
|
||||||
|
"num_inference_steps": 50,
|
||||||
|
"seed": 42
|
||||||
|
}' | jq -r '.data[0].b64_json' | base64 -d > cat.png
|
||||||
|
```
|
||||||
|
|
||||||
|
To create a request-level batch, send multiple independent requests
|
||||||
|
concurrently. Do not put several prompts into one `prompt` field:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
for prompt in "a red fox" "a blue bird" "a green frog"; do
|
||||||
|
curl -s -X POST http://localhost:8091/v1/images/generations \
|
||||||
|
-H "Content-Type: application/json" \
|
||||||
|
-d "{\"prompt\":\"${prompt}\",\"size\":\"1024x1024\",\"seed\":42}" \
|
||||||
|
> /dev/null &
|
||||||
|
done
|
||||||
|
wait
|
||||||
|
```
|
||||||
|
|
||||||
|
The scheduler may batch these requests when their sampling parameters are
|
||||||
|
compatible. See the
|
||||||
|
[Image Generation API](../../serving/image_generation_api.md) for response
|
||||||
|
formats and additional client examples. Streaming models can use
|
||||||
|
model-specific streaming endpoints documented by their serving guide.
|
||||||
|
|
||||||
|
## Python API
|
||||||
|
|
||||||
|
The Python arguments mirror the CLI flags:
|
||||||
|
|
||||||
|
```python
|
||||||
|
from vllm_omni import Omni
|
||||||
|
|
||||||
|
omni = Omni(
|
||||||
|
model="Qwen/Qwen-Image",
|
||||||
|
max_num_seqs=4,
|
||||||
|
request_batch_max_wait_ms=20.0,
|
||||||
|
)
|
||||||
|
|
||||||
|
outputs = omni.generate(
|
||||||
|
[
|
||||||
|
"a cup of coffee on a table",
|
||||||
|
"a toy dinosaur on a sandy beach",
|
||||||
|
"a fox waking up in bed and yawning",
|
||||||
|
]
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
Each list item passed to `Omni.generate()` becomes an independent logical
|
||||||
|
request that the scheduler may batch with compatible items.
|
||||||
|
|
||||||
|
For step execution:
|
||||||
|
|
||||||
|
```python
|
||||||
|
from vllm_omni import Omni
|
||||||
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||||
|
|
||||||
|
omni = Omni(
|
||||||
|
model="Qwen/Qwen-Image",
|
||||||
|
step_execution=True,
|
||||||
|
max_num_seqs=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
outputs = omni.generate(
|
||||||
|
"A cat sitting on a windowsill",
|
||||||
|
OmniDiffusionSamplingParams(num_inference_steps=50),
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
## Stage Configuration
|
||||||
|
|
||||||
|
Use the equivalent engine arguments in a deployment YAML:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
stage_args:
|
||||||
|
- stage_id: 0
|
||||||
|
stage_type: diffusion
|
||||||
|
engine_args:
|
||||||
|
step_execution: false
|
||||||
|
max_num_seqs: 4
|
||||||
|
request_batch_max_wait_ms: 20
|
||||||
|
```
|
||||||
|
|
||||||
|
For step execution, set `step_execution: true` and remove
|
||||||
|
`request_batch_max_wait_ms`.
|
||||||
|
|
||||||
|
## CLI Reference
|
||||||
|
|
||||||
|
| Flag | Default | Effect |
|
||||||
|
|---|---:|---|
|
||||||
|
| `--step-execution` | disabled | Select step-wise scheduling |
|
||||||
|
| `--max-num-seqs` | `1` for diffusion stages | Set request- or step-scheduler capacity |
|
||||||
|
| `--request-batch-max-wait-ms` | `0` | Wait for burst coalescing in request mode |
|
||||||
|
| `--diffusion-streaming-output` | disabled | Expose supported intermediate diffusion outputs and require step execution |
|
||||||
|
|
||||||
|
## Limitations and Troubleshooting
|
||||||
|
|
||||||
|
- Requests batch only when their compatibility-sensitive parameters match.
|
||||||
|
- Different LoRA adapters or scales run in separate batches.
|
||||||
|
- FIFO scheduling can cause an incompatible request to block later compatible
|
||||||
|
requests.
|
||||||
|
- All diffusion cache backends are unsupported in step mode. KV transfer is
|
||||||
|
supported for newly admitted step requests; some other request-mode extras
|
||||||
|
remain unsupported.
|
||||||
|
- If request-mode startup reports that the pipeline does not support batching,
|
||||||
|
use `--max-num-seqs 1`.
|
||||||
|
- If step-mode startup mentions `prepare_encode()`, `denoise_step()`,
|
||||||
|
`step_scheduler()`, or `post_decode()`, the pipeline does not implement the
|
||||||
|
required step contract.
|
||||||
|
|
||||||
|
For implementation details and model-author guidance, see
|
||||||
|
[Diffusion Continuous Batching](../../design/feature/diffusion_continuous_batching.md).
|
||||||
@@ -1,119 +0,0 @@
|
|||||||
# Request-Level Batching
|
|
||||||
|
|
||||||
Request-level batching lets diffusion serving combine multiple compatible
|
|
||||||
logical requests into one pipeline `forward()` call. Each prompt remains a
|
|
||||||
separate request with its own `request_id`, sampling parameters, seed, output,
|
|
||||||
error, and abort state. The scheduler decides which requests can run together.
|
|
||||||
|
|
||||||
!!! warning "Prompt List Semantics"
|
|
||||||
|
|
||||||
Diffusion request-level batching does not support a top-level packed
|
|
||||||
list-prompt request. Submit multiple prompts as independent requests and let
|
|
||||||
the scheduler batch compatible in-flight requests. Multimodal payloads stay
|
|
||||||
inside a single prompt dict, for example
|
|
||||||
`{"prompt": "...", "multi_modal_data": {"image": image}}`.
|
|
||||||
|
|
||||||
## Enablement
|
|
||||||
|
|
||||||
Request-level execution is the default diffusion request mode when
|
|
||||||
`step_execution=False`. `max_num_seqs` controls the maximum number of compatible
|
|
||||||
requests in one scheduler wave:
|
|
||||||
|
|
||||||
- `max_num_seqs=1` is the serial request path.
|
|
||||||
- `max_num_seqs>1` allows fused request-level batching for pipelines that
|
|
||||||
declare request-batch support.
|
|
||||||
|
|
||||||
Increase `max_num_seqs` above `1` to allow the request scheduler to keep more
|
|
||||||
than one compatible request active:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve Qwen/Qwen-Image --omni \
|
|
||||||
--port 8091 \
|
|
||||||
--max-num-seqs 4
|
|
||||||
```
|
|
||||||
|
|
||||||
For bursty online traffic, you can also set a small admission wait window. This
|
|
||||||
lets the engine wait briefly before the first `schedule()` of a new wave so
|
|
||||||
nearby compatible requests can arrive and share the same fused forward pass:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve Qwen/Qwen-Image --omni \
|
|
||||||
--port 8091 \
|
|
||||||
--max-num-seqs 4 \
|
|
||||||
--request-batch-max-wait-ms 20
|
|
||||||
```
|
|
||||||
|
|
||||||
`--request-batch-max-wait-ms 0` is the default and disables admission waiting,
|
|
||||||
so there is no added wait latency.
|
|
||||||
|
|
||||||
For deploy YAMLs, configure the diffusion stage engine args:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
stage_args:
|
|
||||||
- stage_id: 0
|
|
||||||
stage_type: diffusion
|
|
||||||
engine_args:
|
|
||||||
max_num_seqs: 4
|
|
||||||
request_batch_max_wait_ms: 20
|
|
||||||
```
|
|
||||||
|
|
||||||
## Compatibility
|
|
||||||
|
|
||||||
Only pipelines that declare request-batch support can use the fused
|
|
||||||
multi-request `pipeline.forward(batch)` path. Pipelines that do not support this
|
|
||||||
contract can still use request-level execution with `max_num_seqs=1`, which is
|
|
||||||
the serial request path. Setting `max_num_seqs>1` for a pipeline without
|
|
||||||
request-batch support fails early during engine initialization.
|
|
||||||
|
|
||||||
The scheduler batches only compatible requests. Compatibility is based on
|
|
||||||
shape-sensitive and guidance-sensitive sampling fields, including resolution,
|
|
||||||
frame count, CFG settings, output count, and LoRA identity. Requests with
|
|
||||||
different LoRA adapters or scales are kept in separate batches so the worker
|
|
||||||
activates one adapter per batch.
|
|
||||||
|
|
||||||
Request-level batching applies when `step_execution=False`. For the separate
|
|
||||||
step-wise runtime, see [Step Execution](step_execution.md).
|
|
||||||
|
|
||||||
## Tuning
|
|
||||||
|
|
||||||
- `max_num_seqs` caps the number of active compatible requests in one scheduler
|
|
||||||
wave. `1` means serial request execution; values above `1` enable fused
|
|
||||||
request-level batching when the pipeline supports it.
|
|
||||||
- `request_batch_max_wait_ms` is an upper bound on extra admission wait before a
|
|
||||||
new wave starts. Keep it small for latency-sensitive serving; values such as
|
|
||||||
`10` to `50` ms are a practical starting range for bursty HTTP ingress.
|
|
||||||
- `0` disables admission waiting and preserves the lowest first-request latency.
|
|
||||||
- FIFO ordering is conservative: an incompatible request at the front of the
|
|
||||||
waiting queue can block later compatible requests from joining the current
|
|
||||||
batch.
|
|
||||||
|
|
||||||
## Python API
|
|
||||||
|
|
||||||
When constructing `Omni`, pass the same engine arguments:
|
|
||||||
|
|
||||||
```python
|
|
||||||
from vllm_omni.entrypoints.omni import Omni
|
|
||||||
|
|
||||||
omni = Omni(
|
|
||||||
model="Qwen/Qwen-Image",
|
|
||||||
max_num_seqs=4,
|
|
||||||
request_batch_max_wait_ms=20.0,
|
|
||||||
)
|
|
||||||
|
|
||||||
outputs = omni.generate(
|
|
||||||
[
|
|
||||||
"a cup of coffee on a table",
|
|
||||||
"a toy dinosaur on a sandy beach",
|
|
||||||
"a fox waking up in bed and yawning",
|
|
||||||
]
|
|
||||||
)
|
|
||||||
```
|
|
||||||
|
|
||||||
`Omni.generate([...])` submits each list item as its own logical diffusion
|
|
||||||
request. The runtime may batch those requests internally when their sampling
|
|
||||||
parameters are compatible.
|
|
||||||
|
|
||||||
## For Contributors
|
|
||||||
|
|
||||||
For implementation details and model-author guidance, see
|
|
||||||
[Request-Level Batching for Diffusion](../../design/feature/diffusion_request_level_batching.md).
|
|
||||||
@@ -1,100 +0,0 @@
|
|||||||
# Step Execution
|
|
||||||
|
|
||||||
Step execution is an opt-in diffusion execution mode enabled with
|
|
||||||
`step_execution=True` when constructing `Omni`.
|
|
||||||
|
|
||||||
It is not a generic diffusion toggle for every pipeline. Only pipelines that
|
|
||||||
implement the stepwise contract support it today.
|
|
||||||
|
|
||||||
## Quick Start
|
|
||||||
|
|
||||||
### Python API
|
|
||||||
|
|
||||||
```python
|
|
||||||
from vllm_omni import Omni
|
|
||||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
||||||
|
|
||||||
omni = Omni(
|
|
||||||
model="Qwen/Qwen-Image",
|
|
||||||
step_execution=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
outputs = omni.generate(
|
|
||||||
"A cat sitting on a windowsill",
|
|
||||||
OmniDiffusionSamplingParams(
|
|
||||||
num_inference_steps=50,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
```
|
|
||||||
|
|
||||||
### Serving
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve Qwen/Qwen-Image --omni \
|
|
||||||
--port 8091 \
|
|
||||||
--step-execution \
|
|
||||||
--max-num-seqs 8
|
|
||||||
```
|
|
||||||
|
|
||||||
For serving, `--step-execution` enables the step-wise runtime. Continuous
|
|
||||||
batching only becomes relevant when `--max-num-seqs > 1`.
|
|
||||||
|
|
||||||
`max_num_seqs` controls the step scheduler capacity:
|
|
||||||
|
|
||||||
- `max_num_seqs=1` runs one request through the step-wise state machine.
|
|
||||||
- `max_num_seqs>1` allows compatible active requests to share step waves.
|
|
||||||
|
|
||||||
## Supported Pipelines
|
|
||||||
|
|
||||||
| Pipeline | Example models | Step execution |
|
|
||||||
|----------|----------------|----------------|
|
|
||||||
| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes |
|
|
||||||
| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No |
|
|
||||||
|
|
||||||
!!! warning "Experimental continuous batching"
|
|
||||||
When `--step-execution` is enabled and `max_num_seqs > 1` is configured,
|
|
||||||
the step-wise path can batch
|
|
||||||
compatible requests together. This is experimental. Requests with
|
|
||||||
incompatible sampling parameters are intentionally kept in separate batches,
|
|
||||||
and `max_num_seqs=1` remains the conservative single-request step path.
|
|
||||||
|
|
||||||
## Current Limitations
|
|
||||||
|
|
||||||
- Continuous batching under `step_execution` is experimental and only batches
|
|
||||||
compatible requests.
|
|
||||||
- `cache_backend` is not supported together with step execution.
|
|
||||||
- Unsupported pipelines fail early during model loading.
|
|
||||||
- Request-mode extras such as KV transfer are not wired into step mode yet.
|
|
||||||
- LoRA is supported in step mode, but each batch must use a single adapter:
|
|
||||||
requests with different `lora_request` or `lora_scale` are scheduled into
|
|
||||||
separate batches.
|
|
||||||
|
|
||||||
## When To Use It
|
|
||||||
|
|
||||||
Use step execution only when you specifically need the pipeline to run through
|
|
||||||
its stepwise request state machine. For normal diffusion inference, leave it
|
|
||||||
disabled unless your workflow depends on this mode.
|
|
||||||
|
|
||||||
For Qwen-Image online serving, the usual progression is:
|
|
||||||
|
|
||||||
- start with `--step-execution --max-num-seqs 1` if you only need the step-wise path
|
|
||||||
- increase `--max-num-seqs` after that if you want the experimental compatible-request batching behavior
|
|
||||||
|
|
||||||
If you are looking for general diffusion speedups, see
|
|
||||||
[Diffusion Features Overview](../diffusion_features.md).
|
|
||||||
|
|
||||||
## Troubleshooting
|
|
||||||
|
|
||||||
If model loading fails with a message mentioning `prepare_encode()`,
|
|
||||||
`denoise_step()`, `step_scheduler()`, and `post_decode()`, the selected
|
|
||||||
pipeline does not support step execution.
|
|
||||||
|
|
||||||
## For Model Authors
|
|
||||||
|
|
||||||
If you want to add step execution support to a new diffusion pipeline, see the
|
|
||||||
implementation guide:
|
|
||||||
[Diffusion Step Execution Design](../../design/feature/diffusion_step_execution.md).
|
|
||||||
|
|
||||||
If you also want that pipeline to participate in the experimental batched
|
|
||||||
step-wise path, see:
|
|
||||||
[Continuous Batching for Step-Wise Diffusion](../../design/feature/diffusion_continuous_batching.md).
|
|
||||||
@@ -59,7 +59,7 @@ Memory optimization methods help reduce GPU memory usage, enabling inference on
|
|||||||
|
|
||||||
| Method | Description | Best For |
|
| Method | Description | Best For |
|
||||||
|--------|-------------|----------|
|
|--------|-------------|----------|
|
||||||
| **[CPU Offload](diffusion/cpu_offload_diffusion.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs |
|
| **[CPU Offload](diffusion/cpu_offload.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs |
|
||||||
| **[Quantization](quantization/overview.md)** | Reduces transformer stages from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss |
|
| **[Quantization](quantization/overview.md)** | Reduces transformer stages from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss |
|
||||||
| **[VAE Parallelism](diffusion/parallelism/vae_parallelism.md)** | Distributes VAE decode work across GPUs | High-resolution generation with reduced VAE memory peak |
|
| **[VAE Parallelism](diffusion/parallelism/vae_parallelism.md)** | Distributes VAE decode work across GPUs | High-resolution generation with reduced VAE memory peak |
|
||||||
|
|
||||||
@@ -80,12 +80,11 @@ denoise steps.
|
|||||||
|
|
||||||
| Method | Description | Best For |
|
| Method | Description | Best For |
|
||||||
|--------|-------------|----------|
|
|--------|-------------|----------|
|
||||||
| **[Request-Level Batching](diffusion/request_batching.md)** | Scheduler batches compatible independent diffusion requests into one pipeline forward pass | Bursty online serving and multi-request throughput |
|
| **[Diffusion Execution Modes](diffusion/execution_modes.md)** | Configures serial requests, request batching, step execution, continuous batching, and streaming output | Matching latency, throughput, cancellation, and output-delivery requirements |
|
||||||
| **[Step Execution](diffusion/step_execution.md)** | Per-step denoise execution with mid-request abort support | Request cancellation between denoise steps, fine-grained execution control |
|
|
||||||
|
|
||||||
**Note:** Request-level batching is available for pipelines that declare the
|
**Note:** Request-level batching and step execution are capability-based.
|
||||||
request-batch forward contract. Step execution is currently supported by
|
Consult the execution guide and the selected pipeline's documentation for
|
||||||
QwenImagePipeline only. See [Supported Models](#supported-models) for details.
|
current support.
|
||||||
|
|
||||||
### Quantization Methods
|
### Quantization Methods
|
||||||
|
|
||||||
@@ -105,8 +104,8 @@ The following tables show which models support each feature:
|
|||||||
|
|
||||||
> Notes:
|
> Notes:
|
||||||
|
|
||||||
> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload_diffusion.md).
|
> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload.md).
|
||||||
> 2. The **💾Quantization** column is collapsed for readability. See [Quantization Overview](quantization/overview.md) for per-method and per-model support details.
|
> 2. The **💾Quantization** column is collapsed for readability. See [Quantization](quantization/overview.md) for per-method and per-model support details.
|
||||||
|
|
||||||
### ImageGen
|
### ImageGen
|
||||||
|
|
||||||
@@ -160,7 +159,7 @@ The following tables show which models support each feature:
|
|||||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||||
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||||
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||||
|
| **MiniMax-H3** | ❌ | ✅ | ✅ | ❌ | ✅ (DiT/TE) | ❌ | ✅ | ✅ | ✅ (tile) | ✅ (DiT) | ❌ |
|
||||||
|
|
||||||
**Frame Interpolation Support**
|
**Frame Interpolation Support**
|
||||||
|
|
||||||
@@ -206,7 +205,7 @@ The following tables show which models support each feature:
|
|||||||
3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible.
|
3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible.
|
||||||
4. CPU Offloading (Layerwise) supports single-card for now.
|
4. CPU Offloading (Layerwise) supports single-card for now.
|
||||||
5. Using FP8-Quant as an example of qunatization methods.
|
5. Using FP8-Quant as an example of qunatization methods.
|
||||||
6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT). LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches).
|
6. Step Execution is not compatible with any diffusion cache backend. LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches).
|
||||||
|
|
||||||
|
|
||||||
## Multi-Thread Weight Loading
|
## Multi-Thread Weight Loading
|
||||||
@@ -275,9 +274,9 @@ Measured on NVIDIA H800:
|
|||||||
|
|
||||||
**Memory Optimization:**
|
**Memory Optimization:**
|
||||||
|
|
||||||
- **[CPU Offload Guide](diffusion/cpu_offload_diffusion.md)** - Offload model components to CPU, reduce GPU memory usage
|
- **[CPU Offload Guide](diffusion/cpu_offload.md)** - Offload model components to CPU, reduce GPU memory usage
|
||||||
- **[VAE Parallelism Guide](diffusion/parallelism/vae_parallelism.md)** - Distribute VAE decode work across GPUs for high-resolution images and videos
|
- **[VAE Parallelism Guide](diffusion/parallelism/vae_parallelism.md)** - Distribute VAE decode work across GPUs for high-resolution images and videos
|
||||||
- **[Quantization Overview](quantization/overview.md)** - Overview of quantization methods for diffusion, multi-stage omni/TTS, and multi-stage diffusion models
|
- **[Quantization](quantization/overview.md)** - Quantization methods for diffusion, multi-stage omni/TTS, and multi-stage diffusion models
|
||||||
|
|
||||||
**Extensions:**
|
**Extensions:**
|
||||||
|
|
||||||
@@ -286,7 +285,7 @@ Measured on NVIDIA H800:
|
|||||||
|
|
||||||
**Execution Modes:**
|
**Execution Modes:**
|
||||||
|
|
||||||
- **[Step Execution Guide](diffusion/step_execution.md)** - Per-step denoise execution with mid-request abort support
|
- **[Diffusion Execution Modes](diffusion/execution_modes.md)** - Configure request batching, step execution, continuous batching, and streaming output
|
||||||
|
|
||||||
**Startup Optimization:**
|
**Startup Optimization:**
|
||||||
|
|
||||||
|
|||||||
@@ -184,7 +184,7 @@ if __name__ == "__main__":
|
|||||||
!!! info
|
!!! info
|
||||||
|
|
||||||
For diffusion request-level batching controls such as `max_num_seqs`, see
|
For diffusion request-level batching controls such as `max_num_seqs`, see
|
||||||
[Request-Level Batching](../../diffusion/request_batching.md).
|
[Diffusion Execution Modes](../../diffusion/execution_modes.md).
|
||||||
|
|
||||||
### Negative Prompts
|
### Negative Prompts
|
||||||
|
|
||||||
|
|||||||
@@ -5,6 +5,9 @@ path. The same configuration entrypoint is used across diffusion-only models,
|
|||||||
multi-stage omni/TTS models, and multi-stage diffusion models, but each model
|
multi-stage omni/TTS models, and multi-stage diffusion models, but each model
|
||||||
type has a different quantization scope.
|
type has a different quantization scope.
|
||||||
|
|
||||||
|
For the internal architecture and backend extension points, see the
|
||||||
|
[quantization design document](../../design/feature/quantization.md).
|
||||||
|
|
||||||
## Quantization Modes
|
## Quantization Modes
|
||||||
|
|
||||||
| Mode | Guide | Description | Methods |
|
| Mode | Guide | Description | Methods |
|
||||||
|
|||||||
@@ -0,0 +1,77 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
"""Offline PersonaPlex inference through the vllm-omni fullduplex backend.
|
||||||
|
|
||||||
|
This is the runnable end-to-end proof for the integration: it drives a user WAV
|
||||||
|
through :class:`PersonaPlexEngine` + :class:`PersonaPlexSession` (our wrappers
|
||||||
|
around the Moshi/Mimi lockstep loop) and writes the agent's audio + inner-monologue
|
||||||
|
text. It mirrors ``python -m moshi.offline`` but goes through our ``FrameStepper``
|
||||||
|
seam, so a green run validates that the session/adapter plumbing is faithful.
|
||||||
|
|
||||||
|
Prereqs: HF access to ``nvidia/personaplex-7b-v1`` (``export HF_TOKEN=...``). The
|
||||||
|
stack is moshi-free (native temporal + depformer + streaming Mimi via
|
||||||
|
``transformers``); decoding is greedy.
|
||||||
|
|
||||||
|
Example:
|
||||||
|
python examples/offline_inference/personaplex/personaplex_offline.py \\
|
||||||
|
--input-wav input_assistant.wav --output-wav out.wav \\
|
||||||
|
--voice-prompt NATF2.pt --persona "You are a wise and friendly teacher."
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import json
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex import (
|
||||||
|
PersonaPlexConfig,
|
||||||
|
PersonaPlexEngine,
|
||||||
|
PersonaPlexSession,
|
||||||
|
)
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.config import DEFAULT_PERSONA
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
parser = argparse.ArgumentParser(description=__doc__)
|
||||||
|
parser.add_argument("--input-wav", required=True, help="User-side input WAV")
|
||||||
|
parser.add_argument("--output-wav", required=True, help="Agent audio output WAV")
|
||||||
|
parser.add_argument("--output-text", default=None, help="Optional JSON for agent text")
|
||||||
|
parser.add_argument("--voice-prompt", default="NATF2.pt", help="Voice clone (bundled .pt or .wav)")
|
||||||
|
parser.add_argument("--persona", default=DEFAULT_PERSONA, help="System role text")
|
||||||
|
parser.add_argument("--hf-repo", default="nvidia/personaplex-7b-v1")
|
||||||
|
parser.add_argument("--device", default="cuda")
|
||||||
|
parser.add_argument("--cpu-offload", action="store_true")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
import sphn # local import: only the runnable driver needs the audio I/O dep
|
||||||
|
|
||||||
|
config = PersonaPlexConfig(
|
||||||
|
hf_repo=args.hf_repo,
|
||||||
|
voice_prompt=args.voice_prompt,
|
||||||
|
persona=args.persona,
|
||||||
|
device=args.device,
|
||||||
|
cpu_offload=args.cpu_offload,
|
||||||
|
)
|
||||||
|
|
||||||
|
engine = PersonaPlexEngine(config).load()
|
||||||
|
session = PersonaPlexSession(engine, config)
|
||||||
|
session.open()
|
||||||
|
|
||||||
|
# Load exactly like the reference (moshi.lm.load_audio): native read then
|
||||||
|
# explicit resample, so greedy runs stay bit-faithful to `python -m moshi.offline`.
|
||||||
|
user_pcm, src_sr = sphn.read(args.input_wav)
|
||||||
|
user_pcm = sphn.resample(user_pcm, src_sample_rate=src_sr, dst_sample_rate=config.sample_rate)
|
||||||
|
audio, text = session.run(np.asarray(user_pcm[0], dtype=np.float32))
|
||||||
|
|
||||||
|
sphn.write_wav(args.output_wav, audio, config.sample_rate)
|
||||||
|
print(f"wrote {audio.shape[0] / config.sample_rate:.2f}s of agent audio to {args.output_wav}")
|
||||||
|
print(f"agent text: {text!r}")
|
||||||
|
if args.output_text:
|
||||||
|
with open(args.output_text, "w") as f:
|
||||||
|
json.dump({"text": text}, f, ensure_ascii=False)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -260,7 +260,7 @@ if __name__ == "__main__":
|
|||||||
For diffusion pipelines, the input list is sliced into single-item requests
|
For diffusion pipelines, the input list is sliced into single-item requests
|
||||||
before feeding into the diffusion pipeline. For request-level batching
|
before feeding into the diffusion pipeline. For request-level batching
|
||||||
controls such as `max_num_seqs`, see
|
controls such as `max_num_seqs`, see
|
||||||
[Request-Level Batching](../../../docs/user_guide/diffusion/request_batching.md).
|
[Diffusion Execution Modes](../../../docs/user_guide/diffusion/execution_modes.md).
|
||||||
|
|
||||||
### Negative Prompts
|
### Negative Prompts
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,84 @@
|
|||||||
|
# PersonaPlex full-duplex online serving
|
||||||
|
|
||||||
|
Serve [`nvidia/personaplex-7b-v1`](https://huggingface.co/nvidia/personaplex-7b-v1)
|
||||||
|
(a Moshi-based full-duplex speech-to-speech model) with the native vLLM-Omni engine.
|
||||||
|
|
||||||
|
The server hosts the **official PersonaPlex web client** at `/` (auto-downloaded from
|
||||||
|
the model repo) and implements its WebSocket protocol at `/api/chat`, so you talk to the
|
||||||
|
model live in the browser — mic in, agent speech + inner-monologue text out.
|
||||||
|
|
||||||
|
> Experimental. Requires a GPU and Hugging Face access to the gated repo
|
||||||
|
> (`HF_TOKEN` with access to `nvidia/personaplex-7b-v1`).
|
||||||
|
|
||||||
|
## Start the server
|
||||||
|
|
||||||
|
```bash
|
||||||
|
HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 bash run_server.sh
|
||||||
|
# or directly:
|
||||||
|
HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 python -m \
|
||||||
|
vllm_omni.experimental.fullduplex.personaplex.serving.server --port 8124
|
||||||
|
```
|
||||||
|
|
||||||
|
## Talk to it
|
||||||
|
|
||||||
|
- **Browser (recommended):** open `http://localhost:8124/` and allow the microphone.
|
||||||
|
Use headphones so the agent does not hear itself.
|
||||||
|
- **Headless:** stream a 24 kHz mono WAV and save the reply:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python duplex_client.py --url ws://localhost:8124/api/chat --input user.wav --out reply.wav
|
||||||
|
```
|
||||||
|
|
||||||
|
## Unified Realtime API
|
||||||
|
|
||||||
|
The default `vllm_omni/deploy/personaplex.yaml` also enables main's engine-owned
|
||||||
|
full-duplex control plane:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
|
||||||
|
/path/to/personaplex-7b-v1 \
|
||||||
|
--omni \
|
||||||
|
--stage-configs-path vllm_omni/deploy/personaplex.yaml
|
||||||
|
```
|
||||||
|
|
||||||
|
Validate the actual `/v1/realtime?duplex=1` scheduler path with paced 24 kHz
|
||||||
|
PCM, two concurrent sessions, overflow admission, per-session slot recycling,
|
||||||
|
and non-silent output:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python tests/e2e/online_serving/personaplex_realtime_duplex.py \
|
||||||
|
--model /path/to/personaplex-7b-v1 \
|
||||||
|
--input-wav /path/to/speech.wav \
|
||||||
|
--output-dir /tmp/personaplex-realtime-duplex
|
||||||
|
```
|
||||||
|
|
||||||
|
This is distinct from the standalone `/api/chat` compatibility server above.
|
||||||
|
The unified endpoint advertises `supports_barge_in=false`: overlapping speech
|
||||||
|
is native model behavior, but destructive output interruption and model-state
|
||||||
|
rewind have not been validated for PersonaPlex.
|
||||||
|
|
||||||
|
## Protocol (`/api/chat`)
|
||||||
|
|
||||||
|
Binary WebSocket messages, first byte is a tag (identical to `moshi.server`):
|
||||||
|
|
||||||
|
| Direction | Tag | Payload |
|
||||||
|
| --- | --- | --- |
|
||||||
|
| server → client | `\x00` | ready handshake (after the system-prompt prefill) |
|
||||||
|
| client → server | `\x01` | Opus-encoded mic audio (24 kHz) |
|
||||||
|
| server → client | `\x01` | Opus-encoded agent audio |
|
||||||
|
| server → client | `\x02` | UTF-8 inner-monologue text |
|
||||||
|
|
||||||
|
Connect with query params `text_prompt` (persona) and `voice_prompt` (voice file).
|
||||||
|
A raw-PCM endpoint (`/v1/audio/duplex`, JSON `open` + float32 frames) is also available
|
||||||
|
for clients that do not want an Opus dependency.
|
||||||
|
|
||||||
|
## Notes
|
||||||
|
|
||||||
|
- **Run the client near the server.** Real-time 80 ms audio is sensitive to network
|
||||||
|
latency/jitter; over a high-latency remote link playback can stutter regardless of
|
||||||
|
engine speed. On localhost it is smooth.
|
||||||
|
- Greedy decoding by default. One conversation per server by default (a second
|
||||||
|
concurrent connection is rejected with WS close 1013 until the first ends);
|
||||||
|
pass `--batch-size N` to host N concurrent conversations on one engine
|
||||||
|
(elastic batching; connections beyond N are rejected until a slot frees).
|
||||||
|
- Full runbook: [`recipes/NVIDIA/PersonaPlex.md`](../../../recipes/NVIDIA/PersonaPlex.md).
|
||||||
@@ -0,0 +1,114 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
"""Headless client for the PersonaPlex duplex server (official Moshi protocol).
|
||||||
|
|
||||||
|
Streams a 24 kHz mono WAV to ``/api/chat`` exactly as the browser client does
|
||||||
|
(Opus over WebSocket, tags ``\\x00`` ready / ``\\x01`` audio / ``\\x02`` text),
|
||||||
|
collects the agent's reply audio + inner-monologue text, and writes a WAV.
|
||||||
|
|
||||||
|
The browser UI at ``http://<host>:<port>/`` is the primary demo; this script is
|
||||||
|
for headless testing / scripting.
|
||||||
|
|
||||||
|
python duplex_client.py --url ws://localhost:8124/api/chat \
|
||||||
|
--input user.wav --out reply.wav
|
||||||
|
|
||||||
|
Deps: ``websockets``, ``sphn``, ``soundfile``, ``numpy``.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import soundfile as sf
|
||||||
|
import sphn
|
||||||
|
|
||||||
|
SR = 24000 # PersonaPlex / Mimi codec rate
|
||||||
|
|
||||||
|
|
||||||
|
async def converse(url: str, pcm: np.ndarray) -> tuple[np.ndarray, str]:
|
||||||
|
import websockets
|
||||||
|
|
||||||
|
reader = sphn.OpusStreamReader(SR) # decode agent audio
|
||||||
|
writer = sphn.OpusStreamWriter(SR) # encode our mic audio
|
||||||
|
text: list[str] = []
|
||||||
|
ready = asyncio.Event()
|
||||||
|
last = {"t": 0.0}
|
||||||
|
|
||||||
|
async with websockets.connect(url, max_size=None) as ws:
|
||||||
|
|
||||||
|
async def receive() -> None:
|
||||||
|
loop = asyncio.get_event_loop()
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
msg = await ws.recv()
|
||||||
|
except Exception:
|
||||||
|
return
|
||||||
|
if not isinstance(msg, (bytes, bytearray)) or not msg:
|
||||||
|
continue
|
||||||
|
tag, payload = msg[0], msg[1:]
|
||||||
|
if tag == 0:
|
||||||
|
ready.set()
|
||||||
|
elif tag == 1:
|
||||||
|
reader.append_bytes(bytes(payload))
|
||||||
|
last["t"] = loop.time()
|
||||||
|
elif tag == 2:
|
||||||
|
text.append(payload.decode("utf8", errors="replace"))
|
||||||
|
|
||||||
|
recv_task = asyncio.create_task(receive())
|
||||||
|
await asyncio.wait_for(ready.wait(), timeout=120) # waits for system-prompt prefill
|
||||||
|
|
||||||
|
frame = 1920 # 80 ms
|
||||||
|
for i in range(0, len(pcm), frame):
|
||||||
|
chunk = pcm[i : i + frame]
|
||||||
|
if len(chunk) < frame:
|
||||||
|
chunk = np.concatenate([chunk, np.zeros(frame - len(chunk), dtype=np.float32)])
|
||||||
|
writer.append_pcm(chunk)
|
||||||
|
data = writer.read_bytes()
|
||||||
|
if data:
|
||||||
|
await ws.send(b"\x01" + data)
|
||||||
|
await asyncio.sleep(frame / SR) # pace at real time
|
||||||
|
|
||||||
|
loop = asyncio.get_event_loop()
|
||||||
|
last["t"] = loop.time()
|
||||||
|
while loop.time() - last["t"] < 2.5: # drain until the agent stops talking
|
||||||
|
await asyncio.sleep(0.1)
|
||||||
|
recv_task.cancel()
|
||||||
|
|
||||||
|
chunks: list[np.ndarray] = []
|
||||||
|
while True:
|
||||||
|
out = reader.read_pcm()
|
||||||
|
if out is None or out.shape[-1] == 0:
|
||||||
|
break
|
||||||
|
chunks.append(np.asarray(out, dtype=np.float32))
|
||||||
|
audio = np.concatenate(chunks) if chunks else np.zeros(0, dtype=np.float32)
|
||||||
|
return audio, "".join(text)
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
ap = argparse.ArgumentParser()
|
||||||
|
ap.add_argument("--url", default="ws://localhost:8124/api/chat")
|
||||||
|
ap.add_argument("--input", required=True, help="24 kHz mono WAV of the user turn")
|
||||||
|
ap.add_argument("--out", default="personaplex_reply.wav")
|
||||||
|
ap.add_argument("--tail", type=float, default=10.0, help="trailing silence (s) for the agent reply")
|
||||||
|
ap.add_argument("--persona", default="")
|
||||||
|
ap.add_argument("--voice", default="NATF2.pt")
|
||||||
|
args = ap.parse_args()
|
||||||
|
|
||||||
|
wav, sr = sf.read(args.input, dtype="float32")
|
||||||
|
if wav.ndim > 1:
|
||||||
|
wav = wav.mean(axis=1)
|
||||||
|
if sr != SR:
|
||||||
|
raise SystemExit(f"input must be {SR} Hz mono; got {sr} Hz. Resample it first.")
|
||||||
|
wav = np.concatenate([wav, np.zeros(int(args.tail * SR), dtype=np.float32)])
|
||||||
|
|
||||||
|
url = f"{args.url}?text_prompt={args.persona}&voice_prompt={args.voice}"
|
||||||
|
audio, text = asyncio.run(converse(url, wav))
|
||||||
|
sf.write(args.out, audio if audio.size else np.zeros(SR, dtype=np.float32), SR)
|
||||||
|
print(f"agent reply: {audio.size / SR:.2f}s -> {args.out}")
|
||||||
|
print(f"inner-monologue: {text.strip()!r}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
+15
@@ -0,0 +1,15 @@
|
|||||||
|
#!/bin/bash
|
||||||
|
# PersonaPlex full-duplex online serving.
|
||||||
|
#
|
||||||
|
# Serves the OFFICIAL PersonaPlex web client at / (auto-downloaded from
|
||||||
|
# nvidia/personaplex-7b-v1) and the Moshi WebSocket protocol at /api/chat,
|
||||||
|
# backed by the native vLLM-Omni engine. Requires HF access to the gated repo.
|
||||||
|
|
||||||
|
HOST="${HOST:-0.0.0.0}"
|
||||||
|
PORT="${PORT:-8124}"
|
||||||
|
|
||||||
|
echo "Starting PersonaPlex duplex server on ${HOST}:${PORT}..."
|
||||||
|
echo "Open http://localhost:${PORT}/ for the official web client (run it near the server)."
|
||||||
|
|
||||||
|
python -m vllm_omni.experimental.fullduplex.personaplex.serving.server \
|
||||||
|
--host "$HOST" --port "$PORT"
|
||||||
@@ -71,6 +71,14 @@ plugins:
|
|||||||
- search
|
- search
|
||||||
- autorefs
|
- autorefs
|
||||||
- awesome-nav
|
- awesome-nav
|
||||||
|
- redirects:
|
||||||
|
redirect_maps:
|
||||||
|
user_guide/diffusion/request_batching.md: user_guide/diffusion/execution_modes.md
|
||||||
|
user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md
|
||||||
|
user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md
|
||||||
|
user_guide/quantization/index.md: user_guide/quantization/overview.md
|
||||||
|
design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md
|
||||||
|
design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md
|
||||||
- glightbox
|
- glightbox
|
||||||
- git-revision-date-localized:
|
- git-revision-date-localized:
|
||||||
# exclude files
|
# exclude files
|
||||||
|
|||||||
+2
-2
@@ -70,8 +70,7 @@ dev = [
|
|||||||
"FlagEmbedding==1.4.0", # for hunyuan_image3 tests
|
"FlagEmbedding==1.4.0", # for hunyuan_image3 tests
|
||||||
"kernels==0.14.1", # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
|
"kernels==0.14.1", # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
|
||||||
"s3tokenizer==0.3.0", # for tts test
|
"s3tokenizer==0.3.0", # for tts test
|
||||||
"step-audio2==1.0.0", # Step-Audio2 Token2Wav (flashcosyvoice / hyperpyyaml)
|
"step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml)
|
||||||
"stepaudio2-minicpmo", # for MiniCPM-o 4.5 talker vocoder (stepaudio2-minicpmo Token2wav)
|
|
||||||
]
|
]
|
||||||
|
|
||||||
demo = [
|
demo = [
|
||||||
@@ -109,6 +108,7 @@ docs = [
|
|||||||
"mkdocstrings-python",
|
"mkdocstrings-python",
|
||||||
"mkdocs-gen-files",
|
"mkdocs-gen-files",
|
||||||
"mkdocs-awesome-nav",
|
"mkdocs-awesome-nav",
|
||||||
|
"mkdocs-redirects",
|
||||||
"mkdocs-glightbox",
|
"mkdocs-glightbox",
|
||||||
"mkdocs-git-revision-date-localized-plugin",
|
"mkdocs-git-revision-date-localized-plugin",
|
||||||
"mkdocs-minify-plugin",
|
"mkdocs-minify-plugin",
|
||||||
|
|||||||
@@ -332,6 +332,6 @@ bundled offline CLI do not currently expose `sigmas`.
|
|||||||
- <https://huggingface.co/Lightricks/LTX-2.3>
|
- <https://huggingface.co/Lightricks/LTX-2.3>
|
||||||
- <https://huggingface.co/diffusers/LTX-2.3-Diffusers>
|
- <https://huggingface.co/diffusers/LTX-2.3-Diffusers>
|
||||||
- [Online video generation](../../docs/user_guide/examples/online_serving/text_to_video.md)
|
- [Online video generation](../../docs/user_guide/examples/online_serving/text_to_video.md)
|
||||||
- [Request batching](../../docs/user_guide/diffusion/request_batching.md)
|
- [Diffusion execution modes](../../docs/user_guide/diffusion/execution_modes.md)
|
||||||
- [T2V offline example](../../examples/offline_inference/text_to_video/text_to_video.md)
|
- [T2V offline example](../../examples/offline_inference/text_to_video/text_to_video.md)
|
||||||
- [I2V offline example](../../examples/offline_inference/image_to_video/README.md)
|
- [I2V offline example](../../examples/offline_inference/image_to_video/README.md)
|
||||||
@@ -0,0 +1,61 @@
|
|||||||
|
# MiniMax-H3 on RTX 5090
|
||||||
|
|
||||||
|
This recipe uses BF16 weights, tiled VAE decode, tensor parallelism where a
|
||||||
|
second GPU is available, and distributed layerwise offload (DLO). It is a
|
||||||
|
memory-first serving configuration; lower resident counts reduce HBM use and
|
||||||
|
increase CPU-to-GPU transfer time.
|
||||||
|
|
||||||
|
## Capacity requirements
|
||||||
|
|
||||||
|
| Resource | One RTX 5090 | Two RTX 5090s |
|
||||||
|
| --- | ---: | ---: |
|
||||||
|
| GPU HBM | 32 GiB | 32 GiB per GPU |
|
||||||
|
| Checkpoint storage | 135 GiB per partition | 135 GiB per partition |
|
||||||
|
| Available system RAM | 200 GiB minimum | 200 GiB minimum |
|
||||||
|
| Recommended system RAM | 384 GiB | 384 GiB |
|
||||||
|
|
||||||
|
`FL2VA` and `Ref2VA` are separate 135 GiB checkpoint partitions. Start one
|
||||||
|
server at a time. DLO keeps rank-local weights in pinned host memory; increasing
|
||||||
|
`--dlo-resident-layers` improves latency but does **not** reduce host RAM in the
|
||||||
|
current implementation because resident layers retain pinned CPU master copies.
|
||||||
|
|
||||||
|
> **Modular H3:** after #5720 lands, preserve this recipe's one-partition
|
||||||
|
> behavior with `--task-type fl2va` or `--task-type ref2va`.
|
||||||
|
|
||||||
|
## One RTX 5090: 1344x768, 5 seconds
|
||||||
|
|
||||||
|
Use 12 resident DiT layers. A 50-step B300 allocation test with this exact
|
||||||
|
single-rank topology peaked at 26.50 GiB; re-measure peak HBM on the target
|
||||||
|
card before increasing the resident count.
|
||||||
|
|
||||||
|
```bash
|
||||||
|
CUDA_VISIBLE_DEVICES=0 vllm serve /path/to/MiniMax-H3/FL2VA \
|
||||||
|
--omni --trust-remote-code --host 0.0.0.0 --port 8000 \
|
||||||
|
--num-gpus 1 --tensor-parallel-size 1 --text-encoder-tp-size 1 \
|
||||||
|
--usp 1 --ring 1 --vae-patch-parallel-size 1 \
|
||||||
|
--vae-parallel-mode tile --vae-use-tiling \
|
||||||
|
--enable-distributed-layerwise-offload --dlo-no-use-allgather \
|
||||||
|
--dlo-resident-layers 12 --enforce-eager \
|
||||||
|
--diffusion-attention-backend CUDNN_ATTN
|
||||||
|
```
|
||||||
|
|
||||||
|
## Two RTX 5090s: 1344x768, 5 seconds
|
||||||
|
|
||||||
|
Use TP2 and 20 resident DiT layers. The two-rank B300 capacity run peaked at
|
||||||
|
27,726 MiB per rank for this shape and 50 steps. This is a memory/correctness
|
||||||
|
proxy, not a consumer-GPU latency claim.
|
||||||
|
|
||||||
|
```bash
|
||||||
|
CUDA_VISIBLE_DEVICES=0,1 vllm serve /path/to/MiniMax-H3/FL2VA \
|
||||||
|
--omni --trust-remote-code --host 0.0.0.0 --port 8000 \
|
||||||
|
--num-gpus 2 --tensor-parallel-size 2 --text-encoder-tp-size 2 \
|
||||||
|
--usp 1 --ring 1 --vae-patch-parallel-size 2 \
|
||||||
|
--vae-parallel-mode tile --vae-use-tiling \
|
||||||
|
--enable-distributed-layerwise-offload --dlo-no-use-allgather \
|
||||||
|
--dlo-resident-layers 20 --enforce-eager \
|
||||||
|
--diffusion-attention-backend CUDNN_ATTN
|
||||||
|
```
|
||||||
|
|
||||||
|
For Ref2VA, stop the FL2VA server and restart the same command with
|
||||||
|
`/path/to/MiniMax-H3/Ref2VA`. Ref2VA reference video count and prompt length
|
||||||
|
can increase activation memory; begin with one request at a time.
|
||||||
@@ -0,0 +1,147 @@
|
|||||||
|
# MiniMax H3 on Moore Threads MUSA
|
||||||
|
|
||||||
|
> Joint video and audio generation on Moore Threads GPUs
|
||||||
|
|
||||||
|
## Summary
|
||||||
|
|
||||||
|
- Vendor: MiniMaxAI
|
||||||
|
- Model: [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3)
|
||||||
|
- Tasks: T2VA, FL2VA, and Ref2VA
|
||||||
|
- Mode: OpenAI-compatible `/v1/videos` HTTP serving
|
||||||
|
- Hardware: 4x Moore Threads MTT S5000 for the validated Ref2VA profile
|
||||||
|
- Maintainer: Community
|
||||||
|
|
||||||
|
This recipe adapts [MiniMax-H3.md](MiniMax-H3.md) for MUSA environments.
|
||||||
|
MiniMax H3 has two independently served checkpoint partitions:
|
||||||
|
|
||||||
|
- `FL2VA` serves text-to-video+audio (`t2va`) and
|
||||||
|
first-frame-to-video+audio (`fl2va`).
|
||||||
|
- `Ref2VA` serves image/audio or video-reference generation (`ref2va`).
|
||||||
|
|
||||||
|
## Prerequisites
|
||||||
|
|
||||||
|
### Checkpoint
|
||||||
|
|
||||||
|
Only download the partition needed by the server. Downloading the whole model
|
||||||
|
repository also retrieves duplicate shared components and the original
|
||||||
|
top-level pipeline layout.
|
||||||
|
|
||||||
|
For T2VA and FL2VA:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python -m pip install modelscope
|
||||||
|
export MODEL_ROOT=/path/to/MiniMax-H3
|
||||||
|
modelscope download MiniMax/MiniMax-H3 \
|
||||||
|
--local_dir "${MODEL_ROOT}" \
|
||||||
|
--max-workers 16 \
|
||||||
|
--include 'FL2VA/**'
|
||||||
|
```
|
||||||
|
|
||||||
|
For Ref2VA, download `Ref2VA/**` instead.
|
||||||
|
|
||||||
|
### Environment
|
||||||
|
|
||||||
|
Install a compatible PyTorch/torch-musa, torchada, MATE, and vLLM-MUSA
|
||||||
|
stack before installing vLLM-Omni. Importing vLLM-Omni should report that the
|
||||||
|
`musa` vLLM platform plugin is active.
|
||||||
|
|
||||||
|
Install vLLM-Omni from a checkout containing MiniMax H3 support:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python -m pip install -e .
|
||||||
|
```
|
||||||
|
|
||||||
|
MiniMax H3 reference inputs and MP4 output require `soundfile`, `ffmpeg`, and
|
||||||
|
`ffprobe`. The Python dependency is installed by vLLM-Omni; install the two
|
||||||
|
executables with the operating system package manager and verify them:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python -c 'import soundfile; print(soundfile.__version__)'
|
||||||
|
ffmpeg -version
|
||||||
|
ffprobe -version
|
||||||
|
```
|
||||||
|
|
||||||
|
When torchaudio cannot load TorchCodec, vLLM-Omni automatically falls back to
|
||||||
|
soundfile. Formats that libsndfile cannot read are demuxed through ffmpeg.
|
||||||
|
|
||||||
|
## Start a server
|
||||||
|
|
||||||
|
One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
|
||||||
|
and FL2VA, or to `Ref2VA` for Ref2VA.
|
||||||
|
|
||||||
|
The validated Ref2VA configuration uses tensor parallelism across four MTT
|
||||||
|
S5000 GPUs and offloads inactive model components to CPU:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
export MODEL="${MODEL_ROOT}/FL2VA"
|
||||||
|
export PORT=8091
|
||||||
|
|
||||||
|
MUSA_VISIBLE_DEVICES=0,1,2,3 \
|
||||||
|
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||||
|
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
|
||||||
|
vllm serve "${MODEL}" \
|
||||||
|
--omni \
|
||||||
|
--host 0.0.0.0 \
|
||||||
|
--port "${PORT}" \
|
||||||
|
--trust-remote-code \
|
||||||
|
--num-gpus 4 \
|
||||||
|
--tensor-parallel-size 4 \
|
||||||
|
--vae-patch-parallel-size 4 \
|
||||||
|
--enable-cpu-offload \
|
||||||
|
--vae-use-tiling \
|
||||||
|
--diffusion-attention-backend FLASH_ATTN
|
||||||
|
```
|
||||||
|
|
||||||
|
On MUSA, `FLASH_ATTN` selects the MATE/FlashAttention-3 implementation. Do
|
||||||
|
not install the CUDA-only `fa4` optional dependency.
|
||||||
|
|
||||||
|
For request examples and parameter definitions, see
|
||||||
|
[MiniMax-H3.md](MiniMax-H3.md#http-api-examples).
|
||||||
|
|
||||||
|
## MUSA validation
|
||||||
|
|
||||||
|
The following weight-independent gates have been validated on one MTT S5000:
|
||||||
|
|
||||||
|
- packed variable-length attention selects MATE FlashAttention-3 and returns
|
||||||
|
finite BF16 output;
|
||||||
|
- MiniMax H3 RoPE casts CPU FP64 position metadata to FP32 before device
|
||||||
|
arithmetic and returns finite output;
|
||||||
|
- the video-VAE seeded RNG context is deterministic on MUSA and restores both
|
||||||
|
CPU and MUSA RNG states;
|
||||||
|
- the soundfile fallback preserves stereo channel layout, FP32 samples, and
|
||||||
|
the native sample rate when torchaudio/TorchCodec is unavailable.
|
||||||
|
|
||||||
|
Ref2VA image+audio generation was also validated end to end on one MTT S5000
|
||||||
|
with CPU offload and VAE tiling. The smoke used a 448x256 reference image, a
|
||||||
|
four-second 32 kHz stereo audio reference, two denoising steps, and seed 42.
|
||||||
|
The synchronous HTTP request completed in 50.6 seconds and returned an MP4
|
||||||
|
with the following decoded properties:
|
||||||
|
|
||||||
|
- 107 video frames at 448x256 and 24 FPS;
|
||||||
|
- non-static video (`pixel_std=13.07`, `temporal_delta=0.49`);
|
||||||
|
- finite stereo audio at 32 kHz (`audio_rms=0.0458`).
|
||||||
|
|
||||||
|
Official Ref2VA video-reference serving was validated with the expanded
|
||||||
|
MiniMax prompt at 1344x768, 5 seconds, seed 0, and 50 steps. A true TP4 run on
|
||||||
|
four MTT S5000 GPUs returned HTTP 200 in 638.6 seconds and produced 124 H.264
|
||||||
|
frames at 24 FPS with AAC stereo audio at 32 kHz.
|
||||||
|
|
||||||
|
Keep `--vae-use-tiling` enabled for this serving profile.
|
||||||
|
|
||||||
|
## Known limitations
|
||||||
|
|
||||||
|
- Keep Ring Attention at degree 1. The current Ring path does not preserve
|
||||||
|
MiniMax H3 packed padding boundaries.
|
||||||
|
- H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
|
||||||
|
- Each server process loads only one checkpoint partition.
|
||||||
|
- H3 currently executes one generation request per diffusion batch.
|
||||||
|
- FP8 quantization has not been enabled for MiniMax H3.
|
||||||
|
- MP3, M4A, MP4, and reference-video audio fallback requires `ffmpeg` on
|
||||||
|
`PATH`; WAV inputs can be read directly through soundfile.
|
||||||
|
|
||||||
|
## Additional resources
|
||||||
|
|
||||||
|
- [MiniMax-H3.md](MiniMax-H3.md) — full model and API guide
|
||||||
|
- [MiniMax-H3-NPU.md](MiniMax-H3-NPU.md) — Ascend NPU deployment guide
|
||||||
|
- [Supported models](../../docs/models/supported_models.md)
|
||||||
|
- [Video API](../../docs/serving/videos_api.md)
|
||||||
@@ -1,7 +1,7 @@
|
|||||||
# MiniMax H3 on Ascend NPU
|
# MiniMax H3 on Ascend NPU
|
||||||
|
|
||||||
> Joint video and audio generation with text, first-frame, image/audio, or
|
> Joint video and audio generation with text, first/last keyframes, and mixed
|
||||||
> multi-video conditions — Ascend NPU deployment guide
|
> image/video/audio references — Ascend NPU deployment guide
|
||||||
|
|
||||||
## Summary
|
## Summary
|
||||||
|
|
||||||
@@ -149,9 +149,9 @@ throughput guarantee.
|
|||||||
- H3 currently executes one generation request per diffusion batch.
|
- H3 currently executes one generation request per diffusion batch.
|
||||||
- The first regional-compile request is a warmup and should not be included
|
- The first regional-compile request is a warmup and should not be included
|
||||||
in steady-state performance measurements.
|
in steady-state performance measurements.
|
||||||
- Image+audio Ref2VA accepts exactly one image and one audio reference.
|
- The official H3 input matrix and media limits are documented in the [GPU
|
||||||
- Video Ref2VA accepts one or more video files, but not an additional
|
recipe](MiniMax-H3.md#official-input-matrix-and-limits); this NPU path uses
|
||||||
standalone audio reference.
|
the same HTTP request contract.
|
||||||
- VAE patch parallelism requires size 1 or the full DiT group size and
|
- VAE patch parallelism requires size 1 or the full DiT group size and
|
||||||
supports the H3 native `tile` mode only.
|
supports the H3 native `tile` mode only.
|
||||||
|
|
||||||
|
|||||||
+212
-18
@@ -1,7 +1,7 @@
|
|||||||
# MiniMax H3
|
# MiniMax H3
|
||||||
|
|
||||||
> Joint video and audio generation with text, first-frame, image/audio, or
|
> Joint video and audio generation with text, first/last keyframes, and
|
||||||
> multi-video conditions
|
> mixed image/video/audio references
|
||||||
|
|
||||||
## Summary
|
## Summary
|
||||||
|
|
||||||
@@ -16,8 +16,8 @@ checkpoint has two independently served partitions:
|
|||||||
|
|
||||||
- `FL2VA`: text-to-video+audio (`t2va`) and first-frame-to-video+audio
|
- `FL2VA`: text-to-video+audio (`t2va`) and first-frame-to-video+audio
|
||||||
(`fl2va`)
|
(`fl2va`)
|
||||||
- `Ref2VA`: image+audio or one-or-more-video reference-to-video+audio
|
- `Ref2VA`: up to 9 images, 3 videos, and 3 audio references in supported
|
||||||
(`ref2va`)
|
image/video/audio combinations (`ref2va`; audio-only is rejected)
|
||||||
|
|
||||||
The generated MP4 contains H.264 video and synchronized stereo audio.
|
The generated MP4 contains H.264 video and synchronized stereo audio.
|
||||||
|
|
||||||
@@ -33,8 +33,10 @@ export MODEL_ROOT=/path/to/MiniMax-H3
|
|||||||
hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}"
|
hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}"
|
||||||
```
|
```
|
||||||
|
|
||||||
Install vLLM-Omni from the checkout containing MiniMax H3 support. On
|
Install vLLM-Omni from the checkout containing MiniMax H3 support. The
|
||||||
Blackwell, install the optional FlashAttention-4 dependency:
|
two-GPU RTX 5090/4090 profiles use cuDNN attention and do not need
|
||||||
|
FlashAttention-4. Install the optional dependency only for the four-GPU
|
||||||
|
B300/GB200 `FLASH_ATTN` profile:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
uv venv
|
uv venv
|
||||||
@@ -50,6 +52,26 @@ reference-video preparation and MP4 output.
|
|||||||
One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
|
One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
|
||||||
and FL2VA requests, or to `Ref2VA` for either Ref2VA request.
|
and FL2VA requests, or to `Ref2VA` for either Ref2VA request.
|
||||||
|
|
||||||
|
### Memory and storage requirements
|
||||||
|
|
||||||
|
Treat GPU HBM, host RAM, and checkpoint storage as separate requirements. Each
|
||||||
|
H3 checkpoint partition (`FL2VA` or `Ref2VA`) contains about **134 GiB** of
|
||||||
|
BF16 safetensors (about **135 GiB** on disk). Keeping both partitions locally
|
||||||
|
therefore needs roughly **270 GiB** of model storage, although the examples
|
||||||
|
serve only one partition at a time.
|
||||||
|
|
||||||
|
CPU offload and distributed layerwise offload reduce GPU residency; they do
|
||||||
|
not make the model weights disappear. With `--dlo-no-use-allgather`, each
|
||||||
|
worker retains its standard-loader rank-local weights in host memory, including
|
||||||
|
pinned CPU buffers used for H2D streaming. Use at least **200 GiB available
|
||||||
|
system RAM** before starting the two-GPU recipe; a **384 GiB host is
|
||||||
|
recommended** to leave room for the OS, CUDA/PyTorch allocations, request
|
||||||
|
inputs, and filesystem cache. Do not run the FL2VA and Ref2VA servers at the
|
||||||
|
same time on a host sized for this minimum.
|
||||||
|
|
||||||
|
The consumer-GPU profiles below are HBM budgets only. They still require the
|
||||||
|
host-RAM budget above.
|
||||||
|
|
||||||
### Single GPU: accuracy and memory first
|
### Single GPU: accuracy and memory first
|
||||||
|
|
||||||
The single-GPU configuration uses model-level CPU offload.
|
The single-GPU configuration uses model-level CPU offload.
|
||||||
@@ -77,6 +99,98 @@ Use a GPU with enough memory for the active H3 component and enough system RAM
|
|||||||
for the offloaded components. CPU offload reduces GPU memory pressure but adds
|
for the offloaded components. CPU offload reduces GPU memory pressure but adds
|
||||||
PCIe/NVLink transfer latency.
|
PCIe/NVLink transfer latency.
|
||||||
|
|
||||||
|
### Two 24/32 GB GPUs: TP2 distributed layerwise offload
|
||||||
|
|
||||||
|
For two PCIe consumer GPUs, combine TP2 with distributed layerwise offload
|
||||||
|
(DLO). The standard loader first creates the rank-local TP shard. DLO keeps
|
||||||
|
that shard in pinned host memory and streams the 30 tail DiT blocks through a
|
||||||
|
shared two-buffer window without a DP AllGather. The first 20 DiT blocks are
|
||||||
|
copied to the GPUs once per denoise stage, reused by every sampling step, and
|
||||||
|
released before VAE decode so the decoder can reuse their HBM.
|
||||||
|
|
||||||
|
```bash
|
||||||
|
export MODEL="${MODEL_ROOT}/FL2VA"
|
||||||
|
export PORT=8091
|
||||||
|
|
||||||
|
CUDA_VISIBLE_DEVICES=0,1 \
|
||||||
|
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||||
|
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
|
||||||
|
vllm serve "${MODEL}" \
|
||||||
|
--omni \
|
||||||
|
--host 0.0.0.0 \
|
||||||
|
--port "${PORT}" \
|
||||||
|
--trust-remote-code \
|
||||||
|
--num-gpus 2 \
|
||||||
|
--tensor-parallel-size 2 \
|
||||||
|
--usp 1 \
|
||||||
|
--ring 1 \
|
||||||
|
--text-encoder-tp-size 2 \
|
||||||
|
--vae-patch-parallel-size 2 \
|
||||||
|
--vae-parallel-mode tile \
|
||||||
|
--vae-use-tiling \
|
||||||
|
--enable-distributed-layerwise-offload \
|
||||||
|
--dlo-no-use-allgather \
|
||||||
|
--dlo-resident-layers 20 \
|
||||||
|
--enforce-eager \
|
||||||
|
--diffusion-attention-backend CUDNN_ATTN
|
||||||
|
```
|
||||||
|
|
||||||
|
Use the profile that matches the per-GPU memory capacity:
|
||||||
|
|
||||||
|
| Profile | GPUs | Starting shape | Resident DiT blocks | Attention | Execution | Status |
|
||||||
|
|---|---:|---:|---:|---|---|---|
|
||||||
|
| `rtx5090` | 2 x 32 GB | 1344x768 | 20 | cuDNN attention | eager | Target-hardware validated |
|
||||||
|
| `rtx4090` | 2 x 24 GB | 1024x576 | 12 | cuDNN attention | eager | Capacity-proxy starting point |
|
||||||
|
|
||||||
|
This topology uses all available parallel capacity: TP2 shards both the DiT
|
||||||
|
and text encoder, `--dlo-no-use-allgather` streams each rank's local TP shard
|
||||||
|
without reconstructing full blocks, and VAE patch parallelism splits tiled
|
||||||
|
decode across both GPUs. cuDNN attention is selected explicitly for the RTX
|
||||||
|
consumer path; the server stays eager to avoid an unqualified compile path.
|
||||||
|
|
||||||
|
The resident count changes placement and transfer frequency only; it does not
|
||||||
|
quantize or change the BF16/FP32 denoise math. Re-measure peak memory before
|
||||||
|
increasing it on a different request shape.
|
||||||
|
|
||||||
|
### RTX 5090 target-hardware validation
|
||||||
|
|
||||||
|
At vLLM-Omni commit `ae6577ea`, one full 50-step T2VA request completed on
|
||||||
|
2 x RTX 5090 without OOM:
|
||||||
|
|
||||||
|
| Shape | Frames | Client E2E | Sampled peak/GPU | Output validation |
|
||||||
|
|---:|---:|---:|---:|---|
|
||||||
|
| 1344x768 | 124 at 24 FPS | 8 min 38 s | approximately 22.6 GiB | H.264 video + 32 kHz stereo AAC; full `ffmpeg` decode passed |
|
||||||
|
|
||||||
|
This is a single end-to-end validation run, not a warmed multi-run latency
|
||||||
|
benchmark. The sampled `nvidia-smi` peak is also not a CUDA allocator
|
||||||
|
high-water mark. The environment used vLLM 0.26.0, vLLM-Omni
|
||||||
|
`0.26.1.dev14+gae6577ea`, and PyTorch 2.11.0+cu130. The
|
||||||
|
[run record](https://github.com/lishunyang12/vllm-omni-rankings/blob/dcd06d7e83cb069842535918c0169ee9f3f29ba0/scripts/%E5%BE%AE%E4%BF%A1%E5%9B%BE%E7%89%87_20260805000034_86_237.png)
|
||||||
|
captures the environment, output contract, elapsed time, and sampled peak.
|
||||||
|
|
||||||
|
Before the target run, both profiles were exercised on two B300 ranks as an
|
||||||
|
allocation and correctness proxy. At 1344x768, 124 frames, and 50 steps, the
|
||||||
|
20-layer profile peaked at 27,726 MiB per rank. At 1024x576, the 12-layer
|
||||||
|
profile peaked at 18,888 MiB per rank in a 5-step capacity run. The resident
|
||||||
|
and fully streamed placements produced identical decoded video-frame and audio
|
||||||
|
hashes for the same shape, step count, prompt, and seed. The B300 result does
|
||||||
|
not establish RTX 4090 PCIe latency; treat the 4090 profile as a conservative
|
||||||
|
starting point until it is measured on that GPU.
|
||||||
|
|
||||||
|
To run T2VA, FL2VA, image+audio Ref2VA, and two-video Ref2VA in order, validate
|
||||||
|
every MP4's H.264/AAC streams, and retain live server and GPU-memory logs:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
RUN_ROOT=/path/to/run-root \
|
||||||
|
MODEL_ROOT=/path/to/MiniMax-H3 \
|
||||||
|
GPU_IDS=0,1 \
|
||||||
|
PROFILE=rtx5090 \
|
||||||
|
bash examples/offline_inference/minimax_h3/run_h3_2gpu_all_tasks.sh
|
||||||
|
```
|
||||||
|
|
||||||
|
The script selects 20 resident layers for `PROFILE=rtx5090` and 12 for
|
||||||
|
`PROFILE=rtx4090`; `DLO_RESIDENT_LAYERS=N` overrides either default.
|
||||||
|
|
||||||
### Four GPUs: measured best-practice throughput
|
### Four GPUs: measured best-practice throughput
|
||||||
|
|
||||||
The best validated four-GPU configuration on four NVIDIA B300 GPUs is:
|
The best validated four-GPU configuration on four NVIDIA B300 GPUs is:
|
||||||
@@ -85,7 +199,8 @@ The best validated four-GPU configuration on four NVIDIA B300 GPUs is:
|
|||||||
- Ulysses sequence parallelism degree 4;
|
- Ulysses sequence parallelism degree 4;
|
||||||
- native tiled VAE patch parallelism degree 4;
|
- native tiled VAE patch parallelism degree 4;
|
||||||
- regional `torch.compile` for the repeated DiT blocks;
|
- regional `torch.compile` for the repeated DiT blocks;
|
||||||
- FlashAttention, with Ring and TP left at 1.
|
- CuTe FlashAttention-4 through the `FLASH_ATTN` backend, with Ring and TP
|
||||||
|
left at 1.
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
export MODEL="${MODEL_ROOT}/FL2VA"
|
export MODEL="${MODEL_ROOT}/FL2VA"
|
||||||
@@ -109,6 +224,9 @@ vllm serve "${MODEL}" \
|
|||||||
--diffusion-attention-backend FLASH_ATTN
|
--diffusion-attention-backend FLASH_ATTN
|
||||||
```
|
```
|
||||||
|
|
||||||
|
On Blackwell, `FLASH_ATTN` selects FA4. Confirm the server log contains
|
||||||
|
`Using CuTe FlashAttention-4 on Blackwell` before recording measurements.
|
||||||
|
|
||||||
Do not add `--enforce-eager` to this performance configuration. The first
|
Do not add `--enforce-eager` to this performance configuration. The first
|
||||||
request includes regional compilation; warm the server once before measuring
|
request includes regional compilation; warm the server once before measuring
|
||||||
steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain
|
steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain
|
||||||
@@ -225,6 +343,7 @@ curl -sS -X POST "${API_URL}" \
|
|||||||
-F 'prompt=In a snowy blue-purple forest, Ori carefully walks past a sleeping giant; footsteps crunch in the snow while the creature breathes and softly snorts.' \
|
-F 'prompt=In a snowy blue-purple forest, Ori carefully walks past a sleeping giant; footsteps crunch in the snow while the creature breathes and softly snorts.' \
|
||||||
-F 'width=1344' \
|
-F 'width=1344' \
|
||||||
-F 'height=768' \
|
-F 'height=768' \
|
||||||
|
-F 'aspect_ratio=16:9' \
|
||||||
-F 'fps=24' \
|
-F 'fps=24' \
|
||||||
-F 'num_inference_steps=50' \
|
-F 'num_inference_steps=50' \
|
||||||
-F 'flow_shift=12' \
|
-F 'flow_shift=12' \
|
||||||
@@ -253,11 +372,31 @@ curl -sS -X POST "${API_URL}" \
|
|||||||
-o fl2va.mp4
|
-o fl2va.mp4
|
||||||
```
|
```
|
||||||
|
|
||||||
### 3. Ref2VA: image and audio reference
|
Use the same `FL2VA` partition for the official tail-keyframe forms. A single
|
||||||
|
image with `frame_indices=[-1]` conditions the last frame; two ordered images
|
||||||
|
with `frame_indices=[0,-1]` condition the first and last frames:
|
||||||
|
|
||||||
Run this request against the `Ref2VA` partition. `audio_reference` accepts an
|
```bash
|
||||||
HTTP(S) URL or a `data:` URL. In one terminal, expose the local reference
|
export LAST_FRAME=/path/to/fl2va_last_frame.png
|
||||||
assets to the serving host:
|
export FIRST_FRAME=/path/to/fl2va_first_frame.png
|
||||||
|
|
||||||
|
curl -sS -X POST "${API_URL}" \
|
||||||
|
-F 'prompt=The subject moves naturally from the first image to the last image.' \
|
||||||
|
-F 'num_inference_steps=50' \
|
||||||
|
-F 'flow_shift=12' \
|
||||||
|
-F 'seed=2102' \
|
||||||
|
-F 'extra_params={"task":"fl2va","duration":8.7,"frame_indices":[0,-1],"audio_flow_shift":3.0}' \
|
||||||
|
-F "input_references=@${FIRST_FRAME};type=image/png" \
|
||||||
|
-F "input_references=@${LAST_FRAME};type=image/png" \
|
||||||
|
-o fl2va_first_last.mp4
|
||||||
|
```
|
||||||
|
|
||||||
|
### 3. Ref2VA: image-only, image/audio, or mixed references
|
||||||
|
|
||||||
|
Run these requests against the `Ref2VA` partition. Image-only Ref2VA omits
|
||||||
|
`audio_reference`; adding one or more audio references is optional. The typed
|
||||||
|
fields accept one object or an ordered JSON list. `audio_reference` accepts an
|
||||||
|
HTTP(S) URL or a `data:` URL.
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
python -m http.server 8092 \
|
python -m http.server 8092 \
|
||||||
@@ -265,7 +404,24 @@ python -m http.server 8092 \
|
|||||||
--directory /path/to/reference_assets
|
--directory /path/to/reference_assets
|
||||||
```
|
```
|
||||||
|
|
||||||
Then submit the image and audio request from another terminal:
|
Then submit an image-only request from another terminal:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
export REF_IMAGE=/path/to/reference_assets/ref2va_image.png
|
||||||
|
|
||||||
|
curl -sS -X POST "${API_URL}" \
|
||||||
|
-F 'prompt=A white cat sits on a beige couch and slowly looks toward the camera.' \
|
||||||
|
-F 'aspect_ratio=adaptive' \
|
||||||
|
-F 'short_edge=768' \
|
||||||
|
-F 'num_inference_steps=50' \
|
||||||
|
-F 'flow_shift=12' \
|
||||||
|
-F 'seed=3100' \
|
||||||
|
-F 'extra_params={"task":"ref2va","duration":8.0,"audio_flow_shift":3.0}' \
|
||||||
|
-F "input_reference=@${REF_IMAGE};type=image/png" \
|
||||||
|
-o ref2va_image_only.mp4
|
||||||
|
```
|
||||||
|
|
||||||
|
An image-plus-audio request is:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
export REF_IMAGE=/path/to/reference_assets/ref2va_image.png
|
export REF_IMAGE=/path/to/reference_assets/ref2va_image.png
|
||||||
@@ -288,7 +444,7 @@ curl -sS -X POST "${API_URL}" \
|
|||||||
The requested duration should cover the complete audio. If `duration` is
|
The requested duration should cover the complete audio. If `duration` is
|
||||||
shorter, the reference soundtrack is truncated to the generated clip.
|
shorter, the reference soundtrack is truncated to the generated clip.
|
||||||
|
|
||||||
### 4. Ref2VA: two video references
|
### 4. Ref2VA: video, separate audio, and mixed references
|
||||||
|
|
||||||
Run this request against the `Ref2VA` partition. Repeat the
|
Run this request against the `Ref2VA` partition. Repeat the
|
||||||
`input_references` multipart field once per source video. H3 consumes the
|
`input_references` multipart field once per source video. H3 consumes the
|
||||||
@@ -313,10 +469,40 @@ curl -sS -X POST "${API_URL}" \
|
|||||||
-o ref2va_video_video.mp4
|
-o ref2va_video_video.mp4
|
||||||
```
|
```
|
||||||
|
|
||||||
The server stores uploaded multi-video references only for the lifetime of the
|
The server stores uploaded references only for the lifetime of the request and
|
||||||
request and deletes the temporary files after generation. Video Ref2VA uses
|
deletes temporary files after generation. A video may use its embedded
|
||||||
the source-video soundtracks and does not accept a separate
|
soundtrack, a separate `audio_reference`, or both. To send a mixed multipart
|
||||||
`audio_reference`.
|
request, repeat `input_references` for each image, video, or audio file; the
|
||||||
|
server classifies them by MIME type and preserves the per-type order.
|
||||||
|
|
||||||
|
Reference videos must be MP4/MOV with H.264/H.265 video, optional AAC/MP3
|
||||||
|
audio, 2–15 seconds each, and at most 15 seconds combined. They may still be
|
||||||
|
longer than the generated clip. Use `start_time_seconds` to select a
|
||||||
|
synchronized segment; for multiple typed video references, pass one value per
|
||||||
|
video in `extra_params.start_time_seconds`.
|
||||||
|
|
||||||
|
Reference images accept JPG/JPEG, PNG, WEBP, HEIC, or HEIF up to 30 MiB. Standalone
|
||||||
|
audio references accept WAV or MP3 up to 15 MiB, with 2–15 seconds per file and
|
||||||
|
at most 15 seconds combined.
|
||||||
|
|
||||||
|
## Official input matrix and limits
|
||||||
|
|
||||||
|
| Task | Supported references | Limits |
|
||||||
|
|------|----------------------|--------|
|
||||||
|
| T2VA | text only | prompt must be non-empty |
|
||||||
|
| FL2VA | first image, last image, or ordered first+last images | at most 2 images; `frame_indices` is `[0]`, `[-1]`, or `[0,-1]` |
|
||||||
|
| Ref2VA | image-only, image+image, image+video, video+audio, and mixed image/video/audio | images ≤9, videos ≤3, audios ≤3, total references ≤12; audio requires a visual reference |
|
||||||
|
|
||||||
|
The H3 output contract is 4–15 seconds at 24 FPS, stereo 32 kHz audio, and a
|
||||||
|
32-pixel canvas multiple. T2VA requires one named output ratio from `21:9`,
|
||||||
|
`16:9`, `4:3`, `1:1`, `3:4`, or `9:16`. FL2VA always follows the first input
|
||||||
|
image's ratio and ignores a generic `aspect_ratio` override. Ref2VA defaults to
|
||||||
|
`16:9`; `adaptive` and SGLang's `auto` spelling are accepted aliases for that
|
||||||
|
default. `short_edge` controls the 768-pixel canvas and must be `768`.
|
||||||
|
`num_outputs_per_prompt` accepts 1–10 and derives each output seed as
|
||||||
|
`seed + output_index`. The asynchronous endpoint returns all
|
||||||
|
outputs; the synchronous raw-MP4 endpoint returns the first output when more
|
||||||
|
than one is requested.
|
||||||
|
|
||||||
## Key parameters
|
## Key parameters
|
||||||
|
|
||||||
@@ -329,7 +515,15 @@ the source-video soundtracks and does not accept a separate
|
|||||||
| `flow_shift` | `12` | Video sigma shift |
|
| `flow_shift` | `12` | Video sigma shift |
|
||||||
| `audio_flow_shift` | `3` | Audio sigma shift, passed in `extra_params` |
|
| `audio_flow_shift` | `3` | Audio sigma shift, passed in `extra_params` |
|
||||||
| `seed` | Task-specific | Use a fixed value for reproducibility |
|
| `seed` | Task-specific | Use a fixed value for reproducibility |
|
||||||
| `width`, `height` | Multiples of 32 | Aspect ratio must be between 1:4 and 4:1 |
|
| `aspect_ratio` | Task-specific | T2VA requires a named ratio; FL2VA follows the input image; Ref2VA defaults to `16:9` |
|
||||||
|
| `short_edge` | `768` | H3 shape policy requires exactly 768 when `width`/`height` are omitted |
|
||||||
|
| `num_outputs_per_prompt` | `1` | 1–10; async API returns every output |
|
||||||
|
| `start_time_seconds` | `0` | Reference-video segment start; use a list in `extra_params` for multiple videos |
|
||||||
|
| `width`, `height` | Multiples of 32 | Output aspect ratio must be between 1:4 and 4:1 |
|
||||||
|
|
||||||
|
## ComfyUI Frontend
|
||||||
|
|
||||||
|
Users can also use a ComfyUI frontend to interact with a hosted MiniMax-H3 service. The ComfyUI frontend can run in a separate environment or machine. Refer to [vLLM-Omni ComfyUI Integration](../../docs/features/comfyui.md) for details.
|
||||||
|
|
||||||
## Validated four-GPU evidence
|
## Validated four-GPU evidence
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,174 @@
|
|||||||
|
# PersonaPlex
|
||||||
|
|
||||||
|
> Live full-duplex speech-to-speech (talk to the model in your browser, it listens
|
||||||
|
> while it speaks) with persona and voice control, on the native vLLM-Omni engine.
|
||||||
|
|
||||||
|
## Summary
|
||||||
|
|
||||||
|
- Vendor: NVIDIA
|
||||||
|
- Model: [`nvidia/personaplex-7b-v1`](https://huggingface.co/nvidia/personaplex-7b-v1)
|
||||||
|
(gated; a Moshi finetune)
|
||||||
|
- Task: Full-duplex speech-to-speech. 24 kHz mic audio in, agent speech +
|
||||||
|
inner-monologue text out, in 80 ms lockstep. Persona (role text) and voice
|
||||||
|
(zero-shot clone) are set per session.
|
||||||
|
- Mode: Live duplex WebSocket server with the official browser client, single or
|
||||||
|
multi-session (`--batch-size`), plus an offline WAV-in / WAV-out example
|
||||||
|
- Maintainer: [`@linyueqian`](https://github.com/linyueqian)
|
||||||
|
|
||||||
|
## When to use this recipe
|
||||||
|
|
||||||
|
Use this recipe to run `nvidia/personaplex-7b-v1` as a live voice agent on one GPU.
|
||||||
|
The integration is **moshi-free**: the temporal transformer, depformer, and
|
||||||
|
streaming Mimi codec are all native vLLM-Omni modules (Mimi runs on
|
||||||
|
`transformers.MimiModel`), so no vendored fork is installed. Decoding is greedy,
|
||||||
|
matching the reference implementation frame for frame on the golden replays used
|
||||||
|
as the acceptance gate.
|
||||||
|
|
||||||
|
PersonaPlex is the first Moshi-class (pure-lockstep) model in the experimental
|
||||||
|
full-duplex framework (`vllm_omni/experimental/fullduplex/`), and the first with
|
||||||
|
elastic multi-session batching: one engine hosts `--batch-size` concurrent
|
||||||
|
conversations, callers join or leave on any 80 ms tick without disturbing live
|
||||||
|
slots (greedy isolation is bit-exact).
|
||||||
|
|
||||||
|
## References
|
||||||
|
|
||||||
|
- Serving package:
|
||||||
|
[`vllm_omni/experimental/fullduplex/personaplex/`](../../vllm_omni/experimental/fullduplex/personaplex/)
|
||||||
|
- Native model modules (temporal / depformer / streaming Mimi):
|
||||||
|
[`vllm_omni/model_executor/models/personaplex/`](../../vllm_omni/model_executor/models/personaplex/)
|
||||||
|
- Online example (server + headless client):
|
||||||
|
[`examples/online_serving/personaplex/`](../../examples/online_serving/personaplex/)
|
||||||
|
- Offline example:
|
||||||
|
[`examples/offline_inference/personaplex/personaplex_offline.py`](../../examples/offline_inference/personaplex/personaplex_offline.py)
|
||||||
|
- Staged offline pipeline (talker -> Mimi code2wav, async-chunk streaming):
|
||||||
|
[`vllm_omni/deploy/personaplex.yaml`](../../vllm_omni/deploy/personaplex.yaml)
|
||||||
|
- Integration PR: [vllm-project/vllm-omni#4771](https://github.com/vllm-project/vllm-omni/pull/4771)
|
||||||
|
- Framework context: [RFC #3745](https://github.com/vllm-project/vllm-omni/issues/3745)
|
||||||
|
(duplex adapter patterns), [#1335](https://github.com/vllm-project/vllm-omni/issues/1335)
|
||||||
|
(full-duplex target), [PR #3907](https://github.com/vllm-project/vllm-omni/pull/3907)
|
||||||
|
(fullduplex core contracts this adapter conforms to)
|
||||||
|
- Upstream model card:
|
||||||
|
[`nvidia/personaplex-7b-v1`](https://huggingface.co/nvidia/personaplex-7b-v1)
|
||||||
|
|
||||||
|
## Why it is different from MiniCPM-o / JoyVL duplex
|
||||||
|
|
||||||
|
| | MiniCPM-o 4.5 (#3907) | JoyVL | **PersonaPlex (this)** |
|
||||||
|
|---|---|---|---|
|
||||||
|
| Cadence | 1 s chunk groups | ~1 fps frames | **80 ms lockstep (12.5 Hz)** |
|
||||||
|
| Turn control | learned `⟨listen⟩`/`⟨speak⟩` | `</silence>`/`</response>` | **none, pure lockstep** |
|
||||||
|
| Per step | variable-length token group | text decision | **1 user frame in -> 1 agent frame + 1 text token** |
|
||||||
|
| Barge-in | at chunk boundary | n/a | **native (model always hears the user)** |
|
||||||
|
| Session state | chunk-group KV | per-tick HTTP | **persistent ring KV + streaming Mimi state** |
|
||||||
|
|
||||||
|
This is the lockstep ("parallel-frame joint") shape of the duplex adapter
|
||||||
|
patterns: the adapter declares `DuplexCapability.continuous = True`, and
|
||||||
|
`core.DuplexRuntime` runs ONE eternal response that consumes input frames as they
|
||||||
|
arrive and drains on close. The flag is a small, default-off lifecycle mode, so
|
||||||
|
turn-based adapters (JoyVL, MiniCPM-o) are unaffected.
|
||||||
|
|
||||||
|
## Architecture (Moshi RQ-Transformer, all native)
|
||||||
|
|
||||||
|
- **Mimi codec** 24 kHz, 12.5 Hz, 1920 samples/frame, 8 active codebooks
|
||||||
|
(card 2048). Streaming encode + decode via `transformers.MimiModel` with the
|
||||||
|
PersonaPlex reference checkpoint weights.
|
||||||
|
- **Temporal transformer** (Helium backbone) 4096-d / 32 layers / 32 heads,
|
||||||
|
sliding window 3000 frames (a 4-minute rolling context) held in a ring KV cache
|
||||||
|
with per-slot recycle.
|
||||||
|
- **Depformer** 1024-d / 6 layers / 16 heads, autoregressive over 16 codebooks
|
||||||
|
per frame (8 vocoded), reset each frame.
|
||||||
|
- Per-frame token column `[B, 17, 1]`: row 0 = inner-monologue text,
|
||||||
|
rows 1-8 = agent audio, rows 9-16 = user audio.
|
||||||
|
- Persona + voice are injected once at session open through the same lockstep
|
||||||
|
step: voice clone forces the agent stream from reference-audio Mimi codes;
|
||||||
|
persona forces the text stream from `<system> ... <system>` tokens.
|
||||||
|
|
||||||
|
## Hardware Support
|
||||||
|
|
||||||
|
Verified on GPU (Hopper-class). The serving path is plain PyTorch eager plus the
|
||||||
|
native modules, so other CUDA GPUs with enough memory are expected to work.
|
||||||
|
|
||||||
|
## GPU
|
||||||
|
|
||||||
|
### 1x 141 GB Hopper-class GPU (H20-class, verified)
|
||||||
|
|
||||||
|
#### Environment
|
||||||
|
|
||||||
|
- OS: Linux
|
||||||
|
- Python: 3.10+ (CI covers 3.11 / 3.12)
|
||||||
|
- vLLM-Omni: PR #4771 branch or current `main` once merged
|
||||||
|
- Extra dependency for the live server: `pip install sphn` (Opus framing;
|
||||||
|
`aiohttp` already ships with vLLM)
|
||||||
|
- `export HF_TOKEN=...` with access to the gated
|
||||||
|
`nvidia/personaplex-7b-v1` repo (accept the license on the model page)
|
||||||
|
|
||||||
|
First run auto-downloads from the model repo: `model.safetensors`,
|
||||||
|
`tokenizer_spm_32k_3.model`, `voices.tgz`, and the Mimi reference checkpoint;
|
||||||
|
`kyutai/mimi` is fetched via `transformers` for the codec module graph.
|
||||||
|
|
||||||
|
#### Command
|
||||||
|
|
||||||
|
Live duplex server (official browser client served at `/`):
|
||||||
|
|
||||||
|
```bash
|
||||||
|
HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 \
|
||||||
|
python -m vllm_omni.experimental.fullduplex.personaplex.serving.server \
|
||||||
|
--port 8124 --voice NATF2.pt
|
||||||
|
```
|
||||||
|
|
||||||
|
Multi-session on the same GPU (elastic batching, callers join/leave any tick):
|
||||||
|
|
||||||
|
```bash
|
||||||
|
HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 \
|
||||||
|
python -m vllm_omni.experimental.fullduplex.personaplex.serving.server \
|
||||||
|
--port 8124 --batch-size 4
|
||||||
|
```
|
||||||
|
|
||||||
|
Then open `http://localhost:8124/` and allow the microphone (use headphones so
|
||||||
|
the agent does not hear itself). Headless alternative:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python examples/online_serving/personaplex/duplex_client.py \
|
||||||
|
--url ws://localhost:8124/api/chat --input user.wav --out reply.wav
|
||||||
|
```
|
||||||
|
|
||||||
|
Offline WAV-in / WAV-out (no server):
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python examples/offline_inference/personaplex/personaplex_offline.py \
|
||||||
|
--input-wav question.wav --output-wav reply.wav --output-text reply.json \
|
||||||
|
--voice-prompt NATF2.pt --persona "You are a wise and friendly teacher."
|
||||||
|
```
|
||||||
|
|
||||||
|
#### Verification
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# GPU-free contract tests (stubbed FrameStepper)
|
||||||
|
pytest tests/e2e/features/fullduplex/test_personaplex_adapter.py \
|
||||||
|
tests/e2e/features/fullduplex/test_personaplex_session.py \
|
||||||
|
tests/e2e/features/fullduplex/test_personaplex_batched.py -q
|
||||||
|
|
||||||
|
# GPU e2e: batched isolation + slot recycle (coherence check needs a real
|
||||||
|
# question WAV via PPLEX_QUESTION_WAV)
|
||||||
|
pytest tests/e2e/features/fullduplex/test_personaplex_elastic.py -q
|
||||||
|
```
|
||||||
|
|
||||||
|
A green offline run producing a coherent spoken reply (and inner-monologue text
|
||||||
|
in `reply.json`) is the quickest end-to-end check.
|
||||||
|
|
||||||
|
#### Notes
|
||||||
|
|
||||||
|
- Realtime budget is 80 ms/frame. Verified eager per-tick latency at
|
||||||
|
`--batch-size 4` is ~70-74 ms on this hardware class, i.e. all four
|
||||||
|
conversations stay realtime; single-session has comfortable headroom.
|
||||||
|
- Decoding is greedy only (temperature/top-k knobs are intentionally not
|
||||||
|
exposed): greedy is what the parity gates pin against the reference
|
||||||
|
implementation, and sampling amplifies sub-bit numeric drift into divergence.
|
||||||
|
- With `--batch-size N`, connections beyond N are rejected until a slot frees;
|
||||||
|
a recycled slot is bit-exact with a fresh engine for the same inputs.
|
||||||
|
- Voice prompts: bundled `voices.tgz` names (`NATF*`/`NATM*` natural,
|
||||||
|
`VARF*`/`VARM*` varied) or a path to your own `.pt`/`.wav`.
|
||||||
|
- Run the browser client near the server. The 80 ms cadence is sensitive to
|
||||||
|
network jitter; over a high-latency link playback can stutter regardless of
|
||||||
|
engine speed. On localhost it is smooth.
|
||||||
|
- Sliding window is 3000 frames (~4 min). The ring KV recycles beyond that;
|
||||||
|
very long sessions keep running with a rolling context.
|
||||||
@@ -46,6 +46,7 @@ recipes/
|
|||||||
| [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
|
| [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
|
||||||
| [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
|
| [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
|
||||||
| [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 |
|
| [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 |
|
||||||
|
| [`MiniMaxAI/MiniMax-H3-MUSA.md`](./MiniMaxAI/MiniMax-H3-MUSA.md) | MiniMax H3 T2VA, FL2VA, and Ref2VA on MUSA | 4x MTT S5000 (TP4 Ref2VA profile) |
|
||||||
| [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
|
| [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
|
||||||
| [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU |
|
| [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU |
|
||||||
| [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |
|
| [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |
|
||||||
|
|||||||
@@ -164,10 +164,11 @@ curl -s http://localhost:8091/v1/chat/completions \
|
|||||||
PR [#2495](https://github.com/vllm-project/vllm-omni/pull/2495) adds
|
PR [#2495](https://github.com/vllm-project/vllm-omni/pull/2495) adds
|
||||||
performance CI configs for the same DiT-only settings. The CI step is
|
performance CI configs for the same DiT-only settings. The CI step is
|
||||||
currently opt-in (gated by `RUN_HUNYUAN_IMAGE3_PERF=1`) with `soft_fail`
|
currently opt-in (gated by `RUN_HUNYUAN_IMAGE3_PERF=1`) with `soft_fail`
|
||||||
enabled, intended for initial data collection. Performance assertions are
|
enabled, intended for initial data collection. Perf runners no longer
|
||||||
skipped (`skip-performance-assertion: true`); the baseline values in the
|
compare metrics against JSON baselines at assert time; the `baseline`
|
||||||
JSON configs are reference-only and will be promoted to regression gates
|
blocks in the configs are reference-only and are copied into result JSON
|
||||||
once enough nightly data has been collected.
|
for reporting. Regression gates can be reintroduced once enough nightly
|
||||||
|
data has been collected.
|
||||||
|
|
||||||
The user-facing equivalent is to launch one of the CLI commands above and
|
The user-facing equivalent is to launch one of the CLI commands above and
|
||||||
generate 1024x1024 images with 50 denoising steps.
|
generate 1024x1024 images with 50 denoising steps.
|
||||||
|
|||||||
@@ -56,10 +56,14 @@ vllm serve /path/to/Cosmos3-Super \
|
|||||||
--data-parallel-size 4
|
--data-parallel-size 4
|
||||||
```
|
```
|
||||||
|
|
||||||
### Without AllGather (full weights per rank, no sharding)
|
### Without DLO AllGather (rank-local weights, no DLO sharding)
|
||||||
|
|
||||||
Use `--dlo-no-use-allgather` when you want each rank to load full weights independently —
|
Use `--dlo-no-use-allgather` when you want each rank to stream the standard
|
||||||
no AllGather synchronization, no concurrent request requirement, but N× host memory.
|
loader's rank-local weights independently — no DLO AllGather synchronization,
|
||||||
|
no concurrent request requirement, but no DLO weight sharding and higher host
|
||||||
|
memory use. In pure DP, this is equivalent to independent full-weight
|
||||||
|
layerwise-offload replicas; with TP or HSDP, each rank retains the standard
|
||||||
|
loader's local tensor layout.
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
vllm serve /path/to/Cosmos3-Nano \
|
vllm serve /path/to/Cosmos3-Nano \
|
||||||
@@ -154,9 +158,9 @@ wait
|
|||||||
| Flag | Description | Default |
|
| Flag | Description | Default |
|
||||||
|------|-------------|---------|
|
|------|-------------|---------|
|
||||||
| `--enable-distributed-layerwise-offload` | Enable DLO with H2D + AllGather overlap | `false` |
|
| `--enable-distributed-layerwise-offload` | Enable DLO with H2D + AllGather overlap | `false` |
|
||||||
| `--data-parallel-size N` | Number of DP ranks (weight sharding + concurrent requests) | `1` |
|
| `--data-parallel-size N` | Number of DP ranks; DLO weight sharding + concurrent requests when AllGather is enabled | `1` |
|
||||||
| `--dlo-use-allgather` | Use shard + AllGather for weight reconstruction (recommended) | `true` |
|
| `--dlo-use-allgather` | Use shard + AllGather for weight reconstruction (recommended) | `true` |
|
||||||
| `--dlo-no-use-allgather` | Each rank loads full weights independently (no sharding, no AllGather) | `false` |
|
| `--dlo-no-use-allgather` | Stream standard-loader rank-local weights independently (no DLO sharding, no DLO AllGather) | `false` |
|
||||||
| `--usp N` | Use SP instead of DP for weight sharding (long sequences) | `1` |
|
| `--usp N` | Use SP instead of DP for weight sharding (long sequences) | `1` |
|
||||||
|
|
||||||
## Mode comparison
|
## Mode comparison
|
||||||
@@ -166,7 +170,7 @@ wait
|
|||||||
| DLO + AllGather (DP4) | `--enable-distributed-layerwise-offload --data-parallel-size 4` | 1/4 model | 2 blocks | **3.3× HSDP** | Max throughput, short sequences |
|
| DLO + AllGather (DP4) | `--enable-distributed-layerwise-offload --data-parallel-size 4` | 1/4 model | 2 blocks | **3.3× HSDP** | Max throughput, short sequences |
|
||||||
| DLO + AllGather (DP2) | `--enable-distributed-layerwise-offload --data-parallel-size 2` | 1/2 model | 2 blocks | 2× HSDP | Balanced throughput + memory |
|
| DLO + AllGather (DP2) | `--enable-distributed-layerwise-offload --data-parallel-size 2` | 1/2 model | 2 blocks | 2× HSDP | Balanced throughput + memory |
|
||||||
| DLO + AllGather (SP4) | `--enable-distributed-layerwise-offload --usp 4` | 1/4 model | 2 blocks | 1× (single req) | Long sequences (720p+) |
|
| DLO + AllGather (SP4) | `--enable-distributed-layerwise-offload --usp 4` | 1/4 model | 2 blocks | 1× (single req) | Long sequences (720p+) |
|
||||||
| DLO no-AllGather | `--enable-distributed-layerwise-offload --dlo-no-use-allgather` | Full model | 2 blocks | 0.4× HSDP | No AllGather overhead, high CPU |
|
| DLO no-AllGather | `--enable-distributed-layerwise-offload --dlo-no-use-allgather` | Standard-loader rank-local weights | 2 blocks | 0.4× HSDP | No DLO AllGather overhead, high CPU |
|
||||||
| HSDP | `--use-hsdp --hsdp-shard-size 4` | 0 | 1/4 model | 1× (baseline) | Weights fit in HBM |
|
| HSDP | `--use-hsdp --hsdp-shard-size 4` | 0 | 1/4 model | 1× (baseline) | Weights fit in HBM |
|
||||||
|
|
||||||
## Memory expectations
|
## Memory expectations
|
||||||
@@ -185,7 +189,11 @@ wait
|
|||||||
|
|
||||||
- `--dlo-use-allgather` (default) requires all concurrent requests to have the same
|
- `--dlo-use-allgather` (default) requires all concurrent requests to have the same
|
||||||
`num_inference_steps` (AllGather is a collective that needs all ranks synchronized)
|
`num_inference_steps` (AllGather is a collective that needs all ranks synchronized)
|
||||||
- Online quantization (FP8) is incompatible with mmap loading — falls back to regular
|
- Online quantization (FP8) is incompatible with the DLO+AllGather mmap path;
|
||||||
`load_weights()` automatically
|
use `--dlo-no-use-allgather` or disable online quantization
|
||||||
- Tensor Parallel is not supported (DLO uses DP-based sharding, a different dimension)
|
- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path. TP with
|
||||||
- HSDP + AllGather is rejected (would double-shard weights)
|
`--dlo-no-use-allgather` retains the standard TP loader and should be treated
|
||||||
|
as an experimental compatibility path
|
||||||
|
- HSDP + AllGather is rejected (would double-shard weights). HSDP with
|
||||||
|
`--dlo-no-use-allgather` is accepted at configuration level, but is not yet a
|
||||||
|
broadly validated production combination
|
||||||
@@ -423,6 +423,156 @@ ffprobe -v error -show_entries stream=codec_type,nb_frames,width,height cosmos3_
|
|||||||
once in `vllm_omni/model_extras/cosmos3.py` and forwarded through `--extra-body`;
|
once in `vllm_omni/model_extras/cosmos3.py` and forwarded through `--extra-body`;
|
||||||
unknown keys for the model are dropped.
|
unknown keys for the model are dropped.
|
||||||
|
|
||||||
|
## ROCm
|
||||||
|
|
||||||
|
### 1x MI350X (Online serving)
|
||||||
|
|
||||||
|
#### Environment
|
||||||
|
|
||||||
|
- OS: Ubuntu 22.04+
|
||||||
|
- Python: 3.12+
|
||||||
|
- Driver / runtime: ROCm, HIP 7.2, **gfx942 or gfx950** (validated on gfx950 / MI350X,
|
||||||
|
which exposes 252 GiB of HBM to PyTorch)
|
||||||
|
- PyTorch: 2.11.0 (ROCm build)
|
||||||
|
- vLLM version: match the repository requirements from your current checkout
|
||||||
|
- vLLM-Omni version or commit: use the commit you are deploying from
|
||||||
|
|
||||||
|
No ROCm-specific patch or flag is needed: Cosmos3 runs on the stock code path.
|
||||||
|
|
||||||
|
#### Command
|
||||||
|
|
||||||
|
The NVIDIA command works unchanged. Guardrails need the gated
|
||||||
|
`nvidia/Cosmos-1.0-Guardrail` repo, so the quick path is to disable them (you are
|
||||||
|
then responsible for license compliance):
|
||||||
|
|
||||||
|
```bash
|
||||||
|
vllm serve nvidia/Cosmos3-Nano \
|
||||||
|
--omni \
|
||||||
|
--no-guardrails \
|
||||||
|
--host 0.0.0.0 --port 8000 \
|
||||||
|
--init-timeout 1800
|
||||||
|
```
|
||||||
|
|
||||||
|
For 720p video, add `--vae-use-tiling`: it cuts peak VRAM by ~68% for ~13%
|
||||||
|
latency, which is the difference between needing a 120 GiB card and a 40 GiB one.
|
||||||
|
`--quantization fp8` and `--enable-layerwise-offload` are also supported; see the
|
||||||
|
Notes for the measured cost of each. For extra GPUs use `--ulysses-degree N` or
|
||||||
|
`--tensor-parallel-size N` (not validated on ROCm yet).
|
||||||
|
|
||||||
|
#### Verification
|
||||||
|
|
||||||
|
Use the same `curl` requests as the CUDA section above. The ROCm images do not
|
||||||
|
ship `ffprobe`, so the CUDA section's `ffprobe` line does not work here. `imageio`
|
||||||
|
and `pyav` are already installed and report the same thing:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print(im.size, im.mode)"
|
||||||
|
python -c "import imageio.v3 as iio; print(iio.improps('cosmos3_t2v.mp4', plugin='pyav').shape)"
|
||||||
|
```
|
||||||
|
|
||||||
|
Expect `(1024, 1024) RGB` and `(189, 720, 1280, 3)` for the requests above.
|
||||||
|
Re-running the same request with the same seed reproduces a byte-identical mp4,
|
||||||
|
so an md5 comparison across two runs also works as a smoke check.
|
||||||
|
|
||||||
|
#### Notes
|
||||||
|
|
||||||
|
- **Warm up once per output shape before timing anything.** The first generation
|
||||||
|
of a shape pays a one-time cost that is cached on disk and reused by later
|
||||||
|
processes in the same container: a ~250 s aiter attention-kernel JIT build plus
|
||||||
|
a ~390 s build of the VAE decode path *for that shape*, so warming 1024² images
|
||||||
|
does nothing for 189-frame video. At 720p / 189 frames / 35 steps this is
|
||||||
|
**540 s cold vs 161 s warm, with byte-identical output**.
|
||||||
|
- **Measured on 1x MI350X (bf16, guardrails off, warm):** T2I 1024² @ 50 steps
|
||||||
|
**~2.7 s**; T2V 1280×720 / 189 frames @ 35 steps **~161 s**, of which ~92% is
|
||||||
|
the DiT denoise loop and ~4% VAE decode, so optimization effort belongs in the
|
||||||
|
denoise loop. The optional flags act on different terms of the memory bill and
|
||||||
|
are therefore complementary rather than redundant:
|
||||||
|
|
||||||
|
| Flag | Latency | Peak reserved | Peak allocated | Acts on |
|
||||||
|
|---|---|---|---|---|
|
||||||
|
| *(none)* | 161 s | 120 GiB | 95 GiB | — |
|
||||||
|
| `--vae-use-tiling` | 183 s (+13.5%) | **38 GiB** (−68%) | **36 GiB** (−62%) | decode activations |
|
||||||
|
| `--enable-layerwise-offload` | 162 s (+0.8%) | 84 GiB (−30%) | 69 GiB (−27%) | weights |
|
||||||
|
| `--quantization fp8` | 150 s (−6.6%) | 107 GiB (−11%) | 82 GiB (−14%) | weights (online, no calibration) |
|
||||||
|
|
||||||
|
Reserved is the caching-allocator high water mark and runs ~2–25 GiB above
|
||||||
|
allocated; allocated is the figure to compare against another platform.
|
||||||
|
- **`--vae-use-tiling` is by far the biggest lever for video.** Of the 95 GiB
|
||||||
|
allocated by default only ~35 GiB is weights and fixed overhead; the remaining
|
||||||
|
**60 GiB is activations from decoding all 189 frames in one piece**, which is
|
||||||
|
why fp8 and offload cannot touch it and why tiling, which leaves under 5 GiB of
|
||||||
|
them, wins so much. `--vae-use-slicing` is a no-op for a single video request:
|
||||||
|
peak memory is bit-identical to the default.
|
||||||
|
- **`--enable-layerwise-offload` is close to free on this part:** per-block DiT
|
||||||
|
compute is long enough to hide the weight transfer, so on a card this size it
|
||||||
|
buys headroom for concurrency or higher resolution rather than the
|
||||||
|
speed-for-memory trade it is on small cards.
|
||||||
|
- **Attention backend:** the numbers above are aiter FlashAttention, not SDPA.
|
||||||
|
ROCm picks aiter when it is installed and the encoded device capability
|
||||||
|
satisfies `90 < major * 10 + minor < 100` (gfx942 and gfx950); otherwise it
|
||||||
|
falls back to `TORCH_SDPA` and says so only at debug level, which is easy to
|
||||||
|
miss. If you are reporting numbers, assert the choice instead of trusting the
|
||||||
|
default: `RocmOmniPlatform.get_diffusion_attn_backend_cls(None, 128)` returns
|
||||||
|
the resolved path.
|
||||||
|
- **Determinism is per-configuration:** same seed and same flags reproduce a
|
||||||
|
byte-identical mp4 across processes, but different configurations do not match
|
||||||
|
each other — including offload, which nominally only relocates weights, so do
|
||||||
|
not use it as a bit-exact regression baseline. Quality was not evaluated for
|
||||||
|
any configuration.
|
||||||
|
|
||||||
|
### 1x MI350X (Offline generation)
|
||||||
|
|
||||||
|
#### Environment
|
||||||
|
|
||||||
|
Same as the online serving section above.
|
||||||
|
|
||||||
|
#### Command
|
||||||
|
|
||||||
|
The standard task examples run unchanged:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Text-to-image
|
||||||
|
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||||
|
--model nvidia/Cosmos3-Nano \
|
||||||
|
--prompt "A photorealistic red sports car at golden hour, cinematic lighting." \
|
||||||
|
--negative-prompt "blurry, distorted, low quality" \
|
||||||
|
--height 1024 --width 1024 --num-inference-steps 50 --guidance-scale 7.0 \
|
||||||
|
--extra-body '{"flow_shift": 3.0, "guardrails": false}' \
|
||||||
|
--output cosmos3_t2i.png
|
||||||
|
|
||||||
|
# Text-to-video
|
||||||
|
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||||
|
--model nvidia/Cosmos3-Nano \
|
||||||
|
--prompt "A robot arm is cleaning a plate in the kitchen." \
|
||||||
|
--negative-prompt "blurry, distorted, low quality" \
|
||||||
|
--height 720 --width 1280 --num-frames 189 --fps 24 \
|
||||||
|
--num-inference-steps 35 --guidance-scale 6.0 --flow-shift 10.0 \
|
||||||
|
--extra-body '{"max_sequence_length": 4096, "guardrails": false,
|
||||||
|
"use_resolution_template": false, "use_duration_template": false}' \
|
||||||
|
--output cosmos3_t2v.mp4
|
||||||
|
```
|
||||||
|
|
||||||
|
Both scripts print `Total generation time` and `Worker peak GPU memory
|
||||||
|
(reserved)`, which is where the latency and reserved figures in the Notes above
|
||||||
|
come from. Add `--vae-use-tiling`, `--quantization fp8` or
|
||||||
|
`--enable-layerwise-offload` to reproduce the other rows. The allocated column
|
||||||
|
needs `VLLM_LOGGING_LEVEL=DEBUG`, which enables a per-request line reporting
|
||||||
|
reserved, allocated and pool overhead together.
|
||||||
|
|
||||||
|
#### Verification
|
||||||
|
|
||||||
|
Same as the online serving section above.
|
||||||
|
|
||||||
|
#### Notes
|
||||||
|
|
||||||
|
- Run each command twice and read the second run; the first pays the one-time
|
||||||
|
compilation described above. If you run in a container, keep the container
|
||||||
|
alive between runs (`docker run --rm` throws the JIT cache away and you pay
|
||||||
|
the ~250 s kernel build every time).
|
||||||
|
- Latency and memory figures are identical to the online serving section — they
|
||||||
|
were taken from this path, because a single process has no client/server
|
||||||
|
timing ambiguity.
|
||||||
|
|
||||||
## NPU
|
## NPU
|
||||||
|
|
||||||
### 1x Ascend 910B / 910C (Atlas A2 / A3) — Online serving
|
### 1x Ascend 910B / 910C (Atlas A2 / A3) — Online serving
|
||||||
|
|||||||
@@ -81,6 +81,21 @@ def _from_pipeline_key(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_duplex_session_capacity_propagates_to_every_structured_stage() -> None:
|
||||||
|
omni_config = _from_pipeline_key("personaplex")
|
||||||
|
|
||||||
|
assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [2, 2]
|
||||||
|
|
||||||
|
|
||||||
|
def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) -> None:
|
||||||
|
deploy_path = tmp_path / "personaplex-turn.yaml"
|
||||||
|
deploy_path.write_text("session_mode: turn\n", encoding="utf-8")
|
||||||
|
|
||||||
|
omni_config = _from_pipeline_key("personaplex", deploy_config_path=str(deploy_path))
|
||||||
|
|
||||||
|
assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [1, 1]
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
|
@pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
|
||||||
def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(model_type: str):
|
def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(model_type: str):
|
||||||
pipeline = _resolve_pipeline_or_skip(model_type)
|
pipeline = _resolve_pipeline_or_skip(model_type)
|
||||||
@@ -105,6 +120,7 @@ def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(mod
|
|||||||
assert omni_stage.runtime_config.num_replicas == legacy_stage.yaml_runtime.get("num_replicas", 1)
|
assert omni_stage.runtime_config.num_replicas == legacy_stage.yaml_runtime.get("num_replicas", 1)
|
||||||
|
|
||||||
engine_args = legacy_stage.yaml_engine_args
|
engine_args = legacy_stage.yaml_engine_args
|
||||||
|
assert omni_stage.model_config.duplex_max_sessions == engine_args.get("duplex_max_sessions", 1)
|
||||||
assert omni_stage.model_config.enforce_eager == engine_args.get("enforce_eager", False)
|
assert omni_stage.model_config.enforce_eager == engine_args.get("enforce_eager", False)
|
||||||
assert omni_stage.load_config.load_format == engine_args.get("load_format", "auto")
|
assert omni_stage.load_config.load_format == engine_args.get("load_format", "auto")
|
||||||
assert omni_stage.load_config.tokenizer_mode == engine_args.get("tokenizer_mode", "auto")
|
assert omni_stage.load_config.tokenizer_mode == engine_args.get("tokenizer_mode", "auto")
|
||||||
@@ -379,6 +395,7 @@ def test_runtime_config_fields_match_rfc_runtime_scope():
|
|||||||
def test_sub_config_fields_match_rfc_scopes():
|
def test_sub_config_fields_match_rfc_scopes():
|
||||||
assert {f.name for f in fields(OmniStageModelConfig)} == {
|
assert {f.name for f in fields(OmniStageModelConfig)} == {
|
||||||
"active_stream_window",
|
"active_stream_window",
|
||||||
|
"duplex_max_sessions",
|
||||||
"enable_sleep_mode",
|
"enable_sleep_mode",
|
||||||
"default_sampling_params",
|
"default_sampling_params",
|
||||||
"subtalker_sampling_params",
|
"subtalker_sampling_params",
|
||||||
|
|||||||
@@ -110,6 +110,9 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
|
|||||||
session.status = RequestStatus.RUNNING
|
session.status = RequestStatus.RUNNING
|
||||||
session.resumable = True
|
session.resumable = True
|
||||||
session.num_computed_tokens = len(session.prompt_token_ids)
|
session.num_computed_tokens = len(session.prompt_token_ids)
|
||||||
|
# schedule() has counted this step's token as in flight; update_from_output()
|
||||||
|
# must settle it before handling the segment boundary.
|
||||||
|
session.num_in_flight_tokens = 1
|
||||||
|
|
||||||
sched = MagicMock()
|
sched = MagicMock()
|
||||||
sched.requests = {session.request_id: session}
|
sched.requests = {session.request_id: session}
|
||||||
@@ -118,12 +121,20 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
|
|||||||
is_done_receiving_chunks=lambda _request_id: True,
|
is_done_receiving_chunks=lambda _request_id: True,
|
||||||
segment_finished_requests={session.request_id},
|
segment_finished_requests={session.request_id},
|
||||||
)
|
)
|
||||||
sched._handle_stopped_request.return_value = False
|
|
||||||
|
def rearm_resumable_request(request: Request) -> bool:
|
||||||
|
request.status = RequestStatus.WAITING
|
||||||
|
return False
|
||||||
|
|
||||||
|
sched._handle_stopped_request.side_effect = rearm_resumable_request
|
||||||
sched.running = [session]
|
sched.running = [session]
|
||||||
sched.waiting = MagicMock()
|
sched.waiting = MagicMock()
|
||||||
sched.skipped_waiting = MagicMock()
|
sched.skipped_waiting = MagicMock()
|
||||||
sched.structured_output_manager.should_advance.return_value = False
|
sched.structured_output_manager.should_advance.return_value = False
|
||||||
sched._pending_finish_reqs = []
|
# Async scheduling can observe the segment boundary in the next schedule()
|
||||||
|
# before this model output is applied and defer the same request for finish.
|
||||||
|
# update_from_output() must not re-arm the resumable request twice.
|
||||||
|
sched._pending_finish_reqs = [session]
|
||||||
sched.recompute_kv_load_failures = False
|
sched.recompute_kv_load_failures = False
|
||||||
sched.connector = None
|
sched.connector = None
|
||||||
sched.kv_cache_manager.take_events.return_value = None
|
sched.kv_cache_manager.take_events.return_value = None
|
||||||
@@ -155,6 +166,8 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
|
|||||||
output = outputs[session.client_index].outputs[0]
|
output = outputs[session.client_index].outputs[0]
|
||||||
assert output.finish_reason is not None
|
assert output.finish_reason is not None
|
||||||
assert output.is_segment_finished is True
|
assert output.is_segment_finished is True
|
||||||
|
sched._handle_stopped_request.assert_called_once_with(session)
|
||||||
|
assert sched._pending_finish_reqs == []
|
||||||
|
|
||||||
|
|
||||||
def test_async_chunk_resumable_stop_rearms_connector_polling() -> None:
|
def test_async_chunk_resumable_stop_rearms_connector_polling() -> None:
|
||||||
|
|||||||
@@ -1,3 +1,6 @@
|
|||||||
|
import math
|
||||||
|
from typing import NamedTuple
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
@@ -317,6 +320,103 @@ def test_deferred_multimodal_cache_write_commits_on_completion():
|
|||||||
assert torch.equal(rows[8:16], expected)
|
assert torch.equal(rows[8:16], expected)
|
||||||
|
|
||||||
|
|
||||||
|
def test_discard_deferred_mm_outputs_drops_pending_chunks():
|
||||||
|
"""discard_deferred_mm_outputs is the abort path: a request that
|
||||||
|
leaves without a cache commit must not leak its staged chunks into
|
||||||
|
the cache if commit_deferred_mm_outputs is later called for it."""
|
||||||
|
cache = get_omni_pcache()
|
||||||
|
mm_key = "codes.audio"
|
||||||
|
block_table = torch.tensor(
|
||||||
|
[
|
||||||
|
[2, 3, 4, 5],
|
||||||
|
[6, 7, 8, 9],
|
||||||
|
],
|
||||||
|
dtype=torch.long,
|
||||||
|
)
|
||||||
|
input_batch = MockInputBatch(
|
||||||
|
num_computed_tokens_cpu=torch.tensor([4, 0], dtype=torch.long),
|
||||||
|
block_table=block_table,
|
||||||
|
)
|
||||||
|
|
||||||
|
cache.update_omni_tensor_prefix_cache(
|
||||||
|
hidden_states=None,
|
||||||
|
multimodal_outputs={mm_key: torch.arange(8, dtype=torch.long).reshape(4, 2)},
|
||||||
|
num_tokens_unpadded=4,
|
||||||
|
slot_mapping=torch.arange(8, 12),
|
||||||
|
skip_mm_cache_keys={mm_key},
|
||||||
|
)
|
||||||
|
|
||||||
|
chunk = torch.arange(8, dtype=torch.long).reshape(4, 2)
|
||||||
|
cache.stage_deferred_mm_outputs(
|
||||||
|
query_start_loc=torch.tensor([0, 4], dtype=torch.long),
|
||||||
|
input_batch=input_batch,
|
||||||
|
multimodal_outputs={mm_key: chunk},
|
||||||
|
num_scheduled_tokens={"req1": 4, "req2": 0},
|
||||||
|
deferred_mm_cache_keys={mm_key},
|
||||||
|
)
|
||||||
|
assert "req1" in cache._deferred_mm_outputs
|
||||||
|
|
||||||
|
cache.discard_deferred_mm_outputs("req1")
|
||||||
|
assert "req1" not in cache._deferred_mm_outputs
|
||||||
|
|
||||||
|
# A later commit call for the now-discarded req_id must be a no-op,
|
||||||
|
# not a resurrection of the dropped chunk.
|
||||||
|
cache.commit_deferred_mm_outputs({"req1"}, input_batch)
|
||||||
|
rows = cache.mm_outputs_cache[mm_key].view(-1, 2)
|
||||||
|
assert torch.all(rows[8:12] == 0)
|
||||||
|
|
||||||
|
|
||||||
|
def test_deferred_mm_outputs_request_id_reused_after_discard():
|
||||||
|
"""A request ID reused after discard_deferred_mm_outputs must start
|
||||||
|
clean: the new occupant's committed data must not be mixed with the
|
||||||
|
previous (discarded) occupant's staged chunks."""
|
||||||
|
cache = get_omni_pcache()
|
||||||
|
mm_key = "codes.audio"
|
||||||
|
block_table = torch.tensor(
|
||||||
|
[
|
||||||
|
[2, 3, 4, 5],
|
||||||
|
[6, 7, 8, 9],
|
||||||
|
],
|
||||||
|
dtype=torch.long,
|
||||||
|
)
|
||||||
|
input_batch = MockInputBatch(
|
||||||
|
num_computed_tokens_cpu=torch.tensor([4, 0], dtype=torch.long),
|
||||||
|
block_table=block_table,
|
||||||
|
)
|
||||||
|
|
||||||
|
cache.update_omni_tensor_prefix_cache(
|
||||||
|
hidden_states=None,
|
||||||
|
multimodal_outputs={mm_key: torch.arange(8, dtype=torch.long).reshape(4, 2)},
|
||||||
|
num_tokens_unpadded=4,
|
||||||
|
slot_mapping=torch.arange(8, 12),
|
||||||
|
skip_mm_cache_keys={mm_key},
|
||||||
|
)
|
||||||
|
|
||||||
|
stale_chunk = torch.full((4, 2), -1, dtype=torch.long)
|
||||||
|
cache.stage_deferred_mm_outputs(
|
||||||
|
query_start_loc=torch.tensor([0, 4], dtype=torch.long),
|
||||||
|
input_batch=input_batch,
|
||||||
|
multimodal_outputs={mm_key: stale_chunk},
|
||||||
|
num_scheduled_tokens={"req1": 4, "req2": 0},
|
||||||
|
deferred_mm_cache_keys={mm_key},
|
||||||
|
)
|
||||||
|
cache.discard_deferred_mm_outputs("req1")
|
||||||
|
|
||||||
|
# "req1" is reused by a brand-new request with different data.
|
||||||
|
fresh_chunk = torch.arange(8, dtype=torch.long).reshape(4, 2)
|
||||||
|
cache.stage_deferred_mm_outputs(
|
||||||
|
query_start_loc=torch.tensor([0, 4], dtype=torch.long),
|
||||||
|
input_batch=input_batch,
|
||||||
|
multimodal_outputs={mm_key: fresh_chunk},
|
||||||
|
num_scheduled_tokens={"req1": 4, "req2": 0},
|
||||||
|
deferred_mm_cache_keys={mm_key},
|
||||||
|
)
|
||||||
|
cache.commit_deferred_mm_outputs({"req1"}, input_batch)
|
||||||
|
|
||||||
|
rows = cache.mm_outputs_cache[mm_key].view(-1, 2)
|
||||||
|
assert torch.equal(rows[8:12], fresh_chunk)
|
||||||
|
|
||||||
|
|
||||||
def test_deferred_multimodal_cache_can_be_merged_on_full_block_hit():
|
def test_deferred_multimodal_cache_can_be_merged_on_full_block_hit():
|
||||||
cache = get_omni_pcache()
|
cache = get_omni_pcache()
|
||||||
mm_key = "codes.audio"
|
mm_key = "codes.audio"
|
||||||
@@ -599,3 +699,219 @@ def test_get_merged_multimodal_outputs(feat_dims, num_tokens_padded, mocker):
|
|||||||
|
|
||||||
assert req2_merged_mm_outputs.shape == torch.Size([num_new_toks_req2, curr_feat_dim])
|
assert req2_merged_mm_outputs.shape == torch.Size([num_new_toks_req2, curr_feat_dim])
|
||||||
assert torch.all(req2_merged_mm_outputs == req2_new_mm_outputs)
|
assert torch.all(req2_merged_mm_outputs == req2_new_mm_outputs)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"passthrough_shape",
|
||||||
|
[(5,), (5, 3)],
|
||||||
|
ids=["1d_token_metadata", "2d_feature_tensor"],
|
||||||
|
)
|
||||||
|
def test_get_merged_multimodal_outputs_slices_passthrough_tensor_per_request(passthrough_shape, mocker):
|
||||||
|
"""A passthrough key that is itself a real tensor (not the string/list
|
||||||
|
passthrough already covered above) must be sliced per request through
|
||||||
|
to_payload_element's tensor branch, under the same mixed cache
|
||||||
|
hit/miss + uneven scheduled-length batch as test_get_merged_multimodal_outputs.
|
||||||
|
|
||||||
|
Covers both a 1D token-aligned metadata tensor (e.g. ref_code_len,
|
||||||
|
codec_streaming -- shape (total_tokens,)) and a 2D per-token feature
|
||||||
|
tensor (shape (total_tokens, feat_dim)), since W2's batching could
|
||||||
|
easily special-case one and silently mishandle the other.
|
||||||
|
"""
|
||||||
|
feat_dims = {"foo": 100, "bar": 100}
|
||||||
|
cache = get_omni_pcache_with_mm_tensors(feat_dims, seq_len=DEFAULT_SEQ_LEN)
|
||||||
|
|
||||||
|
orig_num_tokens_unpadded = 8
|
||||||
|
slot_offset = 8
|
||||||
|
orig_slot_mapping = torch.arange(slot_offset, slot_offset + orig_num_tokens_unpadded)
|
||||||
|
feature_dims = {key: val.shape[-1] for key, val in cache.mm_outputs_cache.items()}
|
||||||
|
orig_mm_outputs = {
|
||||||
|
key: torch.rand((orig_num_tokens_unpadded, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys
|
||||||
|
}
|
||||||
|
|
||||||
|
cache.update_omni_tensor_prefix_cache(
|
||||||
|
hidden_states=None,
|
||||||
|
multimodal_outputs=orig_mm_outputs,
|
||||||
|
num_tokens_unpadded=orig_num_tokens_unpadded,
|
||||||
|
slot_mapping=orig_slot_mapping,
|
||||||
|
num_tokens_padded=None,
|
||||||
|
)
|
||||||
|
|
||||||
|
num_new_toks_req1 = 3
|
||||||
|
num_new_toks_req2 = 2
|
||||||
|
cache.add_prefix_cached_new_req_id("req1")
|
||||||
|
|
||||||
|
num_scheduled_tokens = {
|
||||||
|
"req1": num_new_toks_req1,
|
||||||
|
"req2": num_new_toks_req2,
|
||||||
|
}
|
||||||
|
|
||||||
|
new_mm_outputs = {}
|
||||||
|
for mm_key in cache.mm_cache_keys:
|
||||||
|
new_mm_outputs[mm_key] = torch.rand(
|
||||||
|
(num_new_toks_req1 + num_new_toks_req2, feature_dims[mm_key]),
|
||||||
|
dtype=DTYPE,
|
||||||
|
)
|
||||||
|
total_scheduled_tokens = num_new_toks_req1 + num_new_toks_req2
|
||||||
|
assert passthrough_shape[0] == total_scheduled_tokens
|
||||||
|
passthrough_tensor = torch.arange(int(torch.prod(torch.tensor(passthrough_shape))), dtype=DTYPE).reshape(
|
||||||
|
passthrough_shape
|
||||||
|
)
|
||||||
|
new_mm_outputs["passthrough_tensor"] = passthrough_tensor
|
||||||
|
|
||||||
|
input_batch = MockInputBatch(num_computed_tokens_cpu=torch.Tensor([orig_num_tokens_unpadded, 0]))
|
||||||
|
|
||||||
|
mocker.patch(
|
||||||
|
"vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids",
|
||||||
|
new=fake_get_cached_block_ids,
|
||||||
|
)
|
||||||
|
merged_mm_outputs = cache.get_merged_multimodal_states(
|
||||||
|
query_start_loc=[0, num_new_toks_req1],
|
||||||
|
input_batch=input_batch,
|
||||||
|
multimodal_outputs=new_mm_outputs,
|
||||||
|
num_scheduled_tokens=num_scheduled_tokens,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert "passthrough_tensor" not in cache.mm_cache_keys
|
||||||
|
|
||||||
|
passthrough_out = merged_mm_outputs["passthrough_tensor"]
|
||||||
|
# req1 (cache hit) gets tokens [0:3), req2 (cache miss) gets tokens [3:5) --
|
||||||
|
# a passthrough tensor is sliced by scheduled position, not merged with
|
||||||
|
# any cached history (unlike the mm_cache_keys tensors above).
|
||||||
|
assert torch.equal(passthrough_out["req1"], passthrough_tensor[0:num_new_toks_req1])
|
||||||
|
assert torch.equal(passthrough_out["req2"], passthrough_tensor[num_new_toks_req1:total_scheduled_tokens])
|
||||||
|
|
||||||
|
|
||||||
|
class MergeHarness(NamedTuple):
|
||||||
|
"""A primed cache plus the arguments get_merged_multimodal_states needs."""
|
||||||
|
|
||||||
|
cache: OmniTensorPrefixCache
|
||||||
|
input_batch: MockInputBatch
|
||||||
|
num_scheduled_tokens: dict[str, int]
|
||||||
|
query_start_loc: list[int]
|
||||||
|
total_scheduled_tokens: int
|
||||||
|
new_mm_outputs: dict[str, object]
|
||||||
|
|
||||||
|
|
||||||
|
def build_mixed_hit_miss_harness(mocker, num_new_toks_req1: int = 3, num_new_toks_req2: int = 2) -> MergeHarness:
|
||||||
|
"""Two requests over one cache key: req1 hits 8 cached tokens, req2 misses,
|
||||||
|
and their scheduled lengths differ -- the same batch shape as
|
||||||
|
test_get_merged_multimodal_outputs, which is what makes a per-request
|
||||||
|
boundary error observable. Callers add their passthrough keys to
|
||||||
|
``new_mm_outputs`` before calling get_merged_multimodal_states.
|
||||||
|
"""
|
||||||
|
cache = get_omni_pcache_with_mm_tensors({"foo": 100}, seq_len=DEFAULT_SEQ_LEN)
|
||||||
|
|
||||||
|
orig_num_tokens_unpadded = 8
|
||||||
|
feature_dims = {key: val.shape[-1] for key, val in cache.mm_outputs_cache.items()}
|
||||||
|
cache.update_omni_tensor_prefix_cache(
|
||||||
|
hidden_states=None,
|
||||||
|
multimodal_outputs={
|
||||||
|
key: torch.rand((orig_num_tokens_unpadded, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys
|
||||||
|
},
|
||||||
|
num_tokens_unpadded=orig_num_tokens_unpadded,
|
||||||
|
slot_mapping=torch.arange(8, 8 + orig_num_tokens_unpadded),
|
||||||
|
num_tokens_padded=None,
|
||||||
|
)
|
||||||
|
cache.add_prefix_cached_new_req_id("req1")
|
||||||
|
|
||||||
|
total_scheduled_tokens = num_new_toks_req1 + num_new_toks_req2
|
||||||
|
mocker.patch(
|
||||||
|
"vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids",
|
||||||
|
new=fake_get_cached_block_ids,
|
||||||
|
)
|
||||||
|
return MergeHarness(
|
||||||
|
cache=cache,
|
||||||
|
input_batch=MockInputBatch(num_computed_tokens_cpu=torch.Tensor([orig_num_tokens_unpadded, 0])),
|
||||||
|
num_scheduled_tokens={"req1": num_new_toks_req1, "req2": num_new_toks_req2},
|
||||||
|
query_start_loc=[0, num_new_toks_req1],
|
||||||
|
total_scheduled_tokens=total_scheduled_tokens,
|
||||||
|
new_mm_outputs={
|
||||||
|
key: torch.rand((total_scheduled_tokens, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_merged_multimodal_outputs_passes_per_request_tensor_list_through(mocker):
|
||||||
|
"""The passthrough shape qwen3-tts actually emits: a nested ``codes``/``meta``
|
||||||
|
dict whose ``codes.ref`` holds one tensor per request
|
||||||
|
(qwen3_tts_talker.make_omni_output), carried through to_payload_element's
|
||||||
|
pass_lists_through branch beside 1D int metadata that IS sliced per request.
|
||||||
|
|
||||||
|
Each request must receive the list *whole*. Per-request indexing happens
|
||||||
|
later, in the runner's _unwrap_lists; a W2 that split the list here would
|
||||||
|
make that second index select a different request's reference codes. The
|
||||||
|
existing list coverage passes a list of strings, which cannot catch a
|
||||||
|
tensor-specific regression, and no other test covers an int dtype.
|
||||||
|
"""
|
||||||
|
harness = build_mixed_hit_miss_harness(mocker)
|
||||||
|
|
||||||
|
ref_codes = [torch.tensor([11, 12], dtype=torch.long), torch.tensor([21], dtype=torch.long)]
|
||||||
|
ref_code_len = torch.arange(harness.total_scheduled_tokens, dtype=torch.int32)
|
||||||
|
harness.new_mm_outputs["codes"] = {"ref": ref_codes}
|
||||||
|
harness.new_mm_outputs["meta"] = {"ref_code_len": ref_code_len}
|
||||||
|
|
||||||
|
merged = harness.cache.get_merged_multimodal_states(
|
||||||
|
query_start_loc=harness.query_start_loc,
|
||||||
|
input_batch=harness.input_batch,
|
||||||
|
multimodal_outputs=harness.new_mm_outputs,
|
||||||
|
num_scheduled_tokens=harness.num_scheduled_tokens,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert "codes" not in harness.cache.mm_cache_keys
|
||||||
|
for req_id in ("req1", "req2"):
|
||||||
|
req_refs = merged["codes"][req_id]["ref"]
|
||||||
|
assert len(req_refs) == len(ref_codes)
|
||||||
|
for got, want in zip(req_refs, ref_codes):
|
||||||
|
assert torch.equal(got, want)
|
||||||
|
assert got.dtype == torch.long
|
||||||
|
|
||||||
|
# Per-request clones: the branch exists so one request's downstream mutation
|
||||||
|
# cannot reach another request's payload or the source list.
|
||||||
|
merged["codes"]["req1"]["ref"][0] += 100
|
||||||
|
assert torch.equal(merged["codes"]["req2"]["ref"][0], torch.tensor([11, 12], dtype=torch.long))
|
||||||
|
assert torch.equal(ref_codes[0], torch.tensor([11, 12], dtype=torch.long))
|
||||||
|
|
||||||
|
# Token-aligned 1D metadata is still split per request, unlike the list.
|
||||||
|
req1_len = harness.num_scheduled_tokens["req1"]
|
||||||
|
assert torch.equal(merged["meta"]["req1"]["ref_code_len"], ref_code_len[:req1_len])
|
||||||
|
assert torch.equal(merged["meta"]["req2"]["ref_code_len"], ref_code_len[req1_len:])
|
||||||
|
assert merged["meta"]["req1"]["ref_code_len"].dtype == torch.int32
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"passthrough_shape",
|
||||||
|
[(3, 4), (1, 4), (7,)],
|
||||||
|
ids=["shorter_2d", "single_row_2d", "longer_1d"],
|
||||||
|
)
|
||||||
|
def test_get_merged_multimodal_outputs_broadcasts_non_token_aligned_passthrough(passthrough_shape, mocker):
|
||||||
|
"""A passthrough tensor whose first dimension is not the scheduled-token
|
||||||
|
count is request-invariant: to_payload_element hands every request the whole
|
||||||
|
tensor, cloned, rather than slicing it.
|
||||||
|
|
||||||
|
W2 batching most plausibly breaks by assuming every passthrough tensor is
|
||||||
|
token-aligned; that assumption would slice these and give each request a
|
||||||
|
truncated view, or an empty one once start exceeds the first dimension.
|
||||||
|
"""
|
||||||
|
harness = build_mixed_hit_miss_harness(mocker)
|
||||||
|
assert passthrough_shape[0] != harness.total_scheduled_tokens
|
||||||
|
|
||||||
|
passthrough = torch.arange(math.prod(passthrough_shape), dtype=DTYPE).reshape(passthrough_shape)
|
||||||
|
expected = passthrough.clone()
|
||||||
|
harness.new_mm_outputs["speaker_embedding"] = passthrough
|
||||||
|
|
||||||
|
merged = harness.cache.get_merged_multimodal_states(
|
||||||
|
query_start_loc=harness.query_start_loc,
|
||||||
|
input_batch=harness.input_batch,
|
||||||
|
multimodal_outputs=harness.new_mm_outputs,
|
||||||
|
num_scheduled_tokens=harness.num_scheduled_tokens,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert "speaker_embedding" not in harness.cache.mm_cache_keys
|
||||||
|
out = merged["speaker_embedding"]
|
||||||
|
for req_id in ("req1", "req2"):
|
||||||
|
assert out[req_id].shape == passthrough_shape
|
||||||
|
assert torch.equal(out[req_id], expected)
|
||||||
|
|
||||||
|
out["req1"] += 100
|
||||||
|
assert torch.equal(out["req2"], expected)
|
||||||
|
assert torch.equal(passthrough, expected)
|
||||||
@@ -0,0 +1,119 @@
|
|||||||
|
import pytest
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from vllm_omni.core.prefix_cache import OmniTensorPrefixCache
|
||||||
|
|
||||||
|
# `cpu` is the lane selector, not a hardware claim: the only Buildkite steps that
|
||||||
|
# reach tests/ select `core_model and cpu`, and they run on GPU hardware (l4_1).
|
||||||
|
# Marking this `cuda` collects it nowhere. Real hardware need is the skipif below.
|
||||||
|
pytestmark = [
|
||||||
|
pytest.mark.core_model,
|
||||||
|
pytest.mark.cpu,
|
||||||
|
pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device"),
|
||||||
|
]
|
||||||
|
|
||||||
|
NUM_BLOCKS = 10
|
||||||
|
BLOCK_SIZE = 4
|
||||||
|
HIDDEN_SIZE = 2
|
||||||
|
DTYPE = torch.float32
|
||||||
|
|
||||||
|
|
||||||
|
def get_omni_pcache() -> OmniTensorPrefixCache:
|
||||||
|
"""Build an OmniTensorPrefixCache, but don't init mm tensors."""
|
||||||
|
return OmniTensorPrefixCache(
|
||||||
|
num_blocks=NUM_BLOCKS,
|
||||||
|
block_size=BLOCK_SIZE,
|
||||||
|
hidden_size=HIDDEN_SIZE,
|
||||||
|
hs_dtype=DTYPE,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_drain_ready_async_writes_is_a_noop_before_the_event_is_ready(monkeypatch):
|
||||||
|
"""drain_ready_async_writes() must not touch the cache while the
|
||||||
|
scheduled D2H's event has not completed, and must consume it exactly
|
||||||
|
once after it has.
|
||||||
|
|
||||||
|
Real GPU timing (e.g. via torch.cuda._sleep to stall the copy stream)
|
||||||
|
was tried and found unreliable on this box: cross-stream wait_stream/
|
||||||
|
wait_event did not reliably delay event.query() from reporting ready
|
||||||
|
immediately, likely a virtualization/driver quirk of this cloud GPU.
|
||||||
|
Forcing the event's query() result directly tests the branching logic
|
||||||
|
in drain_ready_async_writes/_consume_pending_write deterministically,
|
||||||
|
independent of real CUDA scheduling behavior.
|
||||||
|
"""
|
||||||
|
cache = get_omni_pcache()
|
||||||
|
mm_key = "codes.audio"
|
||||||
|
cache.maybe_init_missing_mm_cache_keys({mm_key: torch.zeros(4, HIDDEN_SIZE)}, seq_len=4)
|
||||||
|
|
||||||
|
num_tokens = 4
|
||||||
|
slot_mapping_gpu = torch.arange(8, 8 + num_tokens, dtype=torch.int64, device="cuda")
|
||||||
|
mm_gpu = {
|
||||||
|
mm_key: torch.arange(num_tokens * HIDDEN_SIZE, dtype=torch.float32, device="cuda").reshape(
|
||||||
|
num_tokens, HIDDEN_SIZE
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
cache.schedule_async_write(
|
||||||
|
hidden_states_gpu=None,
|
||||||
|
multimodal_outputs_gpu=mm_gpu,
|
||||||
|
slot_mapping_gpu=slot_mapping_gpu,
|
||||||
|
num_tokens_unpadded=num_tokens,
|
||||||
|
num_tokens_padded=num_tokens,
|
||||||
|
)
|
||||||
|
|
||||||
|
pending_event = cache._pending_write.event
|
||||||
|
rows = cache.mm_outputs_cache[mm_key].view(-1, HIDDEN_SIZE)
|
||||||
|
|
||||||
|
monkeypatch.setattr(pending_event, "query", lambda: False)
|
||||||
|
assert cache.drain_ready_async_writes() == 0
|
||||||
|
assert torch.all(rows[8 : 8 + num_tokens] == 0)
|
||||||
|
assert cache._pending_write is not None
|
||||||
|
|
||||||
|
monkeypatch.setattr(pending_event, "query", lambda: True)
|
||||||
|
assert cache.drain_ready_async_writes() == 1
|
||||||
|
assert torch.equal(rows[8 : 8 + num_tokens], mm_gpu[mm_key].cpu())
|
||||||
|
assert cache._pending_write is None
|
||||||
|
|
||||||
|
# Nothing pending anymore: a second drain is a no-op, not a re-scatter.
|
||||||
|
assert cache.drain_ready_async_writes() == 0
|
||||||
|
|
||||||
|
|
||||||
|
def test_schedule_async_write_consumes_previous_pending_write_first():
|
||||||
|
"""Scheduling step N+1's write must first consume (scatter) step N's
|
||||||
|
pending write, even if the caller never explicitly called
|
||||||
|
drain_ready_async_writes() in between -- otherwise step N's data would
|
||||||
|
be silently dropped when overwritten by step N+1's pending write."""
|
||||||
|
cache = get_omni_pcache()
|
||||||
|
mm_key = "codes.audio"
|
||||||
|
cache.maybe_init_missing_mm_cache_keys({mm_key: torch.zeros(4, HIDDEN_SIZE)}, seq_len=4)
|
||||||
|
|
||||||
|
num_tokens = 4
|
||||||
|
first_slots = torch.arange(8, 8 + num_tokens, dtype=torch.int64, device="cuda")
|
||||||
|
first_gpu = {mm_key: torch.full((num_tokens, HIDDEN_SIZE), 1.0, device="cuda")}
|
||||||
|
second_slots = torch.arange(12, 12 + num_tokens, dtype=torch.int64, device="cuda")
|
||||||
|
second_gpu = {mm_key: torch.full((num_tokens, HIDDEN_SIZE), 2.0, device="cuda")}
|
||||||
|
|
||||||
|
cache.schedule_async_write(
|
||||||
|
hidden_states_gpu=None,
|
||||||
|
multimodal_outputs_gpu=first_gpu,
|
||||||
|
slot_mapping_gpu=first_slots,
|
||||||
|
num_tokens_unpadded=num_tokens,
|
||||||
|
num_tokens_padded=num_tokens,
|
||||||
|
)
|
||||||
|
torch.accelerator.synchronize()
|
||||||
|
|
||||||
|
# This call's internal _consume_pending_write() must flush the FIRST
|
||||||
|
# write before the second one is scheduled; the caller never drained.
|
||||||
|
cache.schedule_async_write(
|
||||||
|
hidden_states_gpu=None,
|
||||||
|
multimodal_outputs_gpu=second_gpu,
|
||||||
|
slot_mapping_gpu=second_slots,
|
||||||
|
num_tokens_unpadded=num_tokens,
|
||||||
|
num_tokens_padded=num_tokens,
|
||||||
|
)
|
||||||
|
torch.accelerator.synchronize()
|
||||||
|
cache.drain_ready_async_writes()
|
||||||
|
|
||||||
|
rows = cache.mm_outputs_cache[mm_key].view(-1, HIDDEN_SIZE)
|
||||||
|
assert torch.all(rows[8:12] == 1.0)
|
||||||
|
assert torch.all(rows[12:16] == 2.0)
|
||||||
+8
-146
@@ -1,16 +1,15 @@
|
|||||||
|
import copy
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
import re
|
import re
|
||||||
import subprocess
|
import subprocess
|
||||||
import threading
|
import threading
|
||||||
import time
|
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from tests.dfx.reliability.helpers import list_remote_process_pids_by_pattern, post_chat_completions_raw
|
|
||||||
from tests.helpers.mark import hardware_marks
|
from tests.helpers.mark import hardware_marks
|
||||||
from tests.helpers.runtime import OmniServerParams
|
from tests.helpers.runtime import OmniServerParams
|
||||||
from tests.helpers.stage_config import modify_stage_config
|
from tests.helpers.stage_config import modify_stage_config
|
||||||
@@ -119,32 +118,6 @@ def create_unique_server_pytest_params(
|
|||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
def create_benchmark_pytest_params(
|
|
||||||
benchmark_configs: list[dict[str, Any]],
|
|
||||||
server_to_benchmark_mapping: dict[str, dict],
|
|
||||||
) -> list[Any]:
|
|
||||||
"""Like :func:`create_benchmark_indices`, but wrap each index in ``pytest.param`` with JSON marks."""
|
|
||||||
marks_by_name = _marks_by_test_name(benchmark_configs)
|
|
||||||
params: list[Any] = []
|
|
||||||
seen: set[str] = set()
|
|
||||||
for config in benchmark_configs:
|
|
||||||
test_name = config["test_name"]
|
|
||||||
if test_name in seen:
|
|
||||||
continue
|
|
||||||
seen.add(test_name)
|
|
||||||
params_list = get_benchmark_params_for_server(test_name, server_to_benchmark_mapping)
|
|
||||||
id_suffixes = _unique_benchmark_param_id_suffixes(params_list)
|
|
||||||
for idx, id_suffix in enumerate(id_suffixes):
|
|
||||||
params.append(
|
|
||||||
pytest.param(
|
|
||||||
(test_name, idx),
|
|
||||||
marks=marks_by_name.get(test_name, []),
|
|
||||||
id=f"{test_name}-{id_suffix}",
|
|
||||||
)
|
|
||||||
)
|
|
||||||
return params
|
|
||||||
|
|
||||||
|
|
||||||
def create_paired_benchmark_pytest_params(
|
def create_paired_benchmark_pytest_params(
|
||||||
server_entries: list[tuple[Any, str]],
|
server_entries: list[tuple[Any, str]],
|
||||||
params_by_test_name: dict[str, list[Any]],
|
params_by_test_name: dict[str, list[Any]],
|
||||||
@@ -357,10 +330,6 @@ def supports_video_generation(model_name: str) -> bool:
|
|||||||
return any(key in lower for key in ("wan", "video", "i2v", "t2v"))
|
return any(key in lower for key in ("wan", "video", "i2v", "t2v"))
|
||||||
|
|
||||||
|
|
||||||
def supports_chat_generation(model_name: str) -> bool:
|
|
||||||
return not supports_video_generation(model_name)
|
|
||||||
|
|
||||||
|
|
||||||
def parse_stage_devices(stage_config_path: str) -> str:
|
def parse_stage_devices(stage_config_path: str) -> str:
|
||||||
text = Path(stage_config_path).read_text(encoding="utf-8")
|
text = Path(stage_config_path).read_text(encoding="utf-8")
|
||||||
raw_devices: list[str] = re.findall(r"^\s*devices:\s*\"?([0-9,\s]+)\"?\s*$", text, flags=re.MULTILINE)
|
raw_devices: list[str] = re.findall(r"^\s*devices:\s*\"?([0-9,\s]+)\"?\s*$", text, flags=re.MULTILINE)
|
||||||
@@ -389,48 +358,6 @@ def assert_fault_exception(exc: Exception, error_keywords: tuple[str, ...]) -> N
|
|||||||
assert any(key in text for key in error_keywords), f"unexpected error under fault injection: {exc}"
|
assert any(key in text for key in error_keywords), f"unexpected error under fault injection: {exc}"
|
||||||
|
|
||||||
|
|
||||||
def wait_chat_request_ready(host: str, port: int, model: str, timeout_sec: int = 180) -> None:
|
|
||||||
"""Poll a minimal chat request until success."""
|
|
||||||
deadline = time.time() + timeout_sec
|
|
||||||
payload = json.dumps(
|
|
||||||
{
|
|
||||||
"model": model,
|
|
||||||
"messages": [{"role": "user", "content": "Say hello in one short sentence."}],
|
|
||||||
"stream": False,
|
|
||||||
"modalities": ["text"],
|
|
||||||
}
|
|
||||||
)
|
|
||||||
last_error: str | None = None
|
|
||||||
while time.time() < deadline:
|
|
||||||
try:
|
|
||||||
status, body = post_chat_completions_raw(host, port, payload)
|
|
||||||
if status == 200:
|
|
||||||
return
|
|
||||||
last_error = f"http={status} body={body[:200]!r}"
|
|
||||||
except Exception as exc: # noqa: BLE001
|
|
||||||
last_error = str(exc)
|
|
||||||
time.sleep(2)
|
|
||||||
raise TimeoutError(f"runtime-teardown warmup request did not succeed within {timeout_sec}s: {last_error}")
|
|
||||||
|
|
||||||
|
|
||||||
def assert_no_extra_worker_processes(
|
|
||||||
baseline_pids: set[int],
|
|
||||||
worker_pattern: str,
|
|
||||||
timeout_sec: int = 60,
|
|
||||||
) -> None:
|
|
||||||
"""Ensure no extra worker PIDs remain after teardown."""
|
|
||||||
deadline = time.time() + timeout_sec
|
|
||||||
while time.time() < deadline:
|
|
||||||
current = set(list_remote_process_pids_by_pattern(worker_pattern))
|
|
||||||
extra = current - baseline_pids
|
|
||||||
if not extra:
|
|
||||||
return
|
|
||||||
time.sleep(2)
|
|
||||||
current = set(list_remote_process_pids_by_pattern(worker_pattern))
|
|
||||||
extra = sorted(current - baseline_pids)
|
|
||||||
assert not extra, f"orphan worker processes remain after container teardown: {extra}"
|
|
||||||
|
|
||||||
|
|
||||||
def create_test_parameter_mapping(configs: list[dict[str, Any]]) -> dict[str, dict]:
|
def create_test_parameter_mapping(configs: list[dict[str, Any]]) -> dict[str, dict]:
|
||||||
mapping = {}
|
mapping = {}
|
||||||
for config in configs:
|
for config in configs:
|
||||||
@@ -637,57 +564,6 @@ def extract_configs_resource_label(configs: list[dict[str, Any]]) -> str:
|
|||||||
return get_runtime_resource_label()
|
return get_runtime_resource_label()
|
||||||
|
|
||||||
|
|
||||||
def resolve_baseline_value(
|
|
||||||
baseline_raw: Any,
|
|
||||||
*,
|
|
||||||
sweep_index: int | None,
|
|
||||||
max_concurrency: Any = None,
|
|
||||||
request_rate: Any = None,
|
|
||||||
) -> Any:
|
|
||||||
"""Pick the baseline threshold for this sweep step."""
|
|
||||||
if baseline_raw is None:
|
|
||||||
return 100000
|
|
||||||
if isinstance(baseline_raw, dict):
|
|
||||||
if max_concurrency is not None:
|
|
||||||
for key in (max_concurrency, str(max_concurrency)):
|
|
||||||
if key in baseline_raw:
|
|
||||||
return baseline_raw[key]
|
|
||||||
if request_rate is not None:
|
|
||||||
for key in (request_rate, str(request_rate)):
|
|
||||||
if key in baseline_raw:
|
|
||||||
return baseline_raw[key]
|
|
||||||
raise KeyError(
|
|
||||||
f"baseline dict has no key for max_concurrency={max_concurrency!r} "
|
|
||||||
f"or request_rate={request_rate!r}; keys={list(baseline_raw.keys())!r}"
|
|
||||||
)
|
|
||||||
if isinstance(baseline_raw, (list, tuple)):
|
|
||||||
if sweep_index is None:
|
|
||||||
raise ValueError("list baseline requires sweep_index")
|
|
||||||
if not (0 <= sweep_index < len(baseline_raw)):
|
|
||||||
raise IndexError(f"baseline list len={len(baseline_raw)} has no index {sweep_index}")
|
|
||||||
return baseline_raw[sweep_index]
|
|
||||||
return baseline_raw
|
|
||||||
|
|
||||||
|
|
||||||
def _baseline_thresholds_for_step(
|
|
||||||
baseline_data: dict[str, Any],
|
|
||||||
*,
|
|
||||||
sweep_index: int | None = None,
|
|
||||||
max_concurrency: Any = None,
|
|
||||||
request_rate: Any = None,
|
|
||||||
) -> dict[str, Any]:
|
|
||||||
"""Resolve baseline config to one threshold per metric for this iteration."""
|
|
||||||
return {
|
|
||||||
metric_name: resolve_baseline_value(
|
|
||||||
baseline_raw,
|
|
||||||
sweep_index=sweep_index,
|
|
||||||
max_concurrency=max_concurrency,
|
|
||||||
request_rate=request_rate,
|
|
||||||
)
|
|
||||||
for metric_name, baseline_raw in baseline_data.items()
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
def run_benchmark(
|
def run_benchmark(
|
||||||
args: list[str],
|
args: list[str],
|
||||||
test_name: str,
|
test_name: str,
|
||||||
@@ -696,18 +572,17 @@ def run_benchmark(
|
|||||||
num_prompt: int,
|
num_prompt: int,
|
||||||
*,
|
*,
|
||||||
baseline_config: dict[str, Any] | None = None,
|
baseline_config: dict[str, Any] | None = None,
|
||||||
sweep_index: int | None = None,
|
|
||||||
request_rate: Any | None = None,
|
|
||||||
max_concurrency: Any | None = None,
|
|
||||||
random_input_len: Any | None = None,
|
random_input_len: Any | None = None,
|
||||||
random_output_len: Any | None = None,
|
random_output_len: Any | None = None,
|
||||||
resource_label: str | None = None,
|
resource_label: str | None = None,
|
||||||
) -> dict[str, Any]:
|
) -> dict[str, Any]:
|
||||||
"""Run one ``vllm bench serve --omni`` iteration and return parsed metrics.
|
"""Run one ``vllm bench serve --omni`` iteration and return parsed metrics.
|
||||||
|
|
||||||
After ``vllm bench`` writes the JSON, ``result["baseline"]`` holds the resolved per-metric thresholds
|
After ``vllm bench`` writes the JSON, ``result["baseline"]`` stores the raw
|
||||||
(when ``baseline_config`` is provided). If the benchmark exits without writing a result file,
|
``baseline`` object from the test config (when ``baseline_config`` is provided),
|
||||||
``result_omni_template.json`` is used as a fallback.
|
including any hardware-nested map (e.g. ``{"H100": {...}}``). Downstream report
|
||||||
|
tooling can select a hardware bucket when needed. If the benchmark exits without
|
||||||
|
writing a result file, ``result_omni_template.json`` is used as a fallback.
|
||||||
"""
|
"""
|
||||||
current_dt = datetime.now().strftime("%Y%m%d-%H%M%S")
|
current_dt = datetime.now().strftime("%Y%m%d-%H%M%S")
|
||||||
ri = _safe_filename_token(random_input_len)
|
ri = _safe_filename_token(random_input_len)
|
||||||
@@ -774,12 +649,8 @@ def run_benchmark(
|
|||||||
result = json.load(f)
|
result = json.load(f)
|
||||||
|
|
||||||
if baseline_config:
|
if baseline_config:
|
||||||
result["baseline"] = _baseline_thresholds_for_step(
|
# Persist the JSON baseline as-is (no hardware / sweep resolution).
|
||||||
baseline_config,
|
result["baseline"] = copy.deepcopy(baseline_config)
|
||||||
sweep_index=sweep_index,
|
|
||||||
request_rate=request_rate,
|
|
||||||
max_concurrency=max_concurrency,
|
|
||||||
)
|
|
||||||
else:
|
else:
|
||||||
result["baseline"] = {}
|
result["baseline"] = {}
|
||||||
if random_input_len is not None:
|
if random_input_len is not None:
|
||||||
@@ -800,12 +671,3 @@ def pytest_addoption(parser: pytest.Parser) -> None:
|
|||||||
default=None,
|
default=None,
|
||||||
help=("Path to benchmark config JSON. Example: --test-config-file tests/dfx/perf/tests/test_tts.json"),
|
help=("Path to benchmark config JSON. Example: --test-config-file tests/dfx/perf/tests/test_tts.json"),
|
||||||
)
|
)
|
||||||
parser.addoption(
|
|
||||||
"--assert-baseline",
|
|
||||||
action="store_true",
|
|
||||||
default=False,
|
|
||||||
help=(
|
|
||||||
"When set, omni/diffusion perf runners compare metrics against the baseline block in the JSON config "
|
|
||||||
"(default: off)."
|
|
||||||
),
|
|
||||||
)
|
|
||||||
@@ -15,15 +15,10 @@ from tests.dfx.conftest import (
|
|||||||
get_runtime_resource_label,
|
get_runtime_resource_label,
|
||||||
is_diffusion_perf_config,
|
is_diffusion_perf_config,
|
||||||
load_benchmark_configs,
|
load_benchmark_configs,
|
||||||
resolve_baseline_value,
|
|
||||||
run_benchmark,
|
run_benchmark,
|
||||||
)
|
)
|
||||||
from tests.helpers.runtime import OmniServer
|
from tests.helpers.runtime import OmniServer
|
||||||
|
|
||||||
# Compare metrics to each test JSON ``baseline`` block only when pytest is run with ``--assert-baseline``
|
|
||||||
# (registered in ``tests/dfx/conftest.py``; default: off). ``run_benchmark`` and ``_resolve_baseline_value`` are
|
|
||||||
# defined in the same module.
|
|
||||||
#
|
|
||||||
# Optional JSON field ``mark`` is applied as pytest marks via
|
# Optional JSON field ``mark`` is applied as pytest marks via
|
||||||
# ``create_paired_omni_benchmark_pytest_params`` (e.g. ``"mark": [{"hardware_marks":
|
# ``create_paired_omni_benchmark_pytest_params`` (e.g. ``"mark": [{"hardware_marks":
|
||||||
# {"res": {"cuda": "H100"}, "num_cards": 2}}, "full_model", "omni"]``).
|
# {"res": {"cuda": "H100"}, "num_cards": 2}}, "full_model", "omni"]``).
|
||||||
@@ -162,36 +157,8 @@ def benchmark_params(request):
|
|||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
def assert_result(
|
def assert_result(result, num_prompt) -> None:
|
||||||
result,
|
|
||||||
params,
|
|
||||||
num_prompt,
|
|
||||||
*,
|
|
||||||
assert_baseline: bool,
|
|
||||||
sweep_index: int | None = None,
|
|
||||||
max_concurrency: Any | None = None,
|
|
||||||
request_rate: Any | None = None,
|
|
||||||
) -> None:
|
|
||||||
assert result["completed"] == num_prompt, "Request failures exist"
|
assert result["completed"] == num_prompt, "Request failures exist"
|
||||||
if not assert_baseline:
|
|
||||||
return
|
|
||||||
baseline_data = params.get("baseline", {})
|
|
||||||
for metric_name, baseline_raw in baseline_data.items():
|
|
||||||
current_value = result[metric_name]
|
|
||||||
baseline_value = resolve_baseline_value(
|
|
||||||
baseline_raw,
|
|
||||||
sweep_index=sweep_index,
|
|
||||||
max_concurrency=max_concurrency,
|
|
||||||
request_rate=request_rate,
|
|
||||||
)
|
|
||||||
if "throughput" in metric_name:
|
|
||||||
if current_value <= baseline_value:
|
|
||||||
print(
|
|
||||||
f"ERROR: Throughput test results were below baseline: {metric_name}: {current_value} > {baseline_value}"
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
if current_value >= baseline_value:
|
|
||||||
print(f"ERROR: Test results exceeded baseline: {metric_name}: {current_value} < {baseline_value}")
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.benchmark
|
@pytest.mark.benchmark
|
||||||
@@ -200,7 +167,7 @@ def assert_result(
|
|||||||
paired_benchmark_params,
|
paired_benchmark_params,
|
||||||
indirect=["omni_server", "benchmark_params"],
|
indirect=["omni_server", "benchmark_params"],
|
||||||
)
|
)
|
||||||
def test_performance_benchmark(omni_server, benchmark_params, request):
|
def test_performance_benchmark(omni_server, benchmark_params):
|
||||||
test_name = benchmark_params["test_name"]
|
test_name = benchmark_params["test_name"]
|
||||||
params = benchmark_params["params"]
|
params = benchmark_params["params"]
|
||||||
dataset_name = params.get("dataset_name", "")
|
dataset_name = params.get("dataset_name", "")
|
||||||
@@ -212,7 +179,6 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
|
|||||||
print(f"Running benchmark for model: {model}")
|
print(f"Running benchmark for model: {model}")
|
||||||
print(f"Benchmark parameters: {benchmark_params}")
|
print(f"Benchmark parameters: {benchmark_params}")
|
||||||
|
|
||||||
assert_baseline = request.config.getoption("--assert-baseline", default=False)
|
|
||||||
resource_label = get_runtime_resource_label()
|
resource_label = get_runtime_resource_label()
|
||||||
|
|
||||||
def to_list(value, default=None):
|
def to_list(value, default=None):
|
||||||
@@ -271,7 +237,7 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
|
|||||||
break
|
break
|
||||||
|
|
||||||
# QPS / request-rate sweep
|
# QPS / request-rate sweep
|
||||||
for i, (qps, num_prompt) in enumerate(zip(qps_list, num_prompt_list)):
|
for qps, num_prompt in zip(qps_list, num_prompt_list):
|
||||||
args = args + ["--request-rate", str(qps), "--num-prompts", str(num_prompt)]
|
args = args + ["--request-rate", str(qps), "--num-prompts", str(num_prompt)]
|
||||||
result = run_benchmark(
|
result = run_benchmark(
|
||||||
args=args,
|
args=args,
|
||||||
@@ -280,24 +246,14 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
|
|||||||
dataset_name=dataset_name,
|
dataset_name=dataset_name,
|
||||||
num_prompt=num_prompt,
|
num_prompt=num_prompt,
|
||||||
baseline_config=params.get("baseline"),
|
baseline_config=params.get("baseline"),
|
||||||
sweep_index=i,
|
|
||||||
request_rate=qps,
|
|
||||||
max_concurrency=None,
|
|
||||||
random_input_len=params.get("random_input_len"),
|
random_input_len=params.get("random_input_len"),
|
||||||
random_output_len=params.get("random_output_len"),
|
random_output_len=params.get("random_output_len"),
|
||||||
resource_label=resource_label,
|
resource_label=resource_label,
|
||||||
)
|
)
|
||||||
assert_result(
|
assert_result(result, num_prompt)
|
||||||
result,
|
|
||||||
params,
|
|
||||||
num_prompt,
|
|
||||||
assert_baseline=assert_baseline,
|
|
||||||
sweep_index=i,
|
|
||||||
request_rate=qps,
|
|
||||||
)
|
|
||||||
|
|
||||||
# concurrency test
|
# concurrency test
|
||||||
for i, (concurrency, num_prompt) in enumerate(zip(max_concurrency_list, num_prompt_list)):
|
for concurrency, num_prompt in zip(max_concurrency_list, num_prompt_list):
|
||||||
args = args + ["--max-concurrency", str(concurrency), "--num-prompts", str(num_prompt), "--request-rate", "inf"]
|
args = args + ["--max-concurrency", str(concurrency), "--num-prompts", str(num_prompt), "--request-rate", "inf"]
|
||||||
result = run_benchmark(
|
result = run_benchmark(
|
||||||
args=args,
|
args=args,
|
||||||
@@ -306,18 +262,8 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
|
|||||||
dataset_name=dataset_name,
|
dataset_name=dataset_name,
|
||||||
num_prompt=num_prompt,
|
num_prompt=num_prompt,
|
||||||
baseline_config=params.get("baseline"),
|
baseline_config=params.get("baseline"),
|
||||||
sweep_index=i,
|
|
||||||
request_rate=None,
|
|
||||||
max_concurrency=concurrency,
|
|
||||||
random_input_len=params.get("random_input_len"),
|
random_input_len=params.get("random_input_len"),
|
||||||
random_output_len=params.get("random_output_len"),
|
random_output_len=params.get("random_output_len"),
|
||||||
resource_label=resource_label,
|
resource_label=resource_label,
|
||||||
)
|
)
|
||||||
assert_result(
|
assert_result(result, num_prompt)
|
||||||
result,
|
|
||||||
params,
|
|
||||||
num_prompt,
|
|
||||||
assert_baseline=assert_baseline,
|
|
||||||
sweep_index=i,
|
|
||||||
max_concurrency=concurrency,
|
|
||||||
)
|
|
||||||
@@ -13,8 +13,6 @@ A config JSON file may be passed via --test-config-file. If omitted, every ``*.j
|
|||||||
pytest run_diffusion_benchmark.py -m "diffusion"
|
pytest run_diffusion_benchmark.py -m "diffusion"
|
||||||
pytest run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
|
pytest run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
|
||||||
|
|
||||||
Optional: ``--assert-baseline`` compares metrics to the ``baseline`` block in each benchmark entry (default: off).
|
|
||||||
|
|
||||||
Optional JSON field ``mark`` is applied as pytest marks on that case via
|
Optional JSON field ``mark`` is applied as pytest marks on that case via
|
||||||
``pytest.param`` (e.g. ``"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}}, "full_model", "diffusion"]``).
|
``pytest.param`` (e.g. ``"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}}, "full_model", "diffusion"]``).
|
||||||
|
|
||||||
@@ -25,6 +23,7 @@ of all runs from cases in that file). Bulk load without ``--test-config-file`` u
|
|||||||
the same per-file aggregation; ``-m`` only selects which cases run.
|
the same per-file aggregation; ``-m`` only selects which cases run.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
import copy
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
import socket
|
import socket
|
||||||
@@ -274,11 +273,6 @@ def _write_result_record(record: dict[str, Any]) -> Path:
|
|||||||
return target
|
return target
|
||||||
|
|
||||||
|
|
||||||
def _append_to_aggregated_file(record: dict[str, Any]) -> None:
|
|
||||||
"""Backward-compatible wrapper; prefer :func:`_write_result_record`."""
|
|
||||||
_write_result_record(record)
|
|
||||||
|
|
||||||
|
|
||||||
_server_lock = threading.Lock()
|
_server_lock = threading.Lock()
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
@@ -678,7 +672,7 @@ def run_benchmark(
|
|||||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", prefix="diffusion_bench_tmp_", delete=False) as tmp:
|
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", prefix="diffusion_bench_tmp_", delete=False) as tmp:
|
||||||
tmp_result_file = Path(tmp.name)
|
tmp_result_file = Path(tmp.name)
|
||||||
|
|
||||||
exclude_keys = {"baseline", "dataset", "task", "name", "skip-performance-assertion"}
|
exclude_keys = {"baseline", "dataset", "task", "name"}
|
||||||
|
|
||||||
cmd = [
|
cmd = [
|
||||||
sys.executable,
|
sys.executable,
|
||||||
@@ -772,12 +766,17 @@ def run_benchmark(
|
|||||||
completed = metrics.get("completed_requests", metrics.get("completed", 0))
|
completed = metrics.get("completed_requests", metrics.get("completed", 0))
|
||||||
failed = metrics.get("failed_requests", metrics.get("failed", 0))
|
failed = metrics.get("failed_requests", metrics.get("failed", 0))
|
||||||
|
|
||||||
|
# Persist the JSON baseline as-is (no hardware / sweep resolution).
|
||||||
|
baseline = copy.deepcopy(params.get("baseline") or {})
|
||||||
|
metrics["baseline"] = baseline
|
||||||
|
|
||||||
record: dict[str, Any] = {
|
record: dict[str, Any] = {
|
||||||
"test_name": test_name,
|
"test_name": test_name,
|
||||||
"endpoint": endpoint,
|
"endpoint": endpoint,
|
||||||
"timestamp": timestamp,
|
"timestamp": timestamp,
|
||||||
"server_params": server_cfg.get("server_params"),
|
"server_params": server_cfg.get("server_params"),
|
||||||
"benchmark_params": params,
|
"benchmark_params": params,
|
||||||
|
"baseline": baseline,
|
||||||
"result": metrics,
|
"result": metrics,
|
||||||
"log_file": str(log_file),
|
"log_file": str(log_file),
|
||||||
"Model": model,
|
"Model": model,
|
||||||
@@ -827,70 +826,14 @@ def run_benchmark(
|
|||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
def _resolve_baseline_value(
|
|
||||||
baseline_raw: Any,
|
|
||||||
*,
|
|
||||||
sweep_index: int | None,
|
|
||||||
max_concurrency: Any = None,
|
|
||||||
request_rate: Any = None,
|
|
||||||
) -> Any:
|
|
||||||
"""Pick the baseline threshold for this sweep step."""
|
|
||||||
if baseline_raw is None:
|
|
||||||
return 100000
|
|
||||||
if isinstance(baseline_raw, dict):
|
|
||||||
if max_concurrency is not None:
|
|
||||||
for key in (max_concurrency, str(max_concurrency)):
|
|
||||||
if key in baseline_raw:
|
|
||||||
return baseline_raw[key]
|
|
||||||
if request_rate is not None:
|
|
||||||
for key in (request_rate, str(request_rate)):
|
|
||||||
if key in baseline_raw:
|
|
||||||
return baseline_raw[key]
|
|
||||||
raise KeyError(
|
|
||||||
f"baseline dict has no key for max_concurrency={max_concurrency!r} "
|
|
||||||
f"or request_rate={request_rate!r}; keys={list(baseline_raw.keys())!r}"
|
|
||||||
)
|
|
||||||
if isinstance(baseline_raw, (list, tuple)):
|
|
||||||
if sweep_index is None:
|
|
||||||
raise ValueError("sweep_index is required when baseline is a list or tuple")
|
|
||||||
return baseline_raw[sweep_index]
|
|
||||||
return baseline_raw
|
|
||||||
|
|
||||||
|
|
||||||
def assert_result(
|
def assert_result(
|
||||||
result: dict[str, Any],
|
result: dict[str, Any],
|
||||||
params: dict[str, Any],
|
|
||||||
num_prompts: int,
|
num_prompts: int,
|
||||||
*,
|
|
||||||
sweep_index: int | None = None,
|
|
||||||
max_concurrency: Any = None,
|
|
||||||
request_rate: Any = None,
|
|
||||||
assert_baseline: bool = True,
|
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Assert that benchmark metrics satisfy the configured baselines."""
|
"""Assert that the expected number of requests completed."""
|
||||||
completed = result.get("completed_requests", result.get("completed", 0))
|
completed = result.get("completed_requests", result.get("completed", 0))
|
||||||
assert completed == num_prompts, f"Expected {num_prompts} completed requests, got {completed}"
|
assert completed == num_prompts, f"Expected {num_prompts} completed requests, got {completed}"
|
||||||
|
|
||||||
if not assert_baseline:
|
|
||||||
return
|
|
||||||
if params.get("skip-performance-assertion", False):
|
|
||||||
print("Skipping performance assertions.")
|
|
||||||
return
|
|
||||||
|
|
||||||
for metric, baseline_raw in params.get("baseline", {}).items():
|
|
||||||
current = result.get(metric)
|
|
||||||
assert current is not None, f"Metric '{metric}' not found in result: {list(result.keys())}"
|
|
||||||
threshold = _resolve_baseline_value(
|
|
||||||
baseline_raw,
|
|
||||||
sweep_index=sweep_index,
|
|
||||||
max_concurrency=max_concurrency,
|
|
||||||
request_rate=request_rate,
|
|
||||||
)
|
|
||||||
if "throughput" in metric:
|
|
||||||
assert current >= threshold, f"{metric}: {current:.4f} < baseline {threshold}"
|
|
||||||
else:
|
|
||||||
assert current <= threshold, f"{metric}: {current:.4f} > baseline {threshold}"
|
|
||||||
|
|
||||||
|
|
||||||
def _default_benchmark_endpoint_for_task(task: str) -> str:
|
def _default_benchmark_endpoint_for_task(task: str) -> str:
|
||||||
"""Return the default client-side benchmark endpoint for a diffusion task."""
|
"""Return the default client-side benchmark endpoint for a diffusion task."""
|
||||||
@@ -919,7 +862,6 @@ def _build_run_params(
|
|||||||
params: dict[str, Any],
|
params: dict[str, Any],
|
||||||
*,
|
*,
|
||||||
num_prompts: int,
|
num_prompts: int,
|
||||||
sweep_index: int | None = None,
|
|
||||||
request_rate: Any | None = None,
|
request_rate: Any | None = None,
|
||||||
max_concurrency: Any | None = None,
|
max_concurrency: Any | None = None,
|
||||||
) -> dict[str, Any]:
|
) -> dict[str, Any]:
|
||||||
@@ -931,16 +873,6 @@ def _build_run_params(
|
|||||||
run_params["request-rate"] = request_rate
|
run_params["request-rate"] = request_rate
|
||||||
if max_concurrency is not None:
|
if max_concurrency is not None:
|
||||||
run_params["max-concurrency"] = max_concurrency
|
run_params["max-concurrency"] = max_concurrency
|
||||||
if "baseline" in params:
|
|
||||||
run_params["baseline"] = {
|
|
||||||
metric: _resolve_baseline_value(
|
|
||||||
baseline_raw,
|
|
||||||
sweep_index=sweep_index,
|
|
||||||
max_concurrency=max_concurrency,
|
|
||||||
request_rate=request_rate,
|
|
||||||
)
|
|
||||||
for metric, baseline_raw in params["baseline"].items()
|
|
||||||
}
|
|
||||||
return run_params
|
return run_params
|
||||||
|
|
||||||
|
|
||||||
@@ -963,36 +895,28 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
|
|||||||
|
|
||||||
sweep_runs: list[dict[str, Any]] = []
|
sweep_runs: list[dict[str, Any]] = []
|
||||||
|
|
||||||
for i, (request_rate, num_prompts) in enumerate(zip(request_rate_list, num_prompt_list)):
|
for request_rate, num_prompts in zip(request_rate_list, num_prompt_list):
|
||||||
sweep_runs.append(
|
sweep_runs.append(
|
||||||
{
|
{
|
||||||
"params": _build_run_params(
|
"params": _build_run_params(
|
||||||
params,
|
params,
|
||||||
request_rate=request_rate,
|
request_rate=request_rate,
|
||||||
num_prompts=num_prompts,
|
num_prompts=num_prompts,
|
||||||
sweep_index=i,
|
|
||||||
),
|
),
|
||||||
"num_prompts": num_prompts,
|
"num_prompts": num_prompts,
|
||||||
"sweep_index": i,
|
|
||||||
"request_rate": request_rate,
|
|
||||||
"max_concurrency": None,
|
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
|
|
||||||
for i, (max_concurrency, num_prompts) in enumerate(zip(max_concurrency_list, num_prompt_list)):
|
for max_concurrency, num_prompts in zip(max_concurrency_list, num_prompt_list):
|
||||||
sweep_runs.append(
|
sweep_runs.append(
|
||||||
{
|
{
|
||||||
"params": _build_run_params(
|
"params": _build_run_params(
|
||||||
params,
|
params,
|
||||||
max_concurrency=max_concurrency,
|
max_concurrency=max_concurrency,
|
||||||
num_prompts=num_prompts,
|
num_prompts=num_prompts,
|
||||||
sweep_index=i,
|
|
||||||
request_rate="inf",
|
request_rate="inf",
|
||||||
),
|
),
|
||||||
"num_prompts": num_prompts,
|
"num_prompts": num_prompts,
|
||||||
"sweep_index": i,
|
|
||||||
"request_rate": None,
|
|
||||||
"max_concurrency": max_concurrency,
|
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -1005,9 +929,6 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
|
|||||||
num_prompts=default_num_prompts,
|
num_prompts=default_num_prompts,
|
||||||
),
|
),
|
||||||
"num_prompts": default_num_prompts,
|
"num_prompts": default_num_prompts,
|
||||||
"sweep_index": None,
|
|
||||||
"request_rate": None,
|
|
||||||
"max_concurrency": None,
|
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -1023,14 +944,12 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
|
|||||||
paired_benchmark_params,
|
paired_benchmark_params,
|
||||||
indirect=["diffusion_server", "benchmark_params"],
|
indirect=["diffusion_server", "benchmark_params"],
|
||||||
)
|
)
|
||||||
def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, request):
|
def test_diffusion_performance_benchmark(diffusion_server, benchmark_params):
|
||||||
"""Run the diffusion performance benchmark and verify request completion.
|
"""Run the diffusion performance benchmark and verify request completion.
|
||||||
|
|
||||||
One server is started per unique parallel configuration (module scope).
|
One server is started per unique parallel configuration (module scope).
|
||||||
For each server, all benchmark parameter sets defined in the config JSON
|
For each server, all benchmark parameter sets defined in the config JSON
|
||||||
are executed sequentially; metrics are recorded to the aggregated result file.
|
are executed sequentially; metrics are recorded to the aggregated result file.
|
||||||
|
|
||||||
Pass ``--assert-baseline`` to compare ``throughput_qps``, ``latency_mean``, etc. to the JSON ``baseline`` block.
|
|
||||||
"""
|
"""
|
||||||
test_name = benchmark_params["test_name"]
|
test_name = benchmark_params["test_name"]
|
||||||
params = benchmark_params["params"]
|
params = benchmark_params["params"]
|
||||||
@@ -1073,14 +992,4 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, req
|
|||||||
print(f"\n Aggregated results: {_aggregated_result_file_for_source(str(source))}")
|
print(f"\n Aggregated results: {_aggregated_result_file_for_source(str(source))}")
|
||||||
print("=" * 60)
|
print("=" * 60)
|
||||||
|
|
||||||
assert_baseline = request.config.getoption("--assert-baseline", default=False)
|
assert_result(result, sweep_run["num_prompts"])
|
||||||
|
|
||||||
assert_result(
|
|
||||||
result,
|
|
||||||
params,
|
|
||||||
sweep_run["num_prompts"],
|
|
||||||
sweep_index=sweep_run["sweep_index"],
|
|
||||||
max_concurrency=sweep_run["max_concurrency"],
|
|
||||||
request_rate=sweep_run["request_rate"],
|
|
||||||
assert_baseline=assert_baseline,
|
|
||||||
)
|
|
||||||
@@ -40,10 +40,12 @@
|
|||||||
"max-concurrency": 1,
|
"max-concurrency": 1,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.495,
|
"H100": {
|
||||||
"latency_mean": 2.0238,
|
"throughput_qps": 0.495,
|
||||||
"peak_memory_mb_max": 29393.4286,
|
"latency_mean": 2.0238,
|
||||||
"peak_memory_mb_mean": 29392.7143
|
"peak_memory_mb_max": 29393.4286,
|
||||||
|
"peak_memory_mb_mean": 29392.7143
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
@@ -90,10 +92,12 @@
|
|||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"num-input-images": 1,
|
"num-input-images": 1,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.3451,
|
"H100": {
|
||||||
"latency_mean": 2.9031,
|
"throughput_qps": 0.3451,
|
||||||
"peak_memory_mb_max": 30480.2857,
|
"latency_mean": 2.9031,
|
||||||
"peak_memory_mb_mean": 30479.2857
|
"peak_memory_mb_max": 30480.2857,
|
||||||
|
"peak_memory_mb_mean": 30479.2857
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
@@ -139,10 +143,12 @@
|
|||||||
"max-concurrency": 1,
|
"max-concurrency": 1,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.5048,
|
"H100": {
|
||||||
"latency_mean": 1.987,
|
"throughput_qps": 0.5048,
|
||||||
"peak_memory_mb_max": 29380.0,
|
"latency_mean": 1.987,
|
||||||
"peak_memory_mb_mean": 29380.0
|
"peak_memory_mb_max": 29380.0,
|
||||||
|
"peak_memory_mb_mean": 29380.0
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
@@ -189,10 +195,12 @@
|
|||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"num-input-images": 1,
|
"num-input-images": 1,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.149,
|
"H100": {
|
||||||
"latency_mean": 6.7157,
|
"throughput_qps": 0.149,
|
||||||
"peak_memory_mb_max": 30574.0,
|
"latency_mean": 6.7157,
|
||||||
"peak_memory_mb_mean": 30574.0
|
"peak_memory_mb_max": 30574.0,
|
||||||
|
"peak_memory_mb_mean": 30574.0
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -24,10 +24,12 @@
|
|||||||
"warmup-num-inference-steps": 28,
|
"warmup-num-inference-steps": 28,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": {"1": 0.0303, "2": 0.0302, "4": 0.0302},
|
"H100": {
|
||||||
"latency_mean": {"1": 32.6, "2": 62.2, "4": 111.7},
|
"throughput_qps": {"1": 0.0303, "2": 0.0302, "4": 0.0302},
|
||||||
"peak_memory_mb_max": 40500,
|
"latency_mean": {"1": 32.6, "2": 62.2, "4": 111.7},
|
||||||
"peak_memory_mb_mean": 40500
|
"peak_memory_mb_max": 40500,
|
||||||
|
"peak_memory_mb_mean": 40500
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -24,10 +24,12 @@
|
|||||||
"warmup-num-inference-steps": 28,
|
"warmup-num-inference-steps": 28,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": {"1": 0.0549, "2": 0.0547, "4": 0.0549},
|
"H100": {
|
||||||
"latency_mean": {"1": 18.0, "2": 34.3, "4": 61.4},
|
"throughput_qps": {"1": 0.0549, "2": 0.0547, "4": 0.0549},
|
||||||
"peak_memory_mb_max": 40500,
|
"latency_mean": {"1": 18.0, "2": 34.3, "4": 61.4},
|
||||||
"peak_memory_mb_mean": 40500
|
"peak_memory_mb_max": 40500,
|
||||||
|
"peak_memory_mb_mean": 40500
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -48,9 +48,11 @@
|
|||||||
"guardrails": false
|
"guardrails": false
|
||||||
},
|
},
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 1.107,
|
"H100": {
|
||||||
"latency_mean": 0.8929,
|
"throughput_qps": 1.107,
|
||||||
"peak_memory_mb_mean": 80000.0
|
"latency_mean": 0.8929,
|
||||||
|
"peak_memory_mb_mean": 80000.0
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
@@ -115,9 +117,11 @@
|
|||||||
"use_duration_template": false
|
"use_duration_template": false
|
||||||
},
|
},
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.027,
|
"H100": {
|
||||||
"latency_mean": 36.6824,
|
"throughput_qps": 0.027,
|
||||||
"peak_memory_mb_mean": 30155.4
|
"latency_mean": 36.6824,
|
||||||
|
"peak_memory_mb_mean": 30155.4
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
@@ -181,9 +185,11 @@
|
|||||||
"guardrails": false
|
"guardrails": false
|
||||||
},
|
},
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.027,
|
"H100": {
|
||||||
"latency_mean": 36.968,
|
"throughput_qps": 0.027,
|
||||||
"peak_memory_mb_mean": 30443.6
|
"latency_mean": 36.968,
|
||||||
|
"peak_memory_mb_mean": 30443.6
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
@@ -251,9 +257,11 @@
|
|||||||
"condition_video_keep": "first"
|
"condition_video_keep": "first"
|
||||||
},
|
},
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.027,
|
"H100": {
|
||||||
"latency_mean": 37.2117,
|
"throughput_qps": 0.027,
|
||||||
"peak_memory_mb_mean": 30135.6
|
"latency_mean": 37.2117,
|
||||||
|
"peak_memory_mb_mean": 30135.6
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -46,12 +46,14 @@
|
|||||||
"seed_tts_locale": "en",
|
"seed_tts_locale": "en",
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"median_audio_ttfp_ms": [
|
"H100": {
|
||||||
223.3584
|
"median_audio_ttfp_ms": [
|
||||||
],
|
223.3584
|
||||||
"median_audio_rtf": [
|
],
|
||||||
0.4393
|
"median_audio_rtf": [
|
||||||
]
|
0.4393
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -70,15 +72,17 @@
|
|||||||
"seed_tts_locale": "en",
|
"seed_tts_locale": "en",
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"median_audio_ttfp_ms": [
|
"H100": {
|
||||||
276.575
|
"median_audio_ttfp_ms": [
|
||||||
],
|
276.575
|
||||||
"median_audio_rtf": [
|
],
|
||||||
0.5073
|
"median_audio_rtf": [
|
||||||
],
|
0.5073
|
||||||
"audio_throughput": [
|
],
|
||||||
27.3728
|
"audio_throughput": [
|
||||||
]
|
27.3728
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -133,13 +133,14 @@
|
|||||||
"bot-task": "think_recaption",
|
"bot-task": "think_recaption",
|
||||||
"max-concurrency": 16,
|
"max-concurrency": 16,
|
||||||
"num-prompts": 96,
|
"num-prompts": 96,
|
||||||
"skip-performance-assertion": true,
|
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.1874,
|
"H100": {
|
||||||
"latency_p99": 101.852,
|
"throughput_qps": 0.1874,
|
||||||
"latency_mean": 79.0385,
|
"latency_p99": 101.852,
|
||||||
"peak_memory_mb_max": 70000,
|
"latency_mean": 79.0385,
|
||||||
"peak_memory_mb_mean": 70000
|
"peak_memory_mb_max": 70000,
|
||||||
|
"peak_memory_mb_mean": 70000
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -37,12 +37,13 @@
|
|||||||
"num-inference-steps": 8,
|
"num-inference-steps": 8,
|
||||||
"num-prompts": 10,
|
"num-prompts": 10,
|
||||||
"max-concurrency": 1,
|
"max-concurrency": 1,
|
||||||
"skip-performance-assertion": true,
|
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.21,
|
"H100": {
|
||||||
"latency_p99": 4.7469,
|
"throughput_qps": 0.21,
|
||||||
"peak_memory_mb_max": 101100,
|
"latency_p99": 4.7469,
|
||||||
"peak_memory_mb_mean": 101100
|
"peak_memory_mb_max": 101100,
|
||||||
|
"peak_memory_mb_mean": 101100
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -37,12 +37,13 @@
|
|||||||
"num-inference-steps": 8,
|
"num-inference-steps": 8,
|
||||||
"num-prompts": 10,
|
"num-prompts": 10,
|
||||||
"max-concurrency": 1,
|
"max-concurrency": 1,
|
||||||
"skip-performance-assertion": true,
|
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.20,
|
"H100": {
|
||||||
"latency_p99": 5.1025,
|
"throughput_qps": 0.20,
|
||||||
"peak_memory_mb_max": 97402,
|
"latency_p99": 5.1025,
|
||||||
"peak_memory_mb_mean": 97402
|
"peak_memory_mb_max": 97402,
|
||||||
|
"peak_memory_mb_mean": 97402
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -36,13 +36,14 @@
|
|||||||
"num-inference-steps": 8,
|
"num-inference-steps": 8,
|
||||||
"num-prompts": 10,
|
"num-prompts": 10,
|
||||||
"max-concurrency": 1,
|
"max-concurrency": 1,
|
||||||
"skip-performance-assertion": true,
|
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.213,
|
"H100": {
|
||||||
"latency_mean": 4.6953,
|
"throughput_qps": 0.213,
|
||||||
"latency_p99": 4.8601,
|
"latency_mean": 4.6953,
|
||||||
"peak_memory_mb_max": 56912.0,
|
"latency_p99": 4.8601,
|
||||||
"peak_memory_mb_mean": 56912.0
|
"peak_memory_mb_max": 56912.0,
|
||||||
|
"peak_memory_mb_mean": 56912.0
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -26,8 +26,7 @@
|
|||||||
"warmup-requests": 1,
|
"warmup-requests": 1,
|
||||||
"warmup-concurrency": 1,
|
"warmup-concurrency": 1,
|
||||||
"warmup-num-inference-steps": 2,
|
"warmup-num-inference-steps": 2,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true
|
||||||
"skip-performance-assertion": true
|
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -44,27 +44,29 @@
|
|||||||
"ignore_eos": true,
|
"ignore_eos": true,
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [
|
"H100": {
|
||||||
96.4061,
|
"mean_ttft_ms": [
|
||||||
140.8147,
|
96.4061,
|
||||||
271.9025,
|
140.8147,
|
||||||
362.2916,
|
271.9025,
|
||||||
507.776
|
362.2916,
|
||||||
],
|
507.776
|
||||||
"mean_e2el_ms": [
|
],
|
||||||
18507.4582,
|
"mean_e2el_ms": [
|
||||||
28365.4956,
|
18507.4582,
|
||||||
31907.9613,
|
28365.4956,
|
||||||
48161.5829,
|
31907.9613,
|
||||||
72630.8601
|
48161.5829,
|
||||||
],
|
72630.8601
|
||||||
"mean_audio_rtf": [
|
],
|
||||||
0.1175,
|
"mean_audio_rtf": [
|
||||||
0.18,
|
0.1175,
|
||||||
0.2359,
|
0.18,
|
||||||
0.3278,
|
0.2359,
|
||||||
0.6296
|
0.3278,
|
||||||
]
|
0.6296
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -91,15 +93,17 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [
|
"H100": {
|
||||||
354.0951
|
"mean_ttft_ms": [
|
||||||
],
|
354.0951
|
||||||
"mean_e2el_ms": [
|
],
|
||||||
3952.5615
|
"mean_e2el_ms": [
|
||||||
],
|
3952.5615
|
||||||
"mean_audio_rtf": [
|
],
|
||||||
0.117
|
"mean_audio_rtf": [
|
||||||
]
|
0.117
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -128,15 +132,17 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [
|
"H100": {
|
||||||
259.6317
|
"mean_ttft_ms": [
|
||||||
],
|
259.6317
|
||||||
"mean_e2el_ms": [
|
],
|
||||||
4253.3009
|
"mean_e2el_ms": [
|
||||||
],
|
4253.3009
|
||||||
"mean_audio_rtf": [
|
],
|
||||||
0.1448
|
"mean_audio_rtf": [
|
||||||
]
|
0.1448
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -167,15 +173,17 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [
|
"H100": {
|
||||||
457.4453
|
"mean_ttft_ms": [
|
||||||
],
|
457.4453
|
||||||
"mean_e2el_ms": [
|
],
|
||||||
5912.5623
|
"mean_e2el_ms": [
|
||||||
],
|
5912.5623
|
||||||
"mean_audio_rtf": [
|
],
|
||||||
0.1967
|
"mean_audio_rtf": [
|
||||||
]
|
0.1967
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -206,20 +214,22 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [
|
"H100": {
|
||||||
90.8033,
|
"mean_ttft_ms": [
|
||||||
202.1779,
|
90.8033,
|
||||||
326.5864,
|
202.1779,
|
||||||
544.8914,
|
326.5864,
|
||||||
961.3929
|
544.8914,
|
||||||
],
|
961.3929
|
||||||
"mean_e2el_ms": [
|
],
|
||||||
4517.3411,
|
"mean_e2el_ms": [
|
||||||
7857.9711,
|
4517.3411,
|
||||||
9861.1687,
|
7857.9711,
|
||||||
13283.3363,
|
9861.1687,
|
||||||
19541.0736
|
13283.3363,
|
||||||
]
|
19541.0736
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -251,12 +261,14 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [
|
"H100": {
|
||||||
318.048
|
"mean_ttft_ms": [
|
||||||
],
|
318.048
|
||||||
"mean_e2el_ms": [
|
],
|
||||||
807.6114
|
"mean_e2el_ms": [
|
||||||
]
|
807.6114
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -290,12 +302,14 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [
|
"H100": {
|
||||||
186.168
|
"mean_ttft_ms": [
|
||||||
],
|
186.168
|
||||||
"mean_e2el_ms": [
|
],
|
||||||
726.1137
|
"mean_e2el_ms": [
|
||||||
]
|
726.1137
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -331,12 +345,14 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [
|
"H100": {
|
||||||
375.5282
|
"mean_ttft_ms": [
|
||||||
],
|
375.5282
|
||||||
"mean_e2el_ms": [
|
],
|
||||||
1006.1663
|
"mean_e2el_ms": [
|
||||||
]
|
1006.1663
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -35,10 +35,12 @@
|
|||||||
"ignore_eos": true,
|
"ignore_eos": true,
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [130.5751, 278.6469, 324.7249, 911.6085],
|
"H100": {
|
||||||
"mean_tpot_ms": [7.9854, 11.1209, 16.1507, 20.7485],
|
"mean_ttft_ms": [130.5751, 278.6469, 324.7249, 911.6085],
|
||||||
"mean_audio_ttfp_ms": [486.486, 954.1791, 1289.6514, 2196.8275],
|
"mean_tpot_ms": [7.9854, 11.1209, 16.1507, 20.7485],
|
||||||
"mean_audio_rtf": [0.1896, 0.3962, 0.4098, 0.7322]
|
"mean_audio_ttfp_ms": [486.486, 954.1791, 1289.6514, 2196.8275],
|
||||||
|
"mean_audio_rtf": [0.1896, 0.3962, 0.4098, 0.7322]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -61,9 +63,11 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [340.5122],
|
"H100": {
|
||||||
"mean_audio_ttfp_ms": [463.0752],
|
"mean_ttft_ms": [340.5122],
|
||||||
"mean_audio_rtf": [0.1153]
|
"mean_audio_ttfp_ms": [463.0752],
|
||||||
|
"mean_audio_rtf": [0.1153]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -88,9 +92,11 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [236.926],
|
"H100": {
|
||||||
"mean_audio_ttfp_ms": [395.2182],
|
"mean_ttft_ms": [236.926],
|
||||||
"mean_audio_rtf": [0.1259]
|
"mean_audio_ttfp_ms": [395.2182],
|
||||||
|
"mean_audio_rtf": [0.1259]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -117,9 +123,11 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [470.6126],
|
"H100": {
|
||||||
"mean_audio_ttfp_ms": [665.5474],
|
"mean_ttft_ms": [470.6126],
|
||||||
"mean_audio_rtf": [0.166]
|
"mean_audio_ttfp_ms": [665.5474],
|
||||||
|
"mean_audio_rtf": [0.166]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -44,27 +44,29 @@
|
|||||||
"ignore_eos": true,
|
"ignore_eos": true,
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [
|
"H100": {
|
||||||
115.8348,
|
"mean_ttft_ms": [
|
||||||
173.1823,
|
115.8348,
|
||||||
446.672,
|
173.1823,
|
||||||
464.5985,
|
446.672,
|
||||||
3420.0634
|
464.5985,
|
||||||
],
|
3420.0634
|
||||||
"mean_e2el_ms": [
|
],
|
||||||
25070.9523,
|
"mean_e2el_ms": [
|
||||||
31992.0646,
|
25070.9523,
|
||||||
37619.7669,
|
31992.0646,
|
||||||
50393.8257,
|
37619.7669,
|
||||||
136619.2743
|
50393.8257,
|
||||||
],
|
136619.2743
|
||||||
"mean_audio_rtf": [
|
],
|
||||||
0.169,
|
"mean_audio_rtf": [
|
||||||
0.2111,
|
0.169,
|
||||||
0.2475,
|
0.2111,
|
||||||
0.3551,
|
0.2475,
|
||||||
0.9266
|
0.3551,
|
||||||
]
|
0.9266
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -91,15 +93,17 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [
|
"H100": {
|
||||||
343.1377
|
"mean_ttft_ms": [
|
||||||
],
|
343.1377
|
||||||
"mean_e2el_ms": [
|
],
|
||||||
4811.1123
|
"mean_e2el_ms": [
|
||||||
],
|
4811.1123
|
||||||
"mean_audio_rtf": [
|
],
|
||||||
0.1463
|
"mean_audio_rtf": [
|
||||||
]
|
0.1463
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -128,15 +132,17 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [
|
"H100": {
|
||||||
245.5446
|
"mean_ttft_ms": [
|
||||||
],
|
245.5446
|
||||||
"mean_e2el_ms": [
|
],
|
||||||
5019.1256
|
"mean_e2el_ms": [
|
||||||
],
|
5019.1256
|
||||||
"mean_audio_rtf": [
|
],
|
||||||
0.1713
|
"mean_audio_rtf": [
|
||||||
]
|
0.1713
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -167,15 +173,17 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [
|
"H100": {
|
||||||
506.1407
|
"mean_ttft_ms": [
|
||||||
],
|
506.1407
|
||||||
"mean_e2el_ms": [
|
],
|
||||||
6814.9656
|
"mean_e2el_ms": [
|
||||||
],
|
6814.9656
|
||||||
"mean_audio_rtf": [
|
],
|
||||||
0.2213
|
"mean_audio_rtf": [
|
||||||
]
|
0.2213
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -51,20 +51,22 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [
|
"H100": {
|
||||||
69.8081,
|
"mean_ttft_ms": [
|
||||||
173.0651,
|
69.8081,
|
||||||
271.8108,
|
173.0651,
|
||||||
413.4079,
|
271.8108,
|
||||||
497.2546
|
413.4079,
|
||||||
],
|
497.2546
|
||||||
"mean_e2el_ms": [
|
],
|
||||||
4720.8225,
|
"mean_e2el_ms": [
|
||||||
7937.2532,
|
4720.8225,
|
||||||
9962.3566,
|
7937.2532,
|
||||||
13724.5479,
|
9962.3566,
|
||||||
18324.091
|
13724.5479,
|
||||||
]
|
18324.091
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -97,12 +99,14 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [
|
"H100": {
|
||||||
25.9477
|
"mean_ttft_ms": [
|
||||||
],
|
25.9477
|
||||||
"mean_e2el_ms": [
|
],
|
||||||
517.3155
|
"mean_e2el_ms": [
|
||||||
]
|
517.3155
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -137,12 +141,14 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [
|
"H100": {
|
||||||
26.7056
|
"mean_ttft_ms": [
|
||||||
],
|
26.7056
|
||||||
"mean_e2el_ms": [
|
],
|
||||||
551.308
|
"mean_e2el_ms": [
|
||||||
]
|
551.308
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -179,12 +185,14 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_ttft_ms": [
|
"H100": {
|
||||||
27.6906
|
"mean_ttft_ms": [
|
||||||
],
|
27.6906
|
||||||
"mean_e2el_ms": [
|
],
|
||||||
596.5284
|
"mean_e2el_ms": [
|
||||||
]
|
596.5284
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -34,10 +34,12 @@
|
|||||||
"num-input-images": 2,
|
"num-input-images": 2,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.0721,
|
"H100": {
|
||||||
"latency_mean": 13.8398,
|
"throughput_qps": 0.0721,
|
||||||
"peak_memory_mb_max": 57610.0,
|
"latency_mean": 13.8398,
|
||||||
"peak_memory_mb_mean": 57610.0
|
"peak_memory_mb_max": 57610.0,
|
||||||
|
"peak_memory_mb_mean": 57610.0
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -52,10 +54,12 @@
|
|||||||
"num-input-images": 2,
|
"num-input-images": 2,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.0186,
|
"H100": {
|
||||||
"latency_mean": 53.6942,
|
"throughput_qps": 0.0186,
|
||||||
"peak_memory_mb_max": 67790.0,
|
"latency_mean": 53.6942,
|
||||||
"peak_memory_mb_mean": 67790.0
|
"peak_memory_mb_max": 67790.0,
|
||||||
|
"peak_memory_mb_mean": 67790.0
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
@@ -100,10 +104,12 @@
|
|||||||
"num-input-images": 2,
|
"num-input-images": 2,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.0538,
|
"H100": {
|
||||||
"latency_mean": 18.5591,
|
"throughput_qps": 0.0538,
|
||||||
"peak_memory_mb_max": 56449.4286,
|
"latency_mean": 18.5591,
|
||||||
"peak_memory_mb_mean": 56449.4286
|
"peak_memory_mb_max": 56449.4286,
|
||||||
|
"peak_memory_mb_mean": 56449.4286
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
@@ -157,10 +163,12 @@
|
|||||||
"num-input-images": 2,
|
"num-input-images": 2,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.2776,
|
"H100": {
|
||||||
"latency_mean": 3.5938,
|
"throughput_qps": 0.2776,
|
||||||
"peak_memory_mb_max": 58113.7143,
|
"latency_mean": 3.5938,
|
||||||
"peak_memory_mb_mean": 58075.1429
|
"peak_memory_mb_max": 58113.7143,
|
||||||
|
"peak_memory_mb_mean": 58075.1429
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -175,10 +183,12 @@
|
|||||||
"num-input-images": 2,
|
"num-input-images": 2,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.1051,
|
"H100": {
|
||||||
"latency_mean": 9.4918,
|
"throughput_qps": 0.1051,
|
||||||
"peak_memory_mb_max": 68107.4286,
|
"latency_mean": 9.4918,
|
||||||
"peak_memory_mb_mean": 68107.4286
|
"peak_memory_mb_max": 68107.4286,
|
||||||
|
"peak_memory_mb_mean": 68107.4286
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -33,10 +33,12 @@
|
|||||||
"max-concurrency": 1,
|
"max-concurrency": 1,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.1129,
|
"H100": {
|
||||||
"latency_mean": 8.8352,
|
"throughput_qps": 0.1129,
|
||||||
"peak_memory_mb_max": 58572.0,
|
"latency_mean": 8.8352,
|
||||||
"peak_memory_mb_mean": 58572.0
|
"peak_memory_mb_max": 58572.0,
|
||||||
|
"peak_memory_mb_mean": 58572.0
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -50,10 +52,12 @@
|
|||||||
"max-concurrency": 1,
|
"max-concurrency": 1,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.024,
|
"H100": {
|
||||||
"latency_mean": 41.5686,
|
"throughput_qps": 0.024,
|
||||||
"peak_memory_mb_max": 66368.0,
|
"latency_mean": 41.5686,
|
||||||
"peak_memory_mb_mean": 66368.0
|
"peak_memory_mb_max": 66368.0,
|
||||||
|
"peak_memory_mb_mean": 66368.0
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
@@ -96,10 +100,12 @@
|
|||||||
"max-concurrency": 1,
|
"max-concurrency": 1,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.0672,
|
"H100": {
|
||||||
"latency_mean": 14.8517,
|
"throughput_qps": 0.0672,
|
||||||
"peak_memory_mb_max": 57076.0,
|
"latency_mean": 14.8517,
|
||||||
"peak_memory_mb_mean": 57076.0
|
"peak_memory_mb_max": 57076.0,
|
||||||
|
"peak_memory_mb_mean": 57076.0
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
@@ -152,10 +158,12 @@
|
|||||||
"max-concurrency": 1,
|
"max-concurrency": 1,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.2879,
|
"H100": {
|
||||||
"latency_mean": 3.4701,
|
"throughput_qps": 0.2879,
|
||||||
"peak_memory_mb_max": 58758.0,
|
"latency_mean": 3.4701,
|
||||||
"peak_memory_mb_mean": 58758.0
|
"peak_memory_mb_max": 58758.0,
|
||||||
|
"peak_memory_mb_mean": 58758.0
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -169,10 +177,12 @@
|
|||||||
"max-concurrency": 1,
|
"max-concurrency": 1,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.1113,
|
"H100": {
|
||||||
"latency_mean": 8.9656,
|
"throughput_qps": 0.1113,
|
||||||
"peak_memory_mb_max": 67424.0,
|
"latency_mean": 8.9656,
|
||||||
"peak_memory_mb_mean": 67424.0
|
"peak_memory_mb_max": 67424.0,
|
||||||
|
"peak_memory_mb_mean": 67424.0
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -33,10 +33,12 @@
|
|||||||
"max-concurrency": 1,
|
"max-concurrency": 1,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.0703,
|
"H100": {
|
||||||
"latency_mean": 14.208,
|
"throughput_qps": 0.0703,
|
||||||
"peak_memory_mb_max": 63225.1429,
|
"latency_mean": 14.208,
|
||||||
"peak_memory_mb_mean": 63225.1429
|
"peak_memory_mb_max": 63225.1429,
|
||||||
|
"peak_memory_mb_mean": 63225.1429
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -50,10 +52,12 @@
|
|||||||
"max-concurrency": 1,
|
"max-concurrency": 1,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.0417,
|
"H100": {
|
||||||
"latency_mean": 23.914,
|
"throughput_qps": 0.0417,
|
||||||
"peak_memory_mb_max": 63225.1429,
|
"latency_mean": 23.914,
|
||||||
"peak_memory_mb_mean": 63225.1429
|
"peak_memory_mb_max": 63225.1429,
|
||||||
|
"peak_memory_mb_mean": 63225.1429
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -33,9 +33,11 @@
|
|||||||
"max-concurrency": 1,
|
"max-concurrency": 1,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.4153,
|
"H100": {
|
||||||
"latency_mean": 2.4065,
|
"throughput_qps": 0.4153,
|
||||||
"peak_memory_mb_mean": 55121.4286
|
"latency_mean": 2.4065,
|
||||||
|
"peak_memory_mb_mean": 55121.4286
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -49,9 +51,11 @@
|
|||||||
"max-concurrency": 1,
|
"max-concurrency": 1,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.0423,
|
"H100": {
|
||||||
"latency_mean": 23.5991,
|
"throughput_qps": 0.0423,
|
||||||
"peak_memory_mb_mean": 65118.0
|
"latency_mean": 23.5991,
|
||||||
|
"peak_memory_mb_mean": 65118.0
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
@@ -91,9 +95,11 @@
|
|||||||
"max-concurrency": 1,
|
"max-concurrency": 1,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.4175,
|
"H100": {
|
||||||
"latency_mean": 2.3953,
|
"throughput_qps": 0.4175,
|
||||||
"peak_memory_mb_mean": 55088.5714
|
"latency_mean": 2.3953,
|
||||||
|
"peak_memory_mb_mean": 55088.5714
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -107,9 +113,11 @@
|
|||||||
"max-concurrency": 1,
|
"max-concurrency": 1,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.0423,
|
"H100": {
|
||||||
"latency_mean": 23.5768,
|
"throughput_qps": 0.0423,
|
||||||
"peak_memory_mb_mean": 66012.2857
|
"latency_mean": 23.5768,
|
||||||
|
"peak_memory_mb_mean": 66012.2857
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
@@ -152,9 +160,11 @@
|
|||||||
"max-concurrency": 1,
|
"max-concurrency": 1,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.1214,
|
"H100": {
|
||||||
"latency_mean": 8.2211,
|
"throughput_qps": 0.1214,
|
||||||
"peak_memory_mb_mean": 54501.1429
|
"latency_mean": 8.2211,
|
||||||
|
"peak_memory_mb_mean": 54501.1429
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
@@ -207,9 +217,11 @@
|
|||||||
"max-concurrency": 1,
|
"max-concurrency": 1,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.559,
|
"H100": {
|
||||||
"latency_mean": 1.794,
|
"throughput_qps": 0.559,
|
||||||
"peak_memory_mb_mean": 55290.5714
|
"latency_mean": 1.794,
|
||||||
|
"peak_memory_mb_mean": 55290.5714
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -223,9 +235,11 @@
|
|||||||
"max-concurrency": 1,
|
"max-concurrency": 1,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.1918,
|
"H100": {
|
||||||
"latency_mean": 5.2022,
|
"throughput_qps": 0.1918,
|
||||||
"peak_memory_mb_mean": 64838.8571
|
"latency_mean": 5.2022,
|
||||||
|
"peak_memory_mb_mean": 64838.8571
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
@@ -269,9 +283,11 @@
|
|||||||
"warmup-num-inference-steps": 20,
|
"warmup-num-inference-steps": 20,
|
||||||
"enable-negative-prompt": true,
|
"enable-negative-prompt": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": [0.4125, 0.6714, 0.6674, 0.8159],
|
"H100": {
|
||||||
"latency_mean": [2.4215, 2.992, 5.9293, 9.656],
|
"throughput_qps": [0.4125, 0.6714, 0.6674, 0.8159],
|
||||||
"peak_memory_mb_mean": [55147.4286, 55148.0, 55150.4643, 55152.3821]
|
"latency_mean": [2.4215, 2.992, 5.9293, 9.656],
|
||||||
|
"peak_memory_mb_mean": [55147.4286, 55148.0, 55150.4643, 55152.3821]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
+137
-111
@@ -39,12 +39,14 @@
|
|||||||
"seed_tts_locale": "en",
|
"seed_tts_locale": "en",
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"median_audio_ttfp_ms": [
|
"H100": {
|
||||||
141.2791
|
"median_audio_ttfp_ms": [
|
||||||
],
|
141.2791
|
||||||
"median_audio_rtf": [
|
],
|
||||||
0.1591
|
"median_audio_rtf": [
|
||||||
]
|
0.1591
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -67,21 +69,23 @@
|
|||||||
"seed_tts_locale": "en",
|
"seed_tts_locale": "en",
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"median_audio_ttfp_ms": [
|
"H100": {
|
||||||
226.747,
|
"median_audio_ttfp_ms": [
|
||||||
647.1272,
|
226.747,
|
||||||
4777.2714
|
647.1272,
|
||||||
],
|
4777.2714
|
||||||
"median_audio_rtf": [
|
],
|
||||||
0.3313,
|
"median_audio_rtf": [
|
||||||
0.3642,
|
0.3313,
|
||||||
1.2282
|
0.3642,
|
||||||
],
|
1.2282
|
||||||
"audio_throughput": [
|
],
|
||||||
24.2296,
|
"audio_throughput": [
|
||||||
40.3552,
|
24.2296,
|
||||||
43.2572
|
40.3552,
|
||||||
]
|
43.2572
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -102,9 +106,11 @@
|
|||||||
"seed_tts_wer_eval": true,
|
"seed_tts_wer_eval": true,
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_audio_rtf": [
|
"H100": {
|
||||||
0.45
|
"mean_audio_rtf": [
|
||||||
]
|
0.45
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
@@ -153,12 +159,14 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"median_audio_ttfp_ms": [
|
"H100": {
|
||||||
46.8434
|
"median_audio_ttfp_ms": [
|
||||||
],
|
46.8434
|
||||||
"median_audio_rtf": [
|
],
|
||||||
0.1367
|
"median_audio_rtf": [
|
||||||
]
|
0.1367
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -186,21 +194,23 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"median_audio_ttfp_ms": [
|
"H100": {
|
||||||
75.3341,
|
"median_audio_ttfp_ms": [
|
||||||
713.7727,
|
75.3341,
|
||||||
5941.7179
|
713.7727,
|
||||||
],
|
5941.7179
|
||||||
"median_audio_rtf": [
|
],
|
||||||
0.1859,
|
"median_audio_rtf": [
|
||||||
0.3148,
|
0.1859,
|
||||||
1.1662
|
0.3148,
|
||||||
],
|
1.1662
|
||||||
"audio_throughput": [
|
],
|
||||||
38.3858,
|
"audio_throughput": [
|
||||||
55.3523,
|
38.3858,
|
||||||
68.7032
|
55.3523,
|
||||||
]
|
68.7032
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -226,9 +236,11 @@
|
|||||||
"seed_tts_wer_eval": true,
|
"seed_tts_wer_eval": true,
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_audio_rtf": [
|
"H100": {
|
||||||
0.35
|
"mean_audio_rtf": [
|
||||||
]
|
0.35
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -252,15 +264,17 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"median_audio_ttfp_ms": [
|
"H100": {
|
||||||
56.746
|
"median_audio_ttfp_ms": [
|
||||||
],
|
56.746
|
||||||
"median_audio_rtf": [
|
],
|
||||||
0.1571
|
"median_audio_rtf": [
|
||||||
],
|
0.1571
|
||||||
"audio_throughput": [
|
],
|
||||||
11.1226
|
"audio_throughput": [
|
||||||
]
|
11.1226
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -283,12 +297,14 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"median_audio_ttfp_ms": [
|
"H100": {
|
||||||
46.3479
|
"median_audio_ttfp_ms": [
|
||||||
],
|
46.3479
|
||||||
"median_audio_rtf": [
|
],
|
||||||
0.1352
|
"median_audio_rtf": [
|
||||||
]
|
0.1352
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -315,21 +331,23 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"median_audio_ttfp_ms": [
|
"H100": {
|
||||||
75.5967,
|
"median_audio_ttfp_ms": [
|
||||||
587.5877,
|
75.5967,
|
||||||
4960.3726
|
587.5877,
|
||||||
],
|
4960.3726
|
||||||
"median_audio_rtf": [
|
],
|
||||||
0.1871,
|
"median_audio_rtf": [
|
||||||
0.3143,
|
0.1871,
|
||||||
1.1703
|
0.3143,
|
||||||
],
|
1.1703
|
||||||
"audio_throughput": [
|
],
|
||||||
38.9378,
|
"audio_throughput": [
|
||||||
55.1846,
|
38.9378,
|
||||||
71.5708
|
55.1846,
|
||||||
]
|
71.5708
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -352,15 +370,17 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"median_audio_ttfp_ms": [
|
"H100": {
|
||||||
55.5029
|
"median_audio_ttfp_ms": [
|
||||||
],
|
55.5029
|
||||||
"median_audio_rtf": [
|
],
|
||||||
0.1508
|
"median_audio_rtf": [
|
||||||
],
|
0.1508
|
||||||
"audio_throughput": [
|
],
|
||||||
9.4885
|
"audio_throughput": [
|
||||||
]
|
9.4885
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
@@ -402,12 +422,14 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"median_audio_ttfp_ms": [
|
"H100": {
|
||||||
82.1886
|
"median_audio_ttfp_ms": [
|
||||||
],
|
82.1886
|
||||||
"median_audio_rtf": [
|
],
|
||||||
0.1241
|
"median_audio_rtf": [
|
||||||
]
|
0.1241
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -429,15 +451,17 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"median_audio_ttfp_ms": [
|
"H100": {
|
||||||
160.9168
|
"median_audio_ttfp_ms": [
|
||||||
],
|
160.9168
|
||||||
"median_audio_rtf": [
|
],
|
||||||
0.1907
|
"median_audio_rtf": [
|
||||||
],
|
0.1907
|
||||||
"audio_throughput": [
|
],
|
||||||
39.2332
|
"audio_throughput": [
|
||||||
]
|
39.2332
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -459,15 +483,17 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"median_audio_ttfp_ms": [
|
"H100": {
|
||||||
100.1265
|
"median_audio_ttfp_ms": [
|
||||||
],
|
100.1265
|
||||||
"median_audio_rtf": [
|
],
|
||||||
0.1437
|
"median_audio_rtf": [
|
||||||
],
|
0.1437
|
||||||
"audio_throughput": [
|
],
|
||||||
11.7605
|
"audio_throughput": [
|
||||||
]
|
11.7605
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -37,12 +37,14 @@
|
|||||||
"seed_tts_locale": "en",
|
"seed_tts_locale": "en",
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"median_audio_ttfp_ms": [
|
"H100": {
|
||||||
137.645
|
"median_audio_ttfp_ms": [
|
||||||
],
|
137.645
|
||||||
"median_audio_rtf": [
|
],
|
||||||
0.0497
|
"median_audio_rtf": [
|
||||||
]
|
0.0497
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -61,15 +63,17 @@
|
|||||||
"seed_tts_locale": "en",
|
"seed_tts_locale": "en",
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"median_audio_ttfp_ms": [
|
"H100": {
|
||||||
475.3357
|
"median_audio_ttfp_ms": [
|
||||||
],
|
475.3357
|
||||||
"median_audio_rtf": [
|
],
|
||||||
0.1243
|
"median_audio_rtf": [
|
||||||
],
|
0.1243
|
||||||
"audio_throughput": [
|
],
|
||||||
36.6358
|
"audio_throughput": [
|
||||||
]
|
36.6358
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -90,9 +94,11 @@
|
|||||||
"seed_tts_wer_eval": true,
|
"seed_tts_wer_eval": true,
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_audio_rtf": [
|
"H100": {
|
||||||
0.45
|
"mean_audio_rtf": [
|
||||||
]
|
0.45
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -114,12 +120,14 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"median_audio_ttfp_ms": [
|
"H100": {
|
||||||
79.6927
|
"median_audio_ttfp_ms": [
|
||||||
],
|
79.6927
|
||||||
"median_audio_rtf": [
|
],
|
||||||
0.0423
|
"median_audio_rtf": [
|
||||||
]
|
0.0423
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -141,15 +149,17 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"median_audio_ttfp_ms": [
|
"H100": {
|
||||||
262.6212
|
"median_audio_ttfp_ms": [
|
||||||
],
|
262.6212
|
||||||
"median_audio_rtf": [
|
],
|
||||||
0.0753
|
"median_audio_rtf": [
|
||||||
],
|
0.0753
|
||||||
"audio_throughput": [
|
],
|
||||||
103.5843
|
"audio_throughput": [
|
||||||
]
|
103.5843
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -171,15 +181,17 @@
|
|||||||
},
|
},
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"median_audio_ttfp_ms": [
|
"H100": {
|
||||||
80.1095
|
"median_audio_ttfp_ms": [
|
||||||
],
|
80.1095
|
||||||
"median_audio_rtf": [
|
],
|
||||||
0.0479
|
"median_audio_rtf": [
|
||||||
],
|
0.0479
|
||||||
"audio_throughput": [
|
],
|
||||||
17.2918
|
"audio_throughput": [
|
||||||
]
|
17.2918
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -203,9 +215,11 @@
|
|||||||
"seed_tts_wer_eval": true,
|
"seed_tts_wer_eval": true,
|
||||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"mean_audio_rtf": [
|
"H100": {
|
||||||
0.35
|
"mean_audio_rtf": [
|
||||||
]
|
0.35
|
||||||
|
]
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -42,9 +42,11 @@
|
|||||||
}
|
}
|
||||||
],
|
],
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.0434,
|
"H100": {
|
||||||
"latency_mean": 23.0654,
|
"throughput_qps": 0.0434,
|
||||||
"peak_memory_mb_mean": 76360.0
|
"latency_mean": 23.0654,
|
||||||
|
"peak_memory_mb_mean": 76360.0
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
@@ -97,9 +99,11 @@
|
|||||||
}
|
}
|
||||||
],
|
],
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.0606,
|
"H100": {
|
||||||
"latency_mean": 16.4928,
|
"throughput_qps": 0.0606,
|
||||||
"peak_memory_mb_mean": 47164.8
|
"latency_mean": 16.4928,
|
||||||
|
"peak_memory_mb_mean": 47164.8
|
||||||
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -122,9 +126,11 @@
|
|||||||
}
|
}
|
||||||
],
|
],
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.0124,
|
"H100": {
|
||||||
"latency_mean": 80.7784,
|
"throughput_qps": 0.0124,
|
||||||
"peak_memory_mb_mean": 55374.3429
|
"latency_mean": 80.7784,
|
||||||
|
"peak_memory_mb_mean": 55374.3429
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
# SPDX-License-Identifier: Apache-2.0
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
"""Omni sleep/wakeup JSON validation (live ``Qwen/Qwen-Image`` process)."""
|
"""Omni sleep/wakeup JSON validation and level-2 wake contract (live ``Qwen/Qwen-Image``)."""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
@@ -15,9 +15,17 @@ pytestmark = [pytest.mark.slow, pytest.mark.diffusion, pytest.mark.omni]
|
|||||||
|
|
||||||
_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
|
_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
|
||||||
|
|
||||||
|
# One module-scoped server for this file. ``--enable-sleep-mode`` is required by the
|
||||||
|
# level-2 wake contract test; invalid-JSON cases only hit request validation and are
|
||||||
|
# unaffected. Do not mix ``omni_server`` (module) with ``omni_server_function`` here:
|
||||||
|
# ``iter_omni_server`` holds ``omni_fixture_lock`` across the module yield, so a
|
||||||
|
# function-scoped second server would deadlock / contend for the GPU.
|
||||||
_QWEN_IMAGE = [
|
_QWEN_IMAGE = [
|
||||||
pytest.param(
|
pytest.param(
|
||||||
OmniServerParams(model="Qwen/Qwen-Image"),
|
OmniServerParams(
|
||||||
|
model="Qwen/Qwen-Image",
|
||||||
|
server_args=["--enable-sleep-mode"],
|
||||||
|
),
|
||||||
id="qwen_image",
|
id="qwen_image",
|
||||||
marks=hardware_marks(res={"cuda": "H100"}),
|
marks=hardware_marks(res={"cuda": "H100"}),
|
||||||
),
|
),
|
||||||
@@ -93,3 +101,37 @@ def test_omni_sleep_wakeup_invalid_requests(
|
|||||||
openai_client.send_omni_sleep_http_request(cfg)
|
openai_client.send_omni_sleep_http_request(cfg)
|
||||||
else:
|
else:
|
||||||
openai_client.send_omni_wakeup_http_request(cfg)
|
openai_client.send_omni_wakeup_http_request(cfg)
|
||||||
|
|
||||||
|
|
||||||
|
# Keep this after the invalid-JSON cases: sleep(level=2) discards weights and the
|
||||||
|
# shared module server is not usable afterward; module teardown then stops it.
|
||||||
|
@pytest.mark.parametrize("omni_server", _QWEN_IMAGE, indirect=True)
|
||||||
|
def test_wakeup_after_level2_sleep_fails(
|
||||||
|
omni_server: OmniServer,
|
||||||
|
openai_client: OpenAIClientHandler,
|
||||||
|
) -> None:
|
||||||
|
"""Regression for #4473 Repro A via HTTP: wake after sleep(level=2) → 501.
|
||||||
|
|
||||||
|
Moved from ``tests/entrypoints/test_omni_sleep_mode.py`` (BAGEL cold start).
|
||||||
|
Live serve surfaces ``NotImplementedError`` as OpenAI-style JSON with code 501.
|
||||||
|
"""
|
||||||
|
assert omni_server is not None
|
||||||
|
sleep_resps = openai_client.send_omni_sleep_http_request(
|
||||||
|
{
|
||||||
|
"json": {"stage_ids": [0], "level": 2},
|
||||||
|
"timeout": 180,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
sleep_resp = sleep_resps[0]
|
||||||
|
assert sleep_resp.status_code == 200, sleep_resp
|
||||||
|
assert sleep_resp.json_body is not None
|
||||||
|
assert sleep_resp.json_body.get("status") == "SUCCESS", sleep_resp.json_body
|
||||||
|
|
||||||
|
openai_client.send_omni_wakeup_http_request(
|
||||||
|
{
|
||||||
|
"json": {"stage_ids": [0]},
|
||||||
|
"timeout": 120,
|
||||||
|
"err_code": 501,
|
||||||
|
"err_message": ("sleep(level=2)", "not yet implemented"),
|
||||||
|
}
|
||||||
|
)
|
||||||
@@ -0,0 +1,64 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
|
||||||
|
from contextlib import nullcontext
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||||
|
from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
|
||||||
|
|
||||||
|
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
|
||||||
|
|
||||||
|
|
||||||
|
def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
|
||||||
|
observed = {}
|
||||||
|
|
||||||
|
def fake_sdpa(query, key, value, **kwargs):
|
||||||
|
observed.update(query=query, key=key, value=value, kwargs=kwargs)
|
||||||
|
return query
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"vllm_omni.diffusion.attention.backends.cudnn_attn.sdpa_kernel",
|
||||||
|
lambda _backends: nullcontext(),
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa)
|
||||||
|
impl = CuDNNAttentionImpl(
|
||||||
|
num_heads=2,
|
||||||
|
head_size=4,
|
||||||
|
softmax_scale=0.5,
|
||||||
|
)
|
||||||
|
query = torch.randn(1, 8, 2, 4)
|
||||||
|
key = torch.randn_like(query)
|
||||||
|
value = torch.randn_like(query)
|
||||||
|
|
||||||
|
output = impl.forward_cuda(
|
||||||
|
query,
|
||||||
|
key,
|
||||||
|
value,
|
||||||
|
AttentionMetadata(extra={"valid_kv_length": 5}),
|
||||||
|
)
|
||||||
|
|
||||||
|
assert output.shape == query.shape
|
||||||
|
assert observed["query"].shape == (1, 2, 8, 4)
|
||||||
|
assert observed["key"].shape == (1, 2, 5, 4)
|
||||||
|
assert observed["value"].shape == (1, 2, 5, 4)
|
||||||
|
assert observed["kwargs"]["attn_mask"] is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_cudnn_rejects_invalid_valid_kv_length():
|
||||||
|
impl = CuDNNAttentionImpl(
|
||||||
|
num_heads=2,
|
||||||
|
head_size=4,
|
||||||
|
softmax_scale=0.5,
|
||||||
|
)
|
||||||
|
query = torch.randn(1, 8, 2, 4)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="valid_kv_length"):
|
||||||
|
impl.forward_cuda(
|
||||||
|
query,
|
||||||
|
query,
|
||||||
|
query,
|
||||||
|
AttentionMetadata(extra={"valid_kv_length": 9}),
|
||||||
|
)
|
||||||
@@ -2,6 +2,9 @@
|
|||||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
"""Unit tests for LayerNorm and RMSNorm custom ops in diffusion layers."""
|
"""Unit tests for LayerNorm and RMSNorm custom ops in diffusion layers."""
|
||||||
|
|
||||||
|
import sys
|
||||||
|
import types
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
import torch
|
import torch
|
||||||
from pytest_mock import MockerFixture
|
from pytest_mock import MockerFixture
|
||||||
@@ -299,6 +302,44 @@ def test_rmsnorm_numerical_correctness():
|
|||||||
torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)
|
torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)
|
||||||
|
|
||||||
|
|
||||||
|
def test_rmsnorm_dtype_matches_parameter_and_fp32_accumulation_contract():
|
||||||
|
"""dtype controls gamma storage; reduced-precision inputs still use fp32 accumulation."""
|
||||||
|
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||||
|
|
||||||
|
hidden_size = 64
|
||||||
|
eps = 1e-6
|
||||||
|
norm = RMSNorm(hidden_size, eps=eps, dtype=torch.bfloat16)
|
||||||
|
x = torch.randn(2, 4, hidden_size, dtype=torch.bfloat16)
|
||||||
|
|
||||||
|
assert norm.weight.dtype == torch.bfloat16
|
||||||
|
|
||||||
|
x_fp32 = x.float()
|
||||||
|
expected = (x_fp32 * torch.rsqrt(x_fp32.square().mean(-1, keepdim=True) + eps) * norm.weight.float()).to(x.dtype)
|
||||||
|
|
||||||
|
torch.testing.assert_close(norm.forward_native(x), expected, atol=0, rtol=0)
|
||||||
|
|
||||||
|
|
||||||
|
def test_rmsnorm_npu_receives_gamma_with_configured_dtype(monkeypatch: pytest.MonkeyPatch):
|
||||||
|
"""The NPU fused op receives the bf16 checkpoint gamma used by H3."""
|
||||||
|
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||||
|
|
||||||
|
captured: dict[str, torch.Tensor] = {}
|
||||||
|
|
||||||
|
def npu_rms_norm(x: torch.Tensor, gamma: torch.Tensor, epsilon: float):
|
||||||
|
captured["gamma"] = gamma
|
||||||
|
captured["epsilon"] = torch.tensor(epsilon)
|
||||||
|
return (x,)
|
||||||
|
|
||||||
|
monkeypatch.setitem(sys.modules, "torch_npu", types.SimpleNamespace(npu_rms_norm=npu_rms_norm))
|
||||||
|
norm = RMSNorm(64, eps=1e-5, dtype=torch.bfloat16)
|
||||||
|
x = torch.randn(2, 4, 64, dtype=torch.bfloat16)
|
||||||
|
|
||||||
|
assert norm.forward_npu(x) is x
|
||||||
|
assert captured["gamma"] is norm.weight
|
||||||
|
assert captured["gamma"].dtype == torch.bfloat16
|
||||||
|
assert captured["epsilon"].item() == pytest.approx(1e-5)
|
||||||
|
|
||||||
|
|
||||||
# ── RMSNorm compile-path regression tests ──
|
# ── RMSNorm compile-path regression tests ──
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -10,10 +10,13 @@ forward_hip, and apply_rotary_emb_mindiesd paths.
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import sys
|
||||||
|
import types
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
from vllm_omni.diffusion.layers.rope import RotaryEmbedding
|
from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rotary_emb_mindiesd
|
||||||
|
|
||||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||||
|
|
||||||
@@ -234,3 +237,33 @@ class TestRotaryEmbeddingNativeShapeRegression:
|
|||||||
output = rope.forward_native(x, cos, sin)
|
output = rope.forward_native(x, cos, sin)
|
||||||
assert output.shape == x.shape
|
assert output.shape == x.shape
|
||||||
assert not torch.isnan(output).any()
|
assert not torch.isnan(output).any()
|
||||||
|
|
||||||
|
|
||||||
|
def test_mindie_full_dim_rope_keeps_h3_cos_sin_layout(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
|
"""H3 passes complete rotary dims to MindIE instead of the CUDA half-dim layout."""
|
||||||
|
captured: dict[str, object] = {}
|
||||||
|
|
||||||
|
def rotary_position_embedding(x, cos, sin, **kwargs):
|
||||||
|
captured.update(x=x, cos=cos, sin=sin, kwargs=kwargs)
|
||||||
|
return x
|
||||||
|
|
||||||
|
monkeypatch.setitem(
|
||||||
|
sys.modules,
|
||||||
|
"mindiesd",
|
||||||
|
types.SimpleNamespace(rotary_position_embedding=rotary_position_embedding),
|
||||||
|
)
|
||||||
|
x = torch.randn(11, 3, 96, dtype=torch.bfloat16)
|
||||||
|
cos = torch.randn(11, 96, dtype=torch.bfloat16)
|
||||||
|
sin = torch.randn(11, 96, dtype=torch.bfloat16)
|
||||||
|
|
||||||
|
actual = apply_rotary_emb_mindiesd(x, cos, sin, interleaved=False, half_head_dim=False)
|
||||||
|
|
||||||
|
assert actual is x
|
||||||
|
assert captured["x"] is x
|
||||||
|
assert captured["cos"] is cos
|
||||||
|
assert captured["sin"] is sin
|
||||||
|
assert captured["kwargs"] == {
|
||||||
|
"rotated_mode": "rotated_half",
|
||||||
|
"head_first": False,
|
||||||
|
"fused": True,
|
||||||
|
}
|
||||||
@@ -30,6 +30,13 @@ def test_pipeline_import_registry_and_component_discovery():
|
|||||||
assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"]
|
assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_dlo_offload_plan_includes_token_refiner():
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||||
|
|
||||||
|
assert MiniMaxH3Pipeline._offload_plan.offload_submodules == {"token_refiner": "blocks"}
|
||||||
|
assert MiniMaxH3Pipeline._offload_plan.resident_dit_paths == frozenset({"transformer"})
|
||||||
|
|
||||||
|
|
||||||
def test_joint_postprocess_is_multiprocessing_picklable():
|
def test_joint_postprocess_is_multiprocessing_picklable():
|
||||||
from vllm_omni.diffusion.models.minimax_h3 import (
|
from vllm_omni.diffusion.models.minimax_h3 import (
|
||||||
get_minimax_h3_post_process_func,
|
get_minimax_h3_post_process_func,
|
||||||
@@ -87,6 +94,87 @@ def test_shifted_sigma_schedule_matches_reference_values():
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
|
||||||
|
MiniMaxH3Attention,
|
||||||
|
)
|
||||||
|
|
||||||
|
class FakeBackend:
|
||||||
|
supports_prefix_kv_slicing = True
|
||||||
|
|
||||||
|
class FakeAttention(torch.nn.Module):
|
||||||
|
attn_backend = FakeBackend
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__()
|
||||||
|
self.metadata = None
|
||||||
|
|
||||||
|
def forward(self, query, key, value, metadata):
|
||||||
|
self.metadata = metadata
|
||||||
|
return query
|
||||||
|
|
||||||
|
attention = object.__new__(MiniMaxH3Attention)
|
||||||
|
torch.nn.Module.__init__(attention)
|
||||||
|
attention.attention = FakeAttention()
|
||||||
|
# Model-side Ulysses shards rows before Attention gathers them again.
|
||||||
|
# The Python packed_total must therefore remain global even though q is
|
||||||
|
# local here.
|
||||||
|
q = torch.randn(4, 2, 4)
|
||||||
|
cu_seqlens = torch.tensor([0, 5, 8], dtype=torch.int32)
|
||||||
|
|
||||||
|
output = attention._run_packed_attention(
|
||||||
|
q,
|
||||||
|
q,
|
||||||
|
q,
|
||||||
|
cu_seqlens=cu_seqlens,
|
||||||
|
max_seqlen=5,
|
||||||
|
packed_total=8,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert output.shape == q.shape
|
||||||
|
assert attention.attention.metadata.attn_mask is None
|
||||||
|
assert attention.attention.metadata.extra["valid_kv_length"] == 5
|
||||||
|
|
||||||
|
|
||||||
|
def test_packed_attention_keeps_padding_mask_for_other_backends():
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
|
||||||
|
MiniMaxH3Attention,
|
||||||
|
)
|
||||||
|
|
||||||
|
class FakeBackend:
|
||||||
|
supports_prefix_kv_slicing = False
|
||||||
|
|
||||||
|
class FakeAttention(torch.nn.Module):
|
||||||
|
attn_backend = FakeBackend
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__()
|
||||||
|
self.metadata = None
|
||||||
|
|
||||||
|
def forward(self, query, key, value, metadata):
|
||||||
|
self.metadata = metadata
|
||||||
|
return query
|
||||||
|
|
||||||
|
attention = object.__new__(MiniMaxH3Attention)
|
||||||
|
torch.nn.Module.__init__(attention)
|
||||||
|
attention.attention = FakeAttention()
|
||||||
|
q = torch.randn(8, 2, 4)
|
||||||
|
|
||||||
|
attention._run_packed_attention(
|
||||||
|
q,
|
||||||
|
q,
|
||||||
|
q,
|
||||||
|
cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32),
|
||||||
|
max_seqlen=5,
|
||||||
|
packed_total=8,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert torch.equal(
|
||||||
|
attention.attention.metadata.attn_mask,
|
||||||
|
torch.tensor([[True, True, True, True, True, False, False, False]]),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def test_reference_image_resize_contract():
|
def test_reference_image_resize_contract():
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
|
|
||||||
@@ -102,6 +190,52 @@ def test_reference_image_resize_contract():
|
|||||||
_reference_image_shape(Image.new("RGB", (100, 501)))
|
_reference_image_shape(Image.new("RGB", (100, 501)))
|
||||||
|
|
||||||
|
|
||||||
|
def test_fl2va_supports_first_last_and_explicit_frame_index_contracts():
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||||
|
_resolve_fl2va_keyframe_indices,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert _resolve_fl2va_keyframe_indices({}, 1) == [0]
|
||||||
|
assert _resolve_fl2va_keyframe_indices({}, 2) == [0, -1]
|
||||||
|
assert _resolve_fl2va_keyframe_indices({"frame_index": -1}, 1) == [-1]
|
||||||
|
assert _resolve_fl2va_keyframe_indices({"frame_indices": [0, -1]}, 2) == [0, -1]
|
||||||
|
with pytest.raises(ValueError, match="frame_indices"):
|
||||||
|
_resolve_fl2va_keyframe_indices({"frame_indices": [0, 1]}, 2)
|
||||||
|
|
||||||
|
|
||||||
|
def test_minimax_h3_uses_the_official_output_canvas_policy():
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||||
|
_resolve_output_canvas,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert _resolve_output_canvas(21 / 9, 768) == (672, 1536)
|
||||||
|
assert _resolve_output_canvas(16 / 9, 768) == (768, 1344)
|
||||||
|
assert _resolve_output_canvas(9 / 16, 768) == (1344, 768)
|
||||||
|
with pytest.raises(ValueError, match="short_edge"):
|
||||||
|
_resolve_output_canvas(16 / 9, 720)
|
||||||
|
|
||||||
|
|
||||||
|
def test_minimax_h3_accepts_sglang_auto_aspect_ratio_alias():
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
|
||||||
|
|
||||||
|
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||||
|
sampling = SimpleNamespace(
|
||||||
|
fps=24,
|
||||||
|
num_frames=1,
|
||||||
|
height=None,
|
||||||
|
width=None,
|
||||||
|
extra_args={"duration": 5.0, "aspect_ratio": "auto"},
|
||||||
|
)
|
||||||
|
height, width, *_ = pipeline._resolve_shape("ref2va", sampling, None)
|
||||||
|
assert (height, width) == (768, 1344)
|
||||||
|
|
||||||
|
|
||||||
|
def test_minimax_h3_advertises_the_official_ref2va_image_limit():
|
||||||
|
from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
|
||||||
|
|
||||||
|
assert get_diffusion_model_metadata("MiniMaxH3Pipeline").max_multimodal_image_inputs == 9
|
||||||
|
|
||||||
|
|
||||||
def test_encoder_forward_uses_hook_compatible_encode_entrypoint():
|
def test_encoder_forward_uses_hook_compatible_encode_entrypoint():
|
||||||
from vllm_omni.diffusion.models.minimax_h3.encoder import (
|
from vllm_omni.diffusion.models.minimax_h3.encoder import (
|
||||||
MiniMaxH3Qwen3VLEncoder,
|
MiniMaxH3Qwen3VLEncoder,
|
||||||
@@ -247,6 +381,127 @@ def test_layerwise_offload_releases_text_encoder():
|
|||||||
pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
|
pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
|
||||||
|
|
||||||
|
|
||||||
|
def test_distributed_layerwise_offload_releases_text_encoder():
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||||
|
|
||||||
|
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||||
|
torch.nn.Module.__init__(pipeline)
|
||||||
|
pipeline.od_config = SimpleNamespace(
|
||||||
|
enable_cpu_offload=False,
|
||||||
|
enable_layerwise_offload=False,
|
||||||
|
enable_distributed_layerwise_offload=True,
|
||||||
|
)
|
||||||
|
pipeline.text_encoder = Mock()
|
||||||
|
expected = torch.ones(2, 3)
|
||||||
|
pipeline.text_encoder.encode_ids.return_value = expected
|
||||||
|
|
||||||
|
actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
|
||||||
|
|
||||||
|
assert actual is expected
|
||||||
|
pipeline.text_encoder.load_to_device.assert_called_once_with()
|
||||||
|
pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
|
||||||
|
|
||||||
|
|
||||||
|
def test_distributed_layerwise_offload_stages_vae_component():
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||||
|
|
||||||
|
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||||
|
torch.nn.Module.__init__(pipeline)
|
||||||
|
pipeline.od_config = SimpleNamespace(
|
||||||
|
enable_layerwise_offload=False,
|
||||||
|
enable_distributed_layerwise_offload=True,
|
||||||
|
)
|
||||||
|
component = Mock()
|
||||||
|
|
||||||
|
with pipeline._component_on_device(component):
|
||||||
|
component.load_to_device.assert_called_once_with()
|
||||||
|
component.offload_to_cpu.assert_not_called()
|
||||||
|
|
||||||
|
component.offload_to_cpu.assert_called_once_with()
|
||||||
|
|
||||||
|
|
||||||
|
def test_distributed_layerwise_resident_blocks_are_stage_scoped():
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||||
|
|
||||||
|
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||||
|
torch.nn.Module.__init__(pipeline)
|
||||||
|
controller = Mock()
|
||||||
|
pipeline._dlo_residency_controller = controller
|
||||||
|
|
||||||
|
with pipeline._resident_dit_layers_on_device():
|
||||||
|
controller.load_resident_layers.assert_called_once_with()
|
||||||
|
controller.offload_resident_layers.assert_not_called()
|
||||||
|
|
||||||
|
controller.offload_resident_layers.assert_called_once_with()
|
||||||
|
|
||||||
|
|
||||||
|
def test_distributed_layerwise_resident_blocks_release_on_failure():
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||||
|
|
||||||
|
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||||
|
torch.nn.Module.__init__(pipeline)
|
||||||
|
controller = Mock()
|
||||||
|
pipeline._dlo_residency_controller = controller
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="denoise failed"):
|
||||||
|
with pipeline._resident_dit_layers_on_device():
|
||||||
|
raise RuntimeError("denoise failed")
|
||||||
|
|
||||||
|
controller.load_resident_layers.assert_called_once_with()
|
||||||
|
controller.offload_resident_layers.assert_called_once_with()
|
||||||
|
|
||||||
|
|
||||||
|
def test_encoder_layerwise_offload_keeps_tp_blocks_rank_local(monkeypatch):
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3.encoder import (
|
||||||
|
MiniMaxH3Qwen3VLEncoder,
|
||||||
|
)
|
||||||
|
|
||||||
|
class Stack(torch.nn.Module):
|
||||||
|
def __init__(self, count):
|
||||||
|
super().__init__()
|
||||||
|
self.blocks = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)])
|
||||||
|
|
||||||
|
class TextStack(torch.nn.Module):
|
||||||
|
def __init__(self, count):
|
||||||
|
super().__init__()
|
||||||
|
self.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)])
|
||||||
|
|
||||||
|
hooks = []
|
||||||
|
|
||||||
|
def fake_apply(block, next_block, device, stream, pin_memory):
|
||||||
|
hook = SimpleNamespace(
|
||||||
|
block=block,
|
||||||
|
next_block=next_block,
|
||||||
|
device=device,
|
||||||
|
stream=stream,
|
||||||
|
pin_memory=pin_memory,
|
||||||
|
_prev_hook=None,
|
||||||
|
offload_layer=Mock(),
|
||||||
|
)
|
||||||
|
hooks.append(hook)
|
||||||
|
return hook
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"vllm_omni.diffusion.offloader.layerwise_backend.apply_block_hook",
|
||||||
|
fake_apply,
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"vllm_omni.platforms.current_omni_platform.Stream",
|
||||||
|
Mock(return_value="copy-stream"),
|
||||||
|
)
|
||||||
|
encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
|
||||||
|
torch.nn.Module.__init__(encoder)
|
||||||
|
encoder.device_target = torch.device("cpu")
|
||||||
|
encoder.vision = Stack(2)
|
||||||
|
encoder.text_model = TextStack(3)
|
||||||
|
|
||||||
|
encoder.enable_omni_layerwise_offload(pin_memory=False)
|
||||||
|
|
||||||
|
assert len(hooks) == 5
|
||||||
|
assert all(hook._prev_hook is not None for hook in hooks)
|
||||||
|
assert all(hook.device == torch.device("cpu") for hook in hooks)
|
||||||
|
|
||||||
|
|
||||||
def test_video_vae_keeps_reference_fp32_weights(monkeypatch):
|
def test_video_vae_keeps_reference_fp32_weights(monkeypatch):
|
||||||
from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
|
from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
|
||||||
|
|
||||||
@@ -278,6 +533,82 @@ def test_video_vae_keeps_reference_fp32_weights(monkeypatch):
|
|||||||
assert next(video_vae.parameters()).dtype == torch.float32
|
assert next(video_vae.parameters()).dtype == torch.float32
|
||||||
|
|
||||||
|
|
||||||
|
def test_video_vae_can_load_on_cpu_for_staged_gpu_residency(monkeypatch):
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
|
||||||
|
|
||||||
|
class FakeRemote(torch.nn.Module):
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__()
|
||||||
|
self.model = torch.nn.Linear(1, 1)
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
vae_module,
|
||||||
|
"_load_component_config",
|
||||||
|
lambda _path: {
|
||||||
|
"latent_channels": 1,
|
||||||
|
"latents_mean": [0.0],
|
||||||
|
"latents_std": [1.0],
|
||||||
|
},
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
vae_module,
|
||||||
|
"_load_remote_component",
|
||||||
|
lambda _path, _config: FakeRemote(),
|
||||||
|
)
|
||||||
|
|
||||||
|
video_vae = vae_module.MiniMaxH3VideoVAE(
|
||||||
|
"unused",
|
||||||
|
device=torch.device("meta"),
|
||||||
|
load_device=torch.device("cpu"),
|
||||||
|
)
|
||||||
|
|
||||||
|
assert next(video_vae.parameters()).device.type == "cpu"
|
||||||
|
assert video_vae._device_target.type == "meta"
|
||||||
|
|
||||||
|
|
||||||
|
def test_video_vae_uses_snapshot_stager_for_accelerator_target(monkeypatch):
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
|
||||||
|
|
||||||
|
class FakeRemote(torch.nn.Module):
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__()
|
||||||
|
self.model = torch.nn.Linear(1, 1)
|
||||||
|
|
||||||
|
stager = Mock()
|
||||||
|
stager_cls = Mock(return_value=stager)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
vae_module,
|
||||||
|
"_load_component_config",
|
||||||
|
lambda _path: {
|
||||||
|
"latent_channels": 1,
|
||||||
|
"latents_mean": [0.0],
|
||||||
|
"latents_std": [1.0],
|
||||||
|
},
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
vae_module,
|
||||||
|
"_load_remote_component",
|
||||||
|
lambda _path, _config: FakeRemote(),
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(vae_module, "PinnedModuleStager", stager_cls)
|
||||||
|
|
||||||
|
video_vae = vae_module.MiniMaxH3VideoVAE(
|
||||||
|
"unused",
|
||||||
|
device=torch.device("cuda"),
|
||||||
|
load_device=torch.device("cpu"),
|
||||||
|
)
|
||||||
|
video_vae.load_to_device()
|
||||||
|
video_vae.offload_to_cpu()
|
||||||
|
|
||||||
|
stager_cls.assert_called_once_with(
|
||||||
|
video_vae.remote,
|
||||||
|
torch.device("cuda"),
|
||||||
|
pin_memory=True,
|
||||||
|
)
|
||||||
|
stager.load.assert_called_once_with()
|
||||||
|
stager.offload.assert_called_once_with()
|
||||||
|
|
||||||
|
|
||||||
def test_video_vae_encode_uses_configured_parallel_tiling():
|
def test_video_vae_encode_uses_configured_parallel_tiling():
|
||||||
from vllm_omni.diffusion.models.minimax_h3.vae import (
|
from vllm_omni.diffusion.models.minimax_h3.vae import (
|
||||||
MiniMaxH3VideoVAE,
|
MiniMaxH3VideoVAE,
|
||||||
@@ -373,3 +704,360 @@ def test_distributed_video_vae_encodes_references_sequentially(monkeypatch):
|
|||||||
torch.tensor([[1.0, 1.0], [2.0, 2.0]]),
|
torch.tensor([[1.0, 1.0], [2.0, 2.0]]),
|
||||||
)
|
)
|
||||||
assert shapes == [(1, 2, 3), (2, 2, 3)]
|
assert shapes == [(1, 2, 3), (2, 2, 3)]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("case", "extra", "image_count", "expected"),
|
||||||
|
[
|
||||||
|
("F1", {"frame_index": -1}, 1, [-1]),
|
||||||
|
("F2", {"frame_indices": [0, -1]}, 2, [0, -1]),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_f1_f2_official_fl2va_keyframe_matrix(case, extra, image_count, expected):
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||||
|
_resolve_fl2va_keyframe_indices,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert case in {"F1", "F2"}
|
||||||
|
assert _resolve_fl2va_keyframe_indices(extra, image_count) == expected
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("case", "counts"),
|
||||||
|
[
|
||||||
|
("R1", (3, 0, 0)),
|
||||||
|
("R2", (1, 0, 0)),
|
||||||
|
("R3", (1, 1, 0)),
|
||||||
|
("R4", (0, 1, 1)),
|
||||||
|
("R5", (1, 1, 1)),
|
||||||
|
("R6", (1, 0, 2)),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_r1_r6_ref2va_reference_count_matrix(case, counts):
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||||
|
_validate_ref2va_reference_counts,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert case.startswith("R")
|
||||||
|
_validate_ref2va_reference_counts(*counts)
|
||||||
|
|
||||||
|
|
||||||
|
def test_ref2va_reference_count_validation_preserves_client_error_metadata():
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||||
|
_validate_ref2va_reference_counts,
|
||||||
|
)
|
||||||
|
from vllm_omni.errors import OmniClientError
|
||||||
|
|
||||||
|
with pytest.raises(OmniClientError, match="at least one image or video"):
|
||||||
|
_validate_ref2va_reference_counts(0, 0, 0)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("case", "start_time", "expected_duration"),
|
||||||
|
[
|
||||||
|
("R7", None, 10.0),
|
||||||
|
("R8", 4.0, 6.0),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_time, expected_duration):
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
|
||||||
|
|
||||||
|
source = tmp_path / f"{case}.mp4"
|
||||||
|
source.touch()
|
||||||
|
metadata = {
|
||||||
|
"width": 1280,
|
||||||
|
"height": 720,
|
||||||
|
"fps": 24.0,
|
||||||
|
"frame_count": 240,
|
||||||
|
"duration": 10.0,
|
||||||
|
"format_names": ("mp4",),
|
||||||
|
"video_codec": "h264",
|
||||||
|
"audio_codecs": (),
|
||||||
|
"file_size": 1024,
|
||||||
|
}
|
||||||
|
transcode_calls = []
|
||||||
|
monkeypatch.setattr(reference_video_module, "_probe_video", lambda _path: metadata)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
reference_video_module,
|
||||||
|
"_transcode_reference_video",
|
||||||
|
lambda source, **kwargs: transcode_calls.append((source, kwargs)) or "prepared.mp4",
|
||||||
|
)
|
||||||
|
|
||||||
|
prepared = reference_video_module.prepare_reference_videos(
|
||||||
|
[str(source)],
|
||||||
|
target_frame_count=209,
|
||||||
|
workdir=str(tmp_path / "work"),
|
||||||
|
start_time_seconds=start_time,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert prepared[0]["duration_seconds"] == pytest.approx(expected_duration)
|
||||||
|
assert prepared[0]["start_time_seconds"] == pytest.approx(start_time or 0.0)
|
||||||
|
assert transcode_calls[0][1]["duration_seconds"] == pytest.approx(expected_duration)
|
||||||
|
|
||||||
|
|
||||||
|
def test_ref2va_two_video_recipe_tolerates_container_rounding(monkeypatch, tmp_path):
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
|
||||||
|
|
||||||
|
first = tmp_path / "first.mp4"
|
||||||
|
second = tmp_path / "second.mp4"
|
||||||
|
first.touch()
|
||||||
|
second.touch()
|
||||||
|
metadata = {
|
||||||
|
str(first): {
|
||||||
|
"width": 1280,
|
||||||
|
"height": 720,
|
||||||
|
"fps": 24.0,
|
||||||
|
"frame_count": 241,
|
||||||
|
"duration": 10.041667,
|
||||||
|
"format_names": ("mp4",),
|
||||||
|
"video_codec": "h264",
|
||||||
|
"audio_codecs": (),
|
||||||
|
"file_size": 1024,
|
||||||
|
},
|
||||||
|
str(second): {
|
||||||
|
"width": 1280,
|
||||||
|
"height": 720,
|
||||||
|
"fps": 24.0,
|
||||||
|
"frame_count": 120,
|
||||||
|
"duration": 4.966667,
|
||||||
|
"format_names": ("mp4",),
|
||||||
|
"video_codec": "h264",
|
||||||
|
"audio_codecs": (),
|
||||||
|
"file_size": 1024,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
transcode_calls = []
|
||||||
|
monkeypatch.setattr(reference_video_module, "_probe_video", lambda path: metadata[str(path)])
|
||||||
|
monkeypatch.setattr(
|
||||||
|
reference_video_module,
|
||||||
|
"_transcode_reference_video",
|
||||||
|
lambda source, **kwargs: transcode_calls.append((source, kwargs)) or f"{source}.prepared.mp4",
|
||||||
|
)
|
||||||
|
|
||||||
|
prepared = reference_video_module.prepare_reference_videos(
|
||||||
|
[str(first), str(second)],
|
||||||
|
target_frame_count=124,
|
||||||
|
workdir=str(tmp_path / "work"),
|
||||||
|
)
|
||||||
|
|
||||||
|
assert [item["duration_seconds"] for item in prepared] == pytest.approx([10.041667, 4.958333])
|
||||||
|
assert sum(item["duration_seconds"] for item in prepared) == pytest.approx(15.0)
|
||||||
|
assert transcode_calls[1][1]["duration_seconds"] == pytest.approx(4.958333)
|
||||||
|
|
||||||
|
|
||||||
|
def test_ref2va_two_video_recipe_rejects_real_duration_overflow(monkeypatch, tmp_path):
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
|
||||||
|
from vllm_omni.errors import OmniClientError
|
||||||
|
|
||||||
|
first = tmp_path / "first.mp4"
|
||||||
|
second = tmp_path / "second.mp4"
|
||||||
|
first.touch()
|
||||||
|
second.touch()
|
||||||
|
metadata = {
|
||||||
|
str(first): {
|
||||||
|
"width": 1280,
|
||||||
|
"height": 720,
|
||||||
|
"fps": 24.0,
|
||||||
|
"frame_count": 240,
|
||||||
|
"duration": 10.0,
|
||||||
|
"format_names": ("mp4",),
|
||||||
|
"video_codec": "h264",
|
||||||
|
"audio_codecs": (),
|
||||||
|
"file_size": 1024,
|
||||||
|
},
|
||||||
|
str(second): {
|
||||||
|
"width": 1280,
|
||||||
|
"height": 720,
|
||||||
|
"fps": 24.0,
|
||||||
|
"frame_count": 120,
|
||||||
|
"duration": 5.02,
|
||||||
|
"format_names": ("mp4",),
|
||||||
|
"video_codec": "h264",
|
||||||
|
"audio_codecs": (),
|
||||||
|
"file_size": 1024,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
monkeypatch.setattr(reference_video_module, "_probe_video", lambda path: metadata[str(path)])
|
||||||
|
monkeypatch.setattr(reference_video_module, "_transcode_reference_video", lambda source, **kwargs: "prepared.mp4")
|
||||||
|
|
||||||
|
with pytest.raises(OmniClientError, match="15 seconds"):
|
||||||
|
reference_video_module.prepare_reference_videos(
|
||||||
|
[str(first), str(second)],
|
||||||
|
target_frame_count=124,
|
||||||
|
workdir=str(tmp_path / "work"),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("duration", [4.0, 15.0])
|
||||||
|
def test_g2_output_duration_accepts_official_boundaries(duration):
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||||
|
|
||||||
|
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||||
|
sampling = SimpleNamespace(
|
||||||
|
fps=24,
|
||||||
|
num_frames=1,
|
||||||
|
height=None,
|
||||||
|
width=None,
|
||||||
|
extra_args={"duration": duration},
|
||||||
|
)
|
||||||
|
|
||||||
|
height, width, *_ = pipeline._resolve_shape("ref2va", sampling, None)
|
||||||
|
assert height > 0 and width > 0
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("duration", [3.99, 15.01, float("nan"), "not-a-duration"])
|
||||||
|
def test_g2_output_duration_rejects_out_of_contract_values(duration):
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||||
|
|
||||||
|
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||||
|
sampling = SimpleNamespace(
|
||||||
|
fps=24,
|
||||||
|
num_frames=1,
|
||||||
|
height=None,
|
||||||
|
width=None,
|
||||||
|
extra_args={"duration": duration},
|
||||||
|
)
|
||||||
|
with pytest.raises(ValueError, match="duration"):
|
||||||
|
pipeline._resolve_shape("ref2va", sampling, None)
|
||||||
|
|
||||||
|
|
||||||
|
def test_g1_fanout_uses_incrementing_output_seeds():
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||||
|
_minimax_h3_output_seeds,
|
||||||
|
_resolve_minimax_h3_num_outputs,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert _resolve_minimax_h3_num_outputs(3) == 3
|
||||||
|
assert _minimax_h3_output_seeds(100, 3) == [100, 101, 102]
|
||||||
|
with pytest.raises(ValueError, match="num_outputs_per_prompt"):
|
||||||
|
_resolve_minimax_h3_num_outputs(11)
|
||||||
|
|
||||||
|
|
||||||
|
def test_g3_task_specific_aspect_ratio_policy():
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||||
|
_resolve_minimax_h3_aspect_ratio,
|
||||||
|
)
|
||||||
|
|
||||||
|
image = Image.new("RGB", (1280, 720))
|
||||||
|
assert _resolve_minimax_h3_aspect_ratio("fl2va", "9:16", image) == pytest.approx(16 / 9)
|
||||||
|
assert _resolve_minimax_h3_aspect_ratio("ref2va", None, None) == pytest.approx(16 / 9)
|
||||||
|
assert _resolve_minimax_h3_aspect_ratio("ref2va", "auto", None) == pytest.approx(16 / 9)
|
||||||
|
with pytest.raises(ValueError, match="requires an explicit"):
|
||||||
|
_resolve_minimax_h3_aspect_ratio("t2va", None, None)
|
||||||
|
with pytest.raises(ValueError, match="one of"):
|
||||||
|
_resolve_minimax_h3_aspect_ratio("t2va", "2:1", None)
|
||||||
|
|
||||||
|
|
||||||
|
def test_g3_t2va_shape_requires_a_named_ratio_and_fl2va_ignores_override():
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||||
|
|
||||||
|
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||||
|
t2va_sampling = SimpleNamespace(
|
||||||
|
fps=24,
|
||||||
|
num_frames=1,
|
||||||
|
height=None,
|
||||||
|
width=None,
|
||||||
|
extra_args={"duration": 5.0},
|
||||||
|
)
|
||||||
|
with pytest.raises(ValueError, match="requires an explicit aspect_ratio"):
|
||||||
|
pipeline._resolve_shape("t2va", t2va_sampling, None)
|
||||||
|
|
||||||
|
fl2va_sampling = SimpleNamespace(
|
||||||
|
fps=24,
|
||||||
|
num_frames=1,
|
||||||
|
height=None,
|
||||||
|
width=None,
|
||||||
|
extra_args={"duration": 5.0, "aspect_ratio": "9:16"},
|
||||||
|
)
|
||||||
|
height, width, *_ = pipeline._resolve_shape("fl2va", fl2va_sampling, Image.new("RGB", (1280, 720)))
|
||||||
|
assert (height, width) == (768, 1344)
|
||||||
|
|
||||||
|
|
||||||
|
def test_g4_reference_image_boundaries_and_aspect_ratio():
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||||
|
_validate_reference_image,
|
||||||
|
)
|
||||||
|
|
||||||
|
_validate_reference_image(Image.new("RGB", (256, 256)))
|
||||||
|
_validate_reference_image(Image.new("RGB", (5760, 2304)))
|
||||||
|
with pytest.raises(ValueError, match="dimensions"):
|
||||||
|
_validate_reference_image(Image.new("RGB", (255, 255)))
|
||||||
|
with pytest.raises(ValueError, match="aspect ratio"):
|
||||||
|
_validate_reference_image(Image.new("RGB", (256, 641)))
|
||||||
|
|
||||||
|
|
||||||
|
def test_g4_reference_image_file_format_and_size_contract(tmp_path):
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||||
|
_load_image,
|
||||||
|
)
|
||||||
|
|
||||||
|
valid_path = tmp_path / "reference.png"
|
||||||
|
Image.new("RGB", (256, 256)).save(valid_path)
|
||||||
|
assert _load_image(valid_path).size == (256, 256)
|
||||||
|
|
||||||
|
invalid_path = tmp_path / "reference.bmp"
|
||||||
|
Image.new("RGB", (256, 256)).save(invalid_path)
|
||||||
|
with pytest.raises(ValueError, match="must use"):
|
||||||
|
_load_image(invalid_path)
|
||||||
|
|
||||||
|
|
||||||
|
def test_g4_standalone_audio_duration_and_total_duration_contract():
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3.reference_video import (
|
||||||
|
validate_reference_audio_waveforms,
|
||||||
|
)
|
||||||
|
|
||||||
|
sample_rate = 16000
|
||||||
|
validate_reference_audio_waveforms(
|
||||||
|
[
|
||||||
|
(torch.zeros(1, 2 * sample_rate), sample_rate),
|
||||||
|
(torch.zeros(1, 13 * sample_rate), sample_rate),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
with pytest.raises(ValueError, match="duration"):
|
||||||
|
validate_reference_audio_waveforms([(torch.zeros(1, int(1.9 * sample_rate)), sample_rate)])
|
||||||
|
with pytest.raises(ValueError, match="at most 15 seconds in total"):
|
||||||
|
validate_reference_audio_waveforms(
|
||||||
|
[
|
||||||
|
(torch.zeros(1, 8 * sample_rate), sample_rate),
|
||||||
|
(torch.zeros(1, 8 * sample_rate), sample_rate),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("field", "value", "message"),
|
||||||
|
[
|
||||||
|
("fps", 23.0, "FPS"),
|
||||||
|
("duration", 1.0, "duration"),
|
||||||
|
("format_names", ("avi",), "container"),
|
||||||
|
("video_codec", "vp9", "H.264"),
|
||||||
|
("audio_codecs", ("opus",), "AAC"),
|
||||||
|
("file_size", 50 * 1024 * 1024 + 1, "size"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_g4_reference_video_metadata_validation(field, value, message, tmp_path):
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3.reference_video import (
|
||||||
|
_validate_reference_video_metadata,
|
||||||
|
)
|
||||||
|
|
||||||
|
metadata = {
|
||||||
|
"width": 1280,
|
||||||
|
"height": 720,
|
||||||
|
"fps": 24.0,
|
||||||
|
"duration": 10.0,
|
||||||
|
"format_names": ("mp4",),
|
||||||
|
"video_codec": "h264",
|
||||||
|
"audio_codecs": ("aac",),
|
||||||
|
"file_size": 1024,
|
||||||
|
}
|
||||||
|
metadata[field] = value
|
||||||
|
with pytest.raises(ValueError, match=message):
|
||||||
|
_validate_reference_video_metadata(metadata, index=0, source=str(tmp_path / "reference.mp4"))
|
||||||
@@ -61,6 +61,8 @@ def test_minimax_h3_t2va_ulysses8_smoke():
|
|||||||
output_type="np",
|
output_type="np",
|
||||||
extra_args={
|
extra_args={
|
||||||
"task": "t2va",
|
"task": "t2va",
|
||||||
|
"duration": 4.0,
|
||||||
|
"aspect_ratio": "16:9",
|
||||||
"flow_shift": 12.0,
|
"flow_shift": 12.0,
|
||||||
"audio_flow_shift": 3.0,
|
"audio_flow_shift": 3.0,
|
||||||
},
|
},
|
||||||
@@ -70,7 +72,14 @@ def test_minimax_h3_t2va_ulysses8_smoke():
|
|||||||
finally:
|
finally:
|
||||||
engine.close()
|
engine.close()
|
||||||
|
|
||||||
_assert_joint_output(outputs)
|
assert len(outputs) == 1
|
||||||
|
frames = np.asarray(outputs[0].images[0])
|
||||||
|
assert frames.shape == (107, 256, 448, 3)
|
||||||
|
multimodal = outputs[0].multimodal_output
|
||||||
|
assert multimodal is not None
|
||||||
|
assert np.asarray(multimodal["audio"]).shape[1] == 2
|
||||||
|
assert multimodal["audio_sample_rate"] == 32000
|
||||||
|
assert multimodal["fps"] == 24
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.skipif(
|
@pytest.mark.skipif(
|
||||||
|
|||||||
@@ -144,3 +144,20 @@ def test_condition_noise_is_seeded_and_keeps_clean_anchor_at_timestep_one():
|
|||||||
noise_aug=0.25,
|
noise_aug=0.25,
|
||||||
)
|
)
|
||||||
torch.testing.assert_close(first, second)
|
torch.testing.assert_close(first, second)
|
||||||
|
|
||||||
|
|
||||||
|
def test_condition_noise_accepts_a_reference_video_longer_than_the_target():
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3.condition_noise import (
|
||||||
|
minimax_h3_imgvid_cond_noise_aug_rows,
|
||||||
|
)
|
||||||
|
|
||||||
|
rows = torch.zeros(4 * 2 * 2, 96)
|
||||||
|
result = minimax_h3_imgvid_cond_noise_aug_rows(
|
||||||
|
rows,
|
||||||
|
condition_shapes=[(4, 4, 4)],
|
||||||
|
target_latent_t=2,
|
||||||
|
imgvid_cond_num_frames=1,
|
||||||
|
seed=7,
|
||||||
|
noise_aug=0.5,
|
||||||
|
)
|
||||||
|
assert result.shape == rows.shape
|
||||||
@@ -61,6 +61,33 @@ def test_packed_attention_is_a_regional_compile_boundary():
|
|||||||
assert getattr(MiniMaxH3Attention._run_packed_attention, "_torchdynamo_disable", False)
|
assert getattr(MiniMaxH3Attention._run_packed_attention, "_torchdynamo_disable", False)
|
||||||
|
|
||||||
|
|
||||||
|
def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims():
|
||||||
|
from vllm_omni.diffusion.layers.rope import RotaryEmbedding
|
||||||
|
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
|
||||||
|
MiniMaxH3Attention,
|
||||||
|
)
|
||||||
|
|
||||||
|
attention = object.__new__(MiniMaxH3Attention)
|
||||||
|
nn.Module.__init__(attention)
|
||||||
|
attention.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False)
|
||||||
|
attention.rope._forward_method = attention.rope.forward_native
|
||||||
|
|
||||||
|
x = torch.randn(11, 3, 128, dtype=torch.bfloat16)
|
||||||
|
freqs_half = torch.randn(11, 48)
|
||||||
|
freqs = torch.cat((freqs_half, freqs_half), dim=-1)
|
||||||
|
actual = attention._apply_rope(x, freqs)
|
||||||
|
|
||||||
|
cos = torch.cos(freqs).to(x.dtype).unsqueeze(1)
|
||||||
|
sin = torch.sin(freqs).to(x.dtype).unsqueeze(1)
|
||||||
|
x_rot = x[..., :96]
|
||||||
|
x1, x2 = x_rot.chunk(2, dim=-1)
|
||||||
|
expected_rot = x_rot * cos + torch.cat((-x2, x1), dim=-1) * sin
|
||||||
|
expected = torch.cat((expected_rot, x[..., 96:]), dim=-1)
|
||||||
|
|
||||||
|
torch.testing.assert_close(actual, expected, atol=0, rtol=0)
|
||||||
|
torch.testing.assert_close(actual[..., 96:], x[..., 96:], atol=0, rtol=0)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
("tp_size", "message"),
|
("tp_size", "message"),
|
||||||
[
|
[
|
||||||
|
|||||||
@@ -135,6 +135,7 @@ def _generate_joint_output(omni, config: QualityTestConfig):
|
|||||||
num_frames=config.num_frames,
|
num_frames=config.num_frames,
|
||||||
guidance_scale=config.guidance_scale,
|
guidance_scale=config.guidance_scale,
|
||||||
sigmas=config.sigmas,
|
sigmas=config.sigmas,
|
||||||
|
extra_args={"aspect_ratio": "16:9"},
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
main_process_peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
|
main_process_peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
|
||||||
|
|||||||
@@ -27,7 +27,9 @@ from vllm_omni.diffusion.offloader.block_discovery import (
|
|||||||
from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
|
from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
|
||||||
DistributedLayerwiseOffloadBackend,
|
DistributedLayerwiseOffloadBackend,
|
||||||
DistributedLayerwiseOffloadHook,
|
DistributedLayerwiseOffloadHook,
|
||||||
|
PinnedResidentLayerGroup,
|
||||||
)
|
)
|
||||||
|
from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
|
||||||
from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan, get_offload_plan
|
from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan, get_offload_plan
|
||||||
from vllm_omni.platforms import current_omni_platform
|
from vllm_omni.platforms import current_omni_platform
|
||||||
|
|
||||||
@@ -97,6 +99,7 @@ def patched_offload_runtime(monkeypatch):
|
|||||||
monkeypatch.setattr(dist_backend_module.current_omni_platform, "Event", DummyEvent)
|
monkeypatch.setattr(dist_backend_module.current_omni_platform, "Event", DummyEvent)
|
||||||
monkeypatch.setattr(dist_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
|
monkeypatch.setattr(dist_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
|
||||||
monkeypatch.setattr(dist_backend_module.current_omni_platform, "stream", dummy_stream)
|
monkeypatch.setattr(dist_backend_module.current_omni_platform, "stream", dummy_stream)
|
||||||
|
monkeypatch.setattr(dist_backend_module.current_omni_platform, "synchronize", lambda: None)
|
||||||
|
|
||||||
|
|
||||||
class TinyBlock(nn.Module):
|
class TinyBlock(nn.Module):
|
||||||
@@ -308,6 +311,164 @@ class TestDistributedLayerwiseOffloadHook:
|
|||||||
assert shard1.numel() == 2
|
assert shard1.numel() == 2
|
||||||
assert torch.equal(shard1, torch.tensor([102.0, 0.0]))
|
assert torch.equal(shard1, torch.tensor([102.0, 0.0]))
|
||||||
|
|
||||||
|
def test_mmap_layout_transform_is_applied_before_sharding(self, dist_group, patched_offload_runtime):
|
||||||
|
block = nn.Module()
|
||||||
|
block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
|
||||||
|
next_block = nn.Module()
|
||||||
|
next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32))
|
||||||
|
next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0)
|
||||||
|
next_block.weight.mmap_weight_transform_pending = True
|
||||||
|
|
||||||
|
hook = DistributedLayerwiseOffloadHook(
|
||||||
|
next_block=next_block,
|
||||||
|
device=torch.device("cpu"),
|
||||||
|
dp_group=None,
|
||||||
|
dp_size=2,
|
||||||
|
rank=0,
|
||||||
|
pin_memory=False,
|
||||||
|
)
|
||||||
|
hook.initialize_hook(block)
|
||||||
|
|
||||||
|
assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([3.0, 2.0]))
|
||||||
|
|
||||||
|
def test_regular_loader_weight_does_not_reapply_mmap_transform(self, dist_group, patched_offload_runtime):
|
||||||
|
block = nn.Module()
|
||||||
|
block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
|
||||||
|
next_block = nn.Module()
|
||||||
|
next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32))
|
||||||
|
next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0)
|
||||||
|
|
||||||
|
hook = DistributedLayerwiseOffloadHook(
|
||||||
|
next_block=next_block,
|
||||||
|
device=torch.device("cpu"),
|
||||||
|
dp_group=None,
|
||||||
|
dp_size=2,
|
||||||
|
rank=0,
|
||||||
|
pin_memory=False,
|
||||||
|
)
|
||||||
|
hook.initialize_hook(block)
|
||||||
|
|
||||||
|
assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([0.0, 1.0]))
|
||||||
|
|
||||||
|
|
||||||
|
class TestPinnedResidentLayerGroup:
|
||||||
|
def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
|
||||||
|
blocks = [nn.Linear(2, 2), nn.Linear(2, 2)]
|
||||||
|
expected = []
|
||||||
|
for index, block in enumerate(blocks):
|
||||||
|
block.weight.data.fill_(index + 1)
|
||||||
|
block.bias.data.fill_(index + 3)
|
||||||
|
expected.append((block.weight.detach().clone(), block.bias.detach().clone()))
|
||||||
|
|
||||||
|
group = PinnedResidentLayerGroup(
|
||||||
|
blocks,
|
||||||
|
device=torch.device("cpu"),
|
||||||
|
copy_stream=DummyStream(),
|
||||||
|
pin_memory=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert all(block.weight.numel() == 0 for block in blocks)
|
||||||
|
|
||||||
|
group.load()
|
||||||
|
for block, (weight, bias) in zip(blocks, expected):
|
||||||
|
assert torch.equal(block.weight, weight)
|
||||||
|
assert torch.equal(block.bias, bias)
|
||||||
|
|
||||||
|
# Device-side changes must not overwrite the persistent host master.
|
||||||
|
blocks[0].weight.data.zero_()
|
||||||
|
group.offload()
|
||||||
|
assert all(block.weight.numel() == 0 for block in blocks)
|
||||||
|
|
||||||
|
group.load()
|
||||||
|
for block, (weight, bias) in zip(blocks, expected):
|
||||||
|
assert torch.equal(block.weight, weight)
|
||||||
|
assert torch.equal(block.bias, bias)
|
||||||
|
|
||||||
|
group.offload()
|
||||||
|
|
||||||
|
def test_mmap_loader_attrs_survive_to_empty_parameter_replacement(self):
|
||||||
|
module = nn.Linear(2, 2, bias=False)
|
||||||
|
|
||||||
|
def transform(tensor):
|
||||||
|
return tensor.flip(0)
|
||||||
|
|
||||||
|
module.weight.mmap_weight_transform = transform
|
||||||
|
|
||||||
|
backend = object.__new__(DistributedLayerwiseOffloadBackend)
|
||||||
|
backend._remember_mmap_param_attrs(module)
|
||||||
|
module.to_empty(device="meta")
|
||||||
|
|
||||||
|
replacement = nn.Parameter(torch.arange(4).reshape(2, 2).float())
|
||||||
|
backend._attach_mmap_param_attrs("weight", replacement, module.weight)
|
||||||
|
|
||||||
|
assert replacement.mmap_weight_transform is transform
|
||||||
|
assert replacement.mmap_weight_transform_pending is True
|
||||||
|
|
||||||
|
def test_shared_allgather_output_is_narrowed_to_current_block(
|
||||||
|
self, dist_group, patched_offload_runtime, monkeypatch
|
||||||
|
):
|
||||||
|
block = nn.Module()
|
||||||
|
block.weight = nn.Parameter(torch.zeros(3, dtype=torch.float32))
|
||||||
|
next_block = nn.Module()
|
||||||
|
next_block.weight = nn.Parameter(torch.arange(3, dtype=torch.float32))
|
||||||
|
|
||||||
|
hook = DistributedLayerwiseOffloadHook(
|
||||||
|
next_block=next_block,
|
||||||
|
device=torch.device("cpu"),
|
||||||
|
dp_group=object(),
|
||||||
|
dp_size=2,
|
||||||
|
rank=0,
|
||||||
|
pin_memory=False,
|
||||||
|
)
|
||||||
|
hook.initialize_hook(block)
|
||||||
|
# Simulate unified buffers sized for a larger heterogeneous block.
|
||||||
|
hook.gpu_buffers = [
|
||||||
|
{torch.float32: torch.empty(12)},
|
||||||
|
{torch.float32: torch.empty(12)},
|
||||||
|
]
|
||||||
|
hook.gpu_shard_buffers = [
|
||||||
|
{torch.float32: torch.empty(6)},
|
||||||
|
{torch.float32: torch.empty(6)},
|
||||||
|
]
|
||||||
|
output_sizes = []
|
||||||
|
|
||||||
|
def fake_allgather(output, local_shard, *, group):
|
||||||
|
del group
|
||||||
|
output_sizes.append(output.numel())
|
||||||
|
output.zero_()
|
||||||
|
output[: local_shard.numel()].copy_(local_shard)
|
||||||
|
|
||||||
|
monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
|
||||||
|
|
||||||
|
hook.prefetch_layer(slot=0, non_blocking=False)
|
||||||
|
|
||||||
|
assert output_sizes == [4]
|
||||||
|
|
||||||
|
|
||||||
|
class TestPinnedModuleStager:
|
||||||
|
def test_offload_rebinds_immutable_cpu_master(self, patched_offload_runtime):
|
||||||
|
module = nn.Linear(2, 2)
|
||||||
|
expected_weight = module.weight.detach().clone()
|
||||||
|
expected_bias = module.bias.detach().clone()
|
||||||
|
stager = PinnedModuleStager(
|
||||||
|
module,
|
||||||
|
torch.device("cpu"),
|
||||||
|
pin_memory=False,
|
||||||
|
copy_stream=DummyStream(),
|
||||||
|
)
|
||||||
|
|
||||||
|
stager.load()
|
||||||
|
module.weight.data.zero_()
|
||||||
|
module.bias.data.zero_()
|
||||||
|
stager.offload()
|
||||||
|
|
||||||
|
assert torch.equal(module.weight, expected_weight)
|
||||||
|
assert torch.equal(module.bias, expected_bias)
|
||||||
|
stager.load()
|
||||||
|
assert torch.equal(module.weight, expected_weight)
|
||||||
|
assert torch.equal(module.bias, expected_bias)
|
||||||
|
stager.offload()
|
||||||
|
|
||||||
|
|
||||||
class _DummyBlock(nn.Module):
|
class _DummyBlock(nn.Module):
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
@@ -504,6 +665,81 @@ class TestCrossGroupSharedBuffer:
|
|||||||
assert not sync_called[0], "Should skip sync-prefetch when slot_group matches (non-contaminated)"
|
assert not sync_called[0], "Should skip sync-prefetch when slot_group matches (non-contaminated)"
|
||||||
|
|
||||||
|
|
||||||
|
class TestAllGatherOutputSlicing:
|
||||||
|
"""Regression test: with heterogeneous block sizes (e.g. MiniMax H3's
|
||||||
|
token_refiner with blocks of ~1231MB and ~239MB), the shared AllGather
|
||||||
|
output buffers are sized to the *largest* block across all groups.
|
||||||
|
prefetch_layer must slice the output buffer down to each block's actual
|
||||||
|
AllGather output size, otherwise output.numel() != dp_size * input.numel()
|
||||||
|
and all_gather_into_tensor raises during enable()'s first prefetch.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def test_prefetch_slices_shared_output_buffer_to_block_size(self, dist_group, patched_offload_runtime, monkeypatch):
|
||||||
|
dp_size = 2
|
||||||
|
|
||||||
|
# Heterogeneous block sizes: 8 vs 32 elements.
|
||||||
|
small_block = nn.Module()
|
||||||
|
small_block.weight = nn.Parameter(torch.arange(8, dtype=torch.float32))
|
||||||
|
large_block = nn.Module()
|
||||||
|
large_block.weight = nn.Parameter(torch.arange(100, 132, dtype=torch.float32))
|
||||||
|
|
||||||
|
def make_hook(next_block: nn.Module) -> DistributedLayerwiseOffloadHook:
|
||||||
|
hook = DistributedLayerwiseOffloadHook(
|
||||||
|
next_block=next_block,
|
||||||
|
device=torch.device("cpu"),
|
||||||
|
dp_group=dist.group.WORLD, # non-None -> AllGather path (mocked below)
|
||||||
|
dp_size=dp_size,
|
||||||
|
rank=0,
|
||||||
|
pin_memory=False,
|
||||||
|
)
|
||||||
|
current_block = nn.Module()
|
||||||
|
current_block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
|
||||||
|
hook.initialize_hook(current_block)
|
||||||
|
return hook
|
||||||
|
|
||||||
|
hook_small = make_hook(small_block)
|
||||||
|
hook_large = make_hook(large_block)
|
||||||
|
hooks = [hook_small, hook_large]
|
||||||
|
|
||||||
|
# Assign max-sized shared buffers exactly as backend.enable() does.
|
||||||
|
shared_buffers = DistributedLayerwiseOffloadBackend._allocate_shared_buffers(hooks)
|
||||||
|
shared_shard_buffers = DistributedLayerwiseOffloadBackend._allocate_shared_shard_buffers(hooks)
|
||||||
|
for hook in hooks:
|
||||||
|
hook.gpu_buffers = shared_buffers
|
||||||
|
hook.gpu_shard_buffers = shared_shard_buffers
|
||||||
|
|
||||||
|
# Shared output buffer must be sized to the largest block (32).
|
||||||
|
assert shared_buffers[0][torch.float32].numel() == 32
|
||||||
|
|
||||||
|
# Fake AllGather that enforces the real torch size contract and
|
||||||
|
# emulates reconstruction of the full block from all ranks' shards.
|
||||||
|
full_values = {
|
||||||
|
8: torch.arange(8, dtype=torch.float32),
|
||||||
|
32: torch.arange(100, 132, dtype=torch.float32),
|
||||||
|
}
|
||||||
|
ag_output_numels: list[int] = []
|
||||||
|
|
||||||
|
def fake_all_gather(output, input, group=None):
|
||||||
|
assert output.numel() == dp_size * input.numel(), (
|
||||||
|
f"AllGather contract violated: output.numel()={output.numel()} "
|
||||||
|
f"!= {dp_size} * input.numel()={input.numel()}"
|
||||||
|
)
|
||||||
|
ag_output_numels.append(output.numel())
|
||||||
|
output.copy_(full_values[output.numel()])
|
||||||
|
|
||||||
|
monkeypatch.setattr(dist, "all_gather_into_tensor", fake_all_gather)
|
||||||
|
|
||||||
|
# Small block: output must be sliced to dp_size * shard = 8, not 32.
|
||||||
|
hook_small.prefetch_layer(slot=0, non_blocking=False)
|
||||||
|
assert ag_output_numels[-1] == 8
|
||||||
|
assert torch.equal(small_block.weight, full_values[8])
|
||||||
|
|
||||||
|
# Large block: uses the full max-sized buffer.
|
||||||
|
hook_large.prefetch_layer(slot=1, non_blocking=False)
|
||||||
|
assert ag_output_numels[-1] == 32
|
||||||
|
assert torch.equal(large_block.weight, full_values[32])
|
||||||
|
|
||||||
|
|
||||||
class TestOffloadPlan:
|
class TestOffloadPlan:
|
||||||
"""Test declarative OffloadPlan metadata."""
|
"""Test declarative OffloadPlan metadata."""
|
||||||
|
|
||||||
@@ -517,6 +753,7 @@ class TestOffloadPlan:
|
|||||||
plan = OffloadPlan(
|
plan = OffloadPlan(
|
||||||
block_attrs={"transformer": ("gen_layers",)},
|
block_attrs={"transformer": ("gen_layers",)},
|
||||||
offload_submodules={"context_encoder": "layers"},
|
offload_submodules={"context_encoder": "layers"},
|
||||||
|
resident_dit_paths=frozenset({"transformer"}),
|
||||||
)
|
)
|
||||||
|
|
||||||
class PipelineWithPlan(nn.Module):
|
class PipelineWithPlan(nn.Module):
|
||||||
@@ -527,12 +764,14 @@ class TestOffloadPlan:
|
|||||||
assert result is plan
|
assert result is plan
|
||||||
assert result.block_attrs == {"transformer": ("gen_layers",)}
|
assert result.block_attrs == {"transformer": ("gen_layers",)}
|
||||||
assert result.offload_submodules == {"context_encoder": "layers"}
|
assert result.offload_submodules == {"context_encoder": "layers"}
|
||||||
|
assert result.resident_dit_paths == frozenset({"transformer"})
|
||||||
|
|
||||||
def test_offload_plan_defaults_to_empty(self):
|
def test_offload_plan_defaults_to_empty(self):
|
||||||
"""OffloadPlan with no arguments should have empty dicts."""
|
"""OffloadPlan with no arguments should have empty dicts."""
|
||||||
plan = OffloadPlan()
|
plan = OffloadPlan()
|
||||||
assert plan.block_attrs == {}
|
assert plan.block_attrs == {}
|
||||||
assert plan.offload_submodules == {}
|
assert plan.offload_submodules == {}
|
||||||
|
assert plan.resident_dit_paths == frozenset()
|
||||||
|
|
||||||
def test_offload_plan_is_frozen(self):
|
def test_offload_plan_is_frozen(self):
|
||||||
"""OffloadPlan should be immutable (frozen=True)."""
|
"""OffloadPlan should be immutable (frozen=True)."""
|
||||||
@@ -540,6 +779,50 @@ class TestOffloadPlan:
|
|||||||
with pytest.raises(Exception):
|
with pytest.raises(Exception):
|
||||||
plan.block_attrs = {"x": ("y",)} # type: ignore
|
plan.block_attrs = {"x": ("y",)} # type: ignore
|
||||||
|
|
||||||
|
def test_all_resident_dit_still_prepares_planned_submodules(self, patched_offload_runtime):
|
||||||
|
class TokenRefiner(nn.Module):
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__()
|
||||||
|
self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
|
||||||
|
|
||||||
|
class Transformer(nn.Module):
|
||||||
|
_layerwise_offload_blocks_attrs = ["blocks"]
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__()
|
||||||
|
self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
|
||||||
|
self.token_refiner = TokenRefiner()
|
||||||
|
|
||||||
|
class Pipeline(nn.Module):
|
||||||
|
_offload_plan = OffloadPlan(
|
||||||
|
offload_submodules={"token_refiner": "blocks"},
|
||||||
|
resident_dit_paths=frozenset({"transformer"}),
|
||||||
|
)
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__()
|
||||||
|
self.transformer = Transformer()
|
||||||
|
|
||||||
|
pipeline = Pipeline()
|
||||||
|
backend = DistributedLayerwiseOffloadBackend(
|
||||||
|
OffloadConfig(
|
||||||
|
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
|
||||||
|
pin_cpu_memory=False,
|
||||||
|
dlo_use_allgather=False,
|
||||||
|
dlo_resident_layers=2,
|
||||||
|
),
|
||||||
|
torch.device("cpu"),
|
||||||
|
)
|
||||||
|
|
||||||
|
backend.enable(pipeline)
|
||||||
|
|
||||||
|
assert len(backend._resident_blocks) == 2
|
||||||
|
assert len(backend._all_hook_groups) == 1
|
||||||
|
assert len(backend._all_hook_groups[0]) == 2
|
||||||
|
assert backend.enabled
|
||||||
|
|
||||||
|
backend.disable()
|
||||||
|
|
||||||
|
|
||||||
class TestMmapValidation:
|
class TestMmapValidation:
|
||||||
"""Tests for mmap loader validation: TP rejection, strict check, validate_loaded_weights."""
|
"""Tests for mmap loader validation: TP rejection, strict check, validate_loaded_weights."""
|
||||||
@@ -746,7 +1029,7 @@ class TestConfigValidation:
|
|||||||
OffloadConfig.from_od_config(FakeODConfig())
|
OffloadConfig.from_od_config(FakeODConfig())
|
||||||
|
|
||||||
def test_hsdp_without_allgather_allowed(self):
|
def test_hsdp_without_allgather_allowed(self):
|
||||||
"""HSDP + DLO + no-AllGather should be allowed (full weights per rank)."""
|
"""HSDP + DLO + no-AllGather uses standard-loader rank-local weights."""
|
||||||
from vllm_omni.diffusion.offloader.base import OffloadConfig
|
from vllm_omni.diffusion.offloader.base import OffloadConfig
|
||||||
|
|
||||||
class FakePC:
|
class FakePC:
|
||||||
@@ -768,12 +1051,33 @@ class TestConfigValidation:
|
|||||||
enable_layerwise_offload = False
|
enable_layerwise_offload = False
|
||||||
enable_distributed_layerwise_offload = True
|
enable_distributed_layerwise_offload = True
|
||||||
dlo_use_allgather = False # no AllGather → should be allowed
|
dlo_use_allgather = False # no AllGather → should be allowed
|
||||||
|
dlo_resident_layers = 20
|
||||||
pin_cpu_memory = True
|
pin_cpu_memory = True
|
||||||
parallel_config = FakePC()
|
parallel_config = FakePC()
|
||||||
model = "/fake/path"
|
model = "/fake/path"
|
||||||
|
|
||||||
config = OffloadConfig.from_od_config(FakeODConfig())
|
config = OffloadConfig.from_od_config(FakeODConfig())
|
||||||
assert config.dp_size == 1 # forced to 1 when no AllGather
|
assert config.dp_size == 1 # forced to 1 when no AllGather
|
||||||
|
assert config.dlo_resident_layers == 20
|
||||||
|
|
||||||
|
def test_resident_layers_with_allgather_rejected(self):
|
||||||
|
class FakePC:
|
||||||
|
data_parallel_size = 2
|
||||||
|
use_hsdp = False
|
||||||
|
sequence_parallel_size = 1
|
||||||
|
|
||||||
|
class FakeODConfig:
|
||||||
|
enable_cpu_offload = False
|
||||||
|
enable_layerwise_offload = False
|
||||||
|
enable_distributed_layerwise_offload = True
|
||||||
|
dlo_use_allgather = True
|
||||||
|
dlo_resident_layers = 20
|
||||||
|
pin_cpu_memory = True
|
||||||
|
parallel_config = FakePC()
|
||||||
|
model = "/fake/path"
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="requires --dlo-no-use-allgather"):
|
||||||
|
OffloadConfig.from_od_config(FakeODConfig())
|
||||||
|
|
||||||
def test_num_inference_steps_none_rejected(self):
|
def test_num_inference_steps_none_rejected(self):
|
||||||
"""DP multi-concurrency should reject None num_inference_steps."""
|
"""DP multi-concurrency should reject None num_inference_steps."""
|
||||||
|
|||||||
@@ -247,6 +247,56 @@ class TestRequestModeDispatch:
|
|||||||
results = {ro.request_id: ro.result.error for ro in out.runner_outputs}
|
results = {ro.request_id: ro.result.error for ro in out.runner_outputs}
|
||||||
assert results == {request_id: f"result_for_{request_id}" for request_id in request_ids}
|
assert results == {request_id: f"result_for_{request_id}" for request_id in request_ids}
|
||||||
|
|
||||||
|
def test_dlo_dp_routes_multiple_requests_without_pipeline_batching(self):
|
||||||
|
executor, _, _ = _make_executor(num_gpus=2)
|
||||||
|
executor.od_config = SimpleNamespace(
|
||||||
|
step_execution=False,
|
||||||
|
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||||||
|
enable_distributed_layerwise_offload=True,
|
||||||
|
dlo_use_allgather=True,
|
||||||
|
)
|
||||||
|
executor.execute_request = Mock(return_value="dlo-dp")
|
||||||
|
executor.collective_rpc = Mock()
|
||||||
|
scheduler_output = _make_sched_output("A", "B")
|
||||||
|
|
||||||
|
result = executor.execute_batch(scheduler_output)
|
||||||
|
|
||||||
|
assert result == "dlo-dp"
|
||||||
|
executor.execute_request.assert_called_once_with(scheduler_output)
|
||||||
|
executor.collective_rpc.assert_not_called()
|
||||||
|
|
||||||
|
def test_dlo_dp_multi_rank_reply_uses_synchronous_rpc_collection(self):
|
||||||
|
executor, req_q, res_q = _make_executor(num_gpus=2)
|
||||||
|
executor.od_config = SimpleNamespace(
|
||||||
|
step_execution=False,
|
||||||
|
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||||||
|
)
|
||||||
|
executor._sync_result_buffer = res_q
|
||||||
|
|
||||||
|
def worker():
|
||||||
|
request = req_q.get(timeout=2)
|
||||||
|
assert "rpc_id" not in request
|
||||||
|
wave_id = request["wave_id"]
|
||||||
|
for dp_rank in range(2):
|
||||||
|
res_q.put(
|
||||||
|
{
|
||||||
|
"dp_rank": dp_rank,
|
||||||
|
"output": _tagged_output(str(dp_rank)),
|
||||||
|
"wave_id": wave_id,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
thread = threading.Thread(target=worker, daemon=True)
|
||||||
|
thread.start()
|
||||||
|
results = executor.collective_rpc(
|
||||||
|
"execute_model",
|
||||||
|
unique_reply_rank=None,
|
||||||
|
exec_all_ranks=True,
|
||||||
|
)
|
||||||
|
thread.join(timeout=2)
|
||||||
|
|
||||||
|
assert [result.error for result in results] == ["0", "1"]
|
||||||
|
|
||||||
|
|
||||||
# ───────────────── concurrent collective RPC ─────────────────
|
# ───────────────── concurrent collective RPC ─────────────────
|
||||||
|
|
||||||
|
|||||||
@@ -390,6 +390,105 @@ def test_send_single_request_respects_processor_receiver_boundary(build_adapter,
|
|||||||
assert sent_payload.meta.is_segment_finished.item() is False
|
assert sent_payload.meta.is_segment_finished.item() is False
|
||||||
|
|
||||||
|
|
||||||
|
def test_send_single_request_personaplex_pending_frame_is_not_segment_boundary(
|
||||||
|
build_adapter,
|
||||||
|
):
|
||||||
|
from vllm_omni.model_executor.stage_input_processors.personaplex import (
|
||||||
|
talker2code2wav_async_chunk,
|
||||||
|
)
|
||||||
|
|
||||||
|
adapter, connector = build_adapter(
|
||||||
|
stage_id=0,
|
||||||
|
connector_extra={
|
||||||
|
"initial_codec_chunk_frames": 1,
|
||||||
|
"codec_chunk_frames": 5,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
request = _req(
|
||||||
|
"req-personaplex",
|
||||||
|
RequestStatus.WAITING,
|
||||||
|
external_req_id="ext-personaplex",
|
||||||
|
)
|
||||||
|
request.resumable = True
|
||||||
|
request.additional_information = {
|
||||||
|
"codes": {
|
||||||
|
"audio": torch.arange(8, dtype=torch.long).reshape(1, 8),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
adapter.custom_process_next_stage_input_func = talker2code2wav_async_chunk
|
||||||
|
|
||||||
|
adapter._send_single_request(
|
||||||
|
{
|
||||||
|
"multimodal_output": None,
|
||||||
|
"request": request,
|
||||||
|
"is_finished": False,
|
||||||
|
"is_segment_finished": True,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
sent_payload = connector.put.call_args.kwargs["data"]
|
||||||
|
assert sent_payload.codes is None
|
||||||
|
assert sent_payload.meta.finished.item() is False
|
||||||
|
assert sent_payload.meta.is_segment_finished.item() is False
|
||||||
|
|
||||||
|
|
||||||
|
def test_personaplex_sender_cleanup_drops_delayed_frame_state(build_adapter):
|
||||||
|
from vllm_omni.model_executor.stage_input_processors.personaplex import (
|
||||||
|
talker2code2wav_async_chunk,
|
||||||
|
)
|
||||||
|
|
||||||
|
adapter, _ = build_adapter(
|
||||||
|
stage_id=0,
|
||||||
|
connector_extra={
|
||||||
|
"initial_codec_chunk_frames": 1,
|
||||||
|
"codec_chunk_frames": 5,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
request = _req(
|
||||||
|
"req-personaplex-first",
|
||||||
|
RequestStatus.WAITING,
|
||||||
|
external_req_id="ext-personaplex-reused",
|
||||||
|
)
|
||||||
|
request.resumable = True
|
||||||
|
request.additional_information = {
|
||||||
|
"codes": {
|
||||||
|
"audio": torch.arange(8, dtype=torch.long).reshape(1, 8),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
first = talker2code2wav_async_chunk(
|
||||||
|
adapter,
|
||||||
|
multimodal_output=None,
|
||||||
|
request=request,
|
||||||
|
is_finished=True,
|
||||||
|
)
|
||||||
|
assert first is not None
|
||||||
|
assert first.codes is None
|
||||||
|
|
||||||
|
adapter.cleanup_sender(request.external_req_id)
|
||||||
|
|
||||||
|
replacement = _req(
|
||||||
|
"req-personaplex-replacement",
|
||||||
|
RequestStatus.WAITING,
|
||||||
|
external_req_id=request.external_req_id,
|
||||||
|
)
|
||||||
|
replacement.resumable = True
|
||||||
|
replacement.additional_information = {
|
||||||
|
"codes": {
|
||||||
|
"audio": torch.arange(8, 16, dtype=torch.long).reshape(1, 8),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
second = talker2code2wav_async_chunk(
|
||||||
|
adapter,
|
||||||
|
multimodal_output=None,
|
||||||
|
request=replacement,
|
||||||
|
is_finished=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert second is not None
|
||||||
|
assert second.codes is None
|
||||||
|
|
||||||
|
|
||||||
def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapter):
|
def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapter):
|
||||||
adapter, _ = build_adapter(stage_id=1)
|
adapter, _ = build_adapter(stage_id=1)
|
||||||
request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
|
request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
|
||||||
|
|||||||
@@ -0,0 +1 @@
|
|||||||
|
"""End-to-end accuracy tests."""
|
||||||
@@ -3,6 +3,7 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import hashlib
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
import subprocess
|
import subprocess
|
||||||
@@ -26,9 +27,13 @@ from tests.helpers.runtime import OmniServer
|
|||||||
pytestmark = [pytest.mark.benchmark, pytest.mark.diffusion, pytest.mark.full_model]
|
pytestmark = [pytest.mark.benchmark, pytest.mark.diffusion, pytest.mark.full_model]
|
||||||
|
|
||||||
MODEL_REPO_ID = "MiniMaxAI/MiniMax-H3"
|
MODEL_REPO_ID = "MiniMaxAI/MiniMax-H3"
|
||||||
MODEL_REVISION = "73372e6cf53e414edd3ab03e357717fb0602e758"
|
# The previously pinned revision was removed from the Hugging Face repository.
|
||||||
|
MODEL_REVISION = "main"
|
||||||
MODEL_ENV_VAR = "VLLM_TEST_MINIMAX_H3_FL2VA_MODEL"
|
MODEL_ENV_VAR = "VLLM_TEST_MINIMAX_H3_FL2VA_MODEL"
|
||||||
REFERENCE_VIDEO_URL = f"https://huggingface.co/{MODEL_REPO_ID}/resolve/main/assets/t2va.mp4"
|
# Keep the official asset link visible in the test report. Hugging Face's
|
||||||
|
# ``blob`` URL is an HTML page, so resolve it to the raw asset for download.
|
||||||
|
REFERENCE_VIDEO_URL = "https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/main/assets/t2va.mp4"
|
||||||
|
REFERENCE_VIDEO_SHA256 = "d66903241362e224085cc93f7a5e70fba6ab378d0ac6ff6af83acf2559849a42"
|
||||||
|
|
||||||
# This is the official reproducible 768p T2VA prompt associated with
|
# This is the official reproducible 768p T2VA prompt associated with
|
||||||
# ``assets/t2va.mp4`` in the MiniMax H3 repository.
|
# ``assets/t2va.mp4`` in the MiniMax H3 repository.
|
||||||
@@ -93,9 +98,15 @@ def _server_args() -> list[str]:
|
|||||||
|
|
||||||
def _download_reference_video(output_dir: Path) -> Path:
|
def _download_reference_video(output_dir: Path) -> Path:
|
||||||
reference_path = output_dir / "reference.mp4"
|
reference_path = output_dir / "reference.mp4"
|
||||||
response = requests.get(REFERENCE_VIDEO_URL, timeout=120)
|
download_url = REFERENCE_VIDEO_URL.replace("/blob/", "/resolve/", 1)
|
||||||
|
response = requests.get(download_url, timeout=120)
|
||||||
response.raise_for_status()
|
response.raise_for_status()
|
||||||
reference_path.write_bytes(response.content)
|
content = response.content
|
||||||
|
digest = hashlib.sha256(content).hexdigest()
|
||||||
|
assert digest == REFERENCE_VIDEO_SHA256, (
|
||||||
|
f"Official MiniMax H3 T2VA reference changed: got {digest}, expected {REFERENCE_VIDEO_SHA256}"
|
||||||
|
)
|
||||||
|
reference_path.write_bytes(content)
|
||||||
return reference_path
|
return reference_path
|
||||||
|
|
||||||
|
|
||||||
@@ -159,6 +170,7 @@ def test_minimax_h3_t2va_matches_official_reference(
|
|||||||
{
|
{
|
||||||
"task": "t2va",
|
"task": "t2va",
|
||||||
"duration": DURATION_SECONDS,
|
"duration": DURATION_SECONDS,
|
||||||
|
"aspect_ratio": "16:9",
|
||||||
"audio_flow_shift": AUDIO_FLOW_SHIFT,
|
"audio_flow_shift": AUDIO_FLOW_SHIFT,
|
||||||
}
|
}
|
||||||
),
|
),
|
||||||
|
|||||||
@@ -0,0 +1,30 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
|
||||||
|
import subprocess
|
||||||
|
import sys
|
||||||
|
|
||||||
|
|
||||||
|
def test_batched_runtime_import_does_not_require_optional_opus_dependency() -> None:
|
||||||
|
script = """
|
||||||
|
import builtins
|
||||||
|
|
||||||
|
real_import = builtins.__import__
|
||||||
|
|
||||||
|
def guarded_import(name, *args, **kwargs):
|
||||||
|
if name == "sphn":
|
||||||
|
raise ModuleNotFoundError("blocked optional dependency: sphn")
|
||||||
|
return real_import(name, *args, **kwargs)
|
||||||
|
|
||||||
|
builtins.__import__ = guarded_import
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.serving.batched import BatchedSessionManager
|
||||||
|
assert BatchedSessionManager is not None
|
||||||
|
"""
|
||||||
|
result = subprocess.run(
|
||||||
|
[sys.executable, "-c", script],
|
||||||
|
check=False,
|
||||||
|
capture_output=True,
|
||||||
|
text=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.returncode == 0, result.stderr
|
||||||
@@ -0,0 +1,303 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
|
||||||
|
import base64
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import pytest
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.policy import (
|
||||||
|
SILENCE_TOKENS,
|
||||||
|
SINE_TOKENS,
|
||||||
|
)
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.stage0 import (
|
||||||
|
PersonaPlexStage0DuplexRuntime,
|
||||||
|
)
|
||||||
|
from vllm_omni.model_executor.models.personaplex.personaplex_talker import (
|
||||||
|
PersonaPlexTalkerForConditionalGeneration,
|
||||||
|
)
|
||||||
|
|
||||||
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeCodec:
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.encode_calls = 0
|
||||||
|
self.reset_calls = 0
|
||||||
|
|
||||||
|
def streaming_init(self, batch_size: int) -> None:
|
||||||
|
assert batch_size == 1
|
||||||
|
|
||||||
|
def encode_frame(self, pcm: torch.Tensor) -> torch.Tensor:
|
||||||
|
assert pcm.shape == (1, 1920)
|
||||||
|
self.encode_calls += 1
|
||||||
|
return torch.full((1, 8), self.encode_calls, dtype=torch.long)
|
||||||
|
|
||||||
|
def reset_streaming(self) -> None:
|
||||||
|
self.reset_calls += 1
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeTalker:
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.device = torch.device("cpu")
|
||||||
|
self.dtype = torch.float32
|
||||||
|
self.frame_calls: list[dict[str, torch.Tensor | None]] = []
|
||||||
|
|
||||||
|
def _build_prefill_embed(
|
||||||
|
self,
|
||||||
|
tokens,
|
||||||
|
offset,
|
||||||
|
span,
|
||||||
|
device,
|
||||||
|
silence=None,
|
||||||
|
user_sine=None,
|
||||||
|
):
|
||||||
|
del offset, silence, user_sine
|
||||||
|
values = tokens[:span].to(device=device, dtype=torch.float32)
|
||||||
|
return values[:, None].expand(-1, 4).contiguous()
|
||||||
|
|
||||||
|
def _build_frame_embed(
|
||||||
|
self,
|
||||||
|
text_token,
|
||||||
|
last_agent,
|
||||||
|
prev_agent,
|
||||||
|
device,
|
||||||
|
user_d0=None,
|
||||||
|
user_d1=None,
|
||||||
|
):
|
||||||
|
self.frame_calls.append(
|
||||||
|
{
|
||||||
|
"text_token": text_token.clone(),
|
||||||
|
"last_agent": None if last_agent is None else last_agent.clone(),
|
||||||
|
"prev_agent": None if prev_agent is None else prev_agent.clone(),
|
||||||
|
"user_d0": None if user_d0 is None else user_d0.clone(),
|
||||||
|
"user_d1": None if user_d1 is None else user_d1.clone(),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
value = float(user_d0[0].item()) if user_d0 is not None else -1.0
|
||||||
|
return torch.full((1, 4), value, device=device)
|
||||||
|
|
||||||
|
|
||||||
|
def _duplex_info(*, seq: int, session_id: str = "session", incarnation: int = 1):
|
||||||
|
pcm = np.zeros(1920, dtype="<f4")
|
||||||
|
return {
|
||||||
|
"data_plane": True,
|
||||||
|
"session_id": session_id,
|
||||||
|
"incarnation": incarnation,
|
||||||
|
"epoch": 0,
|
||||||
|
"seq": seq,
|
||||||
|
"payload": {
|
||||||
|
"format": "pcm_f32le",
|
||||||
|
"sample_rate_hz": 24000,
|
||||||
|
"audio": base64.b64encode(pcm.tobytes()).decode("ascii"),
|
||||||
|
},
|
||||||
|
"runtime_config": {
|
||||||
|
"personaplex_model_path": "/unused",
|
||||||
|
"personaplex_voice_prompt": "NATF2.pt",
|
||||||
|
"personaplex_persona": "Be concise.",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _runtime(*codecs: _FakeCodec) -> PersonaPlexStage0DuplexRuntime:
|
||||||
|
voice_embeddings = torch.arange(8, dtype=torch.float32).reshape(2, 1, 1, 4)
|
||||||
|
codec_args = {"codec": codecs[0]}
|
||||||
|
if len(codecs) > 1:
|
||||||
|
available = iter(codecs)
|
||||||
|
codec_args = {
|
||||||
|
"codec_factory": lambda: next(available),
|
||||||
|
"max_sessions": len(codecs),
|
||||||
|
}
|
||||||
|
return PersonaPlexStage0DuplexRuntime(
|
||||||
|
_FakeTalker(),
|
||||||
|
model_path="/unused",
|
||||||
|
device="cpu",
|
||||||
|
**codec_args,
|
||||||
|
tokenizer=lambda _text: [7, 8, 9],
|
||||||
|
voice_loader=lambda _voice: {
|
||||||
|
"embeddings": voice_embeddings,
|
||||||
|
"cache": torch.zeros((1, 17, 4), dtype=torch.long),
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_first_append_prepends_voice_and_persona_once() -> None:
|
||||||
|
codec = _FakeCodec()
|
||||||
|
runtime = _runtime(codec)
|
||||||
|
|
||||||
|
first = runtime.prepare_append(_duplex_info(seq=1), prompt_len=18)
|
||||||
|
second = runtime.prepare_append(_duplex_info(seq=2), prompt_len=18)
|
||||||
|
|
||||||
|
assert first.prefill_applied is True
|
||||||
|
assert first.prompt_offset == 0
|
||||||
|
assert first.user_codes.shape == (1, 8)
|
||||||
|
assert first.inputs_embeds.shape == (18, 4)
|
||||||
|
assert second.prefill_applied is False
|
||||||
|
assert second.prompt_offset == 17
|
||||||
|
assert second.user_codes.shape == (2, 8)
|
||||||
|
assert second.inputs_embeds.shape == (1, 4)
|
||||||
|
assert codec.encode_calls == 2
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("append_count", "expected_tokens", "expected_provided"),
|
||||||
|
[
|
||||||
|
(
|
||||||
|
1,
|
||||||
|
[*SILENCE_TOKENS, 1, *SINE_TOKENS[1:]],
|
||||||
|
[False, *([True] * 15)],
|
||||||
|
),
|
||||||
|
(
|
||||||
|
2,
|
||||||
|
[*SILENCE_TOKENS, 2, *([1] * 7)],
|
||||||
|
[*([False] * 8), *([True] * 8)],
|
||||||
|
),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_prepare_append_preserves_native_depformer_teacher_forcing(
|
||||||
|
append_count: int,
|
||||||
|
expected_tokens: list[int],
|
||||||
|
expected_provided: list[bool],
|
||||||
|
) -> None:
|
||||||
|
runtime = _runtime(_FakeCodec())
|
||||||
|
|
||||||
|
result = None
|
||||||
|
for seq in range(1, append_count + 1):
|
||||||
|
result = runtime.prepare_append(_duplex_info(seq=seq), prompt_len=18)
|
||||||
|
|
||||||
|
assert result is not None
|
||||||
|
assert result.info_update["pplex_depformer_audio_tokens"].tolist() == expected_tokens
|
||||||
|
assert result.info_update["pplex_depformer_audio_provided"].tolist() == expected_provided
|
||||||
|
|
||||||
|
|
||||||
|
def test_frame_embed_uses_previous_effective_agent_frame() -> None:
|
||||||
|
fake_talker = SimpleNamespace(
|
||||||
|
config=SimpleNamespace(
|
||||||
|
num_audio_codebooks=16,
|
||||||
|
audio_vocab_size=2048,
|
||||||
|
),
|
||||||
|
input_embeddings=lambda stack: stack,
|
||||||
|
)
|
||||||
|
last_agent = torch.arange(10, 18)
|
||||||
|
|
||||||
|
embeds = PersonaPlexTalkerForConditionalGeneration._build_frame_embed(
|
||||||
|
fake_talker,
|
||||||
|
torch.tensor([3]),
|
||||||
|
last_agent,
|
||||||
|
torch.arange(8),
|
||||||
|
torch.device("cpu"),
|
||||||
|
)
|
||||||
|
|
||||||
|
assert embeds.shape == (1, 17)
|
||||||
|
assert torch.equal(embeds[0, 1:9], last_agent)
|
||||||
|
|
||||||
|
|
||||||
|
def test_repeated_append_identity_does_not_advance_codec() -> None:
|
||||||
|
codec = _FakeCodec()
|
||||||
|
runtime = _runtime(codec)
|
||||||
|
|
||||||
|
info = _duplex_info(seq=1)
|
||||||
|
first = runtime.prepare_append(info, prompt_len=18, request_id="req")
|
||||||
|
retry = runtime.prepare_append(info, prompt_len=18, request_id="req")
|
||||||
|
|
||||||
|
assert codec.encode_calls == 1
|
||||||
|
assert torch.equal(first.inputs_embeds, retry.inputs_embeds)
|
||||||
|
assert torch.equal(first.user_codes, retry.user_codes)
|
||||||
|
|
||||||
|
|
||||||
|
def test_next_append_uses_prior_sample_and_causally_delayed_user_frame() -> None:
|
||||||
|
runtime = _runtime(_FakeCodec())
|
||||||
|
|
||||||
|
runtime.prepare_append(_duplex_info(seq=1), prompt_len=18, request_id="req")
|
||||||
|
first_agent = torch.arange(8, dtype=torch.long)
|
||||||
|
runtime.record_sample(request_id="req", text_token=torch.tensor(101), agent_codes=first_agent)
|
||||||
|
runtime.prepare_append(_duplex_info(seq=2), prompt_len=19, request_id="req")
|
||||||
|
second_agent = torch.arange(10, 18, dtype=torch.long)
|
||||||
|
runtime.record_sample(request_id="req", text_token=torch.tensor(102), agent_codes=second_agent)
|
||||||
|
runtime.prepare_append(_duplex_info(seq=3), prompt_len=20, request_id="req")
|
||||||
|
|
||||||
|
silence = torch.tensor(SILENCE_TOKENS, dtype=torch.long)
|
||||||
|
sine = torch.tensor(SINE_TOKENS, dtype=torch.long)
|
||||||
|
first_call = runtime.stage_model.frame_calls[0]
|
||||||
|
assert all(torch.equal(first_call[key], silence) for key in ("last_agent", "prev_agent"))
|
||||||
|
assert all(torch.equal(first_call[key], sine) for key in ("user_d0", "user_d1"))
|
||||||
|
|
||||||
|
second_call = runtime.stage_model.frame_calls[1]
|
||||||
|
assert second_call["text_token"].tolist() == [101]
|
||||||
|
expected_first_effective = torch.cat([first_agent[:1], silence[1:]])
|
||||||
|
assert torch.equal(second_call["last_agent"], expected_first_effective)
|
||||||
|
assert torch.equal(second_call["prev_agent"], expected_first_effective)
|
||||||
|
assert torch.equal(second_call["user_d0"], torch.full((8,), 1, dtype=torch.long))
|
||||||
|
assert torch.equal(second_call["user_d1"], sine)
|
||||||
|
|
||||||
|
third_call = runtime.stage_model.frame_calls[2]
|
||||||
|
assert third_call["text_token"].tolist() == [102]
|
||||||
|
assert torch.equal(third_call["last_agent"], second_agent)
|
||||||
|
assert torch.equal(third_call["prev_agent"], second_agent)
|
||||||
|
assert torch.equal(third_call["user_d0"], torch.full((8,), 2, dtype=torch.long))
|
||||||
|
assert torch.equal(third_call["user_d1"], torch.full((8,), 1, dtype=torch.long))
|
||||||
|
|
||||||
|
|
||||||
|
def test_decoded_pcm_is_writable_for_torch_zero_copy() -> None:
|
||||||
|
pcm = PersonaPlexStage0DuplexRuntime._decode_pcm(_duplex_info(seq=1)["payload"])
|
||||||
|
|
||||||
|
assert pcm.flags.writeable
|
||||||
|
|
||||||
|
|
||||||
|
def _prepare_two_sessions(
|
||||||
|
runtime: PersonaPlexStage0DuplexRuntime,
|
||||||
|
) -> tuple[object, object, object, object]:
|
||||||
|
first_1 = runtime.prepare_append(_duplex_info(seq=1), prompt_len=18)
|
||||||
|
second_1 = runtime.prepare_append(
|
||||||
|
_duplex_info(seq=1, session_id="other", incarnation=2),
|
||||||
|
prompt_len=18,
|
||||||
|
)
|
||||||
|
first_2 = runtime.prepare_append(_duplex_info(seq=2), prompt_len=19)
|
||||||
|
second_2 = runtime.prepare_append(
|
||||||
|
_duplex_info(seq=2, session_id="other", incarnation=2),
|
||||||
|
prompt_len=19,
|
||||||
|
)
|
||||||
|
return first_1, second_1, first_2, second_2
|
||||||
|
|
||||||
|
|
||||||
|
def test_live_sessions_keep_independent_streaming_encoders() -> None:
|
||||||
|
first_codec = _FakeCodec()
|
||||||
|
second_codec = _FakeCodec()
|
||||||
|
runtime = _runtime(first_codec, second_codec)
|
||||||
|
|
||||||
|
first_1, second_1, first_2, second_2 = _prepare_two_sessions(runtime)
|
||||||
|
|
||||||
|
assert first_1.user_codes[:, 0].tolist() == [1]
|
||||||
|
assert second_1.user_codes[:, 0].tolist() == [1]
|
||||||
|
assert first_2.user_codes[:, 0].tolist() == [1, 2]
|
||||||
|
assert second_2.user_codes[:, 0].tolist() == [1, 2]
|
||||||
|
|
||||||
|
|
||||||
|
def test_stage0_session_capacity_fails_before_codec_state_is_shared() -> None:
|
||||||
|
runtime = _runtime(_FakeCodec(), _FakeCodec())
|
||||||
|
_prepare_two_sessions(runtime)
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="capacity 2"):
|
||||||
|
runtime.prepare_append(
|
||||||
|
_duplex_info(seq=1, session_id="third", incarnation=3),
|
||||||
|
prompt_len=18,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_close_session_resets_and_reuses_released_codec() -> None:
|
||||||
|
first_codec = _FakeCodec()
|
||||||
|
second_codec = _FakeCodec()
|
||||||
|
runtime = _runtime(first_codec, second_codec)
|
||||||
|
_prepare_two_sessions(runtime)
|
||||||
|
|
||||||
|
runtime.close_session("session", 1)
|
||||||
|
|
||||||
|
assert first_codec.reset_calls == 1
|
||||||
|
assert second_codec.reset_calls == 0
|
||||||
|
replacement = runtime.prepare_append(
|
||||||
|
_duplex_info(seq=1, session_id="replacement", incarnation=4),
|
||||||
|
prompt_len=18,
|
||||||
|
)
|
||||||
|
assert replacement.user_codes[:, 0].tolist() == [3]
|
||||||
@@ -0,0 +1,316 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import torch
|
||||||
|
from torch import nn
|
||||||
|
|
||||||
|
from vllm_omni.model_executor.models.personaplex import (
|
||||||
|
personaplex_code2wav,
|
||||||
|
personaplex_mimi,
|
||||||
|
)
|
||||||
|
from vllm_omni.model_executor.models.personaplex.personaplex_code2wav import (
|
||||||
|
PersonaPlexCode2Wav,
|
||||||
|
)
|
||||||
|
|
||||||
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeStreamingMimi(nn.Module):
|
||||||
|
def __init__(self, stream_value: int = 0) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.stream_value = stream_value
|
||||||
|
self.decode_frame_calls = 0
|
||||||
|
self.reset_calls = 0
|
||||||
|
|
||||||
|
def decode(self, codes: torch.Tensor, return_dict: bool = True):
|
||||||
|
assert return_dict
|
||||||
|
frames = int(codes.shape[-1])
|
||||||
|
return SimpleNamespace(audio_values=torch.arange(frames * 4, dtype=torch.float32).reshape(1, 1, -1))
|
||||||
|
|
||||||
|
def streaming_init(self, batch_size: int) -> None:
|
||||||
|
assert batch_size == 1
|
||||||
|
|
||||||
|
def decode_frame(self, codes: torch.Tensor) -> torch.Tensor:
|
||||||
|
assert codes.shape == (1, 2)
|
||||||
|
self.decode_frame_calls += 1
|
||||||
|
value = self.stream_value + self.decode_frame_calls
|
||||||
|
return torch.full((1, 4), float(value), dtype=torch.float32)
|
||||||
|
|
||||||
|
def reset_streaming(self) -> None:
|
||||||
|
self.reset_calls += 1
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeChunkedMimi(_FakeStreamingMimi):
|
||||||
|
def __init__(self) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.decode_frames_calls = 0
|
||||||
|
self.decode_frames_shapes: list[tuple[int, ...]] = []
|
||||||
|
|
||||||
|
def decode_frames(self, codes: torch.Tensor) -> torch.Tensor:
|
||||||
|
self.decode_frames_calls += 1
|
||||||
|
self.decode_frames_shapes.append(tuple(codes.shape))
|
||||||
|
return torch.arange(codes.shape[-1] * 4, dtype=torch.float32).reshape(1, -1)
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeMimiModel(nn.Module):
|
||||||
|
def __init__(self, _config) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.encoder_transformer = nn.Linear(1, 1, bias=False)
|
||||||
|
self.decoder_transformer = nn.Linear(1, 1, bias=False)
|
||||||
|
self.quantizer = nn.Linear(1, 1, bias=False)
|
||||||
|
self.encoder = SimpleNamespace(layers=[])
|
||||||
|
self.downsample = SimpleNamespace(conv=nn.Conv1d(1, 1, 1))
|
||||||
|
self.upsample = SimpleNamespace(conv=nn.ConvTranspose1d(1, 1, 1))
|
||||||
|
self.decoder = SimpleNamespace(layers=[])
|
||||||
|
|
||||||
|
def load_state_dict(self, _state_dict, strict: bool = True):
|
||||||
|
assert not strict
|
||||||
|
return SimpleNamespace(
|
||||||
|
missing_keys=[
|
||||||
|
"encoder_transformer.weight",
|
||||||
|
"decoder_transformer.weight",
|
||||||
|
],
|
||||||
|
unexpected_keys=[],
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeMimiTransformer(nn.Module):
|
||||||
|
def __init__(self) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.weight = nn.Parameter(torch.zeros(1))
|
||||||
|
|
||||||
|
def load_weights(self, _state_dict, _prefix: str) -> int:
|
||||||
|
return 80
|
||||||
|
|
||||||
|
|
||||||
|
def test_moshi_checkpoint_mapping_drops_replaced_transformer_weights(
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
mapper = getattr(personaplex_mimi, "_map_moshi_codec_weights", None)
|
||||||
|
assert callable(mapper), "PersonaPlex Mimi must load its bundled codec weights without Hugging Face"
|
||||||
|
|
||||||
|
source = {
|
||||||
|
"encoder.model.3.conv.conv.weight": torch.ones(1),
|
||||||
|
"downsample.conv.conv.conv.weight": torch.ones(1),
|
||||||
|
"quantizer.rvq_first.vq.layers.0._codebook.embedding_sum": torch.ones(1),
|
||||||
|
"quantizer.rvq_rest.vq.layers.3._codebook._initialized": torch.ones(1),
|
||||||
|
"encoder_transformer.transformer.layers.0.norm1.weight": torch.ones(1),
|
||||||
|
}
|
||||||
|
|
||||||
|
mapped = mapper(source)
|
||||||
|
|
||||||
|
assert set(mapped) == {
|
||||||
|
"encoder.layers.3.conv.weight",
|
||||||
|
"downsample.conv.weight",
|
||||||
|
"quantizer.semantic_residual_vector_quantizer.layers.0.codebook.embed_sum",
|
||||||
|
"quantizer.acoustic_residual_vector_quantizer.layers.3.codebook.initialized",
|
||||||
|
}
|
||||||
|
assert mapped["encoder.layers.3.conv.weight"] is source["encoder.model.3.conv.conv.weight"]
|
||||||
|
monkeypatch.setattr("transformers.MimiConfig", lambda: object())
|
||||||
|
monkeypatch.setattr("transformers.MimiModel", _FakeMimiModel)
|
||||||
|
monkeypatch.setattr("safetensors.torch.load_file", lambda *_args, **_kwargs: {})
|
||||||
|
monkeypatch.setattr(personaplex_mimi, "_MimiStreamingTransformer", _FakeMimiTransformer)
|
||||||
|
|
||||||
|
codec = personaplex_mimi.PersonaPlexMimiCodec(
|
||||||
|
checkpoint="/unused",
|
||||||
|
device="cpu",
|
||||||
|
)
|
||||||
|
|
||||||
|
assert not hasattr(codec.model, "encoder_transformer")
|
||||||
|
assert not hasattr(codec.model, "decoder_transformer")
|
||||||
|
assert not any(
|
||||||
|
name.startswith(("model.encoder_transformer.", "model.decoder_transformer."))
|
||||||
|
for name, _ in codec.named_parameters()
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_mimi_full_stream_reset_reuses_all_state_storage_and_clears_offsets() -> None:
|
||||||
|
codec = personaplex_mimi.PersonaPlexMimiCodec.__new__(personaplex_mimi.PersonaPlexMimiCodec)
|
||||||
|
nn.Module.__init__(codec)
|
||||||
|
codec.device = torch.device("cpu")
|
||||||
|
codec.dtype = torch.float32
|
||||||
|
codec._enc_stages = [("conv", personaplex_mimi._StreamConv1d(nn.Conv1d(1, 1, 3)))]
|
||||||
|
codec._dec_stages = []
|
||||||
|
codec._downsample = personaplex_mimi._StreamConv1d(nn.Conv1d(1, 1, 3))
|
||||||
|
codec._upsample = personaplex_mimi._StreamConvTr1d(nn.ConvTranspose1d(1, 1, 3))
|
||||||
|
codec.encoder_transformer = personaplex_mimi._MimiStreamingTransformer(
|
||||||
|
num_layers=1,
|
||||||
|
dim=4,
|
||||||
|
num_heads=1,
|
||||||
|
context=3,
|
||||||
|
)
|
||||||
|
codec.decoder_transformer = personaplex_mimi._MimiStreamingTransformer(
|
||||||
|
num_layers=1,
|
||||||
|
dim=4,
|
||||||
|
num_heads=1,
|
||||||
|
context=3,
|
||||||
|
)
|
||||||
|
codec.streaming_init(batch_size=2)
|
||||||
|
|
||||||
|
conv_states = list(codec._conv_states())
|
||||||
|
conv_buffers = [
|
||||||
|
state.prev if isinstance(state, personaplex_mimi._StreamConv1d) else state.partial for state in conv_states
|
||||||
|
]
|
||||||
|
transformers = [codec.encoder_transformer, codec.decoder_transformer]
|
||||||
|
kv_states = [kv for transformer in transformers for kv in transformer._kv]
|
||||||
|
state_tensors = [
|
||||||
|
*conv_buffers,
|
||||||
|
*(kv.cache for kv in kv_states),
|
||||||
|
*(kv.end_offset for kv in kv_states),
|
||||||
|
*(kv.start_offset for kv in kv_states),
|
||||||
|
*(transformer._offset for transformer in transformers),
|
||||||
|
]
|
||||||
|
for tensor in state_tensors:
|
||||||
|
tensor.fill_(1)
|
||||||
|
for state in conv_states:
|
||||||
|
state._fresh.zero_()
|
||||||
|
storage = [tensor.data_ptr() for tensor in state_tensors]
|
||||||
|
|
||||||
|
codec.reset_streaming()
|
||||||
|
|
||||||
|
assert storage == [tensor.data_ptr() for tensor in state_tensors]
|
||||||
|
cleared = [
|
||||||
|
*conv_buffers,
|
||||||
|
*(kv.end_offset for kv in kv_states),
|
||||||
|
*(kv.start_offset for kv in kv_states),
|
||||||
|
*(transformer._offset for transformer in transformers),
|
||||||
|
]
|
||||||
|
assert all(not tensor.any() for tensor in cleared)
|
||||||
|
assert all(state._fresh.all() for state in conv_states)
|
||||||
|
|
||||||
|
|
||||||
|
def _model(*, max_sessions: int = 1) -> tuple[PersonaPlexCode2Wav, _FakeStreamingMimi]:
|
||||||
|
mimi_config = SimpleNamespace(num_codebooks=2, sample_rate=24000, samples_per_frame=4, mimi_name=None)
|
||||||
|
config = SimpleNamespace(mimi_config=mimi_config, mimi_name=None)
|
||||||
|
vllm_config = SimpleNamespace(
|
||||||
|
model_config=SimpleNamespace(
|
||||||
|
model="/unused",
|
||||||
|
hf_config=config,
|
||||||
|
duplex_max_sessions=max_sessions,
|
||||||
|
),
|
||||||
|
device_config=SimpleNamespace(device="cpu"),
|
||||||
|
)
|
||||||
|
model = PersonaPlexCode2Wav(vllm_config=vllm_config)
|
||||||
|
mimi = _FakeStreamingMimi()
|
||||||
|
model.mimi = mimi
|
||||||
|
model._mimi_device = torch.device("cpu")
|
||||||
|
return model, mimi
|
||||||
|
|
||||||
|
|
||||||
|
def _codes(frames: int, *, start: int = 0) -> torch.Tensor:
|
||||||
|
return torch.stack([torch.arange(start + offset, start + offset + frames) for offset in (0, 100)]).reshape(-1)
|
||||||
|
|
||||||
|
|
||||||
|
def _audio(output) -> torch.Tensor:
|
||||||
|
return output.multimodal_outputs["model_outputs"][0]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("second_codes", [_codes(3), _codes(1, start=100)])
|
||||||
|
def test_resumable_codes_emit_only_new_pcm(second_codes: torch.Tensor) -> None:
|
||||||
|
model, mimi = _model()
|
||||||
|
|
||||||
|
first = model(input_ids=_codes(2), request_ids=["req"])
|
||||||
|
second = model(input_ids=second_codes, request_ids=["req"])
|
||||||
|
|
||||||
|
assert _audio(first).numel() == 8
|
||||||
|
assert _audio(second).numel() == 4
|
||||||
|
assert mimi.decode_frame_calls == 3
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("frames", "expected_shapes"),
|
||||||
|
[
|
||||||
|
(5, [(1, 2, 5)]),
|
||||||
|
(12, [(1, 2, 5), (1, 2, 5), (1, 2, 2)]),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_streaming_decode_uses_bounded_multi_frame_calls(
|
||||||
|
frames: int,
|
||||||
|
expected_shapes: list[tuple[int, ...]],
|
||||||
|
) -> None:
|
||||||
|
model, _ = _model()
|
||||||
|
mimi = _FakeChunkedMimi()
|
||||||
|
model.mimi = mimi
|
||||||
|
|
||||||
|
output = model(input_ids=_codes(frames), request_ids=["req"])
|
||||||
|
|
||||||
|
assert _audio(output).numel() == frames * 4
|
||||||
|
assert mimi.decode_frames_shapes == expected_shapes
|
||||||
|
assert mimi.decode_frame_calls == 0
|
||||||
|
|
||||||
|
|
||||||
|
def test_profile_inputs_skip_decode_while_malformed_online_input_warns(
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
model, mimi = _model()
|
||||||
|
runtime_info = [{"meta": {"personaplex_dummy_profile": True}}]
|
||||||
|
warnings: list[tuple] = []
|
||||||
|
monkeypatch.setattr(personaplex_code2wav.logger, "warning", lambda *args: warnings.append(args))
|
||||||
|
|
||||||
|
output = model(
|
||||||
|
input_ids=torch.arange(3),
|
||||||
|
runtime_additional_information=runtime_info,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert _audio(output).numel() == 0
|
||||||
|
assert mimi.decode_frame_calls == 0
|
||||||
|
assert warnings == []
|
||||||
|
|
||||||
|
output = model(input_ids=torch.arange(3))
|
||||||
|
|
||||||
|
assert _audio(output).numel() == 0
|
||||||
|
assert mimi.decode_frame_calls == 0
|
||||||
|
assert len(warnings) == 1
|
||||||
|
assert "not divisible by" in warnings[0][0]
|
||||||
|
assert model.get_dummy_runtime_additional_information(2) == [
|
||||||
|
{"meta": {"personaplex_dummy_profile": True}},
|
||||||
|
{"meta": {"personaplex_dummy_profile": True}},
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def test_request_id_falls_back_to_runtime_information() -> None:
|
||||||
|
model, _ = _model()
|
||||||
|
info = [{"request_id": "runtime-req"}]
|
||||||
|
|
||||||
|
model(input_ids=_codes(2), runtime_additional_information=info)
|
||||||
|
second = model(input_ids=_codes(3), runtime_additional_information=info)
|
||||||
|
|
||||||
|
assert _audio(second).numel() == 4
|
||||||
|
|
||||||
|
|
||||||
|
def test_decoder_slot_lifecycle_isolated_capacity_and_reuse() -> None:
|
||||||
|
model, _ = _model(max_sessions=2)
|
||||||
|
first = _FakeStreamingMimi(stream_value=10)
|
||||||
|
second = _FakeStreamingMimi(stream_value=20)
|
||||||
|
model._set_mimi_codecs([first, second])
|
||||||
|
assert model._max_codec_sessions == 2
|
||||||
|
|
||||||
|
initial = model(
|
||||||
|
input_ids=torch.cat([_codes(1), _codes(1)]),
|
||||||
|
request_ids=["first", "second"],
|
||||||
|
seq_token_counts=[2, 2],
|
||||||
|
)
|
||||||
|
continued = model(
|
||||||
|
input_ids=torch.cat([_codes(1, start=10), _codes(1, start=20)]),
|
||||||
|
request_ids=["first", "second"],
|
||||||
|
seq_token_counts=[2, 2],
|
||||||
|
)
|
||||||
|
|
||||||
|
initial_audio = initial.multimodal_outputs["model_outputs"]
|
||||||
|
continued_audio = continued.multimodal_outputs["model_outputs"]
|
||||||
|
assert initial_audio[0].tolist() == [11.0] * 4
|
||||||
|
assert initial_audio[1].tolist() == [21.0] * 4
|
||||||
|
assert continued_audio[0].tolist() == [12.0] * 4
|
||||||
|
assert continued_audio[1].tolist() == [22.0] * 4
|
||||||
|
assert model._request_codec_slots == {"first": 0, "second": 1}
|
||||||
|
with pytest.raises(RuntimeError, match="decoder capacity 2 is exhausted"):
|
||||||
|
model._codec_for_request("overflow")
|
||||||
|
|
||||||
|
model.on_requests_finished(["first"])
|
||||||
|
assert first.reset_calls == 1
|
||||||
|
assert second.reset_calls == 0
|
||||||
|
replacement = model(input_ids=_codes(1), request_ids=["replacement"])
|
||||||
|
assert _audio(replacement).tolist() == [13.0] * 4
|
||||||
@@ -0,0 +1,364 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
|
||||||
|
import base64
|
||||||
|
from pathlib import Path
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import pytest
|
||||||
|
from vllm.sampling_params import SamplingParams
|
||||||
|
|
||||||
|
from tests.e2e.online_serving import personaplex_realtime_duplex as e2e_driver
|
||||||
|
from vllm_omni.config.stage_config import (
|
||||||
|
load_deploy_config,
|
||||||
|
merge_pipeline_deploy,
|
||||||
|
)
|
||||||
|
from vllm_omni.experimental.fullduplex.engine.duplex_runtime import (
|
||||||
|
DuplexInputMode,
|
||||||
|
)
|
||||||
|
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.data_plane import (
|
||||||
|
PersonaPlexDataPlaneContext,
|
||||||
|
PersonaPlexDataPlaneSession,
|
||||||
|
)
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.input import (
|
||||||
|
PersonaPlexPcmAppendBuffer,
|
||||||
|
)
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.runtime_extension import (
|
||||||
|
PersonaPlexDuplexRuntimeExtension,
|
||||||
|
)
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.serving_adapter import (
|
||||||
|
PersonaPlexServingRuntimeAdapter,
|
||||||
|
)
|
||||||
|
from vllm_omni.model_executor.models.personaplex.pipeline import (
|
||||||
|
PERSONAPLEX_PIPELINE,
|
||||||
|
)
|
||||||
|
|
||||||
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||||
|
|
||||||
|
|
||||||
|
def _pcm_payload(samples: np.ndarray, *, sample_rate_hz: int = 24000) -> dict[str, object]:
|
||||||
|
samples = np.ascontiguousarray(samples, dtype="<f4")
|
||||||
|
return {
|
||||||
|
"type": "audio",
|
||||||
|
"format": "pcm_f32le",
|
||||||
|
"sample_rate_hz": sample_rate_hz,
|
||||||
|
"audio": base64.b64encode(samples.tobytes()).decode("ascii"),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _audio_client(
|
||||||
|
frame_count: int = 10,
|
||||||
|
*,
|
||||||
|
voiced_frames: int = 0,
|
||||||
|
) -> SimpleNamespace:
|
||||||
|
frames = np.full((frame_count, e2e_driver.FRAME_SAMPLES), 7, dtype="<i2")
|
||||||
|
frames[:voiced_frames] = 1000
|
||||||
|
raw = frames.tobytes()
|
||||||
|
event = {
|
||||||
|
"type": "response.audio.delta",
|
||||||
|
"response_id": "response-1",
|
||||||
|
"sample_rate_hz": e2e_driver.SAMPLE_RATE_HZ,
|
||||||
|
"metadata": {
|
||||||
|
"vllm_omni": {
|
||||||
|
"runtime_impl": "scheduler_data_plane",
|
||||||
|
"uses_model_runner_scheduler": True,
|
||||||
|
"runner_kv_backed": True,
|
||||||
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
|
events = SimpleNamespace(
|
||||||
|
events=[event],
|
||||||
|
response_audio={"response-1": [raw]},
|
||||||
|
audio_bytes=lambda: raw,
|
||||||
|
)
|
||||||
|
return SimpleNamespace(events=events)
|
||||||
|
|
||||||
|
|
||||||
|
def test_realtime_audio_frame_stats_separate_voiced_and_silent_frames() -> None:
|
||||||
|
silent = np.zeros(e2e_driver.FRAME_SAMPLES, dtype="<i2")
|
||||||
|
voiced = np.full(e2e_driver.FRAME_SAMPLES, 1000, dtype="<i2")
|
||||||
|
|
||||||
|
stats = e2e_driver._audio_frame_stats(
|
||||||
|
silent.tobytes() + voiced.tobytes(),
|
||||||
|
input_frames=3,
|
||||||
|
voiced_frame_rms_threshold=1e-4,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert stats == {
|
||||||
|
"output_frames": 2,
|
||||||
|
"voiced_frames": 1,
|
||||||
|
"silent_frames": 1,
|
||||||
|
"frame_deficit": 1,
|
||||||
|
"frame_coverage_ratio": pytest.approx(2 / 3),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def test_realtime_audio_frame_stats_reject_partial_codec_frame() -> None:
|
||||||
|
partial_frame = np.zeros(e2e_driver.FRAME_SAMPLES - 1, dtype="<i2")
|
||||||
|
|
||||||
|
with pytest.raises(AssertionError, match="whole 80 ms codec frames"):
|
||||||
|
e2e_driver._audio_frame_stats(
|
||||||
|
partial_frame.tobytes(),
|
||||||
|
input_frames=1,
|
||||||
|
voiced_frame_rms_threshold=1e-4,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("frame_count", [10, 40])
|
||||||
|
def test_e2e_driver_rejects_inaudible_sessions(frame_count: int) -> None:
|
||||||
|
args = SimpleNamespace(
|
||||||
|
max_frame_deficit=4,
|
||||||
|
voiced_frame_rms_threshold=1e-3,
|
||||||
|
min_voiced_frames=5,
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(AssertionError, match="audible speech"):
|
||||||
|
e2e_driver._session_result(
|
||||||
|
_audio_client(frame_count),
|
||||||
|
input_frames=frame_count,
|
||||||
|
args=args,
|
||||||
|
minimum_chunks=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_e2e_driver_uses_absolute_audible_floor() -> None:
|
||||||
|
args = SimpleNamespace(
|
||||||
|
max_frame_deficit=4,
|
||||||
|
voiced_frame_rms_threshold=1e-3,
|
||||||
|
min_voiced_frames=5,
|
||||||
|
)
|
||||||
|
|
||||||
|
e2e_driver._session_result(
|
||||||
|
_audio_client(200, voiced_frames=5),
|
||||||
|
input_frames=200,
|
||||||
|
args=args,
|
||||||
|
minimum_chunks=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_e2e_driver_defaults_replacement_to_full_input() -> None:
|
||||||
|
args = e2e_driver.parse_args(
|
||||||
|
[
|
||||||
|
"--model",
|
||||||
|
"/model",
|
||||||
|
"--input-wav",
|
||||||
|
"/input.wav",
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
assert args.replacement_frames == 0
|
||||||
|
assert args.min_voiced_frames == 5
|
||||||
|
|
||||||
|
|
||||||
|
def test_e2e_driver_records_and_validates_input_sha256(tmp_path: Path) -> None:
|
||||||
|
input_wav = tmp_path / "input.wav"
|
||||||
|
input_wav.write_bytes(b"fixture-wav")
|
||||||
|
|
||||||
|
identity = e2e_driver._input_identity(
|
||||||
|
input_wav,
|
||||||
|
expected_sha256="bade4b3c163edde390ff391207d34a887257d8a9cc3b621cc8c618b6e6761304",
|
||||||
|
)
|
||||||
|
|
||||||
|
assert identity == {
|
||||||
|
"path": str(input_wav.resolve()),
|
||||||
|
"sha256": "bade4b3c163edde390ff391207d34a887257d8a9cc3b621cc8c618b6e6761304",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def test_e2e_driver_rejects_unexpected_input_sha256(tmp_path: Path) -> None:
|
||||||
|
input_wav = tmp_path / "input.wav"
|
||||||
|
input_wav.write_bytes(b"fixture-wav")
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="input WAV SHA-256 mismatch"):
|
||||||
|
e2e_driver._input_identity(input_wav, expected_sha256="0" * 64)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("attribute", "expected"),
|
||||||
|
[
|
||||||
|
(
|
||||||
|
"duplex_runtime_extension",
|
||||||
|
"vllm_omni.experimental.fullduplex.personaplex.runtime_extension.PersonaPlexDuplexRuntimeExtension",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"duplex_serving_adapter",
|
||||||
|
"vllm_omni.experimental.fullduplex.personaplex.serving_adapter.PersonaPlexServingRuntimeAdapter",
|
||||||
|
),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_personaplex_pipeline_registers_unified_duplex_components(
|
||||||
|
attribute: str,
|
||||||
|
expected: str,
|
||||||
|
) -> None:
|
||||||
|
assert PERSONAPLEX_PIPELINE.duplex_control_enabled is True
|
||||||
|
assert getattr(PERSONAPLEX_PIPELINE, attribute) == expected
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("engine_arg", "expected"),
|
||||||
|
[
|
||||||
|
("skip_tokenizer_init", True),
|
||||||
|
("enable_prefix_caching", False),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_personaplex_stage_engine_args(
|
||||||
|
engine_arg: str,
|
||||||
|
expected: bool,
|
||||||
|
) -> None:
|
||||||
|
deploy_path = Path(__file__).parents[5] / "vllm_omni" / "deploy" / "personaplex.yaml"
|
||||||
|
stages = merge_pipeline_deploy(
|
||||||
|
PERSONAPLEX_PIPELINE,
|
||||||
|
load_deploy_config(deploy_path),
|
||||||
|
)
|
||||||
|
|
||||||
|
assert [stage.yaml_engine_args.get(engine_arg) for stage in stages] == [
|
||||||
|
expected,
|
||||||
|
expected,
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def test_personaplex_duplex_capacity_is_propagated_to_all_model_stages() -> None:
|
||||||
|
deploy_path = Path(__file__).parents[5] / "vllm_omni" / "deploy" / "personaplex.yaml"
|
||||||
|
deploy = load_deploy_config(deploy_path)
|
||||||
|
stages = merge_pipeline_deploy(PERSONAPLEX_PIPELINE, deploy)
|
||||||
|
|
||||||
|
assert deploy.duplex_session.max_sessions == 2
|
||||||
|
assert [stage.yaml_engine_args.get("duplex_max_sessions") for stage in stages] == [2, 2]
|
||||||
|
assert "personaplex_codec_max_sessions" not in deploy.connectors["connector_of_shared_memory"]["extra"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_personaplex_capabilities_are_honest() -> None:
|
||||||
|
single = PersonaPlexServingRuntimeAdapter(lambda *_: None).capabilities(max_sessions=1)
|
||||||
|
multi = PersonaPlexServingRuntimeAdapter(lambda *_: None).capabilities(max_sessions=2)
|
||||||
|
|
||||||
|
assert multi.implementation_level == "model_native_duplex"
|
||||||
|
assert multi.input_modes == ["append_audio_chunk"]
|
||||||
|
assert multi.chunk_period_ms == 80
|
||||||
|
assert single.supports_multi_session is False
|
||||||
|
assert multi.supports_multi_session is True
|
||||||
|
assert multi.supports_multi_session_same_replica is True
|
||||||
|
assert multi.supports_barge_in is False
|
||||||
|
|
||||||
|
|
||||||
|
def test_pcm_buffer_emits_one_80ms_frame_transactionally() -> None:
|
||||||
|
buffer = PersonaPlexPcmAppendBuffer()
|
||||||
|
reservation = buffer.prepare_append(
|
||||||
|
_pcm_payload(np.arange(1920, dtype=np.float32)),
|
||||||
|
operation_id="op-1",
|
||||||
|
chunk_period_ms=80,
|
||||||
|
allow_emit=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert reservation is not None
|
||||||
|
assert reservation.byte_count == 1920 * 4
|
||||||
|
assert buffer.pending_byte_count == 0
|
||||||
|
reservation.rollback()
|
||||||
|
assert buffer.pending_byte_count == 1920 * 4
|
||||||
|
|
||||||
|
|
||||||
|
def test_pcm_buffer_rejects_changed_sample_rate() -> None:
|
||||||
|
buffer = PersonaPlexPcmAppendBuffer()
|
||||||
|
assert (
|
||||||
|
buffer.prepare_append(
|
||||||
|
_pcm_payload(np.zeros(100, np.float32)),
|
||||||
|
operation_id="op-1",
|
||||||
|
chunk_period_ms=80,
|
||||||
|
allow_emit=False,
|
||||||
|
)
|
||||||
|
is None
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="sample_rate"):
|
||||||
|
buffer.prepare_append(
|
||||||
|
_pcm_payload(np.zeros(100, np.float32), sample_rate_hz=16000),
|
||||||
|
operation_id="op-2",
|
||||||
|
chunk_period_ms=80,
|
||||||
|
allow_emit=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_pcm_buffer_rejects_non_finite_samples() -> None:
|
||||||
|
buffer = PersonaPlexPcmAppendBuffer()
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="finite"):
|
||||||
|
buffer.prepare_append(
|
||||||
|
_pcm_payload(np.array([0.0, np.nan], np.float32)),
|
||||||
|
operation_id="op-1",
|
||||||
|
chunk_period_ms=80,
|
||||||
|
allow_emit=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_runtime_extension_builds_one_frame_scheduler_append() -> None:
|
||||||
|
extension = PersonaPlexDuplexRuntimeExtension()
|
||||||
|
defaults = (
|
||||||
|
SamplingParams(temperature=0.8, top_k=10, max_tokens=10),
|
||||||
|
SamplingParams(max_tokens=1024),
|
||||||
|
)
|
||||||
|
|
||||||
|
configured = extension.configure_sampling_params(runtime_config={}, defaults=defaults)
|
||||||
|
assert configured[0].temperature == 0.0
|
||||||
|
assert configured[0].top_k == 1
|
||||||
|
assert configured[0].max_tokens == 1
|
||||||
|
assert configured[1] is defaults[1]
|
||||||
|
|
||||||
|
plan = extension.plan_append(
|
||||||
|
request_id="req",
|
||||||
|
fence=DuplexFence("session", incarnation=2, epoch=3),
|
||||||
|
session_config={"instructions": "Be concise."},
|
||||||
|
runtime_config={"personaplex_prefill_slots": 4},
|
||||||
|
seq=1,
|
||||||
|
turn_seq=1,
|
||||||
|
mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
|
||||||
|
payload=_pcm_payload(np.zeros(1920, np.float32)),
|
||||||
|
final=False,
|
||||||
|
sampling_params=configured[0],
|
||||||
|
)
|
||||||
|
assert len(plan.prompt["prompt_token_ids"]) == 5
|
||||||
|
duplex = plan.prompt["model_intermediate_buffer"]["duplex"]
|
||||||
|
assert duplex["session_id"] == "session"
|
||||||
|
assert duplex["incarnation"] == 2
|
||||||
|
assert duplex["epoch"] == 3
|
||||||
|
assert duplex["seq"] == 1
|
||||||
|
assert duplex["data_plane"] is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_projector_emits_only_cumulative_audio_and_text_suffixes() -> None:
|
||||||
|
encoded_sizes: list[int] = []
|
||||||
|
|
||||||
|
def encode_audio(audio, _sample_rate, _response_format, _speed):
|
||||||
|
size = int(np.asarray(audio, dtype=np.float32).size)
|
||||||
|
encoded_sizes.append(size)
|
||||||
|
return f"audio-{size}"
|
||||||
|
|
||||||
|
projector = PersonaPlexDataPlaneSession(encode_audio)
|
||||||
|
projector.begin_request("req")
|
||||||
|
context = PersonaPlexDataPlaneContext(response_format="wav")
|
||||||
|
|
||||||
|
first = SimpleNamespace(
|
||||||
|
request_id="req",
|
||||||
|
outputs=[SimpleNamespace(text="he", multimodal_output={})],
|
||||||
|
multimodal_output={"model_outputs": np.arange(4, dtype=np.float32), "sr": 24000},
|
||||||
|
finished=False,
|
||||||
|
)
|
||||||
|
second = SimpleNamespace(
|
||||||
|
request_id="req",
|
||||||
|
outputs=[SimpleNamespace(text="hello", multimodal_output={})],
|
||||||
|
multimodal_output={"model_outputs": np.arange(6, dtype=np.float32), "sr": 24000},
|
||||||
|
finished=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
projected = list(
|
||||||
|
projector.project(
|
||||||
|
{"data_plane_outputs": [first, second]},
|
||||||
|
context=context,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
assert encoded_sizes == [4, 2]
|
||||||
|
assert [item["text"] for item in projected] == ["he", "llo"]
|
||||||
|
assert [item["audio_data"] for item in projected] == ["audio-4", "audio-2"]
|
||||||
|
assert all(item["sample_rate_hz"] == 24000 for item in projected)
|
||||||
|
assert all(item["end_of_turn"] is False for item in projected)
|
||||||
@@ -0,0 +1,122 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
"""Contract tests for the fullduplex core (PR #3907 alignment).
|
||||||
|
|
||||||
|
Covers the DuplexFence identity value, the session identity rules (epoch only
|
||||||
|
advances on barge-in, response_index once per begin_response), stale-output
|
||||||
|
dropping in the turn lifecycle. The lockstep (single-eternal-response)
|
||||||
|
lifecycle is model-owned (see
|
||||||
|
tests/e2e/features/fullduplex/test_personaplex_adapter.py).
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import dataclasses
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from vllm_omni.experimental.fullduplex.core import protocol as ev
|
||||||
|
from vllm_omni.experimental.fullduplex.core.adapter import (
|
||||||
|
DuplexAdapter,
|
||||||
|
DuplexCapability,
|
||||||
|
OutputChunk,
|
||||||
|
)
|
||||||
|
from vllm_omni.experimental.fullduplex.core.runtime import DuplexRuntime
|
||||||
|
from vllm_omni.experimental.fullduplex.core.session import (
|
||||||
|
DuplexSession,
|
||||||
|
DuplexSessionConfig,
|
||||||
|
DuplexState,
|
||||||
|
)
|
||||||
|
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
|
||||||
|
|
||||||
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||||
|
|
||||||
|
|
||||||
|
def test_fence_is_frozen_value_identity():
|
||||||
|
f = DuplexFence(session_id="s1")
|
||||||
|
assert (f.epoch, f.turn_id, f.response_seq) == (0, 0, 0)
|
||||||
|
with pytest.raises(dataclasses.FrozenInstanceError):
|
||||||
|
f.epoch = 1 # type: ignore[misc]
|
||||||
|
assert DuplexFence("s1", 1, 2, 3) == DuplexFence("s1", 1, 2, 3)
|
||||||
|
assert DuplexFence("s1", 1, 2, 3) != DuplexFence("s1", 2, 2, 3)
|
||||||
|
|
||||||
|
|
||||||
|
def test_session_identity_rules():
|
||||||
|
s = DuplexSession(session_id="test", config=DuplexSessionConfig())
|
||||||
|
r1, e1 = s.begin_response()
|
||||||
|
r2, e2 = s.begin_response()
|
||||||
|
assert r2 == r1 + 1, "response index advances exactly once per begin"
|
||||||
|
assert e1 == e2 == 0, "epoch does not move on response begin"
|
||||||
|
s.end_response()
|
||||||
|
assert s.epoch == 0, "epoch does not move on response end"
|
||||||
|
new_epoch = s.barge_in()
|
||||||
|
assert new_epoch == 1, "epoch advances only on barge-in"
|
||||||
|
assert s.is_stale(e1) and not s.is_stale(new_epoch)
|
||||||
|
|
||||||
|
|
||||||
|
class _TurnAdapter(DuplexAdapter):
|
||||||
|
"""Emits chunks until cancelled; counts respond() lifecycles."""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
self.responses = 0
|
||||||
|
self.barge_ins = 0
|
||||||
|
|
||||||
|
def capabilities(self) -> DuplexCapability:
|
||||||
|
return DuplexCapability(
|
||||||
|
input_modalities=frozenset({"audio"}),
|
||||||
|
output_modalities=frozenset({"audio"}),
|
||||||
|
proactive=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def on_input(self, session, modality, data) -> None:
|
||||||
|
pass
|
||||||
|
|
||||||
|
def should_respond(self, session) -> bool:
|
||||||
|
return True
|
||||||
|
|
||||||
|
async def on_barge_in(self, session) -> None:
|
||||||
|
self.barge_ins += 1
|
||||||
|
|
||||||
|
async def respond(self, session):
|
||||||
|
self.responses += 1
|
||||||
|
for i in range(100):
|
||||||
|
await asyncio.sleep(0.001)
|
||||||
|
yield OutputChunk("audio", i)
|
||||||
|
|
||||||
|
|
||||||
|
async def _drive(runtime, events):
|
||||||
|
async def gen():
|
||||||
|
for e in events:
|
||||||
|
await asyncio.sleep(0.005)
|
||||||
|
yield e
|
||||||
|
|
||||||
|
out = []
|
||||||
|
|
||||||
|
async def emit(e):
|
||||||
|
out.append(e)
|
||||||
|
|
||||||
|
await runtime.run(gen(), emit)
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
def test_turn_mode_supersede_and_stale_drop():
|
||||||
|
adapter = _TurnAdapter()
|
||||||
|
session = DuplexSession(session_id="test", config=DuplexSessionConfig(proactive=True))
|
||||||
|
runtime = DuplexRuntime(session, adapter)
|
||||||
|
out = asyncio.run(
|
||||||
|
_drive(
|
||||||
|
runtime,
|
||||||
|
[
|
||||||
|
{"type": ev.RESPONSE_CREATE},
|
||||||
|
{"type": ev.RESPONSE_CREATE}, # supersedes the first
|
||||||
|
{"type": ev.CLOSE},
|
||||||
|
],
|
||||||
|
)
|
||||||
|
)
|
||||||
|
assert adapter.responses == 2
|
||||||
|
assert adapter.barge_ins == 1, "superseding an in-flight response barges in"
|
||||||
|
created = [e for e in out if e["type"] == ev.RESPONSE_CREATED]
|
||||||
|
done = [e for e in out if e["type"] == ev.RESPONSE_DONE]
|
||||||
|
assert len(created) == 2
|
||||||
|
assert len(done) == 1, "stale first response must not emit done"
|
||||||
|
assert done[0]["response_index"] == created[1]["response_index"]
|
||||||
|
assert session.state is DuplexState.CLOSED
|
||||||
@@ -0,0 +1,262 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
"""PersonaPlex adapter driven through the model-owned lockstep runtime.
|
||||||
|
|
||||||
|
PersonaPlexDuplexRuntime starts ONE response for the whole session, streams one
|
||||||
|
agent frame per user frame, and drains on close (core stays verbatim upstream;
|
||||||
|
the lockstep lifecycle is model policy) — all with a GPU-free stub stepper.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from vllm_omni.experimental.fullduplex.core import protocol as ev
|
||||||
|
from vllm_omni.experimental.fullduplex.core.session import (
|
||||||
|
DuplexSession,
|
||||||
|
DuplexSessionConfig,
|
||||||
|
DuplexState,
|
||||||
|
)
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.adapter import (
|
||||||
|
PersonaPlexDuplexAdapter,
|
||||||
|
PersonaPlexDuplexRuntime,
|
||||||
|
)
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.config import FRAME_SIZE
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.engine import FrameOutput
|
||||||
|
|
||||||
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||||
|
|
||||||
|
|
||||||
|
class StubStepper:
|
||||||
|
"""Deterministic FrameStepper: frame i -> audio filled with i, text 't{i}' on odds."""
|
||||||
|
|
||||||
|
sample_rate = 24000
|
||||||
|
frame_size = FRAME_SIZE
|
||||||
|
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.opened: tuple[str | None, str | None] | None = None
|
||||||
|
self.steps = 0
|
||||||
|
|
||||||
|
def open_session(self, voice_prompt=None, persona=None) -> None:
|
||||||
|
self.opened = (voice_prompt, persona)
|
||||||
|
self.steps = 0
|
||||||
|
|
||||||
|
def step(self, user_pcm):
|
||||||
|
assert user_pcm.shape[0] == self.frame_size
|
||||||
|
self.steps += 1
|
||||||
|
i = self.steps
|
||||||
|
return FrameOutput(
|
||||||
|
audio=np.full(self.frame_size, float(i), dtype=np.float32),
|
||||||
|
text=(f"t{i}" if i % 2 else None),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class RecordingAdapter(PersonaPlexDuplexAdapter):
|
||||||
|
def __init__(self, stepper: StubStepper) -> None:
|
||||||
|
super().__init__(stepper)
|
||||||
|
self.closed = False
|
||||||
|
|
||||||
|
async def on_close(self, session: DuplexSession) -> None:
|
||||||
|
self.closed = True
|
||||||
|
await super().on_close(session)
|
||||||
|
|
||||||
|
|
||||||
|
def _collector():
|
||||||
|
out: list[dict] = []
|
||||||
|
|
||||||
|
async def emit(event: dict) -> None:
|
||||||
|
out.append(event)
|
||||||
|
|
||||||
|
return out, emit
|
||||||
|
|
||||||
|
|
||||||
|
async def _feed(events):
|
||||||
|
for e in events:
|
||||||
|
yield e
|
||||||
|
|
||||||
|
|
||||||
|
async def _raising_feed(frame):
|
||||||
|
yield {"type": ev.INPUT_APPEND, "modality": "audio", "data": frame}
|
||||||
|
await asyncio.sleep(0)
|
||||||
|
raise RuntimeError("input failed")
|
||||||
|
|
||||||
|
|
||||||
|
def _respond_tasks() -> list[asyncio.Task]:
|
||||||
|
current = asyncio.current_task()
|
||||||
|
return [
|
||||||
|
task
|
||||||
|
for task in asyncio.all_tasks()
|
||||||
|
if task is not current and "_respond_once" in getattr(task.get_coro(), "__qualname__", "")
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def _session():
|
||||||
|
cfg = DuplexSessionConfig(
|
||||||
|
input_modalities=("audio",),
|
||||||
|
output_modalities=("audio", "text"),
|
||||||
|
)
|
||||||
|
return DuplexSession("s", cfg)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_lockstep_one_agent_frame_per_user_frame():
|
||||||
|
stub = StubStepper()
|
||||||
|
adapter = PersonaPlexDuplexAdapter(stub, voice_prompt="NATF2.pt", persona="be terse")
|
||||||
|
rt = PersonaPlexDuplexRuntime(_session(), adapter)
|
||||||
|
out, emit = _collector()
|
||||||
|
|
||||||
|
frame = np.zeros(FRAME_SIZE, dtype=np.float32)
|
||||||
|
events = [{"type": ev.INPUT_APPEND, "modality": "audio", "data": frame} for _ in range(3)]
|
||||||
|
events.append({"type": ev.CLOSE})
|
||||||
|
await rt.run(_feed(events), emit)
|
||||||
|
|
||||||
|
types = [e["type"] for e in out]
|
||||||
|
# exactly one response for the whole session, drained cleanly on close
|
||||||
|
assert types.count(ev.RESPONSE_CREATED) == 1
|
||||||
|
assert types.count(ev.RESPONSE_DONE) == 1
|
||||||
|
audio = [e for e in out if e["type"] == ev.RESPONSE_DELTA and e["modality"] == "audio"]
|
||||||
|
text = [e["data"] for e in out if e["type"] == ev.RESPONSE_DELTA and e["modality"] == "text"]
|
||||||
|
assert len(audio) == 3 # one agent frame per user frame
|
||||||
|
assert [int(e["data"][0]) for e in audio] == [1, 2, 3] # lockstep order
|
||||||
|
assert text == ["t1", "t3"]
|
||||||
|
assert stub.steps == 3
|
||||||
|
assert stub.opened == ("NATF2.pt", "be terse")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_unaligned_chunks_are_reframed_to_80ms():
|
||||||
|
stub = StubStepper()
|
||||||
|
adapter = PersonaPlexDuplexAdapter(stub)
|
||||||
|
rt = PersonaPlexDuplexRuntime(_session(), adapter)
|
||||||
|
out, emit = _collector()
|
||||||
|
|
||||||
|
# 2.5 frames in one chunk + 0.5 frame in another => 3 whole frames
|
||||||
|
events = [
|
||||||
|
{"type": ev.INPUT_APPEND, "modality": "audio", "data": np.zeros(FRAME_SIZE * 5 // 2, dtype=np.float32)},
|
||||||
|
{"type": ev.INPUT_APPEND, "modality": "audio", "data": np.zeros(FRAME_SIZE // 2, dtype=np.float32)},
|
||||||
|
{"type": ev.CLOSE},
|
||||||
|
]
|
||||||
|
await rt.run(_feed(events), emit)
|
||||||
|
audio = [e for e in out if e["type"] == ev.RESPONSE_DELTA and e["modality"] == "audio"]
|
||||||
|
assert len(audio) == 3
|
||||||
|
assert stub.steps == 3
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_close_without_input_is_clean():
|
||||||
|
adapter = PersonaPlexDuplexAdapter(StubStepper())
|
||||||
|
rt = PersonaPlexDuplexRuntime(_session(), adapter)
|
||||||
|
out, emit = _collector()
|
||||||
|
await rt.run(_feed([{"type": ev.CLOSE}]), emit)
|
||||||
|
# no input -> response never started -> no created/done, no error
|
||||||
|
assert not any(e["type"] == ev.ERROR for e in out)
|
||||||
|
assert not any(e["type"] == ev.RESPONSE_CREATED for e in out)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_iterator_eof_without_close_drains_and_stops():
|
||||||
|
# Transport drop: the input iterator ends WITHOUT an explicit CLOSE event.
|
||||||
|
# The runtime must still signal on_close so the eternal response drains
|
||||||
|
# instead of blocking on its inbox forever.
|
||||||
|
stub = StubStepper()
|
||||||
|
adapter = PersonaPlexDuplexAdapter(stub)
|
||||||
|
rt = PersonaPlexDuplexRuntime(_session(), adapter)
|
||||||
|
out, emit = _collector()
|
||||||
|
frame = np.zeros(FRAME_SIZE, dtype=np.float32)
|
||||||
|
events = [{"type": ev.INPUT_APPEND, "modality": "audio", "data": frame}]
|
||||||
|
await asyncio.wait_for(rt.run(_feed(events), emit), timeout=5.0)
|
||||||
|
types = [e["type"] for e in out]
|
||||||
|
assert types.count(ev.RESPONSE_DONE) == 1
|
||||||
|
assert stub.steps == 1
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_input_exception_closes_response_and_runtime() -> None:
|
||||||
|
stub = StubStepper()
|
||||||
|
adapter = RecordingAdapter(stub)
|
||||||
|
rt = PersonaPlexDuplexRuntime(_session(), adapter)
|
||||||
|
out, emit = _collector()
|
||||||
|
frame = np.zeros(FRAME_SIZE, dtype=np.float32)
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="input failed"):
|
||||||
|
await rt.run(_raising_feed(frame), emit)
|
||||||
|
|
||||||
|
leaked = _respond_tasks()
|
||||||
|
for task in leaked:
|
||||||
|
task.cancel()
|
||||||
|
if leaked:
|
||||||
|
await asyncio.gather(*leaked, return_exceptions=True)
|
||||||
|
|
||||||
|
assert adapter.closed is True
|
||||||
|
assert rt.session.state is DuplexState.CLOSED
|
||||||
|
assert not leaked
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_default_session_config_gets_lockstep_lifecycle():
|
||||||
|
# The lifecycle comes from the model-owned runtime class, not a session
|
||||||
|
# flag: a plain default-constructed session config runs lockstep.
|
||||||
|
stub = StubStepper()
|
||||||
|
adapter = PersonaPlexDuplexAdapter(stub)
|
||||||
|
session = DuplexSession(
|
||||||
|
"s",
|
||||||
|
DuplexSessionConfig(input_modalities=("audio",), output_modalities=("audio", "text")),
|
||||||
|
)
|
||||||
|
rt = PersonaPlexDuplexRuntime(session, adapter)
|
||||||
|
out, emit = _collector()
|
||||||
|
frame = np.zeros(FRAME_SIZE, dtype=np.float32)
|
||||||
|
events = [{"type": ev.INPUT_APPEND, "modality": "audio", "data": frame}, {"type": ev.CLOSE}]
|
||||||
|
await rt.run(_feed(events), emit)
|
||||||
|
types = [e["type"] for e in out]
|
||||||
|
assert types.count(ev.RESPONSE_CREATED) == 1
|
||||||
|
assert types.count(ev.RESPONSE_DONE) == 1
|
||||||
|
assert stub.steps == 1
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_partial_tail_is_padded_and_stepped_on_close():
|
||||||
|
# A trailing sub-frame chunk is zero-padded and stepped at close (same
|
||||||
|
# contract as PersonaPlexSession.flush) rather than silently dropped.
|
||||||
|
stub = StubStepper()
|
||||||
|
adapter = PersonaPlexDuplexAdapter(stub)
|
||||||
|
rt = PersonaPlexDuplexRuntime(_session(), adapter)
|
||||||
|
out, emit = _collector()
|
||||||
|
events = [
|
||||||
|
{"type": ev.INPUT_APPEND, "modality": "audio", "data": np.ones(FRAME_SIZE // 2, dtype=np.float32)},
|
||||||
|
{"type": ev.CLOSE},
|
||||||
|
]
|
||||||
|
await rt.run(_feed(events), emit)
|
||||||
|
audio = [e for e in out if e["type"] == ev.RESPONSE_DELTA and e["modality"] == "audio"]
|
||||||
|
assert len(audio) == 1
|
||||||
|
assert stub.steps == 1
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_single_eternal_response_across_barge_in():
|
||||||
|
# Moved from the core contract tests when the lockstep lifecycle became
|
||||||
|
# model-owned: ONE response for the whole session, RESPONSE_CANCEL bumps the
|
||||||
|
# epoch fence (playback cursor) but the stream keeps flowing, and close
|
||||||
|
# drains via on_close.
|
||||||
|
stub = StubStepper()
|
||||||
|
adapter = PersonaPlexDuplexAdapter(stub)
|
||||||
|
session = _session()
|
||||||
|
rt = PersonaPlexDuplexRuntime(session, adapter)
|
||||||
|
out, emit = _collector()
|
||||||
|
|
||||||
|
frame = np.zeros(FRAME_SIZE, dtype=np.float32)
|
||||||
|
events = [
|
||||||
|
{"type": ev.INPUT_APPEND, "modality": "audio", "data": frame},
|
||||||
|
{"type": ev.RESPONSE_CANCEL}, # barge-in: epoch bumps, stream continues
|
||||||
|
{"type": ev.INPUT_APPEND, "modality": "audio", "data": frame},
|
||||||
|
{"type": ev.CLOSE},
|
||||||
|
]
|
||||||
|
await rt.run(_feed(events), emit)
|
||||||
|
|
||||||
|
types = [e["type"] for e in out]
|
||||||
|
assert types.count(ev.RESPONSE_CREATED) == 1, "exactly one eternal response"
|
||||||
|
assert types.count(ev.RESPONSE_DONE) == 1
|
||||||
|
audio = [e for e in out if e["type"] == ev.RESPONSE_DELTA and e["modality"] == "audio"]
|
||||||
|
assert len(audio) == 2, "stream keeps flowing across barge-in"
|
||||||
|
assert session.epoch == 1, "cancel still advanced the epoch fence"
|
||||||
|
assert stub.steps == 2
|
||||||
@@ -0,0 +1,178 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
"""GPU-free tests for the batched PersonaPlex slot manager (stub engine)."""
|
||||||
|
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.engine import FrameOutput, PrefillStep
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.serving.batched import (
|
||||||
|
BatchedSessionManager,
|
||||||
|
SlotPhase,
|
||||||
|
)
|
||||||
|
|
||||||
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||||
|
|
||||||
|
FRAME = 4
|
||||||
|
|
||||||
|
|
||||||
|
class StubEngine:
|
||||||
|
"""Records the slot-API call sequence; echoes each row's pcm as its output."""
|
||||||
|
|
||||||
|
frame_size = FRAME
|
||||||
|
|
||||||
|
def __init__(self, batch_size: int = 4, prefill_len: int = 2) -> None:
|
||||||
|
self.config = SimpleNamespace(batch_size=batch_size)
|
||||||
|
self.prefill_len = prefill_len
|
||||||
|
self.recycled: list[int] = []
|
||||||
|
self.voice_ended: list[int] = []
|
||||||
|
self.prefill_done: list[int] = []
|
||||||
|
self.steps: list[tuple[np.ndarray, dict | None]] = []
|
||||||
|
|
||||||
|
def prefill_steps(self, persona=None):
|
||||||
|
steps = [PrefillStep("sacrifice"), PrefillStep("voice", embedding=None), PrefillStep("voice_end")]
|
||||||
|
steps += [PrefillStep("tokens", moshi_tokens=None, text_token=3, user_sine=True)] * self.prefill_len
|
||||||
|
return steps
|
||||||
|
|
||||||
|
def recycle_slot(self, b):
|
||||||
|
self.recycled.append(b)
|
||||||
|
|
||||||
|
def voice_end(self, b):
|
||||||
|
self.voice_ended.append(b)
|
||||||
|
|
||||||
|
def end_prefill(self, b):
|
||||||
|
self.prefill_done.append(b)
|
||||||
|
|
||||||
|
def step_batch(self, pcm, prefill=None):
|
||||||
|
self.steps.append((pcm.copy(), dict(prefill) if prefill else None))
|
||||||
|
return [FrameOutput(audio=pcm[b].copy(), text=None) for b in range(self.config.batch_size)]
|
||||||
|
|
||||||
|
|
||||||
|
def test_all_idle_does_not_step():
|
||||||
|
eng = StubEngine()
|
||||||
|
mgr = BatchedSessionManager(eng)
|
||||||
|
assert mgr.tick() is False
|
||||||
|
assert eng.steps == []
|
||||||
|
|
||||||
|
|
||||||
|
def test_pristine_slot_is_live_immediately():
|
||||||
|
eng = StubEngine()
|
||||||
|
mgr = BatchedSessionManager(eng)
|
||||||
|
slot, ready_now = mgr.acquire()
|
||||||
|
assert ready_now is True
|
||||||
|
assert slot.phase is SlotPhase.LIVE
|
||||||
|
assert eng.recycled == []
|
||||||
|
|
||||||
|
|
||||||
|
def test_custom_persona_prefills_and_recycles_on_tick_thread():
|
||||||
|
eng = StubEngine(prefill_len=2)
|
||||||
|
mgr = BatchedSessionManager(eng)
|
||||||
|
slot, ready_now = mgr.acquire(persona="custom")
|
||||||
|
assert ready_now is False
|
||||||
|
assert slot.phase is SlotPhase.PREFILL
|
||||||
|
assert eng.recycled == [] # deferred: reset must run between engine steps
|
||||||
|
|
||||||
|
ready = []
|
||||||
|
slot.on_ready = lambda: ready.append(True)
|
||||||
|
|
||||||
|
mgr.tick() # recycle + sacrifice
|
||||||
|
assert eng.recycled == [slot.index]
|
||||||
|
assert eng.steps[-1][1][slot.index].kind == "sacrifice"
|
||||||
|
|
||||||
|
mgr.tick() # voice frame
|
||||||
|
assert eng.steps[-1][1][slot.index].kind == "voice"
|
||||||
|
assert eng.voice_ended == []
|
||||||
|
|
||||||
|
mgr.tick() # voice_end executed BETWEEN ticks, then first tokens frame
|
||||||
|
assert eng.voice_ended == [slot.index]
|
||||||
|
assert eng.steps[-1][1][slot.index].kind == "tokens"
|
||||||
|
|
||||||
|
mgr.tick() # second tokens frame
|
||||||
|
assert not ready
|
||||||
|
mgr.tick() # prefill exhausted -> live
|
||||||
|
assert eng.prefill_done == [slot.index]
|
||||||
|
assert slot.phase is SlotPhase.LIVE
|
||||||
|
assert ready == [True]
|
||||||
|
|
||||||
|
|
||||||
|
def test_pristine_slot_after_ticking_requires_recycle():
|
||||||
|
# step_batch advances EVERY row each tick, so a pristine idle row that has
|
||||||
|
# lived through ticks has drifted from its boot prefill and must be
|
||||||
|
# recycled before a client may join it.
|
||||||
|
eng = StubEngine()
|
||||||
|
mgr = BatchedSessionManager(eng)
|
||||||
|
first, ready_now = mgr.acquire()
|
||||||
|
assert ready_now is True # claimed before any tick: truly pristine
|
||||||
|
mgr.tick()
|
||||||
|
second, ready_now2 = mgr.acquire()
|
||||||
|
assert second.index != first.index
|
||||||
|
assert ready_now2 is False
|
||||||
|
assert second.phase is SlotPhase.PREFILL
|
||||||
|
|
||||||
|
|
||||||
|
def test_used_slot_requires_recycle_even_with_default_persona():
|
||||||
|
eng = StubEngine()
|
||||||
|
mgr = BatchedSessionManager(eng)
|
||||||
|
slot, ready_now = mgr.acquire()
|
||||||
|
assert ready_now
|
||||||
|
mgr.release(slot)
|
||||||
|
slot2, ready_now2 = mgr.acquire()
|
||||||
|
assert slot2.index == slot.index
|
||||||
|
assert ready_now2 is False
|
||||||
|
assert slot2.phase is SlotPhase.PREFILL
|
||||||
|
|
||||||
|
|
||||||
|
def test_acquire_returns_none_when_full():
|
||||||
|
eng = StubEngine(batch_size=2)
|
||||||
|
mgr = BatchedSessionManager(eng)
|
||||||
|
assert mgr.acquire() is not None
|
||||||
|
assert mgr.acquire() is not None
|
||||||
|
assert mgr.acquire() is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_feed_frames_epoch_guard_and_backlog_cap():
|
||||||
|
eng = StubEngine()
|
||||||
|
mgr = BatchedSessionManager(eng, max_buffered_frames=3)
|
||||||
|
slot, _ = mgr.acquire()
|
||||||
|
mgr.feed(slot, np.arange(FRAME * 2, dtype=np.float32), slot.epoch)
|
||||||
|
assert len(slot.inq) == 2
|
||||||
|
mgr.feed(slot, np.zeros(3, dtype=np.float32), slot.epoch) # partial tail buffered
|
||||||
|
assert len(slot.inq) == 2 and slot.pending.shape[0] == 3
|
||||||
|
mgr.feed(slot, np.zeros(FRAME * 5, dtype=np.float32), slot.epoch)
|
||||||
|
assert len(slot.inq) == 3 # oldest dropped beyond the cap
|
||||||
|
mgr.feed(slot, np.zeros(FRAME, dtype=np.float32), slot.epoch - 1) # stale epoch ignored
|
||||||
|
assert len(slot.inq) == 3
|
||||||
|
|
||||||
|
|
||||||
|
def test_live_slot_consumes_input_and_emits_output():
|
||||||
|
eng = StubEngine()
|
||||||
|
mgr = BatchedSessionManager(eng)
|
||||||
|
slot, _ = mgr.acquire()
|
||||||
|
got = []
|
||||||
|
slot.on_output = got.append
|
||||||
|
frame = np.full(FRAME, 0.5, dtype=np.float32)
|
||||||
|
mgr.feed(slot, frame, slot.epoch)
|
||||||
|
mgr.tick()
|
||||||
|
assert len(got) == 1
|
||||||
|
np.testing.assert_array_equal(got[0].audio, frame)
|
||||||
|
# queue drained: next tick feeds zeros
|
||||||
|
mgr.tick()
|
||||||
|
np.testing.assert_array_equal(got[1].audio, np.zeros(FRAME, dtype=np.float32))
|
||||||
|
|
||||||
|
|
||||||
|
def test_release_mid_flight_stops_emission():
|
||||||
|
eng = StubEngine()
|
||||||
|
mgr = BatchedSessionManager(eng)
|
||||||
|
slot, _ = mgr.acquire()
|
||||||
|
got = []
|
||||||
|
slot.on_output = got.append
|
||||||
|
mgr.tick()
|
||||||
|
assert len(got) == 1
|
||||||
|
mgr.release(slot)
|
||||||
|
# another slot keeps the batch non-idle
|
||||||
|
other, _ = mgr.acquire() # same slot index, now PREFILL (used)
|
||||||
|
assert other.phase is SlotPhase.PREFILL
|
||||||
|
mgr.tick()
|
||||||
|
assert len(got) == 1 # no further output for the released epoch
|
||||||
@@ -0,0 +1,240 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
"""GPU e2e for PersonaPlex elastic batching (bit-exact isolation + recycle).
|
||||||
|
|
||||||
|
Greedy decoding makes every cross-slot leak observable as a bit difference:
|
||||||
|
|
||||||
|
1. Isolation (bit-exact): recycling one slot mid-flight leaves every other
|
||||||
|
slot's output bit-identical to a run that never recycled.
|
||||||
|
2. Replay exactness: a recycled slot's opening live ticks match a pristine boot
|
||||||
|
at numerical-noise level (a broken KV mask / staging roll garbles frame 0).
|
||||||
|
3. Coherence (behavioral): a recycled slot fed a real spoken question produces
|
||||||
|
a real reply in its inner-monologue text stream.
|
||||||
|
|
||||||
|
Run (moshi-free native stack; ~30 GB GPU + HF access to the gated repo; set
|
||||||
|
PPLEX_QUESTION_WAV to a spoken-question wav to enable the coherence test):
|
||||||
|
pytest tests/e2e/features/fullduplex/test_personaplex_elastic.py -v -s
|
||||||
|
|
||||||
|
Bench (B sweep, ms/tick vs the 80 ms budget):
|
||||||
|
python tests/e2e/features/fullduplex/test_personaplex_elastic.py --bench 2 4 8 16
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
torch = pytest.importorskip("torch")
|
||||||
|
|
||||||
|
# slow + H100 + omni routes this module into the weekly Buildkite E2E group
|
||||||
|
# ("Omni · H100" in .buildkite/cuda/test-weekly.yml: -m "slow and H100 and
|
||||||
|
# omni" --run-level full_model); the ready/merge CPU sweeps exclude it by mark.
|
||||||
|
pytestmark = [pytest.mark.full_model, pytest.mark.omni, pytest.mark.slow, pytest.mark.H100]
|
||||||
|
|
||||||
|
if not torch.cuda.is_available(): # noqa: E402
|
||||||
|
pytest.skip("needs CUDA", allow_module_level=True)
|
||||||
|
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.config import PersonaPlexConfig # noqa: E402
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.runtime import ( # noqa: E402
|
||||||
|
PersonaPlexEngine,
|
||||||
|
)
|
||||||
|
|
||||||
|
B = 4
|
||||||
|
LIVE_TICKS = 40 # live frames compared per claim
|
||||||
|
RECYCLE_AT = 10
|
||||||
|
|
||||||
|
|
||||||
|
def _make_engine() -> PersonaPlexEngine:
|
||||||
|
cfg = PersonaPlexConfig(batch_size=B) # native stepper is greedy-only
|
||||||
|
try:
|
||||||
|
return PersonaPlexEngine(cfg).load()
|
||||||
|
except Exception as exc: # gated-repo access is environment-dependent
|
||||||
|
name = type(exc).__name__
|
||||||
|
if "Gated" in name or "401" in str(exc) or "403" in str(exc):
|
||||||
|
pytest.skip(f"no access to the gated nvidia/personaplex-7b-v1 repo: {name}")
|
||||||
|
raise
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(scope="module")
|
||||||
|
def engine() -> PersonaPlexEngine:
|
||||||
|
return _make_engine()
|
||||||
|
|
||||||
|
|
||||||
|
def _slot_stream(slot: int, tick: int, frame: int) -> np.ndarray:
|
||||||
|
"""Deterministic per-slot input: distinct low-amplitude tones."""
|
||||||
|
sr = 24000
|
||||||
|
freq = 180.0 + 90.0 * slot
|
||||||
|
t0 = tick * frame
|
||||||
|
t = (np.arange(frame) + t0) / sr
|
||||||
|
return (0.12 * np.sin(2 * np.pi * freq * t)).astype(np.float32)
|
||||||
|
|
||||||
|
|
||||||
|
def _run(engine, total_ticks: int, recycle: bool):
|
||||||
|
"""Run the lockstep loop; optionally recycle slot 1 at RECYCLE_AT.
|
||||||
|
|
||||||
|
Returns per-slot lists of (audio | None) per live tick index, where the
|
||||||
|
recycled slot's list restarts its input stream at its post-prefill tick.
|
||||||
|
"""
|
||||||
|
frame = engine.frame_size
|
||||||
|
outs: list[list[np.ndarray | None]] = [[] for _ in range(B)]
|
||||||
|
stream_tick = [0] * B # per-slot conversation clock driving its input
|
||||||
|
prefill = None
|
||||||
|
tick = 0
|
||||||
|
while tick < total_ticks:
|
||||||
|
if recycle and tick == RECYCLE_AT and prefill is None:
|
||||||
|
engine.recycle_slot(1)
|
||||||
|
prefill = list(engine.prefill_steps(None))
|
||||||
|
pcm = np.zeros((B, frame), dtype=np.float32)
|
||||||
|
step_for_1 = None
|
||||||
|
if prefill:
|
||||||
|
while prefill and prefill[0].kind == "voice_end":
|
||||||
|
engine.voice_end(1)
|
||||||
|
prefill.pop(0)
|
||||||
|
if prefill:
|
||||||
|
step_for_1 = prefill.pop(0)
|
||||||
|
else:
|
||||||
|
engine.end_prefill(1)
|
||||||
|
prefill = None
|
||||||
|
stream_tick[1] = 0 # restart slot 1's input stream from its own t=0
|
||||||
|
for b in range(B):
|
||||||
|
if b == 1 and step_for_1 is not None:
|
||||||
|
continue # engine overrides the user channel for prefill rows
|
||||||
|
pcm[b] = _slot_stream(b, stream_tick[b], frame)
|
||||||
|
stream_tick[b] += 1
|
||||||
|
frame_outs = engine.step_batch(pcm, {1: step_for_1} if step_for_1 is not None else None)
|
||||||
|
for b in range(B):
|
||||||
|
if b == 1 and step_for_1 is not None:
|
||||||
|
continue # prefill tick: nothing conversational to record
|
||||||
|
frame_out = frame_outs[b]
|
||||||
|
outs[b].append(None if frame_out is None else frame_out.audio)
|
||||||
|
tick += 1
|
||||||
|
return outs
|
||||||
|
|
||||||
|
|
||||||
|
def test_recycle_isolation_and_fresh_equivalence(engine):
|
||||||
|
engine.open_batch()
|
||||||
|
prefill_len = sum(1 for s in engine.prefill_steps(None) if s.kind != "voice_end")
|
||||||
|
total = RECYCLE_AT + prefill_len + LIVE_TICKS
|
||||||
|
|
||||||
|
base = _run(engine, total, recycle=False)
|
||||||
|
|
||||||
|
engine.open_batch() # deterministic re-boot (greedy, fixed voice/persona)
|
||||||
|
rec = _run(engine, total, recycle=True)
|
||||||
|
|
||||||
|
# Claim 1 -- isolation: untouched slots bit-identical to the no-recycle run.
|
||||||
|
for b in (0, 2, 3):
|
||||||
|
assert len(base[b]) == total and len(rec[b]) == total
|
||||||
|
for i, (x, y) in enumerate(zip(base[b], rec[b])):
|
||||||
|
assert (x is None) == (y is None), f"slot {b} tick {i} validity diverged"
|
||||||
|
if x is not None:
|
||||||
|
assert np.array_equal(x, y), f"slot {b} tick {i} audio diverged (recycle leaked)"
|
||||||
|
|
||||||
|
# Claim 2 -- replay exactness: the recycled slot's opening live ticks match a
|
||||||
|
# pristine boot at numerical-noise level. A broken KV mask / staging-ring
|
||||||
|
# roll / Mimi reset garbles the audio from the very first frame (diffs of
|
||||||
|
# 0.1-1.0); a correct replay differs only by decoder-state rounding (~1e-4).
|
||||||
|
#
|
||||||
|
# Full-length bitwise identity is NOT expected, for two structural reasons:
|
||||||
|
# the replayed system prompt sits at a different ABSOLUTE RoPE offset (bf16
|
||||||
|
# rounds differently there), and the recycled Mimi ENCODER cannot restore a
|
||||||
|
# fresh stream's start padding (a batch-shared scalar), so from the first
|
||||||
|
# non-silent user frame it emits different-but-equally-valid codes for the
|
||||||
|
# same audio. Once any near-tie flips a token, the two equally valid
|
||||||
|
# trajectories legitimately diverge; coherence of the recycled slot is
|
||||||
|
# asserted behaviorally in test_recycled_slot_converses.
|
||||||
|
fresh = [a for a in base[1] if a is not None][:LIVE_TICKS]
|
||||||
|
reborn = [a for a in rec[1][RECYCLE_AT:] if a is not None][:LIVE_TICKS]
|
||||||
|
assert len(reborn) >= LIVE_TICKS // 2, "recycled slot produced too few live frames"
|
||||||
|
diffs = [float(np.max(np.abs(x - y))) for x, y in zip(fresh, reborn)]
|
||||||
|
first_flip = next((i for i, d in enumerate(diffs) if d > 5e-2), len(diffs))
|
||||||
|
print(
|
||||||
|
f"[elastic] replay exactness: {first_flip}/{len(diffs)} opening ticks at noise level "
|
||||||
|
f"(head diffs: {[f'{d:.1e}' for d in diffs[:6]]})"
|
||||||
|
)
|
||||||
|
assert first_flip >= 4, (
|
||||||
|
f"recycled slot diverged from fresh behavior at tick {first_flip} "
|
||||||
|
f"(diffs {diffs[: first_flip + 2]}); the system-prompt replay is broken"
|
||||||
|
)
|
||||||
|
assert max(diffs[:4]) < 2e-2, f"reborn audio not at noise level: {max(diffs[:4])}"
|
||||||
|
|
||||||
|
|
||||||
|
def test_recycled_slot_converses(engine):
|
||||||
|
"""Behavioral guarantee: a recycled slot holds a coherent conversation.
|
||||||
|
|
||||||
|
Feed a real spoken question into a slot that was recycled mid-flight (while
|
||||||
|
other slots keep running); the model's own inner-monologue text stream must
|
||||||
|
produce a real reply. This is the end-user property the encoder-state
|
||||||
|
caveats above cannot break.
|
||||||
|
"""
|
||||||
|
import os
|
||||||
|
|
||||||
|
import sphn
|
||||||
|
|
||||||
|
wav_path = os.environ.get("PPLEX_QUESTION_WAV")
|
||||||
|
if not wav_path:
|
||||||
|
pytest.skip("set PPLEX_QUESTION_WAV to a spoken-question wav to run")
|
||||||
|
wav, _ = sphn.read(wav_path, sample_rate=24000)
|
||||||
|
wav = wav[0].astype(np.float32)
|
||||||
|
frame = engine.frame_size
|
||||||
|
pad = (-len(wav)) % frame
|
||||||
|
wav = np.concatenate([wav, np.zeros(pad, dtype=np.float32)])
|
||||||
|
q_frames = [wav[i : i + frame] for i in range(0, len(wav), frame)]
|
||||||
|
|
||||||
|
engine.open_batch()
|
||||||
|
pcm = np.zeros((B, frame), dtype=np.float32)
|
||||||
|
for _ in range(5): # everyone live briefly
|
||||||
|
engine.step_batch(pcm)
|
||||||
|
engine.recycle_slot(1)
|
||||||
|
prefill = list(engine.prefill_steps(None))
|
||||||
|
while prefill:
|
||||||
|
while prefill and prefill[0].kind == "voice_end":
|
||||||
|
engine.voice_end(1)
|
||||||
|
prefill.pop(0)
|
||||||
|
if not prefill:
|
||||||
|
break
|
||||||
|
engine.step_batch(pcm, {1: prefill.pop(0)})
|
||||||
|
engine.end_prefill(1)
|
||||||
|
|
||||||
|
text = ""
|
||||||
|
total = len(q_frames) + 120 # question + ~10 s to answer
|
||||||
|
for i in range(total):
|
||||||
|
pcm = np.zeros((B, frame), dtype=np.float32)
|
||||||
|
if i < len(q_frames):
|
||||||
|
pcm[1] = q_frames[i]
|
||||||
|
outs = engine.step_batch(pcm)
|
||||||
|
if outs[1] is not None and outs[1].text:
|
||||||
|
text += outs[1].text
|
||||||
|
alpha = sum(c.isalpha() for c in text)
|
||||||
|
print(f"[elastic] recycled slot inner monologue: {text!r}")
|
||||||
|
assert alpha >= 15, f"recycled slot produced no coherent reply: {text!r}"
|
||||||
|
|
||||||
|
|
||||||
|
def _bench(batch_sizes: list[int]) -> None:
|
||||||
|
import time
|
||||||
|
|
||||||
|
for bs in batch_sizes:
|
||||||
|
cfg = PersonaPlexConfig(batch_size=max(bs, 2))
|
||||||
|
eng = PersonaPlexEngine(cfg).load()
|
||||||
|
eng.open_batch()
|
||||||
|
frame = eng.frame_size
|
||||||
|
pcm = np.zeros((cfg.batch_size, frame), dtype=np.float32)
|
||||||
|
for _ in range(10): # warm
|
||||||
|
eng.step_batch(pcm)
|
||||||
|
torch.cuda.synchronize()
|
||||||
|
t0 = time.monotonic()
|
||||||
|
n = 100
|
||||||
|
for _ in range(n):
|
||||||
|
eng.step_batch(pcm)
|
||||||
|
torch.cuda.synchronize()
|
||||||
|
ms = (time.monotonic() - t0) / n * 1000
|
||||||
|
mem = torch.cuda.max_memory_allocated() / 2**30
|
||||||
|
print(f"B={cfg.batch_size:3d} {ms:6.1f} ms/tick RTF/stream {ms / 80:.3f} {mem:.1f} GiB")
|
||||||
|
del eng
|
||||||
|
torch.cuda.empty_cache()
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
import sys
|
||||||
|
|
||||||
|
sizes = [int(x) for x in sys.argv[2:]] if len(sys.argv) > 2 and sys.argv[1] == "--bench" else [2, 4, 8]
|
||||||
|
_bench(sizes)
|
||||||
@@ -0,0 +1,101 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
"""Single-session lease teardown: a cancelled engine call must be drained.
|
||||||
|
|
||||||
|
Cancelling an asyncio task blocked on a worker-thread engine call returns
|
||||||
|
immediately while the thread keeps running (executor futures cannot be
|
||||||
|
interrupted). The server therefore tracks the in-flight future on the session
|
||||||
|
state and waits it out before releasing the capacity-one lease; otherwise a
|
||||||
|
fast reconnect could call ``open()`` on the shared engine while the previous
|
||||||
|
connection's final ``step()`` is still mutating it.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import contextlib
|
||||||
|
import sys
|
||||||
|
import threading
|
||||||
|
import types
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
try:
|
||||||
|
import sphn # noqa: F401
|
||||||
|
except ImportError:
|
||||||
|
# These lease tests never encode/decode Opus. Keep the optional serving
|
||||||
|
# dependency from hiding dependency-free teardown regressions.
|
||||||
|
sys.modules["sphn"] = types.ModuleType("sphn")
|
||||||
|
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.config import PersonaPlexConfig
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.serving.server import DuplexServer
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.session import PersonaPlexServingSessionState
|
||||||
|
|
||||||
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_cancelled_engine_call_is_drained_before_lease_release():
|
||||||
|
server = DuplexServer(PersonaPlexConfig()) # engine never loaded: stub calls only
|
||||||
|
state = PersonaPlexServingSessionState()
|
||||||
|
|
||||||
|
started = threading.Event()
|
||||||
|
release = threading.Event()
|
||||||
|
finished = threading.Event()
|
||||||
|
|
||||||
|
def slow_engine_op():
|
||||||
|
started.set()
|
||||||
|
release.wait(timeout=5.0)
|
||||||
|
finished.set()
|
||||||
|
return []
|
||||||
|
|
||||||
|
task = asyncio.create_task(server._engine_call(state, slow_engine_op))
|
||||||
|
await asyncio.to_thread(started.wait, 5.0)
|
||||||
|
|
||||||
|
# Cancel while the worker thread is mid-call: the await returns at once,
|
||||||
|
# the thread keeps running, and its completion event stays tracked on the
|
||||||
|
# state (the asyncio future itself reads as cancelled/done here, which is
|
||||||
|
# exactly why the drain must not rely on it).
|
||||||
|
task.cancel()
|
||||||
|
with contextlib.suppress(asyncio.CancelledError):
|
||||||
|
await task
|
||||||
|
assert state.inflight is not None
|
||||||
|
assert not state.inflight.is_set()
|
||||||
|
|
||||||
|
# The drain must genuinely wait for the worker, not just return.
|
||||||
|
drain = asyncio.create_task(server._drain_inflight(state))
|
||||||
|
await asyncio.sleep(0.05)
|
||||||
|
assert not drain.done()
|
||||||
|
|
||||||
|
release.set()
|
||||||
|
await asyncio.wait_for(drain, timeout=5.0)
|
||||||
|
assert finished.is_set()
|
||||||
|
assert state.inflight.is_set()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_completed_engine_call_clears_inflight():
|
||||||
|
server = DuplexServer(PersonaPlexConfig())
|
||||||
|
state = PersonaPlexServingSessionState()
|
||||||
|
result = await server._engine_call(state, lambda: "ok")
|
||||||
|
assert result == "ok"
|
||||||
|
assert state.inflight is None
|
||||||
|
await server._drain_inflight(state) # no-op on a clean state
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_drain_timeout_keeps_lease_until_worker_finishes():
|
||||||
|
server = DuplexServer(PersonaPlexConfig())
|
||||||
|
server._active = True
|
||||||
|
state = PersonaPlexServingSessionState()
|
||||||
|
state.inflight = threading.Event()
|
||||||
|
|
||||||
|
drained = await server._release_lease_after_drain(state, timeout=0.01)
|
||||||
|
|
||||||
|
assert drained is False
|
||||||
|
assert server._active is True
|
||||||
|
|
||||||
|
state.inflight.set()
|
||||||
|
for _ in range(100):
|
||||||
|
if not server._active:
|
||||||
|
break
|
||||||
|
await asyncio.sleep(0.01)
|
||||||
|
assert server._active is False
|
||||||
@@ -0,0 +1,97 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
"""Lockstep-contract tests for the PersonaPlex session driver (no GPU / no moshi).
|
||||||
|
|
||||||
|
A stub :class:`FrameStepper` stands in for the real Moshi backend so we can assert
|
||||||
|
the framing + 1-frame-in / 1-frame-out invariants without weights.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.config import FRAME_SIZE, PersonaPlexConfig
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.engine import FrameOutput, FrameStepper
|
||||||
|
from vllm_omni.experimental.fullduplex.personaplex.session import PersonaPlexSession
|
||||||
|
|
||||||
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||||
|
|
||||||
|
|
||||||
|
class StubStepper:
|
||||||
|
"""Deterministic FrameStepper: frame i -> audio filled with i, text 't{i}' on odds."""
|
||||||
|
|
||||||
|
sample_rate = 24000
|
||||||
|
frame_size = FRAME_SIZE
|
||||||
|
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.opened: tuple[str | None, str | None] | None = None
|
||||||
|
self.steps = 0
|
||||||
|
|
||||||
|
def open_session(self, voice_prompt=None, persona=None) -> None:
|
||||||
|
self.opened = (voice_prompt, persona)
|
||||||
|
self.steps = 0
|
||||||
|
|
||||||
|
def step(self, user_pcm):
|
||||||
|
assert user_pcm.shape[0] == self.frame_size
|
||||||
|
self.steps += 1
|
||||||
|
i = self.steps
|
||||||
|
return FrameOutput(
|
||||||
|
audio=np.full(self.frame_size, float(i), dtype=np.float32),
|
||||||
|
text=(f"t{i}" if i % 2 else None),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_stub_satisfies_protocol():
|
||||||
|
assert isinstance(StubStepper(), FrameStepper)
|
||||||
|
|
||||||
|
|
||||||
|
def test_feed_emits_one_output_per_complete_frame():
|
||||||
|
s = PersonaPlexSession(StubStepper(), PersonaPlexConfig())
|
||||||
|
s.open(voice_prompt="NATF2.pt", persona="be terse")
|
||||||
|
out = s.feed(np.zeros(FRAME_SIZE * 3, dtype=np.float32))
|
||||||
|
assert len(out) == 3
|
||||||
|
# lockstep order preserved: frame i -> audio filled with i
|
||||||
|
assert [int(o.audio[0]) for o in out] == [1, 2, 3]
|
||||||
|
assert [o.text for o in out] == ["t1", None, "t3"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_partial_frames_are_buffered_until_complete():
|
||||||
|
stub = StubStepper()
|
||||||
|
s = PersonaPlexSession(stub, PersonaPlexConfig())
|
||||||
|
s.open()
|
||||||
|
assert s.feed(np.zeros(FRAME_SIZE // 2, dtype=np.float32)) == [] # nothing yet
|
||||||
|
assert stub.steps == 0
|
||||||
|
out = s.feed(np.zeros(FRAME_SIZE, dtype=np.float32)) # 1.5 frames total -> 1 frame
|
||||||
|
assert len(out) == 1
|
||||||
|
assert stub.steps == 1
|
||||||
|
|
||||||
|
|
||||||
|
def test_flush_pads_the_tail():
|
||||||
|
stub = StubStepper()
|
||||||
|
s = PersonaPlexSession(stub, PersonaPlexConfig())
|
||||||
|
s.open()
|
||||||
|
s.feed(np.ones(FRAME_SIZE + 10, dtype=np.float32)) # 1 frame + 10-sample tail
|
||||||
|
assert stub.steps == 1
|
||||||
|
tail = s.flush()
|
||||||
|
assert len(tail) == 1 and stub.steps == 2
|
||||||
|
assert s.flush() == [] # nothing left
|
||||||
|
|
||||||
|
|
||||||
|
def test_run_concatenates_audio_and_joins_text():
|
||||||
|
s = PersonaPlexSession(StubStepper(), PersonaPlexConfig())
|
||||||
|
s.open()
|
||||||
|
audio, text = s.run(np.zeros(FRAME_SIZE * 4, dtype=np.float32))
|
||||||
|
assert audio.shape[0] == FRAME_SIZE * 4
|
||||||
|
assert text == "t1t3" # frames 1..4, text on odds
|
||||||
|
|
||||||
|
|
||||||
|
def test_open_forwards_voice_and_persona():
|
||||||
|
stub = StubStepper()
|
||||||
|
s = PersonaPlexSession(stub, PersonaPlexConfig())
|
||||||
|
s.open(voice_prompt="NATM1.pt", persona="pirate")
|
||||||
|
assert stub.opened == ("NATM1.pt", "pirate")
|
||||||
|
|
||||||
|
|
||||||
|
def test_feed_before_open_raises():
|
||||||
|
s = PersonaPlexSession(StubStepper(), PersonaPlexConfig())
|
||||||
|
with pytest.raises(RuntimeError):
|
||||||
|
s.feed(np.zeros(FRAME_SIZE, dtype=np.float32))
|
||||||
@@ -0,0 +1,492 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
"""Strict PersonaPlex validation for the unified OpenAI Realtime duplex path."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import asyncio
|
||||||
|
import base64
|
||||||
|
import hashlib
|
||||||
|
import json
|
||||||
|
import math
|
||||||
|
import uuid
|
||||||
|
import wave
|
||||||
|
from collections.abc import Sequence
|
||||||
|
from pathlib import Path
|
||||||
|
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
from vllm_omni.experimental.fullduplex.client import (
|
||||||
|
RealtimeDuplexClient,
|
||||||
|
wait_for,
|
||||||
|
write_pcm16_wav,
|
||||||
|
)
|
||||||
|
|
||||||
|
SAMPLE_RATE_HZ = 24_000
|
||||||
|
FRAME_SAMPLES = 1_920
|
||||||
|
FRAME_PERIOD_S = FRAME_SAMPLES / SAMPLE_RATE_HZ
|
||||||
|
|
||||||
|
|
||||||
|
def _input_identity(
|
||||||
|
path: Path,
|
||||||
|
*,
|
||||||
|
expected_sha256: str | None,
|
||||||
|
) -> dict[str, str]:
|
||||||
|
resolved = path.resolve()
|
||||||
|
digest = hashlib.sha256()
|
||||||
|
with resolved.open("rb") as input_file:
|
||||||
|
for chunk in iter(lambda: input_file.read(1024 * 1024), b""):
|
||||||
|
digest.update(chunk)
|
||||||
|
actual = digest.hexdigest()
|
||||||
|
if expected_sha256 is not None and actual != expected_sha256.lower():
|
||||||
|
raise ValueError(
|
||||||
|
f"input WAV SHA-256 mismatch: expected={expected_sha256.lower()}, actual={actual}, path={resolved}"
|
||||||
|
)
|
||||||
|
return {"path": str(resolved), "sha256": actual}
|
||||||
|
|
||||||
|
|
||||||
|
def _realtime_url(base_url: str, model: str, session_id: str) -> str:
|
||||||
|
parts = urlsplit(base_url)
|
||||||
|
query = dict(parse_qsl(parts.query, keep_blank_values=True))
|
||||||
|
query.update(duplex="1", model=model, autostart="0", session_id=session_id)
|
||||||
|
return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment))
|
||||||
|
|
||||||
|
|
||||||
|
def _read_wav_as_float32(path: Path) -> np.ndarray:
|
||||||
|
with wave.open(str(path), "rb") as wav_file:
|
||||||
|
channels = wav_file.getnchannels()
|
||||||
|
sample_width = wav_file.getsampwidth()
|
||||||
|
source_rate = wav_file.getframerate()
|
||||||
|
frames = wav_file.readframes(wav_file.getnframes())
|
||||||
|
dtypes = {1: np.uint8, 2: np.dtype("<i2"), 4: np.dtype("<i4")}
|
||||||
|
if sample_width not in dtypes:
|
||||||
|
raise ValueError(f"unsupported WAV sample width: {sample_width}")
|
||||||
|
pcm = np.frombuffer(frames, dtype=dtypes[sample_width]).astype(np.float32)
|
||||||
|
if sample_width == 1:
|
||||||
|
pcm = (pcm - 128.0) / 128.0
|
||||||
|
else:
|
||||||
|
pcm /= float(1 << (sample_width * 8 - 1))
|
||||||
|
if channels > 1:
|
||||||
|
pcm = pcm.reshape(-1, channels).mean(axis=1)
|
||||||
|
if source_rate != SAMPLE_RATE_HZ and pcm.size:
|
||||||
|
output_samples = max(1, round(pcm.size * SAMPLE_RATE_HZ / source_rate))
|
||||||
|
pcm = np.interp(
|
||||||
|
np.linspace(0, pcm.size - 1, output_samples),
|
||||||
|
np.arange(pcm.size),
|
||||||
|
pcm,
|
||||||
|
).astype(np.float32)
|
||||||
|
pcm = np.ascontiguousarray(pcm, dtype="<f4")
|
||||||
|
if not pcm.size or not np.isfinite(pcm).all():
|
||||||
|
raise ValueError("input WAV must contain finite audio")
|
||||||
|
return pcm
|
||||||
|
|
||||||
|
|
||||||
|
def _events(client: RealtimeDuplexClient, event_type: str) -> list[dict[str, object]]:
|
||||||
|
return [event for event in client.events.events if event.get("type") == event_type]
|
||||||
|
|
||||||
|
|
||||||
|
async def _open_session(
|
||||||
|
args: argparse.Namespace,
|
||||||
|
*,
|
||||||
|
session_id: str,
|
||||||
|
persona: str,
|
||||||
|
expect_error: bool = False,
|
||||||
|
) -> tuple[RealtimeDuplexClient, dict[str, object]]:
|
||||||
|
client = RealtimeDuplexClient(_realtime_url(args.url, args.model, session_id))
|
||||||
|
await client.__aenter__()
|
||||||
|
await client.send(
|
||||||
|
{
|
||||||
|
"type": "session.update",
|
||||||
|
"session": {
|
||||||
|
"session_id": session_id,
|
||||||
|
"model": args.model,
|
||||||
|
"modalities": ["audio", "text"],
|
||||||
|
"input_audio_format": "pcm_f32le",
|
||||||
|
"output_audio_format": "pcm16",
|
||||||
|
"voice": args.voice,
|
||||||
|
"instructions": persona,
|
||||||
|
"turn_detection": None,
|
||||||
|
"extra_body": {"auto_response": True},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
)
|
||||||
|
event_type = "error" if expect_error else "session.created"
|
||||||
|
await wait_for(
|
||||||
|
lambda: client.events.count(event_type) > 0,
|
||||||
|
timeout_s=args.timeout_s,
|
||||||
|
label=event_type,
|
||||||
|
)
|
||||||
|
return client, _events(client, event_type)[-1]
|
||||||
|
|
||||||
|
|
||||||
|
async def _close_session(client: RealtimeDuplexClient, *, timeout_s: float) -> None:
|
||||||
|
await client.send({"type": "session.close"})
|
||||||
|
await wait_for(
|
||||||
|
lambda: client.events.count("session.closed") > 0,
|
||||||
|
timeout_s=timeout_s,
|
||||||
|
label="session.closed",
|
||||||
|
)
|
||||||
|
await client.__aexit__(None, None, None)
|
||||||
|
|
||||||
|
|
||||||
|
async def _stream_frames(
|
||||||
|
client: RealtimeDuplexClient,
|
||||||
|
pcm: np.ndarray,
|
||||||
|
*,
|
||||||
|
max_frames: int | None = None,
|
||||||
|
) -> int:
|
||||||
|
frame_count = math.ceil(pcm.size / FRAME_SAMPLES)
|
||||||
|
if max_frames is not None:
|
||||||
|
if max_frames < 0:
|
||||||
|
raise ValueError(f"max_frames must be non-negative, got {max_frames}")
|
||||||
|
if max_frames > 0:
|
||||||
|
frame_count = min(frame_count, max_frames)
|
||||||
|
for seq in range(frame_count):
|
||||||
|
frame = pcm[seq * FRAME_SAMPLES : (seq + 1) * FRAME_SAMPLES]
|
||||||
|
frame = np.pad(frame, (0, FRAME_SAMPLES - frame.size))
|
||||||
|
frame = np.ascontiguousarray(frame, dtype="<f4")
|
||||||
|
await client.send(
|
||||||
|
{
|
||||||
|
"type": "input_audio_buffer.append",
|
||||||
|
"audio": base64.b64encode(frame).decode("ascii"),
|
||||||
|
"format": "pcm_f32le",
|
||||||
|
"sample_rate_hz": SAMPLE_RATE_HZ,
|
||||||
|
"duration_ms": 80,
|
||||||
|
"audio_end_ms": (seq + 1) * 80,
|
||||||
|
"is_speech": bool(np.sqrt(np.mean(np.square(frame))) >= 1e-4),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
await asyncio.sleep(FRAME_PERIOD_S)
|
||||||
|
return frame_count
|
||||||
|
|
||||||
|
|
||||||
|
def _audio_frame_stats(
|
||||||
|
raw: bytes,
|
||||||
|
*,
|
||||||
|
input_frames: int,
|
||||||
|
voiced_frame_rms_threshold: float,
|
||||||
|
) -> dict[str, int | float]:
|
||||||
|
pcm = np.frombuffer(raw, dtype="<i2")
|
||||||
|
if pcm.size % FRAME_SAMPLES:
|
||||||
|
raise AssertionError(
|
||||||
|
"PersonaPlex output must contain whole 80 ms codec frames: "
|
||||||
|
f"samples={pcm.size}, samples_per_frame={FRAME_SAMPLES}"
|
||||||
|
)
|
||||||
|
output_frames = pcm.size // FRAME_SAMPLES
|
||||||
|
frames = pcm.astype(np.float32).reshape(output_frames, FRAME_SAMPLES) / 32768.0
|
||||||
|
voiced_frames = int(np.count_nonzero(np.sqrt(np.mean(np.square(frames), axis=1)) >= voiced_frame_rms_threshold))
|
||||||
|
return {
|
||||||
|
"output_frames": int(output_frames),
|
||||||
|
"voiced_frames": voiced_frames,
|
||||||
|
"silent_frames": int(output_frames - voiced_frames),
|
||||||
|
"frame_deficit": int(input_frames - output_frames),
|
||||||
|
"frame_coverage_ratio": float(output_frames / input_frames),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _response_ids(client: RealtimeDuplexClient) -> set[str]:
|
||||||
|
return {
|
||||||
|
str(event["response_id"])
|
||||||
|
for event in _events(client, "response.audio.delta")
|
||||||
|
if isinstance(event.get("response_id"), str)
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _session_result(
|
||||||
|
client: RealtimeDuplexClient,
|
||||||
|
*,
|
||||||
|
input_frames: int,
|
||||||
|
args: argparse.Namespace,
|
||||||
|
minimum_chunks: int,
|
||||||
|
) -> tuple[bytes, set[str], dict[str, object]]:
|
||||||
|
chunks = sum(len(items) for items in client.events.response_audio.values())
|
||||||
|
raw = client.events.audio_bytes()
|
||||||
|
if chunks < minimum_chunks or not raw or len(raw) % 2:
|
||||||
|
raise AssertionError(f"invalid PCM16 output: chunks={chunks}, bytes={len(raw)}")
|
||||||
|
pcm = np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0
|
||||||
|
rms = float(np.sqrt(np.mean(np.square(pcm))))
|
||||||
|
# This low floor only catches corrupt, non-finite, or effectively zero PCM.
|
||||||
|
# Audible speech is checked per frame below with the higher 1e-3 default.
|
||||||
|
if not np.isfinite(pcm).all() or rms <= 1e-5:
|
||||||
|
raise AssertionError(f"output audio is non-finite or silent: samples={pcm.size}, rms={rms}")
|
||||||
|
rates = {
|
||||||
|
int(event["sample_rate_hz"])
|
||||||
|
for event in _events(client, "response.audio.delta")
|
||||||
|
if isinstance(event.get("sample_rate_hz"), int)
|
||||||
|
}
|
||||||
|
if rates != {SAMPLE_RATE_HZ}:
|
||||||
|
raise AssertionError(f"unexpected output sample rates: {sorted(rates)}")
|
||||||
|
stats = _audio_frame_stats(
|
||||||
|
raw,
|
||||||
|
input_frames=input_frames,
|
||||||
|
voiced_frame_rms_threshold=args.voiced_frame_rms_threshold,
|
||||||
|
)
|
||||||
|
if not 0 <= int(stats["frame_deficit"]) <= args.max_frame_deficit:
|
||||||
|
raise AssertionError(f"invalid synchronous frame accounting: input={input_frames}, stats={stats}")
|
||||||
|
min_voiced_frames = int(args.min_voiced_frames)
|
||||||
|
if min_voiced_frames < 1:
|
||||||
|
raise ValueError(f"min_voiced_frames must be positive, got {min_voiced_frames}")
|
||||||
|
if int(stats["voiced_frames"]) < min_voiced_frames:
|
||||||
|
raise AssertionError(
|
||||||
|
f"output contained insufficient audible speech: minimum_voiced_frames={min_voiced_frames}, stats={stats}"
|
||||||
|
)
|
||||||
|
runtime_metadata = [
|
||||||
|
metadata["vllm_omni"]
|
||||||
|
for event in _events(client, "response.audio.delta")
|
||||||
|
if isinstance((metadata := event.get("metadata")), dict) and isinstance(metadata.get("vllm_omni"), dict)
|
||||||
|
]
|
||||||
|
if not any(
|
||||||
|
item.get("runtime_impl") == "scheduler_data_plane"
|
||||||
|
and item.get("uses_model_runner_scheduler") is True
|
||||||
|
and item.get("runner_kv_backed") is True
|
||||||
|
for item in runtime_metadata
|
||||||
|
):
|
||||||
|
raise AssertionError(f"audio events did not prove the scheduler data plane: {runtime_metadata}")
|
||||||
|
response_ids = _response_ids(client)
|
||||||
|
if len(response_ids) != 1:
|
||||||
|
raise AssertionError(f"continuous stream used multiple responses: {sorted(response_ids)}")
|
||||||
|
return (
|
||||||
|
raw,
|
||||||
|
response_ids,
|
||||||
|
{
|
||||||
|
**stats,
|
||||||
|
"audio_chunks": chunks,
|
||||||
|
"audio_samples": len(raw) // 2,
|
||||||
|
"audio_rms": rms,
|
||||||
|
"response_ids": sorted(response_ids),
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _capabilities(created: dict[str, object]) -> dict[str, object]:
|
||||||
|
session = created.get("session")
|
||||||
|
capabilities = session.get("capabilities") if isinstance(session, dict) else None
|
||||||
|
if not isinstance(capabilities, dict):
|
||||||
|
raise AssertionError(f"session.created omitted capabilities: {created}")
|
||||||
|
return capabilities
|
||||||
|
|
||||||
|
|
||||||
|
def _save(output_dir: Path, name: str, client: RealtimeDuplexClient, raw: bytes) -> None:
|
||||||
|
(output_dir / f"{name}-events.jsonl").write_text(
|
||||||
|
"".join(json.dumps(event, ensure_ascii=False) + "\n" for event in client.events.events),
|
||||||
|
encoding="utf-8",
|
||||||
|
)
|
||||||
|
write_pcm16_wav(output_dir / f"{name}-output.wav", raw, sample_rate_hz=SAMPLE_RATE_HZ)
|
||||||
|
|
||||||
|
|
||||||
|
async def run(args: argparse.Namespace) -> dict[str, object]:
|
||||||
|
input_path = Path(args.input_wav)
|
||||||
|
input_identity = _input_identity(
|
||||||
|
input_path,
|
||||||
|
expected_sha256=args.expected_input_sha256,
|
||||||
|
)
|
||||||
|
pcm = _read_wav_as_float32(input_path)
|
||||||
|
if args.tail_s > 0:
|
||||||
|
pcm = np.concatenate([pcm, np.zeros(round(args.tail_s * SAMPLE_RATE_HZ), dtype=np.float32)])
|
||||||
|
output_dir = Path(args.output_dir)
|
||||||
|
output_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
ids = {name: f"personaplex-{name}-{uuid.uuid4().hex}" for name in ("primary", "secondary", "replacement")}
|
||||||
|
primary, created = await _open_session(args, session_id=ids["primary"], persona=args.persona)
|
||||||
|
secondary, secondary_created = await _open_session(
|
||||||
|
args,
|
||||||
|
session_id=ids["secondary"],
|
||||||
|
persona=args.secondary_persona,
|
||||||
|
)
|
||||||
|
capabilities = _capabilities(created)
|
||||||
|
if _capabilities(secondary_created) != capabilities:
|
||||||
|
raise AssertionError("concurrent sessions returned different capabilities")
|
||||||
|
expected_capabilities = {
|
||||||
|
"implementation_level": "model_native_duplex",
|
||||||
|
"chunk_period_ms": 80,
|
||||||
|
"supports_multi_session": True,
|
||||||
|
"supports_multi_session_same_replica": True,
|
||||||
|
"supports_barge_in": False,
|
||||||
|
}
|
||||||
|
if any(capabilities.get(key) != value for key, value in expected_capabilities.items()):
|
||||||
|
raise AssertionError(f"unexpected PersonaPlex capabilities: {capabilities}")
|
||||||
|
|
||||||
|
overflow, error = await _open_session(
|
||||||
|
args,
|
||||||
|
session_id=f"personaplex-overflow-{uuid.uuid4().hex}",
|
||||||
|
persona=args.persona,
|
||||||
|
expect_error=True,
|
||||||
|
)
|
||||||
|
error_body = error.get("error")
|
||||||
|
overflow_code = error_body.get("code") if isinstance(error_body, dict) else error.get("code")
|
||||||
|
await overflow.__aexit__(None, None, None)
|
||||||
|
if overflow_code != "resource_exhausted":
|
||||||
|
raise AssertionError(f"overflow session was not rejected: {error}")
|
||||||
|
|
||||||
|
primary_frames, secondary_frames = await asyncio.gather(
|
||||||
|
_stream_frames(primary, pcm),
|
||||||
|
_stream_frames(secondary, pcm),
|
||||||
|
)
|
||||||
|
await asyncio.gather(
|
||||||
|
wait_for(
|
||||||
|
lambda: len(primary.events.audio_bytes()) // (2 * FRAME_SAMPLES) >= primary_frames - args.max_frame_deficit,
|
||||||
|
timeout_s=args.timeout_s,
|
||||||
|
label="primary frame coverage",
|
||||||
|
),
|
||||||
|
wait_for(
|
||||||
|
lambda: (
|
||||||
|
len(secondary.events.audio_bytes()) // (2 * FRAME_SAMPLES) >= secondary_frames - args.max_frame_deficit
|
||||||
|
),
|
||||||
|
timeout_s=args.timeout_s,
|
||||||
|
label="secondary frame coverage",
|
||||||
|
),
|
||||||
|
)
|
||||||
|
await asyncio.sleep(args.drain_s)
|
||||||
|
primary_audio, primary_response_ids, primary_stats = _session_result(
|
||||||
|
primary,
|
||||||
|
input_frames=primary_frames,
|
||||||
|
args=args,
|
||||||
|
minimum_chunks=args.minimum_audio_chunks,
|
||||||
|
)
|
||||||
|
secondary_audio, secondary_response_ids, _ = _session_result(
|
||||||
|
secondary,
|
||||||
|
input_frames=secondary_frames,
|
||||||
|
args=args,
|
||||||
|
minimum_chunks=args.minimum_audio_chunks,
|
||||||
|
)
|
||||||
|
if primary_response_ids & secondary_response_ids:
|
||||||
|
raise AssertionError("concurrent sessions shared a response id")
|
||||||
|
_save(output_dir, "primary", primary, primary_audio)
|
||||||
|
await _close_session(primary, timeout_s=args.timeout_s)
|
||||||
|
|
||||||
|
replacement, replacement_created = await _open_session(
|
||||||
|
args,
|
||||||
|
session_id=ids["replacement"],
|
||||||
|
persona=args.replacement_persona,
|
||||||
|
)
|
||||||
|
if _capabilities(replacement_created) != capabilities:
|
||||||
|
raise AssertionError("replacement session returned different capabilities")
|
||||||
|
continuation_frames, replacement_frames = await asyncio.gather(
|
||||||
|
_stream_frames(secondary, pcm, max_frames=args.replacement_frames),
|
||||||
|
_stream_frames(replacement, pcm, max_frames=args.replacement_frames),
|
||||||
|
)
|
||||||
|
secondary_total_frames = secondary_frames + continuation_frames
|
||||||
|
await asyncio.gather(
|
||||||
|
wait_for(
|
||||||
|
lambda: (
|
||||||
|
len(secondary.events.audio_bytes()) // (2 * FRAME_SAMPLES)
|
||||||
|
>= secondary_total_frames - args.max_frame_deficit
|
||||||
|
),
|
||||||
|
timeout_s=args.timeout_s,
|
||||||
|
label="survivor frame coverage after slot recycle",
|
||||||
|
),
|
||||||
|
wait_for(
|
||||||
|
lambda: (
|
||||||
|
len(replacement.events.audio_bytes()) // (2 * FRAME_SAMPLES)
|
||||||
|
>= replacement_frames - args.max_frame_deficit
|
||||||
|
),
|
||||||
|
timeout_s=args.timeout_s,
|
||||||
|
label="replacement frame coverage",
|
||||||
|
),
|
||||||
|
)
|
||||||
|
await asyncio.sleep(args.drain_s)
|
||||||
|
secondary_audio, secondary_after_ids, secondary_stats = _session_result(
|
||||||
|
secondary,
|
||||||
|
input_frames=secondary_total_frames,
|
||||||
|
args=args,
|
||||||
|
minimum_chunks=args.minimum_audio_chunks + 1,
|
||||||
|
)
|
||||||
|
replacement_audio, replacement_response_ids, replacement_stats = _session_result(
|
||||||
|
replacement,
|
||||||
|
input_frames=replacement_frames,
|
||||||
|
args=args,
|
||||||
|
minimum_chunks=1,
|
||||||
|
)
|
||||||
|
if secondary_after_ids != secondary_response_ids:
|
||||||
|
raise AssertionError("survivor response changed while another session slot was recycled")
|
||||||
|
if replacement_response_ids & (primary_response_ids | secondary_response_ids):
|
||||||
|
raise AssertionError("replacement session reused another session's response id")
|
||||||
|
_save(output_dir, "secondary", secondary, secondary_audio)
|
||||||
|
_save(output_dir, "replacement", replacement, replacement_audio)
|
||||||
|
await asyncio.gather(
|
||||||
|
_close_session(secondary, timeout_s=args.timeout_s),
|
||||||
|
_close_session(replacement, timeout_s=args.timeout_s),
|
||||||
|
)
|
||||||
|
|
||||||
|
errors = primary.events.errors() + secondary.events.errors() + replacement.events.errors()
|
||||||
|
result = {
|
||||||
|
"ok": not errors,
|
||||||
|
"model": args.model,
|
||||||
|
"input": {**input_identity, "tail_s": args.tail_s},
|
||||||
|
"capabilities": capabilities,
|
||||||
|
"overflow_error_code": overflow_code,
|
||||||
|
"primary": {"session_id": ids["primary"], "input_frames": primary_frames, **primary_stats},
|
||||||
|
"secondary": {
|
||||||
|
"session_id": ids["secondary"],
|
||||||
|
"input_frames": secondary_total_frames,
|
||||||
|
"continuation_frames_after_recycle": continuation_frames,
|
||||||
|
**secondary_stats,
|
||||||
|
},
|
||||||
|
"replacement": {
|
||||||
|
"session_id": ids["replacement"],
|
||||||
|
"input_frames": replacement_frames,
|
||||||
|
**replacement_stats,
|
||||||
|
},
|
||||||
|
"errors": errors,
|
||||||
|
"output_dir": str(output_dir),
|
||||||
|
}
|
||||||
|
(output_dir / "result.json").write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||||
|
if errors:
|
||||||
|
raise AssertionError(f"Realtime sessions emitted errors: {errors}")
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
|
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
||||||
|
parser = argparse.ArgumentParser()
|
||||||
|
parser.add_argument("--url", default="ws://127.0.0.1:8099/v1/realtime?duplex=1")
|
||||||
|
parser.add_argument("--model", required=True)
|
||||||
|
parser.add_argument("--input-wav", required=True)
|
||||||
|
parser.add_argument(
|
||||||
|
"--expected-input-sha256",
|
||||||
|
help="Fail before opening sessions when the input WAV does not match this SHA-256.",
|
||||||
|
)
|
||||||
|
parser.add_argument("--output-dir", default="/tmp/personaplex-realtime-duplex")
|
||||||
|
parser.add_argument("--voice", default="NATF2.pt")
|
||||||
|
parser.add_argument("--persona", default="You are a concise and helpful assistant.")
|
||||||
|
parser.add_argument("--secondary-persona", default="You are a calm and precise assistant.")
|
||||||
|
parser.add_argument("--replacement-persona", default="You are a cheerful travel assistant.")
|
||||||
|
parser.add_argument("--tail-s", type=float, default=2.0)
|
||||||
|
parser.add_argument("--drain-s", type=float, default=2.0)
|
||||||
|
parser.add_argument("--timeout-s", type=float, default=180.0)
|
||||||
|
parser.add_argument(
|
||||||
|
"--max-frame-deficit",
|
||||||
|
type=int,
|
||||||
|
default=4,
|
||||||
|
help=(
|
||||||
|
"Maximum unflushed tail frames. PersonaPlex Code2Wav emits five-frame "
|
||||||
|
"chunks, so a session cutoff can leave at most four pending frames."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
parser.add_argument("--voiced-frame-rms-threshold", type=float, default=1e-3)
|
||||||
|
parser.add_argument(
|
||||||
|
"--min-voiced-frames",
|
||||||
|
type=int,
|
||||||
|
default=5,
|
||||||
|
help=(
|
||||||
|
"Require at least one normal five-frame Code2Wav chunk (400 ms) above "
|
||||||
|
"the per-frame RMS threshold; this detects inaudible output, not speech quality."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
parser.add_argument("--minimum-audio-chunks", type=int, default=2)
|
||||||
|
parser.add_argument(
|
||||||
|
"--replacement-frames",
|
||||||
|
type=int,
|
||||||
|
default=0,
|
||||||
|
help="Frames sent after slot reuse; 0 replays the full input so the replacement must produce audible output.",
|
||||||
|
)
|
||||||
|
return parser.parse_args(argv)
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
print(json.dumps(asyncio.run(run(parse_args())), ensure_ascii=False, indent=2))
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -83,7 +83,11 @@ class FakeDeadLLMStageClient(FakeStageClient):
|
|||||||
|
|
||||||
async def get_output_async(self):
|
async def get_output_async(self):
|
||||||
if not self.add_request_calls:
|
if not self.add_request_calls:
|
||||||
return SimpleNamespace(outputs=[])
|
return SimpleNamespace(
|
||||||
|
outputs=[],
|
||||||
|
scheduler_stats=None,
|
||||||
|
finished_requests=None,
|
||||||
|
)
|
||||||
raise EngineDeadError("Stage-0 engine core is dead")
|
raise EngineDeadError("Stage-0 engine core is dead")
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -175,6 +175,26 @@ def test_wakeup_engine_not_support(sleep_incapable_engine):
|
|||||||
assert "wake_up" in response.json()["detail"]
|
assert "wake_up" in response.json()["detail"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_wakeup_route_propagates_not_implemented_error(sleep_capable_engine, mocker):
|
||||||
|
"""``/v1/omni/wakeup`` must not swallow ``NotImplementedError`` from the engine.
|
||||||
|
|
||||||
|
Route exception-propagation only: the mock authors the exception on purpose.
|
||||||
|
The real level-2 discarded-weight path and live HTTP 501 mapping are covered by
|
||||||
|
``test_wakeup_after_level2_sleep_fails``.
|
||||||
|
"""
|
||||||
|
sleep_capable_engine.wake_up = mocker.AsyncMock(side_effect=NotImplementedError("engine wake_up refused"))
|
||||||
|
app = _make_app(sleep_capable_engine)
|
||||||
|
app.state.sleeping_stages = {0}
|
||||||
|
# Re-raise in-process; default TestClient would wrap this as a 500 response.
|
||||||
|
client = TestClient(app, raise_server_exceptions=True)
|
||||||
|
|
||||||
|
with pytest.raises(NotImplementedError, match="engine wake_up refused"):
|
||||||
|
client.post("/v1/omni/wakeup", json={"stage_ids": [0]})
|
||||||
|
|
||||||
|
# Failed wake must not clear the sleeping set.
|
||||||
|
assert app.state.sleeping_stages == {0}
|
||||||
|
|
||||||
|
|
||||||
def test_wakeup_removes_only_requested_stages(sleep_capable_engine):
|
def test_wakeup_removes_only_requested_stages(sleep_capable_engine):
|
||||||
app = _make_app(sleep_capable_engine)
|
app = _make_app(sleep_capable_engine)
|
||||||
app.state.sleeping_stages = {0, 1, 2}
|
app.state.sleeping_stages = {0, 1, 2}
|
||||||
|
|||||||
@@ -17,7 +17,7 @@ from types import SimpleNamespace
|
|||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import pytest
|
import pytest
|
||||||
from fastapi import FastAPI
|
from fastapi import FastAPI, HTTPException
|
||||||
from fastapi.testclient import TestClient
|
from fastapi.testclient import TestClient
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
from pytest_mock import MockerFixture
|
from pytest_mock import MockerFixture
|
||||||
@@ -62,10 +62,14 @@ class FakeAsyncOmni:
|
|||||||
def __init__(self):
|
def __init__(self):
|
||||||
self.stage_configs = [SimpleNamespace(stage_type="diffusion")]
|
self.stage_configs = [SimpleNamespace(stage_type="diffusion")]
|
||||||
self.default_sampling_params_list = [OmniDiffusionSamplingParams()]
|
self.default_sampling_params_list = [OmniDiffusionSamplingParams()]
|
||||||
|
self.model_class_name = "WanPipeline"
|
||||||
self.captured_prompt = None
|
self.captured_prompt = None
|
||||||
self.captured_reference_video_bytes = None
|
self.captured_reference_video_bytes = None
|
||||||
self.captured_sampling_params_list = None
|
self.captured_sampling_params_list = None
|
||||||
|
|
||||||
|
def get_diffusion_od_config(self):
|
||||||
|
return SimpleNamespace(model_class_name=self.model_class_name)
|
||||||
|
|
||||||
async def generate(self, prompt, request_id, sampling_params_list):
|
async def generate(self, prompt, request_id, sampling_params_list):
|
||||||
self.captured_prompt = prompt
|
self.captured_prompt = prompt
|
||||||
self.captured_sampling_params_list = sampling_params_list
|
self.captured_sampling_params_list = sampling_params_list
|
||||||
@@ -487,7 +491,7 @@ def test_v2v_video_generation_form(test_client, mocker: MockerFixture):
|
|||||||
assert input_video[0].size == (32, 24)
|
assert input_video[0].size == (32, 24)
|
||||||
|
|
||||||
|
|
||||||
def test_v2v_video_generation_with_video_reference_form(test_client, mocker: MockerFixture):
|
def test_r9_typed_video_reference_form(test_client, mocker: MockerFixture):
|
||||||
mocker.patch(
|
mocker.patch(
|
||||||
"vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
|
"vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
|
||||||
return_value=b"fake-video",
|
return_value=b"fake-video",
|
||||||
@@ -506,8 +510,53 @@ def test_v2v_video_generation_with_video_reference_form(test_client, mocker: Moc
|
|||||||
|
|
||||||
engine = test_client.app.state.openai_serving_video._engine_client
|
engine = test_client.app.state.openai_serving_video._engine_client
|
||||||
input_video = engine.captured_prompt["multi_modal_data"]["video"]
|
input_video = engine.captured_prompt["multi_modal_data"]["video"]
|
||||||
assert len(input_video) == 2
|
assert len(input_video) == 1
|
||||||
assert input_video[0].size == (32, 24)
|
assert all(isinstance(item, str) for item in input_video)
|
||||||
|
assert engine.captured_reference_video_bytes is not None
|
||||||
|
assert len(engine.captured_reference_video_bytes) == 1
|
||||||
|
assert b"ftyp" in engine.captured_reference_video_bytes[0][:32]
|
||||||
|
|
||||||
|
|
||||||
|
def test_r5_mixed_image_video_audio_references_and_fanout(test_client, mocker: MockerFixture):
|
||||||
|
mocker.patch(
|
||||||
|
"vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
|
||||||
|
return_value=b"fake-video",
|
||||||
|
)
|
||||||
|
test_client.app.state.openai_serving_video._engine_client.model_class_name = "MiniMaxH3Pipeline"
|
||||||
|
response = test_client.post(
|
||||||
|
"/v1/videos",
|
||||||
|
data={
|
||||||
|
"prompt": "Use both references.",
|
||||||
|
"image_reference": json.dumps(
|
||||||
|
[
|
||||||
|
{"image_url": _make_test_image_data_url((40, 24))},
|
||||||
|
{"image_url": _make_test_image_data_url((40, 24))},
|
||||||
|
]
|
||||||
|
),
|
||||||
|
"video_reference": json.dumps(
|
||||||
|
{"video_url": _make_test_video_data_url((32, 24), 2)},
|
||||||
|
),
|
||||||
|
"audio_reference": json.dumps(
|
||||||
|
[{"audio_url": "data:audio/mp3;base64,ZmFrZQ=="}],
|
||||||
|
),
|
||||||
|
"extra_params": json.dumps({"task": "ref2va", "frame_indices": [0, -1], "start_time_seconds": 1.5}),
|
||||||
|
"num_outputs_per_prompt": "2",
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
video_id = response.json()["id"]
|
||||||
|
_wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
|
||||||
|
engine = test_client.app.state.openai_serving_video._engine_client
|
||||||
|
multi_modal_data = engine.captured_prompt["multi_modal_data"]
|
||||||
|
assert len(multi_modal_data["image"]) == 2
|
||||||
|
assert len(multi_modal_data["video"]) == 1
|
||||||
|
assert all(isinstance(item, str) for item in multi_modal_data["video"])
|
||||||
|
assert engine.captured_reference_video_bytes is not None
|
||||||
|
assert len(engine.captured_reference_video_bytes) == 1
|
||||||
|
assert b"ftyp" in engine.captured_reference_video_bytes[0][:32]
|
||||||
|
assert isinstance(multi_modal_data["audio"], str)
|
||||||
|
assert engine.captured_sampling_params_list[0].num_outputs_per_prompt == 2
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("endpoint", ["/v1/videos", "/v1/videos/sync"])
|
@pytest.mark.parametrize("endpoint", ["/v1/videos", "/v1/videos/sync"])
|
||||||
@@ -520,6 +569,7 @@ def test_multi_video_generation_preserves_uploaded_files_until_generation(
|
|||||||
"vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
|
"vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
|
||||||
return_value=b"fake-video",
|
return_value=b"fake-video",
|
||||||
)
|
)
|
||||||
|
test_client.app.state.openai_serving_video._engine_client.model_class_name = "MiniMaxH3Pipeline"
|
||||||
response = test_client.post(
|
response = test_client.post(
|
||||||
endpoint,
|
endpoint,
|
||||||
data={
|
data={
|
||||||
@@ -527,8 +577,10 @@ def test_multi_video_generation_preserves_uploaded_files_until_generation(
|
|||||||
"extra_params": json.dumps({"task": "ref2va", "duration": 15.0}),
|
"extra_params": json.dumps({"task": "ref2va", "duration": 15.0}),
|
||||||
},
|
},
|
||||||
files=[
|
files=[
|
||||||
|
("input_references", ("subject.png", _make_test_image_bytes((40, 24)), "image/png")),
|
||||||
("input_references", ("subject.mp4", b"subject-video", "video/mp4")),
|
("input_references", ("subject.mp4", b"subject-video", "video/mp4")),
|
||||||
("input_references", ("background.mov", b"background-video", "video/quicktime")),
|
("input_references", ("background.mov", b"background-video", "video/quicktime")),
|
||||||
|
("input_references", ("voice.wav", b"reference-audio", "audio/wav")),
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -544,10 +596,29 @@ def test_multi_video_generation_preserves_uploaded_files_until_generation(
|
|||||||
assert engine.captured_reference_video_bytes == [b"subject-video", b"background-video"]
|
assert engine.captured_reference_video_bytes == [b"subject-video", b"background-video"]
|
||||||
assert len(input_videos) == 2
|
assert len(input_videos) == 2
|
||||||
assert all(not Path(path).exists() for path in input_videos)
|
assert all(not Path(path).exists() for path in input_videos)
|
||||||
|
assert isinstance(engine.captured_prompt["multi_modal_data"]["image"], Image.Image)
|
||||||
|
assert isinstance(engine.captured_prompt["multi_modal_data"]["audio"], str)
|
||||||
assert engine.captured_sampling_params_list[0].extra_args["task"] == "ref2va"
|
assert engine.captured_sampling_params_list[0].extra_args["task"] == "ref2va"
|
||||||
assert engine.captured_sampling_params_list[0].extra_args["duration"] == 15.0
|
assert engine.captured_sampling_params_list[0].extra_args["duration"] == 15.0
|
||||||
|
|
||||||
|
|
||||||
|
def test_decode_video_bytes_can_keep_first_frames():
|
||||||
|
from vllm_omni.entrypoints.openai.video_api_utils import _decode_video_bytes
|
||||||
|
|
||||||
|
frames = _decode_video_bytes(
|
||||||
|
_make_test_video_bytes((32, 24), num_frames=6),
|
||||||
|
source="input_reference",
|
||||||
|
max_frames=2,
|
||||||
|
keep="first",
|
||||||
|
)
|
||||||
|
|
||||||
|
assert len(frames) == 2
|
||||||
|
assert frames.fps == pytest.approx(8.0)
|
||||||
|
red_means = [np.asarray(frame)[:, :, 0].mean() for frame in frames]
|
||||||
|
assert red_means[0] < red_means[1]
|
||||||
|
assert red_means[1] < 100
|
||||||
|
|
||||||
|
|
||||||
def test_decode_video_bytes_can_keep_last_frames():
|
def test_decode_video_bytes_can_keep_last_frames():
|
||||||
from vllm_omni.entrypoints.openai.video_api_utils import _decode_video_bytes
|
from vllm_omni.entrypoints.openai.video_api_utils import _decode_video_bytes
|
||||||
|
|
||||||
@@ -1183,17 +1254,67 @@ def test_rejects_input_reference_and_image_reference_together(test_client):
|
|||||||
assert "only one of input_reference, image_reference, or video_reference" in response.json()["detail"].lower()
|
assert "only one of input_reference, image_reference, or video_reference" in response.json()["detail"].lower()
|
||||||
|
|
||||||
|
|
||||||
def test_rejects_image_reference_and_video_reference_together(test_client):
|
def test_r10_typed_image_and_video_reference_form(test_client, mocker: MockerFixture):
|
||||||
|
mocker.patch(
|
||||||
|
"vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
|
||||||
|
return_value=b"fake-video",
|
||||||
|
)
|
||||||
|
test_client.app.state.openai_serving_video._engine_client.model_class_name = "MiniMaxH3Pipeline"
|
||||||
response = test_client.post(
|
response = test_client.post(
|
||||||
"/v1/videos",
|
"/v1/videos",
|
||||||
data={
|
data={
|
||||||
"prompt": "bad refs",
|
"prompt": "use both references",
|
||||||
"image_reference": '{"image_url": "https://example.com/cat.png"}',
|
"image_reference": json.dumps({"image_url": _make_test_image_data_url((40, 24))}),
|
||||||
"video_reference": '{"video_url": "https://example.com/cat.mp4"}',
|
"video_reference": json.dumps({"video_url": _make_test_video_data_url((32, 24), 2)}),
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
assert response.status_code == 200
|
||||||
|
_wait_for_status(test_client, response.json()["id"], VideoGenerationStatus.COMPLETED.value)
|
||||||
|
|
||||||
|
|
||||||
|
def test_generic_video_model_rejects_mixed_image_and_video_references(test_client):
|
||||||
|
response = test_client.post(
|
||||||
|
"/v1/videos/sync",
|
||||||
|
data={
|
||||||
|
"prompt": "use both references",
|
||||||
|
"image_reference": json.dumps({"image_url": _make_test_image_data_url((40, 24))}),
|
||||||
|
"video_reference": json.dumps({"video_url": _make_test_video_data_url((32, 24), 2)}),
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
assert response.status_code == 400
|
assert response.status_code == 400
|
||||||
assert "only one of input_reference, image_reference, or video_reference" in response.json()["detail"].lower()
|
assert "does not support mixed image and video" in response.json()["detail"].lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_h3_multipart_rejects_bmp_image_reference(test_client):
|
||||||
|
image = Image.new("RGB", (64, 64), color="blue")
|
||||||
|
image_buffer = io.BytesIO()
|
||||||
|
image.save(image_buffer, format="BMP")
|
||||||
|
|
||||||
|
test_client.app.state.openai_serving_video._engine_client.model_class_name = "MiniMaxH3Pipeline"
|
||||||
|
response = test_client.post(
|
||||||
|
"/v1/videos/sync",
|
||||||
|
data={"prompt": "reject unsupported image container", "extra_params": '{"task":"ref2va"}'},
|
||||||
|
files=[("input_references", ("reference.bmp", image_buffer.getvalue(), "image/bmp"))],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response.status_code == 400
|
||||||
|
assert "must use jpg" in response.json()["detail"].lower()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_h3_upload_limit_checks_declared_size_before_read():
|
||||||
|
class OversizedUpload:
|
||||||
|
size = api_server.MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES + 1
|
||||||
|
|
||||||
|
async def read(self, _size):
|
||||||
|
raise AssertionError("the oversized upload must be rejected before reading")
|
||||||
|
|
||||||
|
with pytest.raises(HTTPException, match="size limit"):
|
||||||
|
await api_server._read_upload_limited(
|
||||||
|
OversizedUpload(),
|
||||||
|
max_bytes=api_server.MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def test_invalid_seconds_returns_422(test_client):
|
def test_invalid_seconds_returns_422(test_client):
|
||||||
|
|||||||
@@ -262,6 +262,36 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag():
|
|||||||
assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True
|
assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_serve_cli_forwards_distributed_offload_residency():
|
||||||
|
"""Ensure the two-GPU DLO placement controls reach the diffusion stage."""
|
||||||
|
parser = TrackingArgumentParser()
|
||||||
|
subparsers = parser.add_subparsers(dest="command")
|
||||||
|
OmniServeCommand().subparser_init(subparsers)
|
||||||
|
|
||||||
|
args = parser.parse_args(
|
||||||
|
[
|
||||||
|
"serve",
|
||||||
|
"MiniMaxAI/MiniMax-H3",
|
||||||
|
"--omni",
|
||||||
|
"--enable-distributed-layerwise-offload",
|
||||||
|
"--dlo-no-use-allgather",
|
||||||
|
"--dlo-resident-layers",
|
||||||
|
"20",
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
explicit_kwargs = args.get_explicit_kwargs_dict()
|
||||||
|
stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
|
||||||
|
engine_args = stage_cfg["engine_args"]
|
||||||
|
|
||||||
|
assert args.enable_distributed_layerwise_offload is True
|
||||||
|
assert args.dlo_use_allgather is False
|
||||||
|
assert args.dlo_resident_layers == 20
|
||||||
|
assert engine_args["enable_distributed_layerwise_offload"] is True
|
||||||
|
assert engine_args["dlo_use_allgather"] is False
|
||||||
|
assert engine_args["dlo_resident_layers"] == 20
|
||||||
|
|
||||||
|
|
||||||
def test_serve_cli_accepts_diffusion_compile_controls():
|
def test_serve_cli_accepts_diffusion_compile_controls():
|
||||||
"""Ensure both compile controls reach the diffusion stage."""
|
"""Ensure both compile controls reach the diffusion stage."""
|
||||||
parser = TrackingArgumentParser()
|
parser = TrackingArgumentParser()
|
||||||
|
|||||||
Loaded 100 of 190 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user